在VLDB 2022可擴充資料科學最佳論文提出的基于分布式緩存的大規模可擴充嵌入模型訓練架構（HET）的啟發下，騰訊廣告基于北京大學與騰訊大資料團隊聯合共建的Angel4.0平台自研出AngelPS技術，并将其作為太極機器學習平台的核心元件落地應用于廣告系統。在AngelPS的助力下，太極機器學習平台單模型處理上限提升至10TB級别，更能實作7X24小時的線上深度學習與推理，為超大規模廣告模型的生産和使用提供強大基建支援。
騰訊廣告進一步履踐“技術提效”理念，基于太極機器學習平台，訓練出兩個千億維大模型——混元AI大模型、廣告大模型，強化了廣告系統的推薦效率與比對精度，讓前沿技術真正在業務場景上發揮效用。
論文一作苗旭鵬博士領驗證書
在2022年國際資料庫與資料管理頂級會議（International Conference on Very Large Databases，簡稱VLDB 2022）中，由北京大學-騰訊協同創新實驗室（以下簡稱聯合實驗室）提出的學術論文“基于分布式緩存的大規模可擴充嵌入模型訓練架構（HET: Scaling out Huge Embedding Model Training via Cache-enabled Distributed Framework）”獲得大會可擴充資料科學最佳論文獎（Best Scalable Data Science Paper）。
聯合實驗室主任崔斌教授也受邀在VLDB大會期間作特邀報告，介紹了課題組河圖團隊圍繞大規模機器學習以及深度學習系統優化，展開的系統研究工作。相關成果已陸續發表于SIGMOD、VLDB、ICDE、ICML、KDD、TKDE等國際頂級會議和期刊，得到了學術界和工業界的廣泛關注。此次獲獎，表明聯合實驗室在可擴充資料科學與機器學習系統領域的研究水準持續處于國際前沿水準。
獲獎證書
獲獎論文提出了一種新穎的基于Embedding緩存的訓練方法，能夠顯著降低稀疏大模型分布式訓練時通信開銷，提升模型訓練整體效率。
HET目前已正式開源：
https://github.com/PKU-DAIR/Hetu
論文連結：
https://www.vldb.org/pvldb/vol15/p312-miao.pdf
一、稀疏大模型日益多見通信瓶頸或成訓練效率“緻命”問題
圖1 深度學習模型規模發展情況
稀疏大模型，是目前重要的深度學習模型類型之一，廣泛應用在搜尋廣告推薦、圖表示學習等場景。近年來，随着資料規模的逐漸增長，工業界稀疏大模型的規模日益龐大，參數量可以達到萬億規模。如圖1所示，Facebook今年提出的ZionEX[詳見注解1]系統所支援的推薦模型（DLRM）大小已經超過10萬億規模，遠遠超過了Google之前釋出的1.6萬億參數的Switch Transformer[詳見注解2]模型。
稀疏模型的參數，即Embedding參數，可以達到總模型參數量的99%以上。相比于其他模型，這類模型有着更低的計算密度和更大的模型規模，這也對分布式深度學習系統帶來了嚴峻的挑戰。近年來，如何提升稀疏大模型的訓練效率逐漸成為了學術界和工業界都在關注的熱點問題。
對于萬億規模的模型，僅模型參數就需要3.7TB的記憶體空間。由于稀疏大模型中的稀疏參數規模極大，是以工業界目前普遍采用基于參數伺服器(Parameter Server)的解決方案，将Embedding均勻地切分到不同伺服器上。在訓練過程中，計算節點采用稀疏通信的形式，動态地從參數伺服器上拉取所需的Embedding向量，完成目前輪次的計算後，再将Embedding的梯度送出回參數伺服器。盡管這種方式可以靈活地擴充模型規模，但是也面臨着嚴重的通信瓶頸。
以主流深度學習架構TensorFlow為例，在實際資料測試中，通信時間甚至會占到總訓練時間的80%以上。目前大多數改進方向是在參數伺服器的工程實作上進行優化，例如充分挖掘硬體性能來提高整個系統的吞吐率。然而并沒有從根本上解決稀疏參數通信量大的問題，通信仍然是系統的核心痛點。是以需要一種從源頭上解決通信問題的方案。
二、HET：
基于Embedding緩存的稀疏大模型訓練系統
核心思路
圖2 三個常用公開資料集上的Embedding通路頻率分布情況
根據來自在業務場景中的觀察，高維稀疏大模型的輸入資料特征往往具有傾斜分布的特性，具有幂律分布（如圖2所示），進而導緻模型在訓練過程中對Embedding向量的不均衡通路。以推薦資料集Criteo為例，約有10%的Embedding向量引發了整個資料集90%的Embedding通路。訓練過程中，這些高頻Embedding會被頻繁地拉取和推送，成為了通信的主要負載。
我們利用這一特性，提出了Embedding緩存的思想：如果能夠在計算節點利用有限的記憶體空間緩存這些高頻Embedding，那就有機會避免大量的遠端Embedding通路，進而緩解通信瓶頸。根據這一思想，我們提出了基于Embedding緩存的新一代稀疏大模型訓練架構HET。
技術點1：支援Embedding參數緩存的混合通信架構
圖3 HET系統架構
針對稀疏大模型參數中同時存在稀疏以及稠密部分的特點，HET整體上采用參數伺服器（Parameter Server，PS）加全局規約（AllReduce）的混合通信架構，以充分發揮兩者優勢，如圖3所示。其中AllReduce适合于稠密參數的同步，可借助NCCL等通信庫充分發揮GPU間帶寬，而參數伺服器則天然支援稀疏通信，并且在同步協定上也具有較高的靈活性。同時，我們還在計算節點上設計了Cache Embedding Table結構，用于緩存高頻通路的Embedding參數。
每個計算節點上采用Cache Embedding Table可節省大量的通信量，但是也帶來了一個新的問題，即對于某個特定的Embedding來說，其副本可能同時存在于多個不同的計算節點緩存當中。如果不考慮副本間的一緻性，可能會導緻模型訓練發散，無法收斂。為此，我們進一步提出了一種基于細粒度Embedding時鐘的有限異步協定，來解決如何在不同的節點間同步這些Embedding副本的問題。
技術點2：基于細粒度Embedding時鐘的有限異步協定
圖4 HET中的Cache Embedding Table結構
一般來說，Embedding參數采用表的方式進行組織以支援稀疏通路。為了衡量Embedding副本間的一緻性，我們對于每個Embedding向量，在正常的key-value資料結構基礎之上，引入了一個重要的Lamport時鐘，用來記錄Embedding向量的狀态。在模型訓練過程中，通過比較Embedding的時鐘，就可以知道該副本的延遲或超前程度。
圖5 HET中的Cache讀寫操作
對于Embedding緩存表，我們既允許讀取較為陳舊的Embedding，也允許延遲寫回緩存上的梯度更新。為了在充分發揮緩存加速效果的同時保證模型的訓練品質，我們限制了每個Embedding副本和全局Embedding間的時鐘相差不超過一個預先設定的門檻值。在這種情況下，Embedding的每個副本都不會過于超前或落後于它的其他副本。
從全局視角來看，整個模型的稀疏和稠密部分分别采用不同的同步模式，稠密參數采用全同步協定進行通信，稀疏參數采用基于細粒度Embedding時鐘的有限異步協定進行通信。經過理論分析，我們進一步證明了，這種基于細粒度Embedding時鐘的有限異步協定可以保證和全同步協定相似的收斂性。（詳見論文原文）
實驗結果
我們将HET和基于傳統參數伺服器架構的TensorFlow以及同樣是參數伺服器加全局規約的混合通信架構的Parallax[詳見注解3]進行了對比，選取的資料集和模型包括：推薦模型Wide&Deep（WDL）、DeepFM（DFM）、Deep&Cross（DCN）和資料集Criteo，擁有三千多萬稀疏特征，當Embedding次元擴大到4K時，模型參數可以達到萬億級别；以及圖學習模型GraphSAGE和資料集Reddit、Amazon、ogbn-mag（OGB也是目前最權威的圖學習基準資料集之一，Open Graph Benchmark）。
端到端對比
圖6 收斂效果對比
圖7 端到端收斂速度對比
結合圖6和圖7，我們可以看出，在時鐘相差門檻值上界設定為100時，相比于TensorFlow和Parallax，HET可以實作6.37-20.68倍的加速，并且不會對模型收斂性造成顯著影響。對于HET本身，細粒度Embedding緩存帶來了4.36-5.14倍的加速，最多可以減少88%的稀疏參數通信。
緩存效果對比
圖8 在不同Cache空間大小情況下Cache失效率情況
從圖8可以看出，隻需要很少的Cache空間，比如15%的總參數量大小，就可以實作幾乎97%的緩存命中率，即97%的Embedding通路都可以通過本地緩存通路，而無需通信。另外我們也注意到，不同的Cache實作政策在效果上也稍有不同，LFU可以捕獲長期通路傾向性，進而比LRU失效率更低。
可擴充性
圖9 在不同參數規模情況下的收斂效果
我們将模型擴充到32節點，Embedding次元設定到4096，此時總參數量已經達到了萬億規模，從圖9可以看到HET執行時間仍然顯著優于其他基線方案，進而說明了HET的有效性。
作者：騰訊廣告
文章來源：公衆号【騰訊廣告算法大賽】
Illustration by IconScout Store from IconScout
-The End-
掃碼觀看！
本周上新！
關于我“門”
▼
将門是一家以專注于發掘、加速及投資技術驅動型創業公司的新型創投機構，旗下涵蓋将門創新服務、将門技術社群以及将門創投基金。
将門成立于2015年底，創始團隊由微軟創投在中國的創始團隊原班人馬建構而成，曾為微軟優選和深度孵化了126家創新的技術型創業公司。
如果您是技術領域的初創企業，不僅想獲得投資，還希望獲得一系列持續性、有價值的投後服務，
歡迎發送或者推薦項目給我“門”:
⤵一鍵送你進入TechBeat快樂星球