簡介

推薦系統(tǒng)是互聯(lián)網(wǎng)行業(yè)的核心系統(tǒng)，如何高效訓(xùn)練推薦系統(tǒng)是各公司關(guān)注的核心問題。目前，推薦系統(tǒng)基本上都是基于深度學(xué)習(xí)的大規(guī)模 ID 類模型，模型包含數(shù)十億甚至數(shù)百億級別的 ID 特征，典型結(jié)構(gòu)如圖 1 所示。

圖 1. 典型 DLRM 模型結(jié)構(gòu)圖

近年來，以 NVIDIA Merlin HugeCTR 和 TorchRec 為代表的 GPU 解決方案，通過將大規(guī)模 ID 類特征的 embedding 存放在 GPU 上，并對 embedding 進(jìn)行模型并行處理，將其分片到不同 GPU 上，利用 GPU 內(nèi)存帶寬優(yōu)勢，大幅加速了深度推薦系統(tǒng)模型的訓(xùn)練，相較于 CPU 方案有顯著提升。

同時，隨著訓(xùn)練集群 GPU 使用數(shù)量增加（從 8 個 GPU 增加到 128 個 GPU），我們也發(fā)現(xiàn)，embedding 部分通信開銷占整個訓(xùn)練開銷比例越來越大。在一些大規(guī)模訓(xùn)練中（比如在 16 節(jié)點上）甚至超過一半（51%）。這主要是因為兩個原因：

隨著集群 GPU 數(shù)量增加，每個節(jié)點上的 embedding table 數(shù)量逐漸減少，導(dǎo)致不同節(jié)點負(fù)載不均衡，降低訓(xùn)練效率。

相比機(jī)內(nèi)帶寬，機(jī)間帶寬小得多，因此 embedding 模型并行需要進(jìn)行機(jī)間通信耗時較長。

為了幫助行業(yè)用戶更好地理解問題、解決問題，NVIDIA HugeCTR 團(tuán)隊于今年的 RecSys 大會上提出了 EMBark，通過支持 3D 的自定義 sharding 策略和結(jié)合不同的通信壓縮策略，能夠細(xì)粒度的優(yōu)化大規(guī)模集群下深度推薦模型訓(xùn)練的負(fù)載不均衡問題，以及減少 embedding 需要的通信時間，其相關(guān)代碼[1]和論文[2]皆已開源。

圖 2. 不同 cluster 配置下 DLRM 各部分訓(xùn)練耗時占比

EMBark 介紹

EMBark 旨在提高 DLRM 訓(xùn)練中 embedding 在不同集群配置下的性能，并加速整體訓(xùn)練吞吐量。EMBark 是在 NVIDIA Merlin HugeCTR 開源推薦系統(tǒng)框架的基礎(chǔ)上實現(xiàn)的，但所描述的技術(shù)也可以應(yīng)用于其他機(jī)器學(xué)習(xí)框架。

EMBark 有三個關(guān)鍵組件：embedding clusters、靈活的 3D 分片方案和分片規(guī)劃器。下圖展示了 EMBark 的整體架構(gòu)。

圖 3. EMBark 架構(gòu)圖

Embedding Clusters

Embedding clusters 旨在通過將具有相似特征的 embedding 進(jìn)行分組并為每個 cluster 應(yīng)用定制的壓縮策略來高效地訓(xùn)練 embedding。每個 cluster 包括 data distributor、embedding storage 和 embedding operators，協(xié)同將 feature ID 轉(zhuǎn)換為 embedding 向量。

有三種類型的 Embedding clusters：Data-parallel（DP）、Reduction-based（Reduction based）和基于 Unique-based（Unique Based）。每種類型在訓(xùn)練過程中采用不同的通信方法，適用于不同的 embedding。

DP cluster 不壓縮通信，因此簡單高效，但是因為會將 embedding table 在每個 GPU 上重復(fù)，因此僅適用于小表格。

RB cluster 使用歸約操作，對于具有池化操作的多 feature 輸入表格壓縮效果顯著。

UB cluster 僅發(fā)送唯一向量，有利于處理具有明顯訪問熱點的 embedding table。

靈活的 3D 分片方案

靈活的 3D 分片方案旨在解決 RB cluster 中的工作負(fù)載不平衡問題。與固定的分片策略比如 row-wise、table-wise、column-wise 不同，EMBark 使用一個 3D 元組（i, j, k）表示每個分片，其中 I 表示表格索引，j 表示行分片索引，k 表示列分片索引。這種方法允許每個 embedding 跨任意數(shù)量的 GPU 進(jìn)行分片，提供靈活性并實現(xiàn)對工作負(fù)載平衡的精確控制。

分片規(guī)劃器

為了找到最佳分片策略，EMBark 提供了一個分片規(guī)劃器——一種成本驅(qū)動的貪婪搜索算法，根據(jù)硬件規(guī)格和 embedding 配置識別最佳分片策略。

Evaluation

所有實驗均在一個集群上進(jìn)行，該集群由 NVIDIA DGX-H100[3] 節(jié)點組成，每個節(jié)點配備 8 張 NVIDIA H100 GPU（總計 640GB HBM，帶寬為每節(jié)點 24TB/s）。在每個節(jié)點內(nèi)，所有 GPU 通過 NVLink（雙向 900GB/s）互連。節(jié)點間通信使用 InfiniBand（8x400Gbps）。

為了展示 EMBark 可以高效訓(xùn)練任何規(guī)模的 DLRM 模型，我們測試了使用 MLPerf DLRM-DCNv2 模型并生成了幾個具有更大嵌入表和不同屬性的合成模型（參見上表）。我們的訓(xùn)練數(shù)據(jù)集表現(xiàn)出 α=1.2 的冪律偏斜。

圖 4. EMBark evaluation 結(jié)果

Baseline 采用串行的 kernel 執(zhí)行順序，固定的 table-row-wise 分片策略，以及全部使用了 RB-clusters。實驗依次使用了三種優(yōu)化：overlap、更靈活的分片策略和更好的 clusters 配置。

在四種代表性 DLRM 變體（DLRM-DCNv2、T180、T200 和 T510）中，EMBark 實現(xiàn)了平均 1.5 倍的端到端訓(xùn)練吞吐量加速，比 baseline 快最多 1.77 倍。更詳細(xì)的實驗結(jié)果和相關(guān)的分析，可以參考論文。

結(jié)論

EMBark 針對大規(guī)模推薦系統(tǒng)模型訓(xùn)練中 embedding 部分耗時占比過高的問題，通過支持 3D 的自定義 sharding 策略和結(jié)合不同的通信壓縮策略，能夠細(xì)粒度的優(yōu)化大規(guī)模集群下深度推薦模型訓(xùn)練的負(fù)載不均衡問題以及減少 embedding 需要的通信時間，提高大規(guī)模推薦系統(tǒng)模型的訓(xùn)練效率，在四種代表性 DLRM 變體（DLRM-DCNv2、T180、T200 和 T510）中，EMBark 實現(xiàn)了平均 1.5 倍的端到端訓(xùn)練吞吐量加速，比 baseline 快最多 1.77 倍。其中，相關(guān)代碼和論文皆已開源，希望我們的工作對大家有所幫助。同時，我們也在積極探索 embedding offloading 相關(guān)技術(shù)和進(jìn)行 TorchRec 相關(guān)優(yōu)化工作，未來也會及時和大家更新相關(guān)進(jìn)展情況，如果您對這部分感興趣，也可以聯(lián)系我們，大家一起交流和探索。

作者簡介

劉仕杰

劉仕杰于 2020 年加入 NVIDIA DevTech，主要專注于在 NVIDIA GPU 上性能優(yōu)化和推薦系統(tǒng)加速。加入 NVIDIA 之后，他主要參與了 Merlin HugeCTR 開發(fā)和 MLPerf DLRM 優(yōu)化等相關(guān)工作。