EMBark：改善大規模推薦系統嵌入訓練的分片與通訊 NEWS DETAIL

當深度學習推薦模型（DLRM）的訓練從少量 GPU 擴展至多節點叢集時，嵌入表的跨節點通訊與負載不平衡，可能成為限制整體訓練效率的主要因素。NVIDIA HugeCTR 團隊在 RecSys 2024 展示的 EMBark，以嵌入叢集、彈性的 3D 分片方式與分片規劃器，針對大規模嵌入訓練的通訊和資源配置問題提供一條可評估的技術路徑。相關程式碼與論文已開源；實際可用性與效益仍應依專案模型、硬體及軟體版本完成驗證。

大規模 DLRM 訓練面臨的問題

推薦系統常使用包含數十億至數百億 ID 特徵的 DLRM。NVIDIA Merlin HugeCTR 與 TorchRec 等 GPU 方案，可將大規模 ID 特徵嵌入存放並平行處理於 GPU，運用 GPU 記憶體頻寬處理訓練工作。

不過，GPU 數量由 8 張擴展到 128 張時，嵌入通訊在總訓練成本中的占比會提高。來源所述的部分大規模案例中，例如 16 個節點的情境，嵌入通訊開銷曾超過總訓練開銷的一半，達到 51%。造成此情況的因素包括：每個節點持有的嵌入表減少，可能使節點間工作量失衡；以及節點間頻寬通常低於節點內頻寬，使嵌入模型平行所需的通訊時間增加。

EMBark 的三項核心能力

EMBark 以 NVIDIA Merlin HugeCTR 開源推薦系統框架實作，目標是改善不同叢集配置中的 DLRM 嵌入訓練效能，並加速訓練吞吐量。來源亦指出，其技術概念可套用至其他機器學習框架，但實際整合方式需要另行確認。

• 嵌入叢集：依嵌入的特性分組，並為各群組採取對應的通訊壓縮策略。叢集包含資料分發器、嵌入儲存與嵌入運算子，用於將特徵 ID 轉換為嵌入向量。
• 彈性 3D 分片：以 (i, j, k) 三元組描述分片，其中 i 為表格索引、j 為列分片索引、k 為欄分片索引。相較於固定的逐列、逐表或逐欄方式，單一嵌入可跨任意數量 GPU 分片。
• 分片規劃器：採用以成本為基礎的貪婪搜尋演算法，根據硬體規格及嵌入配置尋找分片策略。

三種嵌入叢集的適用條件

EMBark 將嵌入叢集分為資料平行（DP）、基於歸約（RB）與基於唯一（UB）三類。選擇不應只看單一通訊策略，而應檢視表格大小、輸入特徵模式、池化操作及實際存取分布。

其中，3D 分片主要用於處理 RB 叢集的工作負載不平衡。這代表導入評估應先辨識各嵌入表的大小、熱點、輸入稀疏度與池化行為，再判斷哪些表適合放入不同叢集。

評估與導入路徑

1. 盤點模型中的嵌入表、ID 特徵數量、表格容量與多熱輸入特性，建立可量測的訓練基準。
2. 量測節點內與節點間的通訊時間、GPU 記憶體使用量及各 GPU 工作量，確認瓶頸是否確實集中於嵌入通訊或分片不平衡。
3. 依表格特性比較 DP、RB 與 UB 叢集配置，並將硬體規格與嵌入配置納入分片規劃器的評估。
4. 在代表性的多節點規模測試吞吐量、尾端步驟時間、通訊量與記憶體餘裕，再決定是否擴展至正式訓練環境。

來源的實驗平台為 NVIDIA DGX H100 節點，每節點配置 8 張 NVIDIA H100 GPU、總計 640 GB HBM，節點內以 NVLink 互連，標示為 900 GB/s 雙向頻寬；節點間使用 8x400Gbps InfiniBand。此配置可作為理解測試背景的參考，不能直接推論其他 GPU、網路拓撲、資料集或框架版本的結果。

使用限制與驗證重點

EMBark 處理的是嵌入訓練中的配置與通訊問題，並不表示所有推薦模型都會得到相同幅度的改善。DP 複製表格會受到 GPU 記憶體容量限制；RB 與 UB 的效果則取決於池化、多熱特徵及存取熱點是否符合其設計前提。採購或導入前，應在具日期的官方產品文件中確認 HugeCTR、相關程式碼及依賴環境的支援範圍，並以完整 SKU／BOM、網路拓撲與專案測試確認相容性。

常見問題

EMBark 是否只能用於 NVIDIA Merlin HugeCTR？

來源指出 EMBark 以 NVIDIA Merlin HugeCTR 開源推薦系統框架實作，但其技術可應用於其他機器學習框架。其他框架是否已有可直接使用的實作、API 相容性及維運方式，應查核對應版本的官方文件與開源專案內容。

只要增加 GPU 數量，就需要採用 3D 分片嗎？

不一定。3D 分片用來處理 RB 叢集中的工作負載不平衡。若模型的嵌入表配置、通訊比例或 GPU 負載沒有形成相同瓶頸，額外的分片複雜度未必帶來相應效益。應先以實測資料比較固定分片與候選配置。

結論

EMBark 為大規模 DLRM 嵌入訓練提供了以叢集分類、3D 分片與成本式規劃協同調整的方案。其價值在於讓團隊能依嵌入表與叢集條件細分通訊策略；最終選型仍應以實際模型特徵、硬體互連與多節點測試結果為準。

圍繞「EMBark：改善大規模推薦系統嵌入訓練的分片與通訊」繼續瞭解 採購與選型問答。