推薦系統是互聯網行業的核心，而高效地訓練這些系統對於各大公司來說是一個關鍵問題。大多數推薦系統是深度學習推薦模型（DLRMs），包含數十億甚至數百億個 ID 特徵。圖 1 示出了一個典型的結構。
近年來， NVIDIA Merlin HugeCTR 和 TorchRec 等 GPU 解決方案通過在 GPU 上儲存大規模 ID 特徵嵌入並對其進行並行處理，顯著加速了 DLRM 的訓練。與 CPU 解決方案相比，使用 GPU 內存頻寬可實現顯著改進。
與此同時，隨著訓練集群中使用的 GPU 數量增加（從 8 個 GPU 增加到 128 個 GPU），我們發現嵌入的通信開銷在總訓練開銷中佔較大比例。在一些大規模訓練場景中（例如在 16 個節點上），它甚至超過了一半（51%）。
這主要有兩個原因：
1、隨著集群中 GPU 數量的增加，每個節點上的嵌入表數量逐漸減少，導致不同節點之間的負載不平衡，並降低訓練效率。
2. 與節點內頻寬相比，節點間頻寬要小得多，因此嵌入模型並行的通信時間更長。
為了幫助行業用戶更好地理解和解決這些問題，NVIDIA HugeCTR 團隊在 RecSys 2024 上展示了 EMBark。EMBark 支援 3D 靈活的分片策略，並結合通信壓縮策略來微調大規模集群深度推薦模型訓練中的負載不平衡問題，並減少嵌入所需的通信時間。 相關代碼 和 論文 均為開源。
EMBark 簡介
EMBark 可提高不同集群配置下 DLRM 訓練中嵌入的性能，並加速訓練吞吐量。它使用 NVIDIA Merlin HugeCTR 開源推薦系統框架實施，但這些技術也可以應用於其他機器學習框架。
EMBark 包含三個關鍵組件：嵌入集群、靈活的 3D 分片方案和分片規劃器。圖 3 顯示了 EMBark 的整體架構。
嵌入集群
嵌入集群可高效訓練嵌入，方法是將嵌入按類似功能分組，並為每個集群應用自定義壓縮策略。這些集群包括資料分布器、嵌入儲存和嵌入運算符，它們共同工作，將特徵 ID 轉換為嵌入向量。
嵌入集群有三種類型：資料並行（DP）、基于歸約（RB）和基於唯一（UB）。每種方法在訓練期間都使用不同的通信方法，並且適用於不同的嵌入。
DP 集群： DP 集群不會壓縮通信，使通信變得簡單高效，但它們會在每個 GPU 上複製嵌入表，因此僅適用於小型表。
RB 集群 ：使用歸約運算，對於使用池化運算壓縮多熱輸入表格具有顯著效果。
UB 集群： 僅發送唯一向量，對於處理具有明顯訪問熱點的嵌入表是有益的。
靈活的 3D 分片方案
靈活的 3D 分片方案可以解決 RB 集群中的工作負載不平衡問題。與固定分片策略（如逐行、逐表或逐列）不同，EMBark 使用 3D 元組 (i, j, k) 表示每個分片，其中 i 表示表格索引，j 表示行分片索引，k 表示列分片索引。這種方法允許每個嵌入跨任意數量的 GPU 進行分片，從而提供對工作負載平衡的靈活性和精確控制。
分片規劃器
為了找到最佳的分片策略，EMBark 提供了一個分片規劃器，這是一種基於成本的貪婪搜索算法，可以根據硬體規格和嵌入配置確定最佳的分片策略。
評估
利用嵌入集群、靈活的 3D 分片方案和高級分片規劃器，EMBark 可顯著提升訓練性能。為了評估其有效性，我們在一組 NVIDIA DGX H100 節點上進行了實驗，每個節點配備 8 個 NVIDIA H100 GPU (總計 640 GB HBM，每個節點的頻寬為 24TB/s)。
在每個節點內，所有 GPU 都通過 NVLink（900 GB/s 雙向）互聯，節點間通信使用 InfiniBand（8x400Gbps）。
WeChat
Profile