Embedding 在深度學習推薦模型中起着關鍵作用。它們被用于将輸入資料中的離散特征映射到向量，以便下遊的神經網絡進行處理。Embedding 通常構成深度學習推薦模型中的大部分參數，大小可以達到 TB 級。在訓練期間，很難将它們放入單個 GPU 的記憶體中。是以，現代推薦系統可能需要模型并行和資料并行的分布式訓練方法組合，以最佳利用 GPU 計算資源來實作最好的訓練性能。
NVIDIA Merlin Distributed-Embeddings，可以友善TensorFlow 2 使用者用短短幾行代碼輕松完成大規模的推薦模型訓練。
背景
在資料并行分布式訓練中，整個模型被複制到每個 GPU 上。在訓練過程中，一批輸入資料在多個 GPU 中分割，每張卡獨立處理其自己的資料分片，進而允許計算擴充到更大批量的資料。在反向傳播期間，計算的梯度通過 reduction 算子（例如， horovod.tensorflow.allreduce ) 來同步更新多個 GPU 間的參數。
另一方面，模型并行分布式訓練中，模型參數被分割到多個 GPU 上。這種方法更适合分布存儲大型 embedding。訓練中，每個 GPU 通過 alltoall 通信算子（例如， horovod.tensorflow.alltoall) 通路不在本機中的參數。
在之前的相關文章中， 用 TensorFlow 2 在 DGX A100 上訓練 100B + 參數的推薦系統 ， Tomasz 讨論了如何将 1130 億參數的 DLRM 模型中的 embedding 分布到多個 NVIDIA GPU 進行訓練，并相比純 CPU 的方案實作 672 倍的性能提升。這一重大突破可以将訓練時間從幾天縮短到幾分鐘！這是通過模型并行 embedding 層和資料并行 MLP 層來實作的。和 CPU 方案相比，這種混合并行的方法能夠有效利用 GPU 的高記憶體帶寬加速記憶體受限的 embedding 查找，并同時利用多個 GPU 的算力加速 MLP 層。作為參考， NVIDIA A100-80GB GPU 具有超過 2 TB / s 的帶寬和 80 GB HBM2 存儲）。
圖 1. 用于訓練大型推薦系統的通用“混合并行”方法
embedding 表可以按表為分割機關（圖中表 0 和 N ），按“列”分割（圖中表 2），或者按”行”分割。MLP 層跨所有 GPU 複制，而數字特征則可以直接輸入 MLP 層。
然而，實作這種複雜的混合并行訓練方法并不簡單，需要領域内專家設計數百行底層代碼來開發和優化。為了使其更普适，NVIDIA Merlin Distributed-Embeddings 提供了一些易于使用的 TensorFlow 2 的封裝，讓所有人都隻需三行 Python 代碼即可輕松實作模型并行。它提供了一些涵蓋并拓展原生 TensorFlow 功能的高性能 embedding 查找算子。在此基礎上，它提供了一個可規模化的模型并行封裝函數，幫助使用者自動将 embedding 分布于多個 GPU 上。下面将展示它如何實作混合并行。
分布式模型并行
NVIDIA Merlin Distributed-Embeddings 提供了
distributed_embeddings.dist_model_parallel 子產品。它有助于在多個 GPU 之間分布embedding而無需任何複雜的代碼來處理跨GPU間的通信（如 all2all ）。下面的代碼示例顯示了此 API 的用法：
要使用 Horovod 以資料并行方式運作 MLP 層，請将 Horovod的 Distributed GradientTape 和 broadcast 方法替換成 NVIDIA Merlin Distributed-Embeddings 裡同等的 API。以下示例直接取自 Horovod 文檔，并進行了相對應修改。
通過這些微小的改變，您就可以使用混合并行訓練了！
我們還提供了以下完整示例： 使用 Criteo 1TB 點選日志資料訓練 DLRM 模型以及擴充到 22.8 TiB 的合成資料模型。
性能
為了展示 NVIDIA Merlin Distributed-Embeddings 的性能，我們在 Criteo 1TB 資料集 DLRM 模型和最高達到 3 TiB embedding 的合成模型上進行了模型訓練的基準測試。
Criteo 資料集上的 DLRM 模型基準測試
測試表明，我們使用更簡單的 API 取得了近似于專家代碼的性能。NVIDIA 深度學習 DLRM TensorFlow 2 示例代碼現已更新為使用 NVIDIA Merlin Distributed-Embeddings 進行分布式混合并行訓練，更多資訊請參閱我們之前的文章， 用 TensorFlow 2 在 DGX A100 上訓練 100B + 參數的推薦系統。README 中的基準測試部分提供了對性能結果的更多詳述。
我們對 1130 億個參數（ 421 個 GiB 大小）的 DLRM 模型在 Criteo TB 點選日志資料集上用三種不同的硬體設定進行了訓練：
- 僅 CPU 的解決方案。
- 單 GPU 解決方案，其中 CPU 記憶體用于存儲最大的 embedding 表。
- 使用 NVIDIA DGX A100-80GB 的 8 GPU 的混合并行解決方案。此方案利用了 NVIDIA Merlin Distributed-Embeddings 裡提供的模型并行 api 和 embedding API 。
表 1. 各種設定的教育訓練吞吐量和加速
我們觀察到， DGX-A100 上的 NVIDIA Merlin Distributed-Embeddings 方案比僅使用 CPU 的解決方案提供了驚人的 683 倍的加速！我們還注意到與單 GPU 方案相比，混合并行的性能也有顯著提升。這是因為在 GPU 顯存中存儲所有 embedding 避免了通過 CPU-GPU 接口查找 embedding 的開銷。
合成模型基準測試
為了進一步示範方案的可規模化，我們建立了不同大小的合成資料以及對應的 DLRM 模型（表 2 ）。有關模型生成方法和訓練腳本的更多資訊，請參見 GitHub NVIDIA-Merlin/distributed-embeddings 代碼庫。
表 2. 合成模型尺寸
每個合成模型使用一個或多個 DGX-A100-80GB 節點進行訓練，全局資料 batch 大小為 65536 ，并使用 Adagrad 優化器。從表 3 中可以看出， NVIDIA Merlin Distributed-Embeddings 可以在數百個 GPU 上輕松訓練 TB 級模型。
表 3. 各種硬體配置下合成模型的訓練步長時間（ ms ）
另一方面，與傳統的資料并行相比，即使對于可以容納在單個 GPU 中的模型，多 GPU 分布式模型并行仍然提供了顯著加速。表 4 顯示了上述 Tiny 模型在 DGX A100-80GB 上的性能對比。
表 4. Tiny模型（ 4.2GiB ）的訓練步長時間（ ms ）比較 NVIDIA Merlin Distributed-Embeddings 模型并行和原生 TensorFlow 資料并行
本實驗使用了 65536 的全局批量和 Adagrad 優化器。
結論
在這篇文章中，我們介紹了 NVIDIA Merlin Distributed-Embeddings，僅需幾行代碼即可在 NVIDIA GPU 上實作基于 embedding 的深度學習模型，并進行可規模化，高效率地模型并行訓練。歡迎嘗試以下使用合成資料的可擴充訓練[1]示例和基于 Criteo 資料訓練 DLRM 模型示例[2]。
[1]
https://github.com/NVIDIA-Merlin/distributed-embeddings/tree/main/examples/benchmarks/synthetic_models
[2] https://github.com/NVIDIA/DeepLearningExamples/tree/master/TensorFlow2/Recommendation/DLRM#quick-start-guide
本文作者
Shashank Verma
Shashank Verma 是 NVIDIA 的一名深入學習的技術營銷工程師。他負責開發和展示各種深度學習架構中以開發人員為中心的内容。他從威斯康星大學麥迪遜分校獲得電氣工程碩士學位，在那裡他專注于計算機視覺、資料科學的安全方面和 HPC 。
高文雯
高文雯是 NVIDIA Merlin 的進階産品經理，擁有 Amazon 和其他技術公司的産品管理經驗，專注于個性化和推薦。她擁有多倫多大學計算機科學學士學位和麻省理工學院斯隆管理學院工商管理碩士學位。
吳昊
吳昊是 NVIDIA 的進階 GPU 計算架構師。他在中國科學院完成博士學位後于 2011 年加入 NVIDIA 計算架構組。近年來， 吳昊的技術重點是将低精度應用于深度神經網絡訓練和推理。
傅德禹
傅德禹是 NVIDIA 深度學習架構團隊的進階開發技術工程師，負責加速軟體堆棧 CUDA 核心、數學、通信、架構和模型的 DL 教育訓練工作。他最近專注于 NVIDIA Merlin Distributed-Embeddings 和推薦系統。
Tomasz Grel
Tomasz Grel 是一名深度學習工程師。在 NVIDIA ，他專注于確定衆多推薦系統的品質和執行速度，包括 NCF 、 VAE-CF 和 DLRM 。