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

深度推薦模型（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開銷。

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，那就需要需要傳輸大量元素。為了防止緩沖區占用過多內存，緩沖區大小被嚴格限制。如果傳輸的數據大于緩沖區，會分為多次完成傳輸。

軟件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策略有些“短視”，只能根據當前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流水預取帶來兩個好處。

1、攤薄Cache索引開銷

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

2、增加CPU-GPU數據移動帶寬

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

便捷使用

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

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模型作為測試目標，用超大數據集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

參考鏈接：
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

觀點為作者所有。