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