僅需1% Embedding參數(shù)，硬件成本降低十倍，開源方案單GPU訓練超大推薦模型
深度推薦模型（DLRMs）已經(jīng)成為深度學習在互聯(lián)網(wǎng)公司應用的最重要技術場景，如視頻推薦、購物搜索、廣告推送等流量變現(xiàn)業(yè)務，極大改善了用戶體驗和業(yè)務商業(yè)價值。但海量的用戶和業(yè)務數(shù)據(jù)，頻繁地迭代更新需求，以及高昂的訓練成本，都對 DLRM 訓練提出了嚴峻挑戰(zhàn)。
在 DLRM 中，需要先在嵌入表（EmbeddingBags）中進行查表（lookup），再完成下游計算。嵌入表常常貢獻 DLRM 中 99% 以上的內(nèi)存需求，卻只貢獻 1% 的計算量。借助于 GPU 片上高速內(nèi)存（High Bandwidth Memory）和強大算力的幫助，GPU 成為 DLRM 訓練的主流硬件。但是，隨著推薦系統(tǒng)研究的深入，日益增長的嵌入表大小和有限的 GPU 顯存形成顯著矛盾。如何讓利用 GPU 高效訓練超大 DLRM 模型，同時突破 GPU 內(nèi)存墻的限制，已成為 DLRM 領域亟待解決的關鍵問題。
Colossal-AI此前已成功利用異構策略將相同硬件上訓練NLP模型的參數(shù)容量提升上百倍，近期成功將其拓展到推薦系統(tǒng)中，通過軟件緩存（Cache）方法在 CPU 和 GPU 內(nèi)存中動態(tài)存儲嵌入表。基于軟件 Cache 設計，Colossal-AI 還添加流水預取，通過觀察未來即將輸入的訓練數(shù)據(jù)，降低軟件 Cache 檢索和數(shù)據(jù)移動開銷。同時，它以同步更新方式在 GPU 上訓練整個 DLRM 模型，結合廣泛使用的混合并行訓練方法，可以擴展到多個 GPU。實驗表明，Colossal-AI 僅需在 GPU 中保留 1% 的嵌入?yún)?shù)，仍能保持優(yōu)秀的端到端訓練速度。相比 PyTorch 其他方案，顯存需求降低一個數(shù)量級，單塊顯卡即可訓練 TB 級推薦模型。成本優(yōu)勢顯著，例如僅需 5GB 顯存即可訓練占據(jù) 91GB 空間 Embedding Bag 的 DLRM，訓練硬件成本從兩張約 20 萬元的 A100，降低十倍至僅需 2000 元左右的 RTX 3050 等入門級顯卡。
開源地址：https://github.com/hpcaitech/ColossalAI
現(xiàn)有的嵌入表擴展技術
嵌入表將離散的整型特征映射成連續(xù)的浮點特征向量，下圖展示了 DLRM 中的嵌入表訓練過程。首先，在嵌入表中對每個特征查找 Embedding Table 對應的行，然后通過規(guī)約操作，比如 max，mean, sum 操作，變成一個特征向量，傳遞給后續(xù)的稠密神經(jīng)網(wǎng)絡。可見，DLRM 的嵌入表訓練過程主要是不規(guī)則的內(nèi)存訪問操作，因此嚴重受限于硬件訪存速度。
而工業(yè)級 DLRM 的嵌入表可能達到數(shù)百 GB 甚至 TB 級別，遠超單 GPU 最高數(shù)十 GB 的顯存容量。突破單 GPU 的內(nèi)存墻來增大 DLRM 的嵌入表規(guī)模有很多方法。根據(jù)下圖展示的 GPU 集群的內(nèi)存層級圖為例，讓我們來分析幾種常見方案的優(yōu)劣。
GPU 模型并行：將嵌入表切分后分布在多個 GPU 的內(nèi)存中，訓練中通過 GPU 之間互聯(lián)網(wǎng)絡同步中間結果。這種方式的缺點首先是嵌入表切分負載并不均勻，擴展性問題難以解決。其次，增加 GPU 的前期硬件成本大，而且 DLRM 訓練時 GPU 的計算能力并沒有被充分利用，而是僅僅利用了它的 HBM 帶寬優(yōu)勢，導致 GPU 使用率不高。
CPU 部分訓練：將嵌入表分割成兩部分，一部分在 GPU 上訓練，另一部分在 CPU 上訓練。通過利用數(shù)據(jù)分布的長尾效應，我們可以讓 CPU 計算比例盡可能少，讓 GPU 計算比例盡可能大。但是，隨著 batch size 增大，讓 mini-batch 的數(shù)據(jù)全部命中 CPU 或者 GPU 很困難，如果同時命中 CPU 或 GPU 這種方法很難處理。另外，由于 DDR 帶寬和 HBM 相差一個數(shù)據(jù)量級，即使 10% 的輸入數(shù)據(jù)在 CPU 上訓練，整個系統(tǒng)也會有至少一半速度下降。此外，CPU 和 GPU 需要傳輸中間結果，這也有不小的通信開銷，進一步拖慢訓練速度。因此，研究人員設計了異步更新等方式來避免這些性能缺陷，但是異步方式會造成訓練結果的不確定性，在實踐中并不是算法工程師的首選方案。
軟件 Cache：保證訓練全部在 GPU 上進行，嵌入表存在 CPU 和 GPU 組成的異構空間中，每次通過軟件 Cache 方式，將需要的部分換入 GPU。這種方式可以廉價擴展存儲資源，滿足嵌入表不斷增大的需求。而且，相比使用 CPU 來計算，這種方式的整個訓練過程完全在 GPU 上完成，充分利用 HBM 帶寬優(yōu)勢。但 Cache 的查詢、數(shù)據(jù)移動會帶來額外性能損耗。
目前已經(jīng)有一些針對嵌入表優(yōu)秀的軟件 Cache 方案實現(xiàn)，但是它們往往使用定制的 EmbeddingBags Kernel 實現(xiàn)，比如 fbgemm，或者借助第三方深度學習框架。而 Colossal-AI 在原生 PyTorch 基礎上不做任何 Kernel 層次改動，提供了一套開箱用的軟件 Cache EmbeddingBags 實現(xiàn)，還進一步針對 DLRM 訓練流程進行優(yōu)化，提出預取流水來進一步降低 Cache 開銷。
Memory Hierarchy
Colossal-AI 的嵌入表軟件 Cache
Colossal-AI 實現(xiàn)了一個軟件 Cache 并封裝成 nn.Module 提供給用戶在自己模型中使用。DLRM 的嵌入表，一般是由多個 Embedding 組成的 EmbeddingBags，駐留在 CPU 內(nèi)存中。這部分內(nèi)存空間被命名為 CPU Weight。而 EmbeddingBags 一小部分數(shù)據(jù)存儲在 GPU 內(nèi)存中，它包括即將被訓練用到的數(shù)據(jù)。這部分內(nèi)存空間被命名為 CUDA Cached Weight。在 DLRM 訓練期間，首先需要確定本次迭代輸入 mini-batch 的數(shù)據(jù)所對應嵌入表的行，如果有的行不在 GPU 中，需要將它們從 CPU Weight 傳輸?shù)?CUDA Cached Weight 中。如果 GPU 中沒有足夠的空間，它會使用 LFU 算法，根據(jù)訪問緩存的歷史頻率來淘汰被使用最少數(shù)據(jù)。
為了實現(xiàn) Cache 的檢索，需要一些輔助數(shù)據(jù)結構幫忙：cached_idx_map 是一維數(shù)組，存儲 CPU Weight 中行號和 CUDA Cached Weight 的行號對應關系，以及對應行在 GPU 被訪問的頻率信息。CUDA Cached Weight 大小與 CPU Weight 大小的比值命名為 cache_ratio，默認為 1.0%。
Cache 在每個迭代 forward 之前運行，以調(diào)整 CUDA Weight 中的數(shù)據(jù)，具體來說分三個步驟。
Step1：CPU 索引：檢索 CPU Weight 中需要被 Cache 的行號
它需要對輸入 mini-batch 的 input_ids 和 cached_idx_map 取交集，找到 CPU Weight 中需要從 CPU 移動到 GPU 的行號。
Step2：GPU 索引：根據(jù)使用頻率找到 CUDA Weight 中可以被驅(qū)逐的行
這需要我們根據(jù)頻率以從低到高順序，對 cache_idx_map 和 input_ids 取差集合之后的部分進行 top-k（取最大值 k 個數(shù)）操作。
Step3：數(shù)據(jù)搬運：
將 CUDA Cached Weight 中的對應行移動到 CPU Weight 中，然后將 CPU Weight 中的對應行移動到 CUDA Weight 中。
數(shù)據(jù)傳輸模塊負責 CUDA Cached Weight 和 CPU Weight 之間的數(shù)據(jù)雙向傳輸。不同于低效的逐行傳輸，它采用先緩存再集中傳輸方式來提升 PCI-e 的帶寬利用率。分散在內(nèi)存中的嵌入行在源設備的本地內(nèi)存中集中為連續(xù)的數(shù)據(jù)塊，然后塊在 CPU 和 GPU 之間傳輸，并分散到目標內(nèi)存的相應位置。以塊為單位移動數(shù)據(jù)可以提高 PCI-e 帶寬利用率，merge 和 scatter 操作只涉及 CPU 和 GPU 的片上內(nèi)存訪問，因此開銷并不是很大。
Colossal-AI 用一個尺寸受限的緩沖區(qū)來傳輸 CPU 和 GPU 之間數(shù)據(jù)。在最壞的情況下，所有輸入 id 都未命中緩存 cache，那就需要需要傳輸大量元素。為了防止緩沖區(qū)占用過多內(nèi)存，緩沖區(qū)大小被嚴格限制。如果傳輸?shù)臄?shù)據(jù)大于緩沖區(qū)，會分為多次完成傳輸。
Cached EmbeddingBag Workflow
軟件 Cache 性能分析
上述 Cache Step1 和 Step2 的操作都是訪存密集的。因此為了能利用 GPU 的 HBM 的帶寬，它們是在 GPU 上運行的，并使用深度學習框架封裝好的 API 來實現(xiàn)。盡管如此，與嵌入表在 GPU 上的訓練操作相比，Cache 操作的開銷尤為突出。
比如在一次總計 199 秒訓練任務中，Cache 操作的開銷為 99 秒，占比總計算時間接近 50%。經(jīng)過分析，Cache 的主要開銷主要是 Step1 和 Step2 引起。下圖 base 位置展示了此時的 Cache 開銷時間分解，Cache 的 step1,2 紅色和橙色兩階段占 Cache 總開銷的 70%。
Cache 操作的時間分解
而上述問題的原因，是因為傳統(tǒng)的 Cache 策略有些“短視”，只能根據(jù)當前 mini-batch 情況調(diào)整 Cache，因此大部分時間浪費在查詢操作上。
Cache 流水預取
為了縮減 Cache 的開銷，Colossal-AI 設計了一套 “高瞻遠矚” 的 Cache 機制。與其只對前 mini-batch 進行 Cache 操作，Colossal-AI 預取后續(xù)將會被使用的若干 mini-batch，統(tǒng)一進行 Cache 查詢操作。
如下圖所示，Colossal-AI 使用預取來合并多個 mini-batch 數(shù)據(jù)統(tǒng)一進行 Cache 操作，同時采用流水線方式來重疊數(shù)據(jù)讀取和計算的開銷。例子中預取 mini-batch 數(shù)量是 2。在開始訓練前，先從磁盤讀取 mini-batch 0,1 數(shù)據(jù)到 GPU 內(nèi)存，隨后開始 Cache 操作，然后執(zhí)行這兩個 mini-batch 的正、反向傳播和參數(shù)更新。與此同時，可以和對 mini-batch 2，3 的開始數(shù)據(jù)讀取，這部分開銷可以和計算重疊。
和 baseline Cache 執(zhí)行方式相比，圖【Cache 操作的時間分解】對比了 prefetch 8 個 mini-batch 和 baseline 的 Cache 時間分解。訓練總時間從 201 秒下降到 120 秒，圖中所示的 Cache 階段操作時間占比也顯著下降。可以看到和每個 mini-batch 獨立進行 Cache 操作相比，各部分時間都減少了，尤其是 Cache 的前兩步操作。
總結起來，Cache 流水預取帶來兩個好處。
a.攤薄 Cache 索引開銷
預取最顯而易見的好處是減少了 Step1 和 Step2 的開銷，使這個兩步操作在總的訓練過程占比小于 5%。如【Cache 操作的時間分解】所示，通過預取 8 個 mini-batch 數(shù)據(jù)，和沒有預取的 baseline 相比，Cache 查詢的開銷顯著降低。
b.增加 CPU-GPU 數(shù)據(jù)移動帶寬
通過集中更多數(shù)據(jù)，提升數(shù)據(jù)傳輸粒度，從而充分利用 CPU-GPU 傳輸帶寬。對于上面例子，CUDA->CPU 帶寬從 860MB/s 提升到 1477 MB/s，CPU->CUDA 帶寬從 1257 MB/s 提升到 2415 MB/s，幾乎帶來了近一倍的性能增益。
便捷使用
和 Pytorch EmbeddingBag 用法一致，在構建推薦模型時，僅需如下數(shù)行代碼進行初始化，即可大幅提升嵌入表容納量，低成本實現(xiàn) TB 級超大推薦模型訓練。
性能測試
在 NVIDIA A100 GPU （80GB）和 AMD EPYC 7543 32-Core Processor (512GB)硬件平臺上，Colossal-AI 以 Meta 的 DLRM 模型作為測試目標，用超大數(shù)據(jù)集 Cretio 1TB 和 Meta 的 dlrm_datasets 生成數(shù)據(jù)集作為測試模型。實驗中采用將嵌入表全部存儲 GPU 上的 PyTorch 訓練速度作為 baseline。
Cretio 1TB
Cretio 1TB嵌入表總共 177944275 行，設置 embedding dim=128，其嵌入表內(nèi)存需求 91.10 GB。想把 EmbeddingBags 全部存儲在單個 GPU 內(nèi)存中，即使是最高端的英偉達 A100 80GB 也無法滿足其內(nèi)存需求。
但使用 Colossal-AI 仍然在單 GPU 上完成訓練，當 cache ratio=0.05，顯存消耗僅為 5.01 GB，直接降低約 18 倍，可進一步擴展到在單張 GPU 上實現(xiàn) TB 級推薦系統(tǒng)模型的訓練。在訓練速度上，如下圖所示，展示了不同 batch size 下訓練 100M 個樣本的延遲。綠色 Prefetch1 是不使用預取，藍色 Prefetch8 是使用預取（prefetch mini-batch=8）的延遲，可見預取流水優(yōu)化對整體性能提升發(fā)揮了重要作用。圖中每個柱子深色部分為 Cache 開銷，使用預取后，Cache 開銷控制在訓練總時間的 15% 范圍內(nèi)。
多 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 訓練擴展性并不是很好。
下圖展示了顯存使用，顯存使用在不同卡上并不相同，這里展示最大顯存數(shù)值。在僅使用一張 GPU 時，只有 Colossal-AI 的軟件 Cache 方法可以訓練，多卡并行的占用內(nèi)存也顯著減少數(shù)倍。
Meta Research 的合成數(shù)據(jù)集 dlrm_datasets 模仿了工業(yè)界嵌入表的訓練訪問行為，因此常在研究中作為推薦系統(tǒng)相關的軟硬件設計的測試參考。選取其中的 5 億行嵌入表項的作為子數(shù)據(jù)集，構造 256GB 和 128GB 大小的兩個 EmbeddingBags 用于測試。
PyTorch 由于顯存內(nèi)存不足無法在單卡 A100 上訓練。作為對比, Colossal-AI 的軟件 cache 將顯著降低 GPU 內(nèi)存需求，足以訓練大至 256GB 的嵌入表，并可進一步擴展至 TB 級別。而且，流水預取也能體現(xiàn)出加速效果，當預取數(shù)為 32 時，相比沒有預取總時間下降 60%，而且對 GPU 的存儲的需求卻沒有增大。
One More Thing
面向大模型時代的通用深度學習系統(tǒng) Colossal-AI，通過多項自研領先技術如高效多維自動并行、異構內(nèi)存管理、大規(guī)模優(yōu)化庫、自適應任務調(diào)度等實現(xiàn)高效快速部署 AI 大模型訓練和推理，降低 AI 大模型應用成本。
Colossal-AI 相關解決方案已成功在自動駕駛、云計算、零售、醫(yī)藥、芯片等行業(yè)知名廠商落地應用，廣受好評。
Colossal-AI 注重開源社區(qū)建設，提供中文教程，開放用戶社群及論壇，對于用戶反饋進行高效交流與迭代更新，不斷添加 PaLM、AlphaFold、OPT 等前沿應用。
自然開源以來，Colossal-AI 已經(jīng)多次在 GitHub 及 Papers With Code 熱榜位列世界第一，與眾多已有數(shù)萬 star 的明星開源項目一起受到海內(nèi)外關注！
項目開源地址：https://github.com/hpcaitech/ColossalAI
