1. Introduction

HugeCTR 能夠高效地利用 GPU 來(lái)進(jìn)行推薦系統(tǒng)的訓(xùn)練，為了使它還能直接被其他 DL 用戶，比如 TensorFlow 所直接使用，我們開(kāi)發(fā)了 SparseOperationKit （SOK），來(lái)將 HugeCTR 中的高級(jí)特性封裝為 TensorFlow 可直接調(diào)用的形式，從而幫助用戶在 TensorFlow 中直接使用 HugeCTR 中的高級(jí)特性來(lái)加速他們的推薦系統(tǒng)。

圖 1. SOK embedding 工作流程

SOK 以數(shù)據(jù)并行的方式接收輸入數(shù)據(jù)，然后在 SOK 內(nèi)部做黑盒式地模型轉(zhuǎn)換，最后將計(jì)算結(jié)果以數(shù)據(jù)并行的方式傳遞給初始 GPU。這種方式可以盡可能少地修改用戶已有的代碼，以更方便、快捷地在多個(gè) GPU 上進(jìn)行擴(kuò)展。

SOK 不僅僅是加速了 TensorFlow 中的算子，而是根據(jù)業(yè)界中的實(shí)際需求提供了對(duì)應(yīng)的新解決方案，比如說(shuō) GPU HashTable。SOK 可以與 TensorFlow 1.15 和 TensorFlow 2.x 兼容使用；既可以使用 TensorFlow 自帶的通信工具，也可以使用 Horovod 等第三方插件來(lái)作為 embedding parameters 以為參數(shù)的通信工具。

2. TF2 Comparison/Performance

使用 MLPerf 的標(biāo)準(zhǔn)模型 DLRM 來(lái)對(duì) SOK 的性能進(jìn)行測(cè)試。

圖 2. SOK性能測(cè)試數(shù)據(jù)

相比于 NVIDIA 的 DeepLearning Examples，使用 SOK 可以獲得更快的訓(xùn)練速度以及更高的吞吐量。

3. API

SOK 提供了簡(jiǎn)潔的、類 TensorFlow 的 API；使用 SOK 的方式非常簡(jiǎn)單、直接；讓用戶通過(guò)修改幾行代碼就可以使用 SOK。

定義模型結(jié)構(gòu)

左側(cè)是使用 TensorFlow 的 API 來(lái)搭建模型，右側(cè)是使用 SOK 的 API 來(lái)搭建相同的模型。使用 SOK 來(lái)搭建模型的時(shí)候，只需要將 TensorFlow 中的 Embedding Layer 替換為 SOK 對(duì)應(yīng)的 API 即可。

使用 Horovod 來(lái)定義 training loop

同樣的，左側(cè)是使用 TensorFlow 來(lái)定義 training loop，右側(cè)是使用 SOK 時(shí)，training loop 的定義方式?？梢钥吹剑褂?SOK 時(shí)，只需要對(duì) Embedding Variables 和 Dense Variables 進(jìn)行分別處理即可。其中，Embedding Variables 部分由 SOK 管理，Dense Variables 由 TensorFlow 管理。

使用 tf.distribute.MirroredStrategy 來(lái)定義 training loop

類似的，還可以使用 TensorFlow 自帶的通信工具來(lái)定義 training loop。

開(kāi)始訓(xùn)練

在開(kāi)始訓(xùn)練過(guò)程時(shí)，使用 SOK 與使用 TensorFlow 時(shí)所用代碼完全一致。

4. 結(jié)語(yǔ)

SOK 將 HugeCTR 中的高級(jí)特性包裝為 TensorFlow 可以直接使用的模塊，通過(guò)修改少數(shù)幾行代碼即可在已有模型代碼中利用上 HugeCTR 的的先進(jìn)設(shè)計(jì)。

關(guān)于作者

NVIDIA GPU 計(jì)算專家研發(fā)團(tuán)隊(duì)成立于2020年7月，主要面向全球用戶開(kāi)發(fā)基于英偉達(dá)通用計(jì)算產(chǎn)品的軟件解決方案。 目前團(tuán)隊(duì)致力于推薦系統(tǒng)訓(xùn)練框架 Merlin HugeCTR的開(kāi)發(fā)。作為Merlin 系統(tǒng)三大核心軟件之一，Merlin HugeCTR 是一套針對(duì)于推薦系統(tǒng)設(shè)計(jì)的模型訓(xùn)練框架。項(xiàng)目啟動(dòng)于2018年，目前產(chǎn)品應(yīng)用于多家國(guó)際互聯(lián)網(wǎng)公司，自2020年以來(lái)，HugeCTR持續(xù)參與并推動(dòng)國(guó)際性競(jìng)賽MLPerf DLRM訓(xùn)練項(xiàng)目，從而實(shí)現(xiàn)了業(yè)界領(lǐng)先的優(yōu)化性能。 NVIDIA GPU 計(jì)算專家研發(fā)團(tuán)隊(duì)是一個(gè)小型國(guó)際化工程團(tuán)隊(duì)，核心成員來(lái)自于多個(gè)國(guó)家（中國(guó)，韓國(guó)，荷蘭，德國(guó)），客戶分布于世界各地。日常工作使用英語(yǔ)溝通。 團(tuán)隊(duì)主要工作圍繞軟件產(chǎn)品開(kāi)發(fā)展開(kāi)，涉及產(chǎn)品維護(hù)、文檔、測(cè)試、發(fā)布、部署，日常工作中與各國(guó)硬件、產(chǎn)品、工程、客戶支持團(tuán)隊(duì)有著緊密的合作。

審核編輯：郭婷