近日,騰訊TEG數(shù)據(jù)平臺部機器學(xué)習(xí)團隊與北京大學(xué)-騰訊協(xié)同創(chuàng)新實驗室,合作研發(fā)了全新的稀疏大模型訓(xùn)練加速解決方案HET,其研究成果《HET: Scaling out Huge Embedding Model Training via Cache-enabled Distributed Framework》已被國際頂會VLDB 2022錄用。HET提出了一種新穎的基于Embedding緩存的訓(xùn)練方法,能夠顯著降低稀疏大模型分布式訓(xùn)練時通信開銷,提升模型訓(xùn)練整體效率。
稀疏大模型日益多見,通信瓶頸或成訓(xùn)練效率“致命”問題
圖1深度學(xué)習(xí)模型規(guī)模發(fā)展情況
稀疏大模型,是目前重要的深度學(xué)習(xí)模型類型之一,廣泛應(yīng)用在搜索廣告推薦、圖表示學(xué)習(xí)等場景。近年來,隨著數(shù)據(jù)規(guī)模的逐步增長,工業(yè)界稀疏大模型的規(guī)模日益龐大,參數(shù)量可以達到萬億規(guī)模。如圖1所示,Facebook今年提出的ZionEX[詳見注解1]系統(tǒng)所支持的推薦模型(DLRM)大小已經(jīng)超過10萬億規(guī)模,遠遠超過了Google之前發(fā)布的1.6萬億參數(shù)的Switch Transformer[詳見注解2]模型。
稀疏模型的參數(shù),即Embedding參數(shù),可以達到總模型參數(shù)量的99%以上。相比于其他模型,這類模型有著更低的計算密度和更大的模型規(guī)模,這也對分布式深度學(xué)習(xí)系統(tǒng)帶來了嚴峻的挑戰(zhàn)。近年來,如何提升稀疏大模型的訓(xùn)練效率逐漸成為了學(xué)術(shù)界和工業(yè)界都在關(guān)注的熱點問題。
對于萬億規(guī)模的模型,僅模型參數(shù)就需要3.7TB的內(nèi)存空間。由于稀疏大模型中的稀疏參數(shù)規(guī)模極大,因此工業(yè)界目前普遍采用基于參數(shù)服務(wù)器(Parameter Server)的解決方案,將Embedding均勻地切分到不同服務(wù)器上。在訓(xùn)練過程中,計算節(jié)點采用稀疏通信的形式,動態(tài)地從參數(shù)服務(wù)器上拉取所需的Embedding向量,完成當前輪次的計算后,再將Embedding的梯度提交回參數(shù)服務(wù)器。盡管這種方式可以靈活地擴展模型規(guī)模,但是也面臨著嚴重的通信瓶頸。以主流深度學(xué)習(xí)框架TensorFlow為例,在實際數(shù)據(jù)測試中,通信時間甚至?xí)嫉娇傆?xùn)練時間的80%以上。目前大多數(shù)改進方向是在參數(shù)服務(wù)器的工程實現(xiàn)上進行優(yōu)化,例如充分挖掘硬件性能來提高整個系統(tǒng)的吞吐率。然而并沒有從根本上解決稀疏參數(shù)通信量大的問題,通信仍然是系統(tǒng)的核心痛點。因此需要一種從源頭上解決通信問題的方案。
HET:基于Embedding緩存的稀疏大模型訓(xùn)練系統(tǒng)
核心思路
圖2三個常用公開數(shù)據(jù)集上的Embedding訪問頻率分布情況
根據(jù)來自在業(yè)務(wù)場景中的觀察,高維稀疏大模型的輸入數(shù)據(jù)特征往往具有傾斜分布的特性,具有冪律分布(如圖2所示),從而導(dǎo)致模型在訓(xùn)練過程中對Embedding向量的不均衡訪問。以推薦數(shù)據(jù)集Criteo為例,約有10%的Embedding向量引發(fā)了整個數(shù)據(jù)集90%的Embedding訪問。訓(xùn)練過程中,這些高頻Embedding會被頻繁地拉取和推送,成為了通信的主要負載。
我們利用這一特性,提出了Embedding緩存的思想:如果能夠在計算節(jié)點利用有限的內(nèi)存空間緩存這些高頻Embedding,那就有機會避免大量的遠程Embedding訪問,從而緩解通信瓶頸。根據(jù)這一思想,我們提出了基于Embedding緩存的新一代稀疏大模型訓(xùn)練框架HET。
技術(shù)點1:支持Embedding參數(shù)緩存的混合通信架構(gòu)
圖3 HET系統(tǒng)架構(gòu)
針對稀疏大模型參數(shù)中同時存在稀疏以及稠密部分的特點,HET整體上采用參數(shù)服務(wù)器(Parameter Server,PS)加全局規(guī)約(AllReduce)的混合通信架構(gòu),以充分發(fā)揮兩者優(yōu)勢,如圖3所示。其中AllReduce適合于稠密參數(shù)的同步,可借助NCCL等通信庫充分發(fā)揮GPU間帶寬,而參數(shù)服務(wù)器則天然支持稀疏通信,并且在同步協(xié)議上也具有較高的靈活性。同時,我們還在計算節(jié)點上設(shè)計了Cache Embedding Table結(jié)構(gòu),用于緩存高頻訪問的Embedding參數(shù)。
每個計算節(jié)點上采用Cache Embedding Table可節(jié)省大量的通信量,但是也帶來了一個新的問題,即對于某個特定的Embedding來說,其副本可能同時存在于多個不同的計算節(jié)點緩存當中。如果不考慮副本間的一致性,可能會導(dǎo)致模型訓(xùn)練發(fā)散,無法收斂。為此,我們進一步提出了一種基于細粒度Embedding時鐘的有限異步協(xié)議,來解決如何在不同的節(jié)點間同步這些Embedding副本的問題。
技術(shù)點2:基于細粒度Embedding時鐘的有限異步協(xié)議
圖4 HET中的Cache Embedding Table結(jié)構(gòu)
一般來說,Embedding參數(shù)采用表的方式進行組織以支持稀疏訪問。為了衡量Embedding副本間的一致性,我們對于每個Embedding向量,在常規(guī)的key-value數(shù)據(jù)結(jié)構(gòu)基礎(chǔ)之上,引入了一個重要的Lamport時鐘,用來記錄Embedding向量的狀態(tài)。在模型訓(xùn)練過程中,通過比較Embedding的時鐘,就可以知道該副本的延遲或超前程度。
圖5 HET中的Cache讀寫操作
對于Embedding緩存表,我們既允許讀取較為陳舊的Embedding,也允許延遲寫回緩存上的梯度更新。為了在充分發(fā)揮緩存加速效果的同時保證模型的訓(xùn)練質(zhì)量,我們限制了每個Embedding副本和全局Embedding間的時鐘相差不超過一個預(yù)先設(shè)定的閾值。在這種情況下,Embedding的每個副本都不會過于超前或落后于它的其他副本。
從全局視角來看,整個模型的稀疏和稠密部分分別采用不同的同步模式,稠密參數(shù)采用全同步協(xié)議進行通信,稀疏參數(shù)采用基于細粒度Embedding時鐘的有限異步協(xié)議進行通信。經(jīng)過理論分析,我們進一步證明了,這種基于細粒度Embedding時鐘的有限異步協(xié)議可以保證和全同步協(xié)議相似的收斂性。(詳情見論文鏈接)
實驗結(jié)果
我們將HET和基于傳統(tǒng)參數(shù)服務(wù)器架構(gòu)的TensorFlow以及同樣是參數(shù)服務(wù)器加全局規(guī)約的混合通信架構(gòu)的Parallax[詳見注解3]進行了對比,選取的數(shù)據(jù)集和模型包括:推薦模型Wide&Deep(WDL)、DeepFM(DFM)、Deep&Cross(DCN)和數(shù)據(jù)集Criteo,擁有三千多萬稀疏特征,當Embedding維度擴大到4K時,模型參數(shù)可以達到萬億級別;以及圖學(xué)習(xí)模型GraphSAGE和數(shù)據(jù)集Reddit、Amazon、ogbn-mag(OGB也是目前最權(quán)威的圖學(xué)習(xí)基準數(shù)據(jù)集之一,Open Graph Benchmark)。
端到端對比
圖6收斂效果對比
圖7端到端收斂速度對比
結(jié)合圖6和圖7,我們可以看出,在時鐘相差閾值上界設(shè)置為100時,相比于TensorFlow和Parallax,HET可以實現(xiàn)6.37-20.68倍的加速,并且不會對模型收斂性造成顯著影響。對于HET本身,細粒度Embedding緩存帶來了4.36-5.14倍的加速,最多可以減少88%的稀疏參數(shù)通信。
緩存效果對比:
圖8在不同Cache空間大小情況下Cache失效率情況
從圖8可以看出,只需要很少的Cache空間,比如15%的總參數(shù)量大小,就可以實現(xiàn)幾乎97%的緩存命中率,即97%的Embedding訪問都可以通過本地緩存訪問,而無需通信。另外我們也注意到,不同的Cache實現(xiàn)策略在效果上也稍有不同,LFU可以捕獲長期訪問傾向性,從而比LRU失效率更低。
可擴展性:
圖9在不同參數(shù)規(guī)模情況下的收斂效果
我們將模型擴展到32節(jié)點,Embedding維度設(shè)置到4096,此時總參數(shù)量已經(jīng)達到了萬億規(guī)模,從圖9可以看到HET執(zhí)行時間仍然顯著優(yōu)于其他基線方案,從而說明了HET的有效性。
騰訊TEG數(shù)據(jù)平臺部機器學(xué)習(xí)團隊:
該團隊致力于研發(fā)騰訊分布式機器學(xué)習(xí)平臺Angel,解決高維模型和稀疏數(shù)據(jù)訓(xùn)練問題。Angel誕生于騰訊大數(shù)據(jù)生態(tài)中,通過融合大數(shù)據(jù)、傳統(tǒng)機器學(xué)習(xí)和深度學(xué)習(xí)生態(tài),建立起一個端到端機器學(xué)習(xí)平臺,功能涵蓋傳統(tǒng)機器學(xué)習(xí)、圖挖掘、圖學(xué)習(xí)、深度學(xué)習(xí)和隱私計算等。在騰訊公司內(nèi)部, Angel已經(jīng)廣泛應(yīng)用于廣告推薦、金融風控、用戶畫像和短視頻推薦等業(yè)務(wù)。除了服務(wù)于公司內(nèi)部業(yè)務(wù)外,Angel于2017年對外開源,是國內(nèi)第一個LF AI基金會頂級項目。
針對模型規(guī)模擴大后對性能和擴展性帶來的挑戰(zhàn)。Angel平臺工程師聯(lián)合北京大學(xué)-騰訊協(xié)同創(chuàng)新實驗室,推出了稀疏大模型訓(xùn)練框架HEAP,并在廣告推薦全鏈路中,依托騰訊內(nèi)部業(yè)務(wù)需求對各種規(guī)模模型的訓(xùn)練提速增效,開展了包含基于Embedding緩存的新一代稀疏大模型訓(xùn)練,基于層次化參數(shù)服務(wù)器的萬億Embedding模型訓(xùn)練、多GPU分布式訓(xùn)練性能優(yōu)化等在內(nèi)的多項前瞻研究,并落地到廣告精排、粗排、預(yù)排序和召回等多個業(yè)務(wù)模型的訓(xùn)練中,在騰訊各個業(yè)務(wù)線條累計獲取GMV提升約4%。本次發(fā)表的研究成果HET也是在該框架下進行的新的探索。
北京大學(xué)-騰訊協(xié)同創(chuàng)新實驗室:
北京大學(xué)-騰訊協(xié)同創(chuàng)新實驗室成立于2017年,主要在人工智能、大數(shù)據(jù)等領(lǐng)域展開前沿探索和人才培養(yǎng),打造國際領(lǐng)先的校企合作科研平臺和產(chǎn)學(xué)研用基地。
實驗室通過合作研究,在理論和技術(shù)創(chuàng)新、系統(tǒng)研發(fā)和產(chǎn)業(yè)應(yīng)用方面取得重要成果和進展,已在國際頂級學(xué)術(shù)會議和期刊發(fā)表學(xué)術(shù)論文20余篇,除合作研發(fā)Angel外,實驗室還自主開發(fā)了多個開源系統(tǒng)。
今年8月,實驗室已經(jīng)宣布,將自研深度學(xué)習(xí)框架“河圖”融入Angel生態(tài),北京大學(xué)與騰訊團隊將聯(lián)合共建Angel4.0——新一代分布式深度學(xué)習(xí)平臺,面向擁有海量訓(xùn)練數(shù)據(jù)、超大模型參數(shù)的深度學(xué)習(xí)訓(xùn)練場景,為產(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.
（免責聲明：本網(wǎng)站內(nèi)容主要來自原創(chuàng)、合作伙伴供稿和第三方自媒體作者投稿，凡在本網(wǎng)站出現(xiàn)的信息，均僅供參考。本網(wǎng)站將盡力確保所提供信息的準確性及可靠性，但不保證有關(guān)資料的準確性及可靠性，讀者在使用前請進一步核實，并對任何自主決定的行為負責。本網(wǎng)站對有關(guān)資料所引致的錯誤、不確或遺漏，概不負任何法律責任。
任何單位或個人認為本網(wǎng)站中的網(wǎng)頁或鏈接內(nèi)容可能涉嫌侵犯其知識產(chǎn)權(quán)或存在不實內(nèi)容時，應(yīng)及時向本網(wǎng)站提出書面權(quán)利通知或不實情況說明，并提供身份證明、權(quán)屬證明及詳細侵權(quán)或不實情況證明。本網(wǎng)站在收到上述法律文件后，將會依法盡快聯(lián)系相關(guān)文章源頭核實，溝通刪除相關(guān)內(nèi)容或斷開相關(guān)鏈接。 ）
