作者 |  石浪、滿神

近年來，隨著稀疏模型對算力日益增長的需求, 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是什么

HybridBackend是阿里云機器學習平臺PAI自研的、面向稀疏模型訓練的高性能同步訓練框架，核心能力是大幅提升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設備訓練的兩個問題。

1.變化的硬件資源瓶頸

從上圖的 Wide and Deep 模型結構可以看出，稀疏訓練主要由Embedding階段、特征交叉(feature interation)階段和多層感知器(MLP)階段組成，Embedding階段在PS范式的訓練下占據了至少50%以上的訓練時間。經過分析發現，Embedding階段的算子主要以訪存密集型(memory access intensive)和通信密集型的算子(communication intensive)為主，主要需要的硬件資源是內存和網絡的帶寬，而后兩個階段的算子則是計算密集型的算子占主導,需要的資源是算力。這意味著在PS的范式訓練下，任何一個階段都有可能存在某一種硬件資源成為瓶頸而其他硬件資源被浪費的現象。以GPU的算力資源為例，我們觀察GPU使用率(SM Util)在不同的訓練階段之間呈現脈沖式變化(pulse)。

2.算子細碎化(fragmentation)

生產實際中的模型往往擁有上百路的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的系統架構設計和技術亮點。

四、HybridBackend的系統架構

傳統的參數服務器(PS)訓練范式體現的是通過擴展硬件數量來適應模型訓練規模的思路。我們的系統則是同時考慮到了硬件和軟件(模型)兩個層面的特點，并做到協同設計。高性能GPU集群的硬件特性決定了基本的訓練范式，而稀疏模型本身的結構特點和數據分布帶來的問題則通過更精細的系統優化手段來解決。

1.利用大 Batch Size 進行同步訓練

因為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之類的基于數據并行的同步訓練任務上取得了很好的效果。

2.使用資源異構而角色同構的訓練單元

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內)來解決算子細碎化的問題,這樣的做法也很難支持模型結構快速迭代的需求,背離了通用編程架構的設計初衷。

據此，我們從軟硬協同的思路出發，設計了如下的幾個系統優化手段:

1.基于數據和算子感知的合并

根據稀疏模型的結構特點，大部分細碎的算子來源于龐大的Embedding特征查詢(lookup)數量，我們設計了D-Packing這一優化技術。

對于每一路查詢，盡管輸入的訓練數據不同，但使用的算子組合是相同的。對于這種具有數據并行特點的模式，具有相同屬性(維度、初始化器、標定特征組等)的Embedding表將被合并為一張新的Embedding表，而后續的訪存查詢算子也可以被合并為一個新的大算子。合并算子可以用多線程的方式有序查詢Embedding，相對于亂序查詢或分成若干小表查詢，能有顯著的性能提升。查詢完畢后，再依原有代碼需要進行反去重和歸位，真正做到了對用戶透明。

此外，通過分析特征查詢階段各個算子在分布式環境下的語義，我們將部分的kernel進行融合K-Packing,比如通過融合shuffle和stitch算子來消除冗余的數據拷貝。

通過數據和算子兩個維度的基于語義的融合，我們既減少了總體的算子數量，降低fragmentation,同時又避免了所有算子融合在一起而丟失了通過算子間穿插遮掩來提升硬件利用率的優化機會。

2.基于硬件資源瓶頸感知的交錯執行

為了消除同時執行相同硬件資源需求的算子而造成的瓶頸, 我們設計了兩種算子穿插遮掩執行(interleaving)的優化手段。

其一，D-Interleaving是通過對訓練數據batch的切分利用pipeline的機制來調度穿插不同資源類型的算子，這樣可以在訓練的任何階段緩解某一種資源的瓶頸。比如在大batch size的訓練場景下，稀疏模型的MLP階段也會產生很高的feature map顯存占用，通過D-Interleaving就可以有效降低單個GPU設備上的峰值顯存占用，從而使得更大的batch size訓練成為可能。

其二，K-Interleaving是在Embedding Layer內部不同的特征查詢路數之間做算子的穿插和遮掩，比如將通信密集的Shuffle操作和內存訪問密集的Gather進行遮掩，可以有效提升這兩種資源的使用率。

3.基于數據頻次感知的參數緩存

在解決Executor內部多個級別的存儲(GPU顯存、DRAM等)之間的帶寬和延遲問題上，我們針對稀疏模型訓練數據的分布特點，提出了一種感知數據訪問頻次分布的caching機制。通過統計訓練數據的ID，將最熱的訪問數據緩存到GPU的顯存中，而冷數據以及哈希表結構則存放在主內存中，主內存中的數據將根據ID的訪問頻率變化，定期將top-k的高頻ID對應的embeddings刷新到GPU顯存上的緩存中。這樣的混合存儲可以同時結合GPU顯存的高帶寬和DRAM的大容量，后續，這套混合存儲的設計還可以擴展到使用 Intel Persistent Memory、Non-volatile Memory 等更多的硬件設備上。

六、應用場景

HybridBackend已經成功在阿里媽媽智能引擎訓練引擎團隊定向廣告業務有了落地。在阿里媽媽CAN模型下HybridBackend相對于上一代的XDL訓練框架具有明顯的性能優勢，在下表中可以看到其在訓練時長等多個指標下獲得的顯著提升。

同時，我們還基于阿里媽媽定向廣告一年累計的訓練數據對模型規模增長下的HybridBackend性能表現做了測試，結果如下表所示。可以看到，在使用128張GPU進行千億規模參數模型的訓練時，同樣是消費1年的數據量，高性能集群上的HybridBackend僅僅需要2天的時間就能完成訓練任務，而普通集群上的XDL-PS模式則需要約1個月的時間。

七、Roadmap

后續我們計劃定期發布Release版本。近期的Roadmap如下：

• v0.6.0 (2022年5月)：支持端到端分布式同步訓練與評估。
• v0.7.0 (2022年9月)：優化 GPU 利用率與顯存占用。
• v0.8.0 (2023年1月)：進一步優化云上訓練性能。

此外，中長期，我們將在訓練策略的演進，新硬件的優化，服務化能力的支持等幾個探索性方向上持續投入精力，也歡迎各種維度的反饋和改進建議以及技術討論，同時我們十分歡迎和期待對開源社區建設感興趣的同行一起參與共建。

開源地址：https://github.com/alibaba/HybridBackend

參考文獻

[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 論文作者: 張遠行、陳浪石(并列一作)、楊斯然、袁滿、易慧民、張杰、王家忙、董建波、許云龍、宋鉞、李永、張迪、林偉、曲琳、鄭波

論文鏈接: https://arxiv.org/abs/2204.04903