一塊RTX3050搞定DLRM訓(xùn)練！僅需1%Embedding參數(shù)，硬件

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

可是，海量的用戶和業(yè)務(wù)數(shù)據(jù)，頻繁的迭代更新需求以及高昂的培訓(xùn)成本給DLRM培訓(xùn)帶來了嚴(yán)峻的挑戰(zhàn)。

在DLRM，在完成下游計算之前，需要查找嵌入式表。

在DLRM，嵌入式表通常貢獻(xiàn)了99%以上的內(nèi)存需求，但只貢獻(xiàn)了1%的計算量。

借助GPU的片上高速內(nèi)存和強大的計算能力，GPU已經(jīng)成為DLRM培訓(xùn)的主流硬件。

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

龐—AI已經(jīng)成功使用異構(gòu)策略，在相同硬件上提高NLP模型訓(xùn)練的參數(shù)容量數(shù)百倍于之前最近，spool—AI成功將其擴(kuò)展到推薦系統(tǒng)，通過軟件緩存的方式，將嵌入式表格動態(tài)存儲在CPU和GPU內(nèi)存中

在軟件緩存的設(shè)計基礎(chǔ)上，龐—AI還增加了流水線預(yù)取，通過觀察未來要輸入的訓(xùn)練數(shù)據(jù)，降低軟件緩存檢索和數(shù)據(jù)移動的成本。

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

實驗表明，龐氏人工智能只需要在GPU中保留1%的嵌入?yún)?shù)，仍然可以保持優(yōu)秀的端到端訓(xùn)練速度。

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

成本是顯而易見的，比如訓(xùn)練占用91GB嵌入包的DLRM只需要5GB顯存，訓(xùn)練硬件的成本從兩個大約20萬元的A100降低到RTX 3050等入門級顯卡的十分之一，只需要2000元左右。

開放源地址:

基于現(xiàn)有嵌入式表格擴(kuò)展技術(shù)

嵌入式表將離散的整數(shù)特征映射為連續(xù)的浮點特征向量下圖顯示了DLRM嵌入式表的培訓(xùn)過程

首先為嵌入表中的每個特征搜索嵌入表的對應(yīng)行，然后通過常規(guī)運算，如最大，均值，和運算，將特征向量轉(zhuǎn)化為特征向量，再傳遞給后續(xù)的密集神經(jīng)網(wǎng)絡(luò)。

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

而工業(yè)DLRM的嵌入式表可能會達(dá)到幾百GB甚至TB級別，遠(yuǎn)遠(yuǎn)超過單個GPU高達(dá)幾十GB的內(nèi)存容量。

有許多方法可以突破單個GPU的內(nèi)存墻來增加DLRM的嵌入式表大小。

GPU模型并行性:

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

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

其次，增加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計算比盡可能大但伴隨著批量的增加，很難讓mini—batch的所有數(shù)據(jù)都命中CPU或GPU，如果同時命中CPU或GPU，這種方法很難處理

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

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

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

軟件緩存:

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

該方法可以低成本地擴(kuò)展存儲資源，滿足日益增長的嵌入式表格需求。

而且與CPU相比，該方法的整個訓(xùn)練過程完全在GPU上完成，充分利用了HBM的帶寬優(yōu)勢但是，緩存的查詢和數(shù)據(jù)移動會帶來額外的性能損失

目前有一些優(yōu)秀的嵌入式表的軟件緩存方案，但往往是通過自定義的EmbeddingBags內(nèi)核來實現(xiàn)，比如fbgemm，或者借助第三方深度學(xué)習(xí)框架來實現(xiàn)。

龐—AI在原生PyTorch的基礎(chǔ)上，不做任何內(nèi)核級的改動，提供了一套開箱即用的軟件Cache EmbeddingBags實現(xiàn)，進(jìn)一步優(yōu)化了DLRM訓(xùn)練過程，并提出預(yù)取流水線，進(jìn)一步降低緩存開銷。

內(nèi)存層次無損AI嵌入式表格軟件緩存

龐—AI實現(xiàn)了一個軟件緩存，并將其封裝為nn模塊，供用戶在自己的模型中使用

DLRM的嵌入式表，通常由多個嵌入式包組成，駐留在CPU內(nèi)存中。

這部分內(nèi)存空間稱為CPU權(quán)重但是，嵌入包的少量數(shù)據(jù)存儲在GPU內(nèi)存中，其中包括用于訓(xùn)練的數(shù)據(jù)

這部分內(nèi)存空間稱為CUDA緩存權(quán)重。

在DLRM訓(xùn)練期間，首先需要確定表中嵌入的行，這些行對應(yīng)于在該迭代中輸入到小批量中的數(shù)據(jù)如果有些行不在GPU中，需要從CPU權(quán)重轉(zhuǎn)移到CUDA緩存權(quán)重

如果GPU空間不夠，它會根據(jù)訪問緩存的歷史頻率，使用LFU算法剔除最少使用的數(shù)據(jù)。

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

dacachedweight大小與CPU權(quán)重大小的比值命名為cache_ratio，默認(rèn)值為1.0%。

每次迭代前運行緩存來調(diào)整CUDA中的數(shù)據(jù)權(quán)重，具體分三步。

第一步:CPU索引

檢索需要緩存的CPU權(quán)重中的行號。

它需要對輸入小批的input_ids和cached_idx_map求交集，找到CPU權(quán)重中需要從CPU移到GPU的行號。

第二步:GPU索引

根據(jù)使用頻率找到CUDA權(quán)重中可以被驅(qū)逐的行。

這就要求我們按照頻率從低到高，對cache_idx_map和input_ids的差集之后的部分進(jìn)行top—k。

第三步:數(shù)據(jù)處理:

將CUDA緩存權(quán)重中的對應(yīng)行移到CPU權(quán)重中，然后將CPU權(quán)重中的對應(yīng)行移到CUDA權(quán)重中。

傳輸模塊負(fù)責(zé)CUDA緩存權(quán)重和CPU權(quán)重之間的數(shù)據(jù)雙向傳輸。

與低效的逐行傳輸不同，它采用先緩存再集中傳輸?shù)姆椒ǎ岣逷CI—E的帶寬利用率。

分散內(nèi)存中的嵌入行在源設(shè)備的本地內(nèi)存中集中成連續(xù)的數(shù)據(jù)塊，然后在CPU和GPU之間傳輸，分散到目標(biāo)內(nèi)存中相應(yīng)的位置分塊移動數(shù)據(jù)可以提高PCI—E的帶寬利用率，合并和分散操作只涉及CPU和GPU的片內(nèi)內(nèi)存訪問，所以開銷不是很高

龐—AI使用有限大小的緩沖區(qū)在CPU和GPU之間傳輸數(shù)據(jù)。

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

緩存嵌入包工作流軟件的緩存性能分析

高速緩存步驟1和步驟2的上述操作是存儲器訪問密集型的。

因此，為了利用GPU的HBM的帶寬，它們運行在GPU上，由深度學(xué)習(xí)框架封裝的API來實現(xiàn)但是，相對于GPU上嵌入式表的訓(xùn)練操作，緩存操作的開銷尤為突出

例如，在一個總共199秒的訓(xùn)練任務(wù)中，緩存操作的開銷為99秒，占總計算時間的近50%。

經(jīng)過分析，緩存的主要開銷主要是由Step1和Step2造成的下圖中的基準(zhǔn)位置顯示了此時的緩存開銷時間偏差緩存步驟1和2的紅色和橙色階段占總緩存開銷的70%

高速緩存操作的時間偏差

產(chǎn)生上述問題的原因是傳統(tǒng)的緩存策略短視，只能根據(jù)當(dāng)前的小批量情況調(diào)整緩存，因此大部分時間浪費在查詢操作上。

緩存預(yù)取

為了降低緩存的開銷，龐—AI設(shè)計了一種前瞻性的緩存機制spool—AI不是只緩存以前的小批量，而是預(yù)取幾個以后要用的小批量，統(tǒng)一進(jìn)行緩存查詢操作

如下圖所示，龐—AI利用預(yù)取合并多個小批量數(shù)據(jù)進(jìn)行統(tǒng)一緩存操作，同時利用流水線技術(shù)重疊數(shù)據(jù)讀取和計算的開銷。

在本例中，預(yù)取小批量的數(shù)量是2在訓(xùn)練之前，將mini—batch 0，1數(shù)據(jù)從磁盤讀入GPU內(nèi)存，然后啟動緩存操作，再對這兩個mini—batch進(jìn)行正反向傳播和參數(shù)更新

同時可以用來讀取Mini—Batch 2，3的數(shù)據(jù)，這部分開銷可以和計算重疊。

與基線緩存的執(zhí)行模式相比，該圖比較了八個小批量預(yù)取和基線的緩存時間偏差。

總訓(xùn)練時間從201秒下降到120秒，圖中顯示的緩存階段的操作時間比例也明顯下降可以看出，與每個小批量獨立執(zhí)行緩存操作相比，每個部分的時間都有所減少，尤其是緩存操作的前兩步

綜上所述，緩存流水線預(yù)取帶來了兩個好處。

1.稀釋緩存索引開銷

預(yù)取最明顯的優(yōu)點是減少了Step1和Step2的開銷，使得這兩步操作占整個訓(xùn)練過程的比例不到5%如所示，通過預(yù)取8個小批量數(shù)據(jù)，與沒有預(yù)取的基線相比，緩存查詢的開銷顯著降低

2.增加CPU—GPU數(shù)據(jù)移動帶寬

通過集中更多的數(shù)據(jù)，提高數(shù)據(jù)傳輸?shù)牧６龋瑥亩浞掷肅PU—GPU的傳輸帶寬上面的例子，CUDA—gt，CPU帶寬從860MB/s提升到1477 MB/s，CPU—gt，CUDA的帶寬從1257 MB/s提升到2415 MB/s，性能增益幾乎翻倍

使用方便

這與Pytorch嵌入包的用法是一致的在構(gòu)建推薦模型時，只需要初始化以下幾行代碼，可以大幅增加嵌入式表的容量，低成本實現(xiàn)TB級超大推薦模型訓(xùn)練

fromcolosalai . nn . parallel . layers . cache _ embedding import cached embedding bag _ module = cached embedding bag，warmup _ ratio = 0.7，cache _ ratio = 0.01。)

在NVIDIA A100 GPU和AMD EPYC 7543 32核處理器(512 GB)硬件平臺上，龐—AI以Meta DLRM模型為測試對象，使用超大型數(shù)據(jù)集Cretio 1TB和Meta dlrm_datasets作為測試模型。

實驗中，以存儲所有嵌入式表的GPU上的PyTorch訓(xùn)練速度為基線。

Cretio 1TB

Cretio 1TB嵌入式表共有177，944，275行設(shè)置嵌入dim=128需要91.10 GB的嵌入式表內(nèi)存

如果想在單個GPU內(nèi)存中存儲所有的EmbeddingBags，即使是最高端的NVIDIA A100 80GB也無法滿足其內(nèi)存需求。

可是，龐—AI仍然用于在單個GPU上完成訓(xùn)練當(dāng)緩存比=0.05時，內(nèi)存消耗僅為5.01 GB，直接降低了18倍左右還可以進(jìn)一步擴(kuò)展，在單個GPU上實現(xiàn)TB級推薦系統(tǒng)模型訓(xùn)練

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

綠色預(yù)取1不使用預(yù)取，藍(lán)色預(yù)取8是預(yù)取的延遲可見預(yù)取流水線優(yōu)化對整體性能提升的重要作用

圖中每列的黑色部分是緩存開銷預(yù)取后，緩存開銷控制在總訓(xùn)練時間的15%以內(nèi)

多GPU可擴(kuò)展性

8192被用作全局批量大小，表式分片被用作在8個GPU卡上嵌入標(biāo)簽的并行模式，以訓(xùn)練DLRM和100M樣本。

此時預(yù)取大小設(shè)置為4，ColossalAI—mem—cr0.05的緩存比是0.05，龐AI—mem—cr0.5是0.5。

下圖顯示了不同GPU情況下的訓(xùn)練延遲除了PyTorch OOM不能在1個GPU上訓(xùn)練之外，PyTorch和龐—AI的訓(xùn)練時間差不多

可以看出，使用4個和8個GPU不會帶來顯著的性能提升，因為:

因此，同步需要巨大的通信開銷。

表方向分片將導(dǎo)致不平衡的拆分負(fù)載也說明用多個GPU來擴(kuò)展嵌入表的訓(xùn)練擴(kuò)展性不是很好

下圖顯示了顯存的使用，不同的卡上使用的顯存是不同的此處顯示了最大視頻內(nèi)存值

在只使用一個GPU的情況下，只需要訓(xùn)練龐—AI的軟件緩存方法，多張卡并行占用的內(nèi)存就可以顯著降低數(shù)倍。

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

其中，選擇表項中嵌入的5億行作為子數(shù)據(jù)集，構(gòu)建256GB和128GB兩個嵌入包進(jìn)行測試。

PyTorch由于顯存不足無法在單卡A100上訓(xùn)練相比之下，龐—AI的軟件緩存將大幅降低GPU內(nèi)存需求，足以訓(xùn)練高達(dá)256GB的嵌入式表格，并可以進(jìn)一步擴(kuò)展到TB級別

而且流水線預(yù)取也能體現(xiàn)加速效果預(yù)取次數(shù)為32時，總時間比不預(yù)取時減少了60%，且不增加GPU的存儲需求

還有一點

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

龐氏人工智能相關(guān)解決方案已成功應(yīng)用于自動駕駛，云計算，零售，醫(yī)藥，芯片等行業(yè)的知名廠商，并獲得好評。

龐—AI注重開源社區(qū)建設(shè)，提供中文教程，開放用戶社區(qū)和論壇，高效交流和迭代更新用戶反饋，不斷增加PaLM，AlphaFold，OPT等前沿應(yīng)用。

自自然開源以來，龐—AI多次在GitHub和有代碼熱榜的論文上排名世界第一，與眾多上萬顆星的明星開源項目一起引起了國內(nèi)外的關(guān)注！

項目的開放源地址:

參考鏈接:

這篇文章引用自: