摘要
1. Introduction
HugeCTR 能夠高效地利用 GPU 來進行推薦系統的訓練,為了使它還能直接被其他 DL 用戶,比如 TensorFlow 所直接使用,我們開發了 SparseOperationKit (SOK),來將 HugeCTR 中的高級特性封裝為 TensorFlow 可直接調用的形式,從而幫助用戶在 TensorFlow 中直接使用 HugeCTR 中的高級特性來加速他們的推薦系統。

圖 1. SOK embedding 工作流程
SOK 以數據并行的方式接收輸入數據,然后在 SOK 內部做黑盒式地模型轉換,最后將計算結果以數據并行的方式傳遞給初始 GPU。這種方式可以盡可能少地修改用戶已有的代碼,以更方便、快捷地在多個 GPU 上進行擴展。
SOK 不僅僅是加速了 TensorFlow 中的算子,而是根據業界中的實際需求提供了對應的新解決方案,比如說 GPU HashTable。SOK 可以與 TensorFlow 1.15 和 TensorFlow 2.x 兼容使用;既可以使用 TensorFlow 自帶的通信工具,也可以使用 Horovod 等第三方插件來作為 embedding parameters 以為參數的通信工具。
2. TF2 Comparison/Performance
使用 MLPerf 的標準模型 DLRM 來對 SOK 的性能進行測試。

圖 2. SOK性能測試數據
相比于 NVIDIA 的 DeepLearning Examples,使用 SOK 可以獲得更快的訓練速度以及更高的吞吐量。
3. API
SOK 提供了簡潔的、類 TensorFlow 的 API;使用 SOK 的方式非常簡單、直接;讓用戶通過修改幾行代碼就可以使用 SOK。
- 定義模型結構

左側是使用 TensorFlow 的 API 來搭建模型,右側是使用 SOK 的 API 來搭建相同的模型。使用 SOK 來搭建模型的時候,只需要將 TensorFlow 中的 Embedding Layer 替換為 SOK 對應的 API 即可。
- 使用 Horovod 來定義 training loop

同樣的,左側是使用 TensorFlow 來定義 training loop,右側是使用 SOK 時,training loop 的定義方式。可以看到,使用 SOK 時,只需要對 Embedding Variables 和 Dense Variables 進行分別處理即可。其中,Embedding Variables 部分由 SOK 管理,Dense Variables 由 TensorFlow 管理。
- 使用 tf.distribute.MirroredStrategy 來定義 training loop

類似的,還可以使用 TensorFlow 自帶的通信工具來定義 training loop。
- 開始訓練

在開始訓練過程時,使用 SOK 與使用 TensorFlow 時所用代碼完全一致。
4. 結語
SOK 將 HugeCTR 中的高級特性包裝為 TensorFlow 可以直接使用的模塊,通過修改少數幾行代碼即可在已有模型代碼中利用上 HugeCTR 的的先進設計。
更多信息,請參考 SOK 官方文檔:
以下是 HugeCTR 的 Github repo 以及其他發布的文章,歡迎感興趣的朋友閱讀和反饋。
Github: https://github.com/NVIDIA-Merlin/HugeCTR (更多文章詳見 README)
除此之外,NVIDIA Merlin HugeCTR 團隊正在積極招募 C++ 以及 CUDA 工程師,歡迎各位有意向的同學發郵件至 jershi@nvidia.com 或 yincanw@nvidia.com 積極申請!JD 請點擊:NVIDIA GPU 計算專家研發團隊招新啦!你將成為 Merlin 應用框架的核心成員!查看詳細介紹!。