一塊RTX3050搞定DLRM訓練！僅需1%Embedding參數，硬件

推薦模型已經成為互聯網公司深度學習應用最重要的技術場景，如視頻推薦，購物搜索，廣告推送等流量變現服務，大大提升了用戶體驗和商業價值。

可是，海量的用戶和業務數據，頻繁的迭代更新需求以及高昂的培訓成本給DLRM培訓帶來了嚴峻的挑戰。

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

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

借助GPU的片上高速內存和強大的計算能力，GPU已經成為DLRM培訓的主流硬件。

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

龐—AI已經成功使用異構策略，在相同硬件上提高NLP模型訓練的參數容量數百倍于之前最近，spool—AI成功將其擴展到推薦系統，通過軟件緩存的方式，將嵌入式表格動態存儲在CPU和GPU內存中

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

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

實驗表明，龐氏人工智能只需要在GPU中保留1%的嵌入參數，仍然可以保持優秀的端到端訓練速度。

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

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

開放源地址:

基于現有嵌入式表格擴展技術

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

首先為嵌入表中的每個特征搜索嵌入表的對應行，然后通過常規運算，如最大，均值，和運算，將特征向量轉化為特征向量，再傳遞給后續的密集神經網絡。

可以看出，DLRM的嵌入式表訓練過程主要是不規則的內存訪問操作，因此受到硬件內存訪問速度的嚴重限制。

而工業DLRM的嵌入式表可能會達到幾百GB甚至TB級別，遠遠超過單個GPU高達幾十GB的內存容量。

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

GPU模型并行性:

嵌入式表被分割分布在多個GPU的內存中，訓練時通過GPU之間的互聯網絡同步中間結果。

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

其次，增加GPU的前期硬件成本較高，DLRM訓練時GPU的計算能力沒有得到充分利用，只利用了其HBM帶寬優勢，導致GPU利用率較低。

部分CPU訓練:

嵌入式表分為兩部分，一部分在GPU上訓練，一部分在CPU上訓練。

利用數據分布的長尾效應，可以讓CPU計算比盡可能小，GPU計算比盡可能大但伴隨著批量的增加，很難讓mini—batch的所有數據都命中CPU或GPU，如果同時命中CPU或GPU，這種方法很難處理

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

此外，CPU和GPU需要傳輸中間結果，這也有很大的通信開銷，進一步拖慢了訓練速度。

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

軟件緩存:

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

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

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

目前有一些優秀的嵌入式表的軟件緩存方案，但往往是通過自定義的EmbeddingBags內核來實現，比如fbgemm，或者借助第三方深度學習框架來實現。

龐—AI在原生PyTorch的基礎上，不做任何內核級的改動，提供了一套開箱即用的軟件Cache EmbeddingBags實現，進一步優化了DLRM訓練過程，并提出預取流水線，進一步降低緩存開銷。

內存層次無損AI嵌入式表格軟件緩存

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

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

這部分內存空間稱為CPU權重但是，嵌入包的少量數據存儲在GPU內存中，其中包括用于訓練的數據

這部分內存空間稱為CUDA緩存權重。

在DLRM訓練期間，首先需要確定表中嵌入的行，這些行對應于在該迭代中輸入到小批量中的數據如果有些行不在GPU中，需要從CPU權重轉移到CUDA緩存權重

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

為了實現緩存檢索，需要一些輔助的數據結構:cached_idx_map是一個一維數組，存儲CPU權重中的行號和CUDA緩存權重的行號的對應關系，以及對應行在GPU中被訪問的頻率信息。

dacachedweight大小與CPU權重大小的比值命名為cache_ratio，默認值為1.0%。

每次迭代前運行緩存來調整CUDA中的數據權重，具體分三步。

第一步:CPU索引

檢索需要緩存的CPU權重中的行號。

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

第二步:GPU索引

根據使用頻率找到CUDA權重中可以被驅逐的行。

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

第三步:數據處理:

將CUDA緩存權重中的對應行移到CPU權重中，然后將CPU權重中的對應行移到CUDA權重中。

傳輸模塊負責CUDA緩存權重和CPU權重之間的數據雙向傳輸。

與低效的逐行傳輸不同，它采用先緩存再集中傳輸的方法，提高PCI—E的帶寬利用率。

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

龐—AI使用有限大小的緩沖區在CPU和GPU之間傳輸數據。

在最壞的情況下，所有的輸入id都沒有命中緩存，因此需要傳輸大量的元素為了防止緩沖區占用過多內存，緩沖區大小受到嚴格限制如果傳輸的數據大于緩沖區，傳輸將分多次完成

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

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

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

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

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

高速緩存操作的時間偏差

產生上述問題的原因是傳統的緩存策略短視，只能根據當前的小批量情況調整緩存，因此大部分時間浪費在查詢操作上。

緩存預取

為了降低緩存的開銷，龐—AI設計了一種前瞻性的緩存機制spool—AI不是只緩存以前的小批量，而是預取幾個以后要用的小批量，統一進行緩存查詢操作

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

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

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

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

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

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

1.稀釋緩存索引開銷

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

2.增加CPU—GPU數據移動帶寬

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

使用方便

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

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模型為測試對象，使用超大型數據集Cretio 1TB和Meta dlrm_datasets作為測試模型。

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

Cretio 1TB

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

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

可是，龐—AI仍然用于在單個GPU上完成訓練當緩存比=0.05時，內存消耗僅為5.01 GB，直接降低了18倍左右還可以進一步擴展，在單個GPU上實現TB級推薦系統模型訓練

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

綠色預取1不使用預取，藍色預取8是預取的延遲可見預取流水線優化對整體性能提升的重要作用

圖中每列的黑色部分是緩存開銷預取后，緩存開銷控制在總訓練時間的15%以內

多GPU可擴展性

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

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

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

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

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

表方向分片將導致不平衡的拆分負載也說明用多個GPU來擴展嵌入表的訓練擴展性不是很好

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

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

Meta Research的合成數據集dlrm_datasets模仿了業內嵌入式表的訓練訪問行為，因此在研究中經常作為推薦系統相關軟硬件設計的測試參考。

其中，選擇表項中嵌入的5億行作為子數據集，構建256GB和128GB兩個嵌入包進行測試。

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

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

還有一點

面向大模型時代的通用深度學習系統panol—AI，通過高效多維自動并行，異構內存管理，大規模優化庫，自適應任務調度等多項自主研發的領先技術，實現AI大模型訓練和推理的高效快速部署，降低了AI大模型的應用成本

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

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

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

項目的開放源地址:

參考鏈接:

這篇文章引用自: