僅需1% Embedding參數，硬體成本降低十倍，開源方案單GPU訓練超大推薦模型

深度推荐模型（DLRMs）已经成为深度学习在互联网公司应用的最重要技术场景，如视频推荐、购物搜索、广告推送等流量变现业务，极大改善了用户体验和业务商业价值。但海量的用户和业务数据，频繁地迭代更新需求，以及高昂的训练成本，都对 DLRM 训练提出了严峻挑战。

在 DLRM 中，需要先在嵌入表（EmbeddingBags）中进行查表（lookup），再完成下游计算。嵌入表常常贡献 DLRM 中 99% 以上的内存需求，却只贡献 1% 的计算量。借助于 GPU 片上高速内存（High Bandwidth Memory）和强大算力的帮助，GPU 成为 DLRM 训练的主流硬件。但是，随着推荐系统研究的深入，日益增长的嵌入表大小和有限的 GPU 显存形成显著矛盾。如何让利用 GPU 高效训练超大 DLRM 模型，同时突破 GPU 内存墙的限制，已成为 DLRM 领域亟待解决的关键问题。

Colossal-AI此前已成功利用异构策略将相同硬件上训练NLP模型的参数容量提升上百倍，近期成功将其拓展到推荐系统中，通过软件缓存（Cache）方法在 CPU 和 GPU 内存中动态存储嵌入表。基于软件 Cache 设计，Colossal-AI 还添加流水预取，通过观察未来即将输入的训练数据，降低软件 Cache 检索和数据移动开销。同时，它以同步更新方式在 GPU 上训练整个 DLRM 模型，结合广泛使用的混合并行训练方法，可以扩展到多个 GPU。实验表明，Colossal-AI 仅需在 GPU 中保留 1% 的嵌入参数，仍能保持优秀的端到端训练速度。相比 PyTorch 其他方案，显存需求降低一个数量级，单块显卡即可训练 TB 级推荐模型。成本优势显著，例如仅需 5GB 显存即可训练占据 91GB 空间 Embedding Bag 的 DLRM，训练硬件成本从两张约 20 万元的 A100，降低十倍至仅需 2000 元左右的 RTX 3050 等入门级显卡。

开源地址：https://github.com/hpcaitech/ColossalAI

现有的嵌入表扩展技术

嵌入表将离散的整型特征映射成连续的浮点特征向量，下图展示了 DLRM 中的嵌入表训练过程。首先，在嵌入表中对每个特征查找 Embedding Table 对应的行，然后通过规约操作，比如 max，mean, sum 操作，变成一个特征向量，传递给后续的稠密神经网络。可见，DLRM 的嵌入表训练过程主要是不规则的内存访问操作，因此严重受限于硬件访存速度。

而工业级 DLRM 的嵌入表可能达到数百 GB 甚至 TB 级别，远超单 GPU 最高数十 GB 的显存容量。突破单 GPU 的内存墙来增大 DLRM 的嵌入表规模有很多方法。根据下图展示的 GPU 集群的内存层级图为例，让我们来分析几种常见方案的优劣。

GPU 模型并行：将嵌入表切分后分布在多个 GPU 的内存中，训练中通过 GPU 之间互联网络同步中间结果。这种方式的缺点首先是嵌入表切分负载并不均匀，扩展性问题难以解决。其次，增加 GPU 的前期硬件成本大，而且 DLRM 训练时 GPU 的计算能力并没有被充分利用，而是仅仅利用了它的 HBM 带宽优势，导致 GPU 使用率不高。

CPU 部分訓練：將嵌入表格分割成兩部分，一部分在 GPU 上訓練，另一部分在 CPU 上訓練。透過利用資料分佈的長尾效應，我們可以讓 CPU 運算比例盡可能少，讓 GPU 運算比例盡可能大。但是，隨著 batch size 增大，讓 mini-batch 的資料全部命中 CPU 或者 GPU 很困難，如果同時命中 CPU 或 GPU 這種方法很難處理。另外，由於 DDR 頻寬和 HBM 相差一個資料量級，即使 10% 的輸入資料在 CPU 上訓練，整個系統也會有至少一半速度下降。此外，CPU 和 GPU 需要傳輸中間結果，這也有不小的通訊開銷，進一步拖慢訓練速度。因此，研究人員設計了非同步更新等方式來避免這些效能缺陷，但是非同步方式會造成訓練結果的不確定性，在實踐中並不是演算法工程師的首選方案。

軟體Cache：保證訓練全部在GPU 上進行，嵌入表存在CPU 和GPU 組成的異質空間中，每次通過軟體Cache 方式，將所需的部分換入GPU。這種方式可以廉價擴充儲存資源，滿足嵌入表不斷增加的需求。而且，相較於使用 CPU 來計算，這種方式的整個訓練過程完全在 GPU 上完成，充分利用 HBM 頻寬優勢。但 Cache 的查詢、資料移動會帶來額外效能損耗。

目前已經有一些針對嵌入表優秀的軟體 Cache 方案實現，但是它們往往使用定制的 EmbeddingBags Kernel 實現，例如 fbgemm，或者藉助第三方深度學習框架。 而Colossal-AI 在原生PyTorch 基礎上不做任何Kernel 層次改動，提供了一套開箱用的軟體Cache EmbeddingBags 實現，也進一步針對DLRM 訓練流程進行最佳化，提出預取流水來進一步降低Cache 開銷。

Memory Hierarchy

Colossal-AI的嵌入表軟體Cache

Colossal-AI 實作了一個軟體Cache 並封裝成nn.Module 提供給使用者在自己模型中使用。 DLRM 的嵌入表，一般是由多個 Embedding 組成的 EmbeddingBags，駐留在 CPU 記憶體中。這部分記憶體空間被命名為 CPU Weight。而 EmbeddingBags 一小部分資料儲存在 GPU 記憶體中，它包括即將被訓練用到的資料。這部分內存空間被命名為 CUDA Cached Weight。在 DLRM 訓練期間，首先需要確定本次迭代輸入 mini-batch 的資料所對應嵌入表的行，如果有的行不在 GPU 中，則需要將它們從 CPU Weight 傳送到 CUDA Cached Weight 中。如果 GPU 中沒有足夠的空間，它會使用 LFU 演算法，根據存取快取的歷史頻率來淘汰被使用最少資料。

為了實現Cache 的檢索，需要一些輔助資料結構幫忙：cached_idx_map 是一維數組，儲存CPU Weight 中行號和CUDA Cached Weight 的行號對應關係，以及對應行在GPU 被存取的頻率資訊。 CUDA Cached Weight 大小與 CPU Weight 大小的比值命名為 cache_ratio，預設為 1.0%。

Cache 在每個迭代 forward 之前運行，以調整 CUDA Weight 中的數據，具體來說分為三個步驟。

Step1：CPU 索引：檢索CPU Weight 中需要被Cache 的行號

它需要對輸入mini -batch 的input_ids 和cached_idx_map 取交集，找出CPU Weight 中需要從CPU 移到GPU 的行號。

Step2：GPU 索引：根據使用頻率找到CUDA Weight 中可以被驅逐的行

這需要我們根據頻率以從低到高順序，對cache_idx_map 和input_ids 取差集合之後的部分進行top-k（取最大值k 個數）操作。

Step3：資料搬運：

#將CUDA Cached Weight 中的對應行移到CPU Weight 中，然後將CPU Weight 中的對應行移到CUDA Weight 中。

資料傳輸模組負責 CUDA Cached Weight 與 CPU Weight 之間的資料雙向傳輸。有別於低效率的逐行傳輸，它採用先快取再集中傳輸方式來提升 PCI-e 的頻寬利用率。分散在記憶體中的嵌入行在來源裝置的本機記憶體中集中為連續的資料區塊，然後區塊在 CPU 和 GPU 之間傳輸，並分散到目標記憶體的相應位置。以區塊為單位移動資料可以提高 PCI-e 頻寬利用率，merge 和 scatter 操作只涉及 CPU 和 GPU 的片上記憶體訪問，因此開銷並不是很大。

Colossal-AI 用一個尺寸受限的緩衝區來傳輸 CPU 和 GPU 之間資料。在最壞的情況下，所有輸入 id 都未命中快取 cache，那就需要需要傳輸大量元素。為了防止緩衝區佔用過多內存，緩衝區大小被嚴格限制。如果傳輸的資料大於緩衝區，則會分成多次完成傳輸。

Cached EmbeddingBag Workflow

軟體Cache 效能分析

上述Cache Step1 和Step2 的操作都是訪問密集的。因此為了能利用 GPU 的 HBM 的頻寬，它們是在 GPU 上運行的，並使用深度學習框架封裝好的 API 來實現。儘管如此，與嵌入表在 GPU 上的訓練操作相比，Cache 操作的開銷尤為突出。

例如在一個總共199 秒訓練任務中，Cache 作業的開銷為99 秒，佔總計算時間接近 50%# 。經過分析，Cache 的主要開銷主要是 Step1 和 Step2 所引起。下圖 base 位置展示了此時的 Cache 開銷時間分解，Cache 的 step1,2 紅色和橙色兩階段佔 Cache 總開銷的 70%。

Cache 操作的時間分解

而上述問題的原因，是因為傳統的Cache 策略有些“短視”，只能根據當前mini-batch 情況調整Cache，因此大部分時間浪費在查詢操作上。

Cache 管線預取

為了縮減Cache 的開銷，Colossal-AI 設計了一套 「高瞻遠矚」 #的Cache 機制。與其只對前 mini-batch 進行 Cache 操作，Colossal-AI 預取後續將會被使用的若干 mini-batch，統一進行 Cache 查詢操作。

如下圖所示，Colossal-AI 使用預取來合併多個 mini-batch 資料統一進行 Cache 操作，同時採用管線方式來重疊資料讀取和計算的開銷。範例中預取 mini-batch 數量是 2。在開始訓練前，先從磁碟讀取 mini-batch 0,1 資料到 GPU 內存，隨後開始 Cache 操作，然後執行這兩個 mini-batch 的正、反向傳播和參數更新。同時，可以和對 mini-batch 2，3 的開始資料讀取，這部分開銷可以和計算重疊。

和baseline Cache 執行方式相比，圖【Cache 操作的時間分解】比較了prefetch 8 個mini-batch 和baseline 的Cache 時間分解。訓練總時間從 201 秒下降到 120 秒，圖中所示的 Cache 階段操作時間佔比也顯著下降。可以看到和每個 mini-batch 獨立進行 Cache 操作相比，各部分時間都減少了，尤其是 Cache 的前兩個步驟。

總結起來，Cache 流水預取帶來兩個好處。

a.攤薄Cache 索引開銷

預先最顯而易見的好處是減少了Step1 和Step2 的開銷，使這個兩步驟操作在總的訓練過程佔比小於5%。如【Cache 操作的時間分解】所示，透過預取 8 個 mini-batch 數據，和沒有預取的 baseline 相比，Cache 查詢的開銷顯著降低。

b.增加CPU-GPU 數據移動頻寬

#透過集中更多數據，提升數據傳輸粒度，從而充分利用CPU-GPU傳輸頻寬。對於上面例子，CUDA->CPU 頻寬從 860MB/s 提升到 1477 MB/s，CPU->CUDA 頻寬從 1257 MB/s 提升到 2415 MB/s，幾乎帶來了近一倍的效能增益。

便捷使用

和Pytorch EmbeddingBag 用法一致，在建立推薦模型時，僅需如下數行程式碼進行初始化，即可大幅提升嵌入表容納量，低成本實現TB 級超大推薦模型訓練。

Bashfrom colossalai.nn.parallel.layers.cache_embedding import CachedEmbeddingBag
emb_module = CachedEmbeddingBag(num_embeddings=num_embeddings,embedding_dim=embedding_dim,mode="sum"include_last_offset=True,sparse=True,_weight=torch.randn(num_embeddings, embedding_dim),warmup_ratio=0.7,cache_ratio = 0.01,)

效能測試

在NVIDIA A100 GPU （80GB）和AMD EPYC 7543 32-Core Processor (512GB)硬體平台上，Colossal -AI 以Meta 的DLRM 模型作為測試目標，以超大資料集Cretio 1TB 和Meta 的dlrm_datasets 產生資料集作為測試模型。實驗中採用將嵌入表全部儲存 GPU 上的 PyTorch 訓練速度作為 baseline。

Cretio 1TB

Cretio 1TB嵌入表總共177944275 行，設定embedding dim=128，其嵌入表內存需求91.10 GB。想把 EmbeddingBags 全部儲存在單一 GPU 記憶體中，即使是最高階的英偉達 A100 80GB 也無法滿足其記憶體需求。

但使用Colossal-AI 仍在單GPU 上完成訓練，當cache ratio=0.05，顯存消耗僅為5.01 GB，直接降低約18 倍，可進一步擴展到在單張GPU 上實現TB 級推薦系統模型的訓練。在訓練速度上，如下圖所示，展示了不同 batch size 下訓練 100M 個樣本的延遲。綠色 Prefetch1 是不使用預取，藍色 Prefetch8 是使用預取（prefetch mini-batch=8）的延遲，可見預取流水優化對整體效能提升發揮了重要作用。圖中每個柱子深色部分為 Cache 開銷，使用預取後，Cache 開銷控制在訓練總時間的 15% 範圍內。

多GPU 擴充功能

用8192 作為全域batch size，在8 張GPU 卡上使用table-wise sharding 作為EmbeddingBags 並行方式訓練DLRM，訓練100M samples。此時設定 Prefetch 大小為 4，ColossalAI-mem-cr0.05 是 cache ratio=0.05，ColossalAI-mem-cr0.5=0.5。下圖展示了不同 GPU 情況下的訓練延遲。除了 1 GPU 時 PyTorch OOM 無法訓練之外，其餘情況 PyTorch 和 Colossal-AI 訓練時間類似。可以觀察到使用 4 和 8 GPU 並沒有帶來明顯效能提升，這是因為，1. 同步結果需要通訊開銷巨大。 2. table-wise sharding 會導致切分負載不平衡。也說明使用多 GPU 來擴充 embedding table 訓練擴充性並不是很好。

下圖展示了顯存使用，顯存使用在不同卡片上並不相同，這裡展示最大顯存數值。 在只使用一張 GPU 時，只有 Colossal-AI 的軟體 Cache 方法可以訓練，多卡並行的佔用記憶體也顯著減少數倍。

Meta Research 的合成資料集dlrm_datasets 模仿了工業界嵌入表的訓練存取行為，因此常在研究中作為推薦系統相關的軟硬體設計的測試參考。選取其中的 5 億行嵌入表項的作為子資料集，建構 256GB 和 128GB 大小的兩個 EmbeddingBags 用於測試。

PyTorch 由於顯存記憶體不足無法在單卡 A100 上訓練。作為對比, Colossal-AI 的軟體 cache 將顯著降低 GPU 記憶體需求，足以訓練大至 256GB 的嵌入表，並可進一步擴展至 TB 等級。而且，流水預取也能反映出加速效果，當預取數為 32 時，相比沒有預取總時間下降 60%，而且對 GPU 的儲存的需求卻沒有增大。

One More Thing

#針對大模型時代的通用深度學習系統Colossal- AI，透過多項自研領先技術如高效多維自動並行、異質記憶體管理、大規模優化庫、自適應任務調度等實現高效快速部署AI 大模型訓練和推理，降低AI 大模型應用成本。

Colossal-AI 相關解決方案已成功在自動駕駛、雲端運算、零售、醫藥、晶片等業界知名廠商落地應用，廣受好評。

Colossal-AI 注重開源社群建設，提供中文教程，開放用戶社群及論壇，對於用戶回饋進行高效交流與迭代更新，不斷添加PaLM、AlphaFold、OPT 等前沿應用。

自然開源以來，Colossal-AI 已經多次在GitHub 及Papers With Code 熱榜位列世界第一，與眾多已有數萬star 的明星開源專案一起受到海內外關注！

專案開源位址：https://github.com/hpcaitech/ColossalAI

##

以上是僅需1% Embedding參數，硬體成本降低十倍，開源方案單GPU訓練超大推薦模型的詳細內容。更多資訊請關注PHP中文網其他相關文章！