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

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

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

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

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

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

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

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

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

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

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

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

開(kāi)放源地址:

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

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

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

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

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

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

GPU模型并行性:

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

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

其次，增加GPU的前期硬件成本較高，DLRM訓(xùn)練時(shí)GPU的計(jì)算能力沒(méi)有得到充分利用，只利用了其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ì)算比盡可能大但伴隨著批量的增加，很難讓mini—batch的所有數(shù)據(jù)都命中CPU或GPU，如果同時(shí)命中CPU或GPU，這種方法很難處理

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

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

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

軟件緩存:

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

該方法可以低成本地?cái)U(kuò)展存儲(chǔ)資源，滿(mǎn)足日益增長(zhǎng)的嵌入式表格需求。

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

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

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

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

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

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

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

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

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

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

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

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

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

第一步:CPU索引

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

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

第二步:GPU索引

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

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

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

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

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

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

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

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

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

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

高速緩存步驟1和步驟2的上述操作是存儲(chǔ)器訪(fǎng)問(wèn)密集型的。

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

例如，在一個(gè)總共199秒的訓(xùn)練任務(wù)中，緩存操作的開(kāi)銷(xiāo)為99秒，占總計(jì)算時(shí)間的近50%。

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

高速緩存操作的時(shí)間偏差

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

緩存預(yù)取

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

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

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

同時(shí)可以用來(lái)讀取Mini—Batch 2，3的數(shù)據(jù)，這部分開(kāi)銷(xiāo)可以和計(jì)算重疊。

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

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

綜上所述，緩存流水線(xiàn)預(yù)取帶來(lái)了兩個(gè)好處。

1.稀釋緩存索引開(kāi)銷(xiāo)

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

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

通過(guò)集中更多的數(shù)據(jù)，提高數(shù)據(jù)傳輸?shù)牧６?，從而充分利用CPU—GPU的傳輸帶寬上面的例子，CUDA—gt，CPU帶寬從860MB/s提升到1477 MB/s，CPU—gt，CUDA的帶寬從1257 MB/s提升到2415 MB/s，性能增益幾乎翻倍

使用方便

這與Pytorch嵌入包的用法是一致的在構(gòu)建推薦模型時(shí)，只需要初始化以下幾行代碼，可以大幅增加嵌入式表的容量，低成本實(shí)現(xiàn)TB級(jí)超大推薦模型訓(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)硬件平臺(tái)上，龐—AI以Meta DLRM模型為測(cè)試對(duì)象，使用超大型數(shù)據(jù)集Cretio 1TB和Meta dlrm_datasets作為測(cè)試模型。

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

Cretio 1TB

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

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

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

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

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

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

多GPU可擴(kuò)展性

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

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

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

可以看出，使用4個(gè)和8個(gè)GPU不會(huì)帶來(lái)顯著的性能提升，因?yàn)?

因此，同步需要巨大的通信開(kāi)銷(xiāo)。

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

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

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

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

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

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

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

還有一點(diǎn)

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

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

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

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

項(xiàng)目的開(kāi)放源地址:

參考鏈接:

這篇文章引用自: