近年來，隨著稀疏模型對算力日益增長的需求, CPU叢集必須不斷擴大叢集規模來滿足訓練的時效需求，這同時也帶來了不斷上升的資源成本以及實驗的除錯成本。
為了解決這一問題，阿里雲機器學習PAI平臺和阿里媽媽智慧引擎訓練引擎團隊合作開發了稀疏模型高效能同步訓練框架HybridBackend，使得在同成本下GPU叢集訓練吞吐較CPU叢集提升至5倍，大幅降低除錯成本，同時 HybridBackend 相關論文 《PICASSO: Unleashing the Potential of GPU-centric Training for Wide-and-deep Recommender Systems》也被 ICDE 22' 所收錄。HybridBackend背後的技術框架如何設計？未來有哪些規劃？今天一起來深入瞭解。
HybridBackend是阿里雲機器學習平臺PA
和阿里媽媽智慧引擎訓練引擎團隊合作開發的、面向稀疏模型訓練的高效能同步訓練框架，核心能力是大幅提升GPU叢集單位成本下的訓練吞吐效能。目前HybridBackend已經在阿里巴巴集團內部有多個業務落地，將阿里媽媽智慧引擎訓練引擎團隊的定向廣告業務年資料訓練任務時間由1個月縮短至2天，同時HybridBackend在公有云多個頭部網際網路企業中也有成功應用。
以搜尋、推薦、廣告業務為主要應用的稀疏模型訓練系統一直是學界和業界研究的熱點之一。相比於計算機視覺(CV)和自然語言處理(NLP)為代表的稠密模型訓練，稀疏模型針對離散型特徵(以 categorical ID 作為訓練資料)使用Embedding特徵表達有著百GB至數十TB級別的記憶體佔用消耗(比普通的CV、NLP模型引數高出一到兩個數量級),從而突破了單機的記憶體容量限制，需要基於分散式系統的訓練方案。
早期的此類分散式任務由於模型結構相對簡單並且更新迭代緩慢,往往採用定製化的引數伺服器(Parameter Server,PS)系統在大規模的CPU叢集上進行訓練。隨著TensorFlow為代表的通用機器學習程式設計框架的出現，以及深度神經網路(DNN)在推薦類模型上的流行(deep recommender systems),業界逐漸轉向基於通用機器學習程式設計框架(TensorFlow、PyTorch等)進行模型的端到端訓練和推理，但是此時依然以引數伺服器(PS)和大規模CPU叢集作為訓練的正規化和基礎設施。
隨著稀疏模型對算力日益增長的需求(比如Attention等結構的加入),CPU叢集
必須不斷擴大叢集規模來滿足訓練的時效需求，這同時也帶來了不斷上升的資源成本以及實驗的除錯成本。
以NVIDIA GPU為代表的加速器(accelerator)彌補了CPU裝置單位成本算力低下
的劣勢，在CV、NLP等算力需求大的訓練任務上的應用已經成為行業共識。然而實踐證明，如只是簡單地將PS訓練正規化中的worker從CPU裝置替換為GPU裝置，並不能有效地提升訓練任務的吞吐，透過 profiling GPU 的使用率，發現大量的GPU算力資源被閒置浪費。這說明，相比於CV、NLP類任務，稀疏模型訓練有著自身的模型結構和訓練資料的特性，使得傳統的PS訓練正規化不能有效地發揮出GPU裝置的優勢。以深度推薦系統經典的 Wide and Deep 模型結構和TensorFlow框架為例,我們分析並總結了在PS架構下使用GPU裝置訓練的兩個問題。
從上圖的 Wide and Deep 模型結構可以看出，稀疏訓練主要由Embedding階段、
特徵交叉(feature interation)階段和多層感知器(MLP)階段組成，Embedding階段在PS正規化的訓練下佔據了至少50%以上的訓練時間。經過分析發現，Embedding階段的運算元主要以訪存密集型(memory access intensive)和通訊密集型的運算元(communication intensive)為主，主要需要的硬體資源是記憶體和網路的頻寬，而後兩個階段的運算元則是計算密集型的運算元佔主導,需要的資源是算力。這意味著在PS的正規化訓練下，任何一個階段都有可能存在某一種硬體資源成為瓶頸而其他硬體資源被浪費的現象。以GPU的算力資源為例，我們觀察GPU使用率(SM Util)在不同的訓練階段之間呈現脈衝式變化(pulse)。
生產實際中的模型往往擁有上百路的Embedding特徵查詢，每一路的特徵查詢在
TensorFlow內都會呼叫數十個運算元操作(operations)。TensorFlow的引擎在排程上千級別的大量的運算元操作需要額外的CPU執行緒開銷;對於GPU裝置來說,過多的 CUDA kernel 提交到流處理器上(TensorFlow下每個GPU裝置只有一個stream抽象)帶來了GPU Stream Multiprocessor (SM)的排程開銷，同時每個運算元處理資料的併發度又不高，從而很難打滿GPU的計算單元。
類似的問題在CV、NLP等稠密模型的訓練中也有涉及，一般採用基於編譯技術的
最佳化手段進行運算元合併。在 Wide and Deep 模型這樣的稀疏場景下，Embedding階段的這些運算元又往往具有 dynamic shape 的特點，在TensorFlow靜態構圖階段無法獲取準確的運算元尺寸進行最佳化，導致類似TensorFlow-XLA等技術在此類場景下沒有明顯的收益。
這些問題說明，想要發揮出GPU等高效能硬體資源的極致價效比，提高單位成本
下的訓練吞吐，就必須設計新的訓練框架。據我們瞭解，擁有大型搜尋、廣告、推薦業務的國內外企業以及硬體廠商都在著手進行新框架的研發，比如NVIDIA的Merlin-HugeCTR[1]等，然而阿里巴巴集團內雲上叢集普遍部署的是通用計算節點，且叢集上需要執行多種異構的任務，換用專用硬體是很昂貴且不切實際的。
基於這種實際需求，我們推出了HybridBackend，能夠同時適應集團內多元化且
不斷演進的稀疏模型技術。下文中我們將簡要介紹HybridBackend的系統架構設計和技術亮點。
傳統的引數伺服器(PS)訓練正規化體現的是透過擴充套件硬體數量來適應模型訓練規
模的思路。我們的系統則是同時考慮到了硬體和軟體(模型)兩個層面的特點，並做到協同設計。高效能GPU叢集的硬體特性決定了基本的訓練正規化，而稀疏模型本身的結構特點和資料分佈帶來的問題則透過更精細的系統最佳化手段來解決。
因為GPU裝置相對於CPU帶來的巨大的算力提升，以往需要上百臺CPU節點的集
群可以用幾十臺機器的GPU叢集來代替。要保持相同的總訓練規模，同時提升單個GPU節點上的資源利用率，提升單個 GPU worker 上的 batch size 成為必然的選項。同時，因為叢集規模的縮小，可以透過同步訓練的方式有效避免過期梯度(staleness)，從而提升模型訓練的精度。
相對於CPU裝置之間透過PCIe以及TCP進行網路通訊，高效能的GPU叢集在單
個節點內的多個GPU裝置之間往往配備了高速的網路互連(NVLink、NVSwitch),這些高速連線的頻寬通常是TCP網路頻寬的數百倍(第一代NVLINK標定達到300GB/s),而在多個機器節點之間也可以配備基於RDMA技術的高速網路裝置，達到100-200Gbps的頻寬。
選擇同步訓練的第二個好處是，可以使用高效能集合通訊運算元庫(NVIDIA NCCL、
阿里自研的ACCL等)來有效利用硬體機器的網路拓撲結構，從而提升通訊的效能。上述通訊庫已經在CV、NLP之類的基於資料並行的同步訓練任務上取得了很好的效果。
PS訓練正規化在系統的邏輯層面會指定不同的訓練角色,比如server、worker、
evaluator等。server節點一般分配具有大記憶體的CPU機器，而worker節點則會被分配到高主頻的計算型CPU硬體上。這樣形成了訓練單元-任務角色-同構資源的耦合，透過增加訓練單元數量來水平擴充套件(scale out)訓練的規模。
而在高效能的GPU叢集上，一個物理的機器節點往往包括多種異構的硬體資源，
如CPU、GPU處理器、GPU之間的高速互連、DRAM（動態隨機存取記憶體）、Non-volatile Memory（非易失性記憶體）等。這樣，除了水平擴充套件節點數量外，還可以透過垂直擴充套件利用多種異構硬體資源來達到擴大訓練規模的目標。
針對這種硬體架構，我們的系統設計中只保留統一的訓練執行單元(Executor),
每個Executor透過內部的異構硬體資源來執行不同的訓練任務角色。一方面，Executor內部任務執行時，可以有效地利用底層硬體資源之間的locality加速訓練；另一方面，Executor內部的硬體資源可以同時滿足不同的分散式訓練正規化所需要的硬體資源，以方便我們在模型結構的不同部分進行混合並行訓練策略。
五 深入最佳化：HybridBackend的技術亮點
因為稀疏模型結構和訓練資料本身的特性, 變化的硬體資源瓶頸和運算元細碎化，上述的系統架構在實際任務中還是會存在一些影響GPU等硬體裝置使用率的問題。
舉例來說，同步訓練正規化下，所有Executor透過集合通訊進行embedding的shuffle時，網路頻寬資源成為瓶頸，而GPU的計算資源被閒置。一種解決思路
是對硬體資源進行定製化，比如增加網路頻寬資源來消除通訊瓶頸，但是這樣的做法會使得硬體的資源配置和特定的模型結構耦合，是專用推薦系統的老思路。
我們的目標還是希望系統可以架構在雲服務上可得的，數量容易水平擴充套件的通用
硬體配置之上(commodity hardware)。某些硬體廠商也嘗試透過 Huge kernel 的形式(將Embedding層所有的計算手工融合到一個kernel內)來解決運算元細碎化的問題,這樣的做法也很難支援模型結構快速迭代的需求,背離了通用程式設計架構的設計初衷。
據此，我們從軟硬協同的思路出發，設計瞭如下的幾個系統最佳化手段:
根據稀疏模型的結構特點，大部分細碎的運算元來源於龐大的Embedding特徵查詢(lookup)數量，我們設計了D-Packing這一最佳化技術。
對於每一路查詢，儘管輸入的訓練資料不同，但使用的運算元組合是相同的。對於
這種具有資料並行特點的模式，具有相同屬性（維度、初始化器、標定特徵組等）的Embedding表將被合併為一張新的Embedding表，而後續的訪存查詢運算元也可以被合併為一個新的大運算元。合併運算元可以用多執行緒的方式有序查詢Embedding，相對於亂序查詢或分成若干小表查詢，能有顯著的效能提升。查詢完畢後，再依原有程式碼需要進行反去重和歸位，真正做到了對使用者透明。
此外，透過分析特徵查詢階段各個運算元在分散式環境下的語義，我們將部分的
kernel進行融合K-Packing,比如透過融合shuffle和stitch運算元來消除冗餘的資料複製。
透過資料和運算元兩個維度的基於語義的融合，我們既減少了總體的運算元數量，降
低fragmentation,同時又避免了所有運算元融合在一起而丟失了透過運算元間穿插遮掩來提升硬體利用率的最佳化機會。
為了消除同時執行相同硬體資源需求的運算元而造成的瓶頸, 我們設計了兩種運算元穿插遮掩執行(interleaving)的最佳化手段。
其一，D-Interleaving是透過對訓練資料batch的切分利用pipeline的機制來排程穿插不同資源型別的運算元，這樣可以在訓練的任何階段緩解某一種資源的瓶頸。比如在大batch size的訓練場景下，稀疏模型的MLP階段也會產生很高的feature map視訊記憶體佔用，透過D-Interleaving就可以有效降低單個GPU裝置上的峰值視訊記憶體佔用，從而使得更大的batch size訓練成為可能。
其二，K-Interleaving是在Embedding Layer內部不同的特徵查詢路數之間做運算元的穿插和遮掩，比如將通訊密集的Shuffle操作和記憶體訪問密集的Gather進行遮掩，可以有效提升這兩種資源的使用率。
在解決Executor內部多個級別的儲存(GPU視訊記憶體、DRAM等)之間的頻寬和延遲問
題上，我們針對稀疏模型訓練資料的分佈特點，提出了一種感知資料訪問頻次分佈的caching機制。透過統計訓練資料的ID，將最熱的訪問資料快取到GPU的視訊記憶體中，而冷資料以及雜湊表結構則存放在主記憶體中，主記憶體中的資料將根據ID的訪問頻率變化，定期將top-k的高頻ID對應的embeddings重新整理到GPU視訊記憶體上的快取中。這樣的混合儲存可以同時結合GPU視訊記憶體的高頻寬和DRAM的大容量，後續，這套混合儲存的設計還可以擴充套件到使用 Intel Persistent Memory、Non-volatile Memory 等更多的硬體裝置上。
HybridBackend已經成功在阿里媽媽智慧引擎訓練引擎團隊定向廣告業務有了落地。在阿里媽媽CAN模型下HybridBackend相對於上一代的訓練框架
(使用 Parameter Server 模式)具有明顯的
效能優勢，在下表中可以看到其在訓練時長等多個指標下獲得的顯著提升。
同時，我們還基於阿里媽媽定向廣告一年累計的訓練資料對模型規模增長下的
HybridBackend效能表現做了測試，結果如下表所示。可以看到，在使用128張GPU進行千億規模引數模型的訓練時，同樣是消費1年的資料量，高效能叢集上的HybridBackend僅僅需要2天的時間就能完成訓練任務，而普通叢集上的
PS模式則需要約1個月的時間。
後續我們計劃定期釋出Release版本。近期的Roadmap如下：
此外，中長期，我們將在訓練策略的演進，新硬體的最佳化，服務化能力的支援等
幾個探索性方向上持續投入精力，也歡迎各種維度的反饋和改進建議以及技術討論，同時我們十分歡迎和期待對開源社群建設感興趣的同行一起參與共建。
[1] Oldridge, Even, Julio Perez, Ben Frederickson, Nicolas Koumchatzky, Minseok Lee, Zehuan Wang, Lei Wu et al. "Merlin: A GPU Accelerated Recommendation Framework." In Proceedings of IRS . 2020.
論文標題：PICASSO: Unleashing the Potential of GPU-centric Training for Wide-and-deep Recommender Systems
論文作者: 張遠行、陳浪石(並列一作)、楊斯然、袁滿、易慧民、張傑、王家忙、董建波、許雲龍、宋鉞、李永、張迪、林偉、曲琳、鄭波
idBac
來自 “ ITPUB部落格 ” ，連結：http://blog.itpub.net/70004426/viewspace-2892970/，如需轉載，請註明出處，否則將追究法律責任。
