一塊RTX3050搞定DLRM訓(xùn)練！僅需1%Embedding參數(shù)，硬件成本降低至百分之一 | 開源

深度推薦模型（DLRMs）已經(jīng)成為深度學(xué)習(xí)在互聯(lián)網(wǎng)公司應(yīng)用的最重要技術(shù)場景，如視頻推薦、購物搜索、廣告推送等流量變現(xiàn)業(yè)務(wù)，極大改善了用戶體驗和業(yè)務(wù)商業(yè)價值。

但海量的用戶和業(yè)務(wù)數(shù)據(jù)，頻繁地迭代更新需求，以及高昂的訓(xùn)練成本，都對DLRM訓(xùn)練提出了嚴(yán)峻挑戰(zhàn)。

在DLRM中，需要先在嵌入表（EmbeddingBags）中進行查表（lookup），再完成下游計算。

嵌入表常常貢獻DLRM中99%以上的內(nèi)存需求，卻只貢獻1%的計算量。

借助于GPU片上高速內(nèi)存（High Bandwidth Memory）和強大算力的幫助，GPU成為DLRM訓(xùn)練的主流硬件。

但是，隨著推薦系統(tǒng)研究的深入，日益增長的嵌入表大小和有限的GPU顯存形成顯著矛盾。如何讓利用GPU高效訓(xùn)練超大DLRM模型，同時突破GPU內(nèi)存墻的限制，已成為DLRM領(lǐng)域亟待解決的關(guān)鍵問題。

Colossal-AI此前已成功利用異構(gòu)策略將相同硬件上訓(xùn)練NLP模型的參數(shù)容量提升上百倍，近期成功將其拓展到推薦系統(tǒng)中，通過軟件緩存（Cache）方法在CPU 和 GPU 內(nèi)存中動態(tài)存儲嵌入表。

基于軟件Cache設(shè)計，Colossal-AI還添加流水預(yù)取，通過觀察未來即將輸入的訓(xùn)練數(shù)據(jù)，降低軟件Cache檢索和數(shù)據(jù)移動開銷。

同時，它以同步更新方式在 GPU 上訓(xùn)練整個 DLRM模型，結(jié)合廣泛使用的混合并行訓(xùn)練方法，可以擴展到多個 GPU。

實驗表明，Colossal-AI僅需在 GPU 中保留 1% 的嵌入?yún)?shù)，仍能保持優(yōu)秀的端到端訓(xùn)練速度。

相比PyTorch其他方案，顯存需求降低一個數(shù)量級，單塊顯卡即可訓(xùn)練TB級推薦模型。

成本優(yōu)勢顯著，例如僅需5GB顯存即可訓(xùn)練占據(jù)91GB空間 Embedding Bag的DLRM，訓(xùn)練硬件成本從兩張約20萬元的A100，降低至百分之一僅需2000元左右的RTX 3050等入門級顯卡。

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

現(xiàn)有的嵌入表擴展技術(shù)

嵌入表將離散的整型特征映射成連續(xù)的浮點特征向量，下圖展示了DLRM中的嵌入表訓(xùn)練過程。

首先，在嵌入表中對每個特征查找Embedding Table對應(yīng)的行，然后通過規(guī)約操作，比如max，mean, sum操作，變成一個特征向量，傳遞給后續(xù)的稠密神經(jīng)網(wǎng)絡(luò)。

可見，DLRM的嵌入表訓(xùn)練過程主要是不規(guī)則的內(nèi)存訪問操作，因此嚴(yán)重受限于硬件訪存速度。

而工業(yè)級DLRM的嵌入表可能達到數(shù)百GB甚至TB級別，遠(yuǎn)超單GPU最高數(shù)十GB的顯存容量。

突破單GPU的內(nèi)存墻來增大DLRM的嵌入表規(guī)模有很多方法。

根據(jù)下圖展示的GPU集群的內(nèi)存層級圖為例，讓我們來分析幾種常見方案的優(yōu)劣。

GPU模型并行：

將嵌入表切分后分布在多個GPU的內(nèi)存中，訓(xùn)練中通過GPU之間互聯(lián)網(wǎng)絡(luò)同步中間結(jié)果。

這種方式的缺點首先是嵌入表切分負(fù)載并不均勻，擴展性問題難以解決。

其次，增加GPU的前期硬件成本大，而且DLRM訓(xùn)練時GPU的計算能力并沒有被充分利用，而是僅僅利用了它的HBM帶寬優(yōu)勢，導(dǎo)致GPU使用率不高。

CPU部分訓(xùn)練：

將嵌入表分割成兩部分，一部分在GPU上訓(xùn)練，另一部分在CPU上訓(xùn)練。

通過利用數(shù)據(jù)分布的長尾效應(yīng)，我們可以讓CPU計算比例盡可能少，讓GPU計算比例盡可能大。但是，隨著batch size增大，讓mini-batch的數(shù)據(jù)全部命中CPU或者GPU很困難，如果同時命中CPU或GPU這種方法很難處理。

另外，由于DDR帶寬和HBM相差一個數(shù)據(jù)量級，即使10%的輸入數(shù)據(jù)在CPU上訓(xùn)練，整個系統(tǒng)也會有至少一半速度下降。

此外，CPU和GPU需要傳輸中間結(jié)果，這也有不小的通信開銷，進一步拖慢訓(xùn)練速度。

因此，研究人員設(shè)計了異步更新等方式來避免這些性能缺陷，但是異步方式會造成訓(xùn)練結(jié)果的不確定性，在實踐中并不是算法工程師的首選方案。

軟件Cache：

保證訓(xùn)練全部在GPU上進行，嵌入表存在CPU和GPU組成的異構(gòu)空間中，每次通過軟件Cache方式，將需要的部分換入GPU。

這種方式可以廉價擴展存儲資源，滿足嵌入表不斷增大的需求。

而且，相比使用CPU來計算，這種方式的整個訓(xùn)練過程完全在GPU上完成，充分利用HBM帶寬優(yōu)勢。但Cache的查詢、數(shù)據(jù)移動會帶來額外性能損耗。

目前已經(jīng)有一些針對嵌入表優(yōu)秀的軟件Cache方案實現(xiàn)，但是它們往往使用定制的EmbeddingBags Kernel實現(xiàn)，比如fbgemm，或者借助第三方深度學(xué)習(xí)框架。

而Colossal-AI在原生PyTorch基礎(chǔ)上不做任何Kernel層次改動，提供了一套開箱用的軟件Cache EmbeddingBags實現(xiàn)，還進一步針對DLRM訓(xùn)練流程進行優(yōu)化，提出預(yù)取流水來進一步降低Cache開銷。

Colossal-AI的嵌入表軟件Cache

Colossal-AI實現(xiàn)了一個軟件Cache并封裝成nn.Module提供給用戶在自己模型中使用。

DLRM的嵌入表，一般是由多個Embedding組成的EmbeddingBags，駐留在 CPU 內(nèi)存中。

這部分內(nèi)存空間被命名為CPU Weight。而EmbeddingBags一小部分?jǐn)?shù)據(jù)存儲在 GPU內(nèi)存中，它包括即將被訓(xùn)練用到的數(shù)據(jù)。

這部分內(nèi)存空間被命名為CUDA Cached Weight。

在 DLRM 訓(xùn)練期間，首先需要確定本次迭代輸入mini-batch的數(shù)據(jù)所對應(yīng)嵌入表的行，如果有的行不在GPU中，需要將它們從CPU Weight傳輸?shù)?CUDA Cached Weight中。

如果GPU中沒有足夠的空間，它會使用LFU算法，根據(jù)訪問緩存的歷史頻率來淘汰被使用最少數(shù)據(jù)。

為了實現(xiàn)Cache的檢索，需要一些輔助數(shù)據(jù)結(jié)構(gòu)幫忙：cached_idx_map是一維數(shù)組，存儲CPU Weight中行號和CUDA Cached Weight的行號對應(yīng)關(guān)系，以及對應(yīng)行在GPU被訪問的頻率信息。

CUDA Cached Weight 大小與 CPU Weight 大小的比值命名為 cache_ratio，默認(rèn)為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中的對應(yīng)行移動到CPU Weight中，然后將CPU Weight中的對應(yī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)位置。以塊為單位移動數(shù)據(jù)可以提高 PCI-e 帶寬利用率，merge和scatter操作只涉及CPU和GPU的片上內(nèi)存訪問，因此開銷并不是很大。

Colossal-AI用一個尺寸受限的緩沖區(qū)來傳輸CPU和GPU之間數(shù)據(jù)。

在最壞的情況下，所有輸入 id 都未命中緩存cache，那就需要需要傳輸大量元素。為了防止緩沖區(qū)占用過多內(nèi)存，緩沖區(qū)大小被嚴(yán)格限制。如果傳輸?shù)臄?shù)據(jù)大于緩沖區(qū)，會分為多次完成傳輸。

軟件Cache性能分析

上述Cache Step1和Step2的操作都是訪存密集的。

因此為了能利用GPU的HBM的帶寬，它們是在GPU上運行的，并使用深度學(xué)習(xí)框架封裝好的API來實現(xiàn)。盡管如此，與嵌入表在GPU上的訓(xùn)練操作相比，Cache操作的開銷尤為突出。

比如在一次總計199秒訓(xùn)練任務(wù)中，Cache操作的開銷為99秒，占比總計算時間接近50%。

經(jīng)過分析，Cache的主要開銷主要是Step1和Step2引起。下圖base位置展示了此時的Cache開銷時間分解，Cache的step1,2 紅色和橙色兩階段占Cache總開銷的70%。

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

Cache流水預(yù)取

為了縮減Cache的開銷，Colossal-AI設(shè)計了一套“高瞻遠(yuǎn)矚”的Cache機制。與其只對前mini-batch進行Cache操作，Colossal-AI預(yù)取后續(xù)將會被使用的若干mini-batch，統(tǒng)一進行Cache查詢操作。

如下圖所示，Colossal-AI使用預(yù)取來合并多個mini-batch數(shù)據(jù)統(tǒng)一進行Cache操作，同時采用流水線方式來重疊數(shù)據(jù)讀取和計算的開銷。

例子中預(yù)取mini-batch數(shù)量是2。在開始訓(xùn)練前，先從磁盤讀取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時間分解。

訓(xùn)練總時間從201秒下降到120秒，圖中所示的Cache階段操作時間占比也顯著下降。可以看到和每個mini-batch獨立進行Cache操作相比，各部分時間都減少了，尤其是Cache的前兩步操作。

總結(jié)起來，Cache流水預(yù)取帶來兩個好處。

1、攤薄Cache索引開銷

預(yù)取最顯而易見的好處是減少了Step1和Step2的開銷，使這個兩步操作在總的訓(xùn)練過程占比小于5%。如【Cache操作的時間分解】所示，通過預(yù)取8個mini-batch數(shù)據(jù)，和沒有預(yù)取的baseline相比，Cache查詢的開銷顯著降低。

2、增加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用法一致，在構(gòu)建推薦模型時，僅需如下數(shù)行代碼進行初始化，即可大幅提升嵌入表容納量，低成本實現(xiàn)TB級超大推薦模型訓(xùn)練。

bash

from 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模型作為測試目標(biāo)，用超大數(shù)據(jù)集Cretio 1TB和Meta的dlrm_datasets生成數(shù)據(jù)集作為測試模型。

實驗中采用將嵌入表全部存儲GPU上的PyTorch訓(xùn)練速度作為baseline。

Cretio 1TB

Cretio 1TB嵌入表總共177944275行，設(shè)置embedding dim=128，其嵌入表內(nèi)存需求91.10 GB。

想把EmbeddingBags全部存儲在單個GPU內(nèi)存中，即使是最高端的英偉達A100 80GB也無法滿足其內(nèi)存需求。

但使用Colossal-AI仍然在單GPU上完成訓(xùn)練，當(dāng)cache ratio=0.05，顯存消耗僅為5.01 GB，直接降低約18倍，可進一步擴展到在單張GPU上實現(xiàn)TB級推薦系統(tǒng)模型的訓(xùn)練。

在訓(xùn)練速度上，如下圖所示，展示了不同batch size下訓(xùn)練100M個樣本的延遲。

綠色Prefetch1是不使用預(yù)取，藍色Prefetch8是使用預(yù)?。╬refetch mini-batch=8）的延遲，可見預(yù)取流水優(yōu)化對整體性能提升發(fā)揮了重要作用。

圖中每個柱子深色部分為Cache開銷，使用預(yù)取后，Cache開銷控制在訓(xùn)練總時間的15%范圍內(nèi)。

多GPU擴展性

用8192作為全局batch size，在8張GPU卡上使用table-wise sharding作為EmbeddingBags并行方式訓(xùn)練DLRM，訓(xùn)練100M samples。

此時設(shè)置Prefetch大小為4，ColossalAI-mem-cr0.05是cache ratio=0.05，ColossalAI-mem-cr0.5=0.5。

下圖展示了不同GPU情況下的訓(xùn)練延遲。除了1 GPU時PyTorch OOM無法訓(xùn)練之外，其余情況PyTorch和Colossal-AI訓(xùn)練時間類似。

可以觀察到使用4和8 GPU并沒有帶來明顯性能提升，這是因為：

1. 同步結(jié)果需要通信開銷巨大。
2. table-wise sharding會導(dǎo)致切分負(fù)載不均衡。也說明使用多GPU來擴展embedding table訓(xùn)練擴展性并不是很好。

下圖展示了顯存使用，顯存使用在不同卡上并不相同，這里展示最大顯存數(shù)值。

在僅使用一張GPU時，只有Colossal-AI的軟件Cache方法可以訓(xùn)練，多卡并行的占用內(nèi)存也顯著減少數(shù)倍。

Meta Research的合成數(shù)據(jù)集dlrm_datasets模仿了工業(yè)界嵌入表的訓(xùn)練訪問行為，因此常在研究中作為推薦系統(tǒng)相關(guān)的軟硬件設(shè)計的測試參考。

選取其中的5億行嵌入表項的作為子數(shù)據(jù)集，構(gòu)造256GB和128GB大小的兩個EmbeddingBags用于測試。

PyTorch由于顯存內(nèi)存不足無法在單卡A100上訓(xùn)練。作為對比, Colossal-AI的軟件cache將顯著降低GPU內(nèi)存需求，足以訓(xùn)練大至256GB的嵌入表，并可進一步擴展至TB級別。

而且，流水預(yù)取也能體現(xiàn)出加速效果，當(dāng)預(yù)取數(shù)為32時，相比沒有預(yù)取總時間下降60%，而且對GPU的存儲的需求卻沒有增大。

One More Thing

面向大模型時代的通用深度學(xué)習(xí)系統(tǒng) Colossal-AI，通過多項自研領(lǐng)先技術(shù)如高效多維自動并行、異構(gòu)內(nèi)存管理、大規(guī)模優(yōu)化庫、自適應(yīng)任務(wù)調(diào)度等實現(xiàn)高效快速部署AI大模型訓(xùn)練和推理，降低AI大模型應(yīng)用成本。

Colossal-AI相關(guān)解決方案已成功在自動駕駛、云計算、零售、醫(yī)藥、芯片等行業(yè)知名廠商落地應(yīng)用，廣受好評。

Colossal-AI注重開源社區(qū)建設(shè)，提供中文教程，開放用戶社群及論壇，對于用戶反饋進行高效交流與迭代更新，不斷添加PaLM、AlphaFold、OPT等前沿應(yīng)用。

自然開源以來，Colossal-AI已經(jīng)多次在GitHub及Papers With Code熱榜位列世界第一，與眾多已有數(shù)萬star的明星開源項目一起受到海內(nèi)外關(guān)注！

項目開源地址：
https://github.com/hpcaitech/ColossalAI

參考鏈接：
https://ai.facebook.com/blog/dlrm-an-advanced-open-source-deep-learning-recommendation-model/

本文參考自：
https://medium.com/@yangyou_berkeley/embedding-training-with-1-gpu-memory-and-10-times-less-budget-an-open-source-solution-for-6b4c3aba07a8

觀點為作者所有。