近日,騰訊TEG數(shù)據(jù)平臺(tái)部機(jī)器學(xué)習(xí)團(tuán)隊(duì)與北京大學(xué)-騰訊協(xié)同創(chuàng)新實(shí)驗(yàn)室,合作研發(fā)了全新的稀疏大模型訓(xùn)練加速解決方案HET,其研究成果《HET: Scaling out Huge Embedding Model Training via Cache-enabled Distributed Framework》已被國(guó)際頂會(huì)VLDB 2022錄用。HET提出了一種新穎的基于Embedding緩存的訓(xùn)練方法,能夠顯著降低稀疏大模型分布式訓(xùn)練時(shí)通信開銷,提升模型訓(xùn)練整體效率。

稀疏大模型日益多見,通信瓶頸或成訓(xùn)練效率“致命”問題

圖1深度學(xué)習(xí)模型規(guī)模發(fā)展情況

稀疏大模型,是目前重要的深度學(xué)習(xí)模型類型之一,廣泛應(yīng)用在搜索廣告推薦、圖表示學(xué)習(xí)等場(chǎng)景。近年來,隨著數(shù)據(jù)規(guī)模的逐步增長(zhǎng),工業(yè)界稀疏大模型的規(guī)模日益龐大,參數(shù)量可以達(dá)到萬(wàn)億規(guī)模。如圖1所示,Facebook今年提出的ZionEX[詳見注解1]系統(tǒng)所支持的推薦模型(DLRM)大小已經(jīng)超過10萬(wàn)億規(guī)模,遠(yuǎn)遠(yuǎn)超過了Google之前發(fā)布的1.6萬(wàn)億參數(shù)的Switch Transformer[詳見注解2]模型。

稀疏模型的參數(shù),即Embedding參數(shù),可以達(dá)到總模型參數(shù)量的99%以上。相比于其他模型,這類模型有著更低的計(jì)算密度和更大的模型規(guī)模,這也對(duì)分布式深度學(xué)習(xí)系統(tǒng)帶來了嚴(yán)峻的挑戰(zhàn)。近年來,如何提升稀疏大模型的訓(xùn)練效率逐漸成為了學(xué)術(shù)界和工業(yè)界都在關(guān)注的熱點(diǎn)問題。

對(duì)于萬(wàn)億規(guī)模的模型,僅模型參數(shù)就需要3.7TB的內(nèi)存空間。由于稀疏大模型中的稀疏參數(shù)規(guī)模極大,因此工業(yè)界目前普遍采用基于參數(shù)服務(wù)器(Parameter Server)的解決方案,將Embedding均勻地切分到不同服務(wù)器上。在訓(xùn)練過程中,計(jì)算節(jié)點(diǎn)采用稀疏通信的形式,動(dòng)態(tài)地從參數(shù)服務(wù)器上拉取所需的Embedding向量,完成當(dāng)前輪次的計(jì)算后,再將Embedding的梯度提交回參數(shù)服務(wù)器。盡管這種方式可以靈活地?cái)U(kuò)展模型規(guī)模,但是也面臨著嚴(yán)重的通信瓶頸。以主流深度學(xué)習(xí)框架TensorFlow為例,在實(shí)際數(shù)據(jù)測(cè)試中,通信時(shí)間甚至?xí)嫉娇傆?xùn)練時(shí)間的80%以上。目前大多數(shù)改進(jìn)方向是在參數(shù)服務(wù)器的工程實(shí)現(xiàn)上進(jìn)行優(yōu)化,例如充分挖掘硬件性能來提高整個(gè)系統(tǒng)的吞吐率。然而并沒有從根本上解決稀疏參數(shù)通信量大的問題,通信仍然是系統(tǒng)的核心痛點(diǎn)。因此需要一種從源頭上解決通信問題的方案。

HET:基于Embedding緩存的稀疏大模型訓(xùn)練系統(tǒng)

核心思路

圖2三個(gè)常用公開數(shù)據(jù)集上的Embedding訪問頻率分布情況

根據(jù)來自在業(yè)務(wù)場(chǎng)景中的觀察,高維稀疏大模型的輸入數(shù)據(jù)特征往往具有傾斜分布的特性,具有冪律分布(如圖2所示),從而導(dǎo)致模型在訓(xùn)練過程中對(duì)Embedding向量的不均衡訪問。以推薦數(shù)據(jù)集Criteo為例,約有10%的Embedding向量引發(fā)了整個(gè)數(shù)據(jù)集90%的Embedding訪問。訓(xùn)練過程中,這些高頻Embedding會(huì)被頻繁地拉取和推送,成為了通信的主要負(fù)載。

我們利用這一特性,提出了Embedding緩存的思想:如果能夠在計(jì)算節(jié)點(diǎn)利用有限的內(nèi)存空間緩存這些高頻Embedding,那就有機(jī)會(huì)避免大量的遠(yuǎn)程Embedding訪問,從而緩解通信瓶頸。根據(jù)這一思想,我們提出了基于Embedding緩存的新一代稀疏大模型訓(xùn)練框架HET。

技術(shù)點(diǎn)1:支持Embedding參數(shù)緩存的混合通信架構(gòu)

圖3 HET系統(tǒng)架構(gòu)

針對(duì)稀疏大模型參數(shù)中同時(shí)存在稀疏以及稠密部分的特點(diǎn),HET整體上采用參數(shù)服務(wù)器(Parameter Server,PS)加全局規(guī)約(AllReduce)的混合通信架構(gòu),以充分發(fā)揮兩者優(yōu)勢(shì),如圖3所示。其中AllReduce適合于稠密參數(shù)的同步,可借助NCCL等通信庫(kù)充分發(fā)揮GPU間帶寬,而參數(shù)服務(wù)器則天然支持稀疏通信,并且在同步協(xié)議上也具有較高的靈活性。同時(shí),我們還在計(jì)算節(jié)點(diǎn)上設(shè)計(jì)了Cache Embedding Table結(jié)構(gòu),用于緩存高頻訪問的Embedding參數(shù)。

每個(gè)計(jì)算節(jié)點(diǎn)上采用Cache Embedding Table可節(jié)省大量的通信量,但是也帶來了一個(gè)新的問題,即對(duì)于某個(gè)特定的Embedding來說,其副本可能同時(shí)存在于多個(gè)不同的計(jì)算節(jié)點(diǎn)緩存當(dāng)中。如果不考慮副本間的一致性,可能會(huì)導(dǎo)致模型訓(xùn)練發(fā)散,無法收斂。為此,我們進(jìn)一步提出了一種基于細(xì)粒度Embedding時(shí)鐘的有限異步協(xié)議,來解決如何在不同的節(jié)點(diǎn)間同步這些Embedding副本的問題。

技術(shù)點(diǎn)2:基于細(xì)粒度Embedding時(shí)鐘的有限異步協(xié)議

圖4 HET中的Cache Embedding Table結(jié)構(gòu)

一般來說,Embedding參數(shù)采用表的方式進(jìn)行組織以支持稀疏訪問。為了衡量Embedding副本間的一致性,我們對(duì)于每個(gè)Embedding向量,在常規(guī)的key-value數(shù)據(jù)結(jié)構(gòu)基礎(chǔ)之上,引入了一個(gè)重要的Lamport時(shí)鐘,用來記錄Embedding向量的狀態(tài)。在模型訓(xùn)練過程中,通過比較Embedding的時(shí)鐘,就可以知道該副本的延遲或超前程度。

圖5 HET中的Cache讀寫操作

對(duì)于Embedding緩存表,我們既允許讀取較為陳舊的Embedding,也允許延遲寫回緩存上的梯度更新。為了在充分發(fā)揮緩存加速效果的同時(shí)保證模型的訓(xùn)練質(zhì)量,我們限制了每個(gè)Embedding副本和全局Embedding間的時(shí)鐘相差不超過一個(gè)預(yù)先設(shè)定的閾值。在這種情況下,Embedding的每個(gè)副本都不會(huì)過于超前或落后于它的其他副本。

從全局視角來看,整個(gè)模型的稀疏和稠密部分分別采用不同的同步模式,稠密參數(shù)采用全同步協(xié)議進(jìn)行通信,稀疏參數(shù)采用基于細(xì)粒度Embedding時(shí)鐘的有限異步協(xié)議進(jìn)行通信。經(jīng)過理論分析,我們進(jìn)一步證明了,這種基于細(xì)粒度Embedding時(shí)鐘的有限異步協(xié)議可以保證和全同步協(xié)議相似的收斂性。(詳情見論文鏈接)

實(shí)驗(yàn)結(jié)果

我們將HET和基于傳統(tǒng)參數(shù)服務(wù)器架構(gòu)的TensorFlow以及同樣是參數(shù)服務(wù)器加全局規(guī)約的混合通信架構(gòu)的Parallax[詳見注解3]進(jìn)行了對(duì)比,選取的數(shù)據(jù)集和模型包括:推薦模型Wide&Deep(WDL)、DeepFM(DFM)、Deep&Cross(DCN)和數(shù)據(jù)集Criteo,擁有三千多萬(wàn)稀疏特征,當(dāng)Embedding維度擴(kuò)大到4K時(shí),模型參數(shù)可以達(dá)到萬(wàn)億級(jí)別;以及圖學(xué)習(xí)模型GraphSAGE和數(shù)據(jù)集Reddit、Amazon、ogbn-mag(OGB也是目前最權(quán)威的圖學(xué)習(xí)基準(zhǔn)數(shù)據(jù)集之一,Open Graph Benchmark)。

端到端對(duì)比

圖6收斂效果對(duì)比

圖7端到端收斂速度對(duì)比

結(jié)合圖6和圖7,我們可以看出,在時(shí)鐘相差閾值上界設(shè)置為100時(shí),相比于TensorFlow和Parallax,HET可以實(shí)現(xiàn)6.37-20.68倍的加速,并且不會(huì)對(duì)模型收斂性造成顯著影響。對(duì)于HET本身,細(xì)粒度Embedding緩存帶來了4.36-5.14倍的加速,最多可以減少88%的稀疏參數(shù)通信。

緩存效果對(duì)比:

圖8在不同Cache空間大小情況下Cache失效率情況

從圖8可以看出,只需要很少的Cache空間,比如15%的總參數(shù)量大小,就可以實(shí)現(xiàn)幾乎97%的緩存命中率,即97%的Embedding訪問都可以通過本地緩存訪問,而無需通信。另外我們也注意到,不同的Cache實(shí)現(xiàn)策略在效果上也稍有不同,LFU可以捕獲長(zhǎng)期訪問傾向性,從而比LRU失效率更低。

可擴(kuò)展性:

圖9在不同參數(shù)規(guī)模情況下的收斂效果

我們將模型擴(kuò)展到32節(jié)點(diǎn),Embedding維度設(shè)置到4096,此時(shí)總參數(shù)量已經(jīng)達(dá)到了萬(wàn)億規(guī)模,從圖9可以看到HET執(zhí)行時(shí)間仍然顯著優(yōu)于其他基線方案,從而說明了HET的有效性。

騰訊TEG數(shù)據(jù)平臺(tái)部機(jī)器學(xué)習(xí)團(tuán)隊(duì):

該團(tuán)隊(duì)致力于研發(fā)騰訊分布式機(jī)器學(xué)習(xí)平臺(tái)Angel,解決高維模型和稀疏數(shù)據(jù)訓(xùn)練問題。Angel誕生于騰訊大數(shù)據(jù)生態(tài)中,通過融合大數(shù)據(jù)、傳統(tǒng)機(jī)器學(xué)習(xí)和深度學(xué)習(xí)生態(tài),建立起一個(gè)端到端機(jī)器學(xué)習(xí)平臺(tái),功能涵蓋傳統(tǒng)機(jī)器學(xué)習(xí)、圖挖掘、圖學(xué)習(xí)、深度學(xué)習(xí)和隱私計(jì)算等。在騰訊公司內(nèi)部, Angel已經(jīng)廣泛應(yīng)用于廣告推薦、金融風(fēng)控、用戶畫像和短視頻推薦等業(yè)務(wù)。除了服務(wù)于公司內(nèi)部業(yè)務(wù)外,Angel于2017年對(duì)外開源,是國(guó)內(nèi)第一個(gè)LF AI基金會(huì)頂級(jí)項(xiàng)目。

針對(duì)模型規(guī)模擴(kuò)大后對(duì)性能和擴(kuò)展性帶來的挑戰(zhàn)。Angel平臺(tái)工程師聯(lián)合北京大學(xué)-騰訊協(xié)同創(chuàng)新實(shí)驗(yàn)室,推出了稀疏大模型訓(xùn)練框架HEAP,并在廣告推薦全鏈路中,依托騰訊內(nèi)部業(yè)務(wù)需求對(duì)各種規(guī)模模型的訓(xùn)練提速增效,開展了包含基于Embedding緩存的新一代稀疏大模型訓(xùn)練,基于層次化參數(shù)服務(wù)器的萬(wàn)億Embedding模型訓(xùn)練、多GPU分布式訓(xùn)練性能優(yōu)化等在內(nèi)的多項(xiàng)前瞻研究,并落地到廣告精排、粗排、預(yù)排序和召回等多個(gè)業(yè)務(wù)模型的訓(xùn)練中,在騰訊各個(gè)業(yè)務(wù)線條累計(jì)獲取GMV提升約4%。本次發(fā)表的研究成果HET也是在該框架下進(jìn)行的新的探索。

北京大學(xué)-騰訊協(xié)同創(chuàng)新實(shí)驗(yàn)室:

北京大學(xué)-騰訊協(xié)同創(chuàng)新實(shí)驗(yàn)室成立于2017年,主要在人工智能、大數(shù)據(jù)等領(lǐng)域展開前沿探索和人才培養(yǎng),打造國(guó)際領(lǐng)先的校企合作科研平臺(tái)和產(chǎn)學(xué)研用基地。

實(shí)驗(yàn)室通過合作研究,在理論和技術(shù)創(chuàng)新、系統(tǒng)研發(fā)和產(chǎn)業(yè)應(yīng)用方面取得重要成果和進(jìn)展,已在國(guó)際頂級(jí)學(xué)術(shù)會(huì)議和期刊發(fā)表學(xué)術(shù)論文20余篇,除合作研發(fā)Angel外,實(shí)驗(yàn)室還自主開發(fā)了多個(gè)開源系統(tǒng)。

今年8月,實(shí)驗(yàn)室已經(jīng)宣布,將自研深度學(xué)習(xí)框架“河圖”融入Angel生態(tài),北京大學(xué)與騰訊團(tuán)隊(duì)將聯(lián)合共建Angel4.0——新一代分布式深度學(xué)習(xí)平臺(tái),面向擁有海量訓(xùn)練數(shù)據(jù)、超大模型參數(shù)的深度學(xué)習(xí)訓(xùn)練場(chǎng)景,為產(chǎn)業(yè)界帶來新的大規(guī)模深度學(xué)習(xí)破局之策。

參考資料:

[1] Mudigere D, Hao Y, Huang J, et al. High-performance, distributed training of large-scale deep learning recommendation models[J]. arXiv preprint arXiv:2104.05158, 2021.

[2] Fedus W, Zoph B, Shazeer N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity[J]. arXiv preprint arXiv:2101.03961, 2021.

[3] Kim S, Yu G I, Park H, et al. Parallax: Sparsity-aware data parallel training of deep neural networks[C]//Proceedings of the Fourteenth EuroSys Conference 2019. 2019: 1-15.

（免責(zé)聲明：本網(wǎng)站內(nèi)容主要來自原創(chuàng)、合作伙伴供稿和第三方自媒體作者投稿，凡在本網(wǎng)站出現(xiàn)的信息，均僅供參考。本網(wǎng)站將盡力確保所提供信息的準(zhǔn)確性及可靠性，但不保證有關(guān)資料的準(zhǔn)確性及可靠性，讀者在使用前請(qǐng)進(jìn)一步核實(shí)，并對(duì)任何自主決定的行為負(fù)責(zé)。本網(wǎng)站對(duì)有關(guān)資料所引致的錯(cuò)誤、不確或遺漏，概不負(fù)任何法律責(zé)任。
任何單位或個(gè)人認(rèn)為本網(wǎng)站中的網(wǎng)頁(yè)或鏈接內(nèi)容可能涉嫌侵犯其知識(shí)產(chǎn)權(quán)或存在不實(shí)內(nèi)容時(shí)，應(yīng)及時(shí)向本網(wǎng)站提出書面權(quán)利通知或不實(shí)情況說明，并提供身份證明、權(quán)屬證明及詳細(xì)侵權(quán)或不實(shí)情況證明。本網(wǎng)站在收到上述法律文件后，將會(huì)依法盡快聯(lián)系相關(guān)文章源頭核實(shí)，溝通刪除相關(guān)內(nèi)容或斷開相關(guān)鏈接。 ）