作者 |  石浪、滿神

近年來(lái)，隨著稀疏模型對(duì)算力日益增長(zhǎng)的需求, CPU集群必須不斷擴(kuò)大集群規(guī)模來(lái)滿足訓(xùn)練的時(shí)效需求，這同時(shí)也帶來(lái)了不斷上升的資源成本以及實(shí)驗(yàn)的調(diào)試成本。

為了解決這一問(wèn)題，阿里云機(jī)器學(xué)習(xí)PAI平臺(tái)開(kāi)源了稀疏模型高性能同步訓(xùn)練框架HybridBackend，使得在同成本下GPU集群訓(xùn)練吞吐較CPU集群提升至5倍，大幅降低調(diào)試成本，同時(shí) HybridBackend 相關(guān)論文 《PICASSO: Unleashing the Potential of GPU-centric Training for Wide-and-deep Recommender Systems》也被 ICDE 22' 所收錄。HybridBackend背后的技術(shù)框架如何設(shè)計(jì)?未來(lái)有哪些規(guī)劃?今天一起來(lái)深入了解。

一、HybridBackend是什么

HybridBackend是阿里云機(jī)器學(xué)習(xí)平臺(tái)PAI自研的、面向稀疏模型訓(xùn)練的高性能同步訓(xùn)練框架，核心能力是大幅提升GPU集群?jiǎn)挝怀杀鞠碌挠?xùn)練吞吐性能。目前HybridBackend已經(jīng)在阿里巴巴集團(tuán)內(nèi)部有多個(gè)業(yè)務(wù)落地，將阿里媽媽智能引擎訓(xùn)練引擎團(tuán)隊(duì)的定向廣告業(yè)務(wù)年數(shù)據(jù)訓(xùn)練任務(wù)時(shí)間由1個(gè)月縮短至2天，同時(shí)HybridBackend在公有云多個(gè)頭部互聯(lián)網(wǎng)企業(yè)中也有成功應(yīng)用。

二、項(xiàng)目背景

以搜索、推薦、廣告業(yè)務(wù)為主要應(yīng)用的稀疏模型訓(xùn)練系統(tǒng)一直是學(xué)界和業(yè)界研究的熱點(diǎn)之一。相比于計(jì)算機(jī)視覺(jué)(CV)和自然語(yǔ)言處理(NLP)為代表的稠密模型訓(xùn)練，稀疏模型針對(duì)離散型特征(以 categorical ID 作為訓(xùn)練數(shù)據(jù))使用Embedding特征表達(dá)有著百GB至數(shù)十TB級(jí)別的內(nèi)存占用消耗(比普通的CV、NLP模型參數(shù)高出一到兩個(gè)數(shù)量級(jí)),從而突破了單機(jī)的內(nèi)存容量限制，需要基于分布式系統(tǒng)的訓(xùn)練方案。 早期的此類分布式任務(wù)由于模型結(jié)構(gòu)相對(duì)簡(jiǎn)單并且更新迭代緩慢,往往采用定制化的參數(shù)服務(wù)器(Parameter Server,PS)系統(tǒng)在大規(guī)模的CPU集群上進(jìn)行訓(xùn)練。隨著TensorFlow為代表的通用機(jī)器學(xué)習(xí)編程框架的出現(xiàn)，以及深度神經(jīng)網(wǎng)絡(luò)(DNN)在推薦類模型上的流行(deep recommender systems),業(yè)界逐漸轉(zhuǎn)向基于通用機(jī)器學(xué)習(xí)編程框架(TensorFlow、PyTorch等)進(jìn)行模型的端到端訓(xùn)練和推理，但是此時(shí)依然以參數(shù)服務(wù)器(PS)和大規(guī)模CPU集群作為訓(xùn)練的范式和基礎(chǔ)設(shè)施。

三、面臨挑戰(zhàn)

隨著稀疏模型對(duì)算力日益增長(zhǎng)的需求(比如Attention等結(jié)構(gòu)的加入),CPU集群必須不斷擴(kuò)大集群規(guī)模來(lái)滿足訓(xùn)練的時(shí)效需求，這同時(shí)也帶來(lái)了不斷上升的資源成本以及實(shí)驗(yàn)的調(diào)試成本。

以NVIDIA GPU為代表的加速器(accelerator)彌補(bǔ)了CPU設(shè)備單位成本算力低下的劣勢(shì)，在CV、NLP等算力需求大的訓(xùn)練任務(wù)上的應(yīng)用已經(jīng)成為行業(yè)共識(shí)。然而實(shí)踐證明，如只是簡(jiǎn)單地將PS訓(xùn)練范式中的worker從CPU設(shè)備替換為GPU設(shè)備，并不能有效地提升訓(xùn)練任務(wù)的吞吐，通過(guò) profiling GPU 的使用率，發(fā)現(xiàn)大量的GPU算力資源被閑置浪費(fèi)。這說(shuō)明，相比于CV、NLP類任務(wù)，稀疏模型訓(xùn)練有著自身的模型結(jié)構(gòu)和訓(xùn)練數(shù)據(jù)的特性，使得傳統(tǒng)的PS訓(xùn)練范式不能有效地發(fā)揮出GPU設(shè)備的優(yōu)勢(shì)。以深度推薦系統(tǒng)經(jīng)典的 Wide and Deep 模型結(jié)構(gòu)和TensorFlow框架為例,我們分析并總結(jié)了在PS架構(gòu)下使用GPU設(shè)備訓(xùn)練的兩個(gè)問(wèn)題。

1.變化的硬件資源瓶頸

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

2.算子細(xì)碎化(fragmentation)

生產(chǎn)實(shí)際中的模型往往擁有上百路的Embedding特征查詢，每一路的特征查詢?cè)赥ensorFlow內(nèi)都會(huì)調(diào)用數(shù)十個(gè)算子操作(operations)。TensorFlow的引擎在調(diào)度上千級(jí)別的大量的算子操作需要額外的CPU線程開(kāi)銷;對(duì)于GPU設(shè)備來(lái)說(shuō),過(guò)多的 CUDA kernel 提交到流處理器上(TensorFlow下每個(gè)GPU設(shè)備只有一個(gè)stream抽象)帶來(lái)了GPU Stream Multiprocessor (SM)的調(diào)度開(kāi)銷，同時(shí)每個(gè)算子處理數(shù)據(jù)的并發(fā)度又不高，從而很難打滿GPU的計(jì)算單元。

類似的問(wèn)題在CV、NLP等稠密模型的訓(xùn)練中也有涉及，一般采用基于編譯技術(shù)的優(yōu)化手段進(jìn)行算子合并。在 Wide and Deep 模型這樣的稀疏場(chǎng)景下，Embedding階段的這些算子又往往具有 dynamic shape 的特點(diǎn)，在TensorFlow靜態(tài)構(gòu)圖階段無(wú)法獲取準(zhǔn)確的算子尺寸進(jìn)行優(yōu)化，導(dǎo)致類似TensorFlow-XLA等技術(shù)在此類場(chǎng)景下沒(méi)有明顯的收益。

這些問(wèn)題說(shuō)明，想要發(fā)揮出GPU等高性能硬件資源的極致性價(jià)比，提高單位成本下的訓(xùn)練吞吐，就必須設(shè)計(jì)新的訓(xùn)練框架。據(jù)我們了解，擁有大型搜索、廣告、推薦業(yè)務(wù)的國(guó)內(nèi)外企業(yè)以及硬件廠商都在著手進(jìn)行新框架的研發(fā)，比如NVIDIA的Merlin-HugeCTR[1]等，然而阿里巴巴集團(tuán)內(nèi)云上集群普遍部署的是通用計(jì)算節(jié)點(diǎn)，且集群上需要執(zhí)行多種異構(gòu)的任務(wù)，換用專用硬件是很昂貴且不切實(shí)際的。

基于這種實(shí)際需求，我們推出了HybridBackend，能夠同時(shí)適應(yīng)集團(tuán)內(nèi)多元化且不斷演進(jìn)的稀疏模型技術(shù)。下文中我們將簡(jiǎn)要介紹HybridBackend的系統(tǒng)架構(gòu)設(shè)計(jì)和技術(shù)亮點(diǎn)。

四、HybridBackend的系統(tǒng)架構(gòu)

傳統(tǒng)的參數(shù)服務(wù)器(PS)訓(xùn)練范式體現(xiàn)的是通過(guò)擴(kuò)展硬件數(shù)量來(lái)適應(yīng)模型訓(xùn)練規(guī)模的思路。我們的系統(tǒng)則是同時(shí)考慮到了硬件和軟件(模型)兩個(gè)層面的特點(diǎn)，并做到協(xié)同設(shè)計(jì)。高性能GPU集群的硬件特性決定了基本的訓(xùn)練范式，而稀疏模型本身的結(jié)構(gòu)特點(diǎn)和數(shù)據(jù)分布帶來(lái)的問(wèn)題則通過(guò)更精細(xì)的系統(tǒng)優(yōu)化手段來(lái)解決。

1.利用大 Batch Size 進(jìn)行同步訓(xùn)練

因?yàn)镚PU設(shè)備相對(duì)于CPU帶來(lái)的巨大的算力提升，以往需要上百臺(tái)CPU節(jié)點(diǎn)的集群可以用幾十臺(tái)機(jī)器的GPU集群來(lái)代替。要保持相同的總訓(xùn)練規(guī)模，同時(shí)提升單個(gè)GPU節(jié)點(diǎn)上的資源利用率，提升單個(gè) GPU worker 上的 batch size 成為必然的選項(xiàng)。同時(shí)，因?yàn)榧阂?guī)模的縮小，可以通過(guò)同步訓(xùn)練的方式有效避免過(guò)期梯度(staleness)，從而提升模型訓(xùn)練的精度。

相對(duì)于CPU設(shè)備之間通過(guò)PCIe以及TCP進(jìn)行網(wǎng)絡(luò)通信，高性能的GPU集群在單個(gè)節(jié)點(diǎn)內(nèi)的多個(gè)GPU設(shè)備之間往往配備了高速的網(wǎng)絡(luò)互連(NVLink、NVSwitch),這些高速連接的帶寬通常是TCP網(wǎng)絡(luò)帶寬的數(shù)百倍(第一代NVLINK標(biāo)定達(dá)到300GB/s),而在多個(gè)機(jī)器節(jié)點(diǎn)之間也可以配備基于RDMA技術(shù)的高速網(wǎng)絡(luò)設(shè)備，達(dá)到100-200Gbps的帶寬。

選擇同步訓(xùn)練的第二個(gè)好處是，可以使用高性能集合通信算子庫(kù)(NVIDIA NCCL、阿里自研的ACCL等)來(lái)有效利用硬件機(jī)器的網(wǎng)絡(luò)拓?fù)浣Y(jié)構(gòu)，從而提升通信的性能。上述通信庫(kù)已經(jīng)在CV、NLP之類的基于數(shù)據(jù)并行的同步訓(xùn)練任務(wù)上取得了很好的效果。

2.使用資源異構(gòu)而角色同構(gòu)的訓(xùn)練單元

PS訓(xùn)練范式在系統(tǒng)的邏輯層面會(huì)指定不同的訓(xùn)練角色,比如server、worker、evaluator等。server節(jié)點(diǎn)一般分配具有大內(nèi)存的CPU機(jī)器，而worker節(jié)點(diǎn)則會(huì)被分配到高主頻的計(jì)算型CPU硬件上。這樣形成了訓(xùn)練單元-任務(wù)角色-同構(gòu)資源的耦合，通過(guò)增加訓(xùn)練單元數(shù)量來(lái)水平擴(kuò)展(scale out)訓(xùn)練的規(guī)模。

而在高性能的GPU集群上，一個(gè)物理的機(jī)器節(jié)點(diǎn)往往包括多種異構(gòu)的硬件資源，如CPU、GPU處理器、GPU之間的高速互連、DRAM(動(dòng)態(tài)隨機(jī)存取內(nèi)存)、Non-volatile Memory(非易失性內(nèi)存)等。這樣，除了水平擴(kuò)展節(jié)點(diǎn)數(shù)量外，還可以通過(guò)垂直擴(kuò)展利用多種異構(gòu)硬件資源來(lái)達(dá)到擴(kuò)大訓(xùn)練規(guī)模的目標(biāo)。

針對(duì)這種硬件架構(gòu)，我們的系統(tǒng)設(shè)計(jì)中只保留統(tǒng)一的訓(xùn)練執(zhí)行單元(Executor),每個(gè)Executor通過(guò)內(nèi)部的異構(gòu)硬件資源來(lái)執(zhí)行不同的訓(xùn)練任務(wù)角色。一方面，Executor內(nèi)部任務(wù)執(zhí)行時(shí)，可以有效地利用底層硬件資源之間的locality加速訓(xùn)練;另一方面，Executor內(nèi)部的硬件資源可以同時(shí)滿足不同的分布式訓(xùn)練范式所需要的硬件資源，以方便我們?cè)谀Ｐ徒Y(jié)構(gòu)的不同部分進(jìn)行混合并行訓(xùn)練策略。

五、深入優(yōu)化：HybridBackend的技術(shù)亮點(diǎn)

因?yàn)橄∈枘Ｐ徒Y(jié)構(gòu)和訓(xùn)練數(shù)據(jù)本身的特性, 變化的硬件資源瓶頸和算子細(xì)碎化，上述的系統(tǒng)架構(gòu)在實(shí)際任務(wù)中還是會(huì)存在一些影響GPU等硬件設(shè)備使用率的問(wèn)題。

舉例來(lái)說(shuō)，同步訓(xùn)練范式下，所有Executor通過(guò)集合通信進(jìn)行embedding的shuffle時(shí)，網(wǎng)絡(luò)帶寬資源成為瓶頸，而GPU的計(jì)算資源被閑置。一種解決思路是對(duì)硬件資源進(jìn)行定制化，比如增加網(wǎng)絡(luò)帶寬資源來(lái)消除通信瓶頸，但是這樣的做法會(huì)使得硬件的資源配置和特定的模型結(jié)構(gòu)耦合，是專用推薦系統(tǒng)的老思路。

我們的目標(biāo)還是希望系統(tǒng)可以架構(gòu)在云服務(wù)上可得的，數(shù)量容易水平擴(kuò)展的通用硬件配置之上(commodity hardware)。某些硬件廠商也嘗試通過(guò) Huge kernel 的形式(將Embedding層所有的計(jì)算手工融合到一個(gè)kernel內(nèi))來(lái)解決算子細(xì)碎化的問(wèn)題,這樣的做法也很難支持模型結(jié)構(gòu)快速迭代的需求,背離了通用編程架構(gòu)的設(shè)計(jì)初衷。

據(jù)此，我們從軟硬協(xié)同的思路出發(fā)，設(shè)計(jì)了如下的幾個(gè)系統(tǒng)優(yōu)化手段:

1.基于數(shù)據(jù)和算子感知的合并

根據(jù)稀疏模型的結(jié)構(gòu)特點(diǎn)，大部分細(xì)碎的算子來(lái)源于龐大的Embedding特征查詢(lookup)數(shù)量，我們?cè)O(shè)計(jì)了D-Packing這一優(yōu)化技術(shù)。

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

此外，通過(guò)分析特征查詢階段各個(gè)算子在分布式環(huán)境下的語(yǔ)義，我們將部分的kernel進(jìn)行融合K-Packing,比如通過(guò)融合shuffle和stitch算子來(lái)消除冗余的數(shù)據(jù)拷貝。

通過(guò)數(shù)據(jù)和算子兩個(gè)維度的基于語(yǔ)義的融合，我們既減少了總體的算子數(shù)量，降低fragmentation,同時(shí)又避免了所有算子融合在一起而丟失了通過(guò)算子間穿插遮掩來(lái)提升硬件利用率的優(yōu)化機(jī)會(huì)。

2.基于硬件資源瓶頸感知的交錯(cuò)執(zhí)行

為了消除同時(shí)執(zhí)行相同硬件資源需求的算子而造成的瓶頸, 我們?cè)O(shè)計(jì)了兩種算子穿插遮掩執(zhí)行(interleaving)的優(yōu)化手段。

其一，D-Interleaving是通過(guò)對(duì)訓(xùn)練數(shù)據(jù)batch的切分利用pipeline的機(jī)制來(lái)調(diào)度穿插不同資源類型的算子，這樣可以在訓(xùn)練的任何階段緩解某一種資源的瓶頸。比如在大batch size的訓(xùn)練場(chǎng)景下，稀疏模型的MLP階段也會(huì)產(chǎn)生很高的feature map顯存占用，通過(guò)D-Interleaving就可以有效降低單個(gè)GPU設(shè)備上的峰值顯存占用，從而使得更大的batch size訓(xùn)練成為可能。

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

3.基于數(shù)據(jù)頻次感知的參數(shù)緩存

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

六、應(yīng)用場(chǎng)景

HybridBackend已經(jīng)成功在阿里媽媽智能引擎訓(xùn)練引擎團(tuán)隊(duì)定向廣告業(yè)務(wù)有了落地。在阿里媽媽CAN模型下HybridBackend相對(duì)于上一代的XDL訓(xùn)練框架具有明顯的性能優(yōu)勢(shì)，在下表中可以看到其在訓(xùn)練時(shí)長(zhǎng)等多個(gè)指標(biāo)下獲得的顯著提升。

同時(shí)，我們還基于阿里媽媽定向廣告一年累計(jì)的訓(xùn)練數(shù)據(jù)對(duì)模型規(guī)模增長(zhǎng)下的HybridBackend性能表現(xiàn)做了測(cè)試，結(jié)果如下表所示?？梢钥吹剑谑褂?28張GPU進(jìn)行千億規(guī)模參數(shù)模型的訓(xùn)練時(shí)，同樣是消費(fèi)1年的數(shù)據(jù)量，高性能集群上的HybridBackend僅僅需要2天的時(shí)間就能完成訓(xùn)練任務(wù)，而普通集群上的XDL-PS模式則需要約1個(gè)月的時(shí)間。

七、Roadmap

后續(xù)我們計(jì)劃定期發(fā)布Release版本。近期的Roadmap如下：

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

此外，中長(zhǎng)期，我們將在訓(xùn)練策略的演進(jìn)，新硬件的優(yōu)化，服務(wù)化能力的支持等幾個(gè)探索性方向上持續(xù)投入精力，也歡迎各種維度的反饋和改進(jìn)建議以及技術(shù)討論，同時(shí)我們十分歡迎和期待對(duì)開(kāi)源社區(qū)建設(shè)感興趣的同行一起參與共建。

開(kāi)源地址：https://github.com/alibaba/HybridBackend

參考文獻(xiàn)

[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.

論文詳情

論文標(biāo)題：PICASSO: Unleashing the Potential of GPU-centric Training for Wide-and-deep Recommender Systems 論文作者: 張遠(yuǎn)行、陳浪石(并列一作)、楊斯然、袁滿、易慧民、張杰、王家忙、董建波、許云龍、宋鉞、李永、張迪、林偉、曲琳、鄭波

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