深度推薦模型（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 CachedEmbeddingBagemb_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] https://ai.facebook.com/blog/dlrm-an-advanced-open-source-deep-learning-recommendation-model/
[2] https://medium.com/@yangyou_berkeley/embedding-training-with-1-gpu-memory-and-10-times-less-budget-an-open-source-solution-for-6b4c3aba07a8
