1. Introduction

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

圖 1. SOK embedding 工作流程

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

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

2. TF2 Comparison/Performance

使用 MLPerf 的標(biāo)準(zhǔn)模型 DLRM 來對 SOK 的性能進行測試。

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

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

3. API

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

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

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

使用 Horovod 來定義 training loop

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

使用 tf.distribute.MirroredStrategy 來定義 training loop

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

開始訓(xùn)練

在開始訓(xùn)練過程時，使用 SOK 與使用 TensorFlow 時所用代碼完全一致。

4. 結(jié)語

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

關(guān)于作者

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

審核編輯：郭婷