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ì)有著緊密的合作。

審核編輯：郭婷