ai移動單位太大怎麼辦
深度推薦模型（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流水預取帶來兩個好處。
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級超大推薦模型訓練。
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並沒有帶來明顯效能提升，這是因為：
同步結果需要通訊開銷巨大。
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
觀點為作者所有。
—
完
—
量子位
QbitAI
‘’ 追蹤AI技術和產品新動態
一鍵三連「分享」「點贊」和「在看」
科技前沿進展日日相見 ~