簡介
推薦系統是互聯網行業的核心系統，如何高效訓練推薦系統是各公司關注的核心問題。目前，推薦系統基本上都是基于深度學習的大規模 ID 類模型，模型包含數十億甚至數百億級別的 ID 特征，典型結構如圖 1 所示。
圖1. 典型 DLRM 模型結構圖
近年來，以 NVIDIA Merlin HugeCTR 和 TorchRec 為代表的 GPU 解決方案，通過將大規模 ID 類特征的 embedding 存放在 GPU 上，并對 embedding 進行模型并行處理，將其分片到不同 GPU 上，利用 GPU 內存帶寬優勢，大幅加速了深度推薦系統模型的訓練，相較于 CPU 方案有顯著提升。
同時，隨著訓練集群 GPU 使用數量增加（從 8 個 GPU 增加到 128 個 GPU），我們也發現，embedding 部分通信開銷占整個訓練開銷比例越來越大。在一些大規模訓練中（比如在 16 節點上）甚至超過一半（51%）。這主要是因為兩個原因：
- 隨著集群 GPU 數量增加，每個節點上的 embedding table 數量逐漸減少，導致不同節點負載不均衡，降低訓練效率。
- 相比機內帶寬，機間帶寬小得多，因此 embedding 模型并行需要進行機間通信耗時較長。
為了幫助行業用戶更好地理解問題、解決問題，NVIDIA HugeCTR 團隊于今年的 RecSys 大會上提出了 EMBark，通過支持 3D 的自定義 sharding 策略和結合不同的通信壓縮策略，能夠細粒度的優化大規模集群下深度推薦模型訓練的負載不均衡問題以及減少 embedding 需要的通信時間，其相關代碼和論文皆已開源。
圖2. 不同 cluster 配置下 DLRM 各部分訓練耗時占比
EMBark 介紹
EMBark 旨在提高 DLRM 訓練中 embedding 在不同集群配置下的性能，并加速整體訓練吞吐量。EMBark是在NVIDIA Merlin HugeCTR開源推薦系統框架的基礎上實現的，但所描述的技術也可以應用于其他機器學習框架。
EMBark有三個關鍵組件：embedding clusers、靈活的3D分片方案和分片規劃器。下圖展示了EMBark的整體架構。
圖3. EMBark 架構圖
Embedding Clusters
Embedding clusters 旨在通過將具有相似特征的 embedding 進行分組并為每個 cluster 應用定制的壓縮策略來高效地訓練 embedding。每個 cluster 包括 data distributor、embedding storage 和 embedding operators，協同將 feature ID 轉換為 embedding 向量。
有三種類型的 Embedding clusters：Data-parallel(DP)、Reduction-based(Reduction based) 和基于Unique-based(Unique Based)。每種類型在訓練過程中采用不同的通信方法，適用于不同的 embedding。
- DP cluster 不壓縮通信，因此簡單高效，但是因為會將 embedding table 在每個 GPU 上重復，因此僅適用于小表格。
- RB cluster使用歸約操作，對于具有池化操作的多 feature 輸入表格壓縮效果顯著。
- UB cluster僅發送唯一向量，有利于處理具有明顯訪問熱點的 embedding table。
| Embedding clusters | Target |
| Data-parallel (DP) cluster | Small tables |
| Reduction-based (RB) cluster | High-hotness tables with pooling |
| Unique-based (UB) cluster | Skewed tables |
靈活的 3D 分片方案
靈活的 3D 分片方案旨在解決 RB cluster 中的工作負載不平衡問題。與固定的分片策略比如 row-wise、table-wise、column-wise 不同，EMBark 使用一個 3D 元組（i, j, k）表示每個分片，其中 I 表示表格索引，j 表示行分片索引，k 表示列分片索引。這種方法允許每個 embedding 跨任意數量的 GPU 進行分片，提供靈活性并實現對工作負載平衡的精確控制。
分片規劃器
為了找到最佳分片策略，EMBark 提供了一個分片規劃器——一種成本驅動的貪婪搜索算法，根據硬件規格和 embedding 配置識別最佳分片策略。
Evaluation
所有實驗均在一個集群上進行，該集群由 NVIDIA DGX-H100[1] 節點組成，每個節點配備 8 張 NVIDIA H100 GPU（總計 640GB HBM，帶寬為每節點 24TB/s）。在每個節點內，所有 GPU 通過 NVLink（雙向 900GB/s）互連。節點間通信使用 InfiniBand（8x400Gbps）。
為了展示 EMBark 可以高效訓練任何規模的 DLRM 模型，我們測試了使用 MLPerf DLRM-DCNv2 模型并生成了幾個具有更大嵌入表和不同屬性的合成模型（參見上表）。我們的訓練數據集表現出 α=1.2 的冪律偏斜。
圖4. EMBark evaluation 結果
Baseline 采用串行的 kernel 執行順序，固定的 table-row-wise 分片策略，以及全部使用了 RB-clusters。實驗依次使用了三種優化：overlap、更靈活的分片策略和更好的 clusters 配置。
在四種代表性 DLRM 變體（DLRM-DCNv2、T180、T200 和 T510）中，EMBark 實現了平均 1.5 倍的端到端訓練吞吐量加速，比 baseline 快最多 1.77 倍。更詳細的實驗結果和相關的分析，可以參考論文。
結論
EMBark 針對大規模推薦系統模型訓練中 embedding 部分耗時占比過高的問題，通過支持 3D 的自定義 sharding 策略和結合不同的通信壓縮策略，能夠細粒度的優化大規模集群下深度推薦模型訓練的負載不均衡問題以及減少 embedding 需要的通信時間，提高大規模推薦系統模型的訓練效率，在四種代表性 DLRM 變體（DLRM-DCNv2、T180、T200 和 T510）中，EMBark 實現了平均 1.5 倍的端到端訓練吞吐量加速，比 baseline 快最多 1.77 倍。其中，相關代碼和論文皆已開源，希望我們的工作對大家有所幫助。同時，我們也在積極探索 embedding offloading 相關技術和進行 torchrec 相關優化工作，未來也會及時和大家更新相關進展情況，如果您對這部分感興趣，也可以聯系我們，大家一起交流和探索。
[1] 此為 NVIDIA 技術團隊所執行的測試環境，僅供學術論文之相關技術討論。
