作者:玄慈、石浪、松岳、滿(mǎn)神

概述

近日,阿里巴巴自研稀疏訓練引擎論文《PICASSO: Unleashing the Potential of GPU-centric Training for Wide-and-deep Recommender Systems》被國際數據工程頂級會(huì )議 ICDE 2022 (International Conference on Data Engineering) 收錄。PICASSO(Packing, Interleaving and Caching Augmented Software System Optimization)是阿里巴巴智能引擎事業(yè)部廣告訓練引擎 XDL 團隊和阿里云計算平臺 PAI 團隊深度合作研發(fā)的成果,在集團內部 PICASSO 作為 XDL2、PAI-Tensorflow 和 PAI-HybridBackend 三個(gè)產(chǎn)品服務(wù)于搜索、推薦、廣告業(yè)務(wù)。本文將對 PICASSO 訓練引擎的設計方案及思路展開(kāi)介紹,歡迎閱讀交流。

搜推廣模型訓練的性能會(huì )直接影響到迭代開(kāi)發(fā)效率及模型時(shí)效性,如今已是一項非常熱門(mén)的研究主題。阿里巴巴在搜推廣稀疏模型訓練優(yōu)化上有著(zhù)深厚的經(jīng)驗與積累,近年來(lái)我們也看到大量相關(guān)設計或是為解決某類(lèi)問(wèn)題的新方案不斷涌現。然而,隨著(zhù)模型復雜度以及數據規模的快速增長(cháng),即便有多種稠密計算及稀疏處理的優(yōu)化,在通用硬件上的訓練任務(wù)仍然不能表現出持續高水平的資源使用率。結合阿里巴巴實(shí)際業(yè)務(wù)場(chǎng)景,我們進(jìn)一步探索了提升資源利用效率的方法,在滿(mǎn)足業(yè)務(wù)質(zhì)量的嚴苛要求的同時(shí),也期望可以起到拋磚引玉的作用,吸引更多從事相關(guān)工作的同學(xué)提出更為高效優(yōu)雅的稀疏訓練解決方案。

PICASSO 技術(shù)轉化

阿里巴巴 XDL2 訓練引擎

XDL 是阿里巴巴自研的 PS 訓練模式的稀疏訓練引擎,通過(guò)高并發(fā)的方式提升在混部集群中搜推廣任務(wù)的訓練效率。自 2018 年開(kāi)源以來(lái),XDL 團隊結合工業(yè)實(shí)踐對 XDL 框架進(jìn)行了大范圍的革新,包括但不限于 I / O、通信、訪(fǎng)存、編譯優(yōu)化、軟硬件計算加速等。XDL2 是綜合 XDL 多年的業(yè)務(wù)經(jīng)驗與 PICASSO 中所介紹技術(shù)的成果,支持在高性能集群中充分利用硬件優(yōu)勢提升訓練效率,為阿里巴巴內部的多條搜推廣業(yè)務(wù)線(xiàn)訓練任務(wù)保駕護航。

阿里云 HybridBackend 高性能稀疏訓練引擎

HybridBackend 是阿里云推出的提升稀疏模型在異構硬件集群訓練效率的訓練框架,服務(wù)阿里云數十家搜推廣業(yè)務(wù)用戶(hù),平均帶來(lái)了 2-3 倍的訓練性能提升。HybridBackend 將 PICASSO 中所介紹的技術(shù)作為 Tensorflow 的擴展,能夠享受到 Tensorflow 社區的最新加速技術(shù)并在異構集群中最大化提升訓練效率,且與基于 Tensorflow 實(shí)現的前沿學(xué)術(shù)或工業(yè)優(yōu)化策略完全正交。

工業(yè)級稀疏訓練

圖 1 典型推薦系統模型結構

稀疏訓練主要由 Embedding 階段, 特征交叉 (feature interation) 階段和多層感知器 (MLP) 階段組成。實(shí)際訓練中,Embedding 階段在 PS 范式的訓練下經(jīng)常能夠占據 50% 以上的訓練時(shí)間。經(jīng)過(guò)分析發(fā)現,Embedding 階段的算子主要以訪(fǎng)存密集型 (memory access intensive) 和通信密集型的算子 (communication intensive) 為主,主要需要的硬件資源是內存和網(wǎng)絡(luò )的帶寬,而后兩個(gè)階段的算子則是計算密集型的算子占主導,需要的資源是算力。這意味著(zhù)在 PS 的范式訓練下,任何一個(gè)階段都有可能存在某一種硬件資源成為瓶頸而其他硬件資源被浪費的現象。以 GPU 的算力資源為例,我們觀(guān)察 GPU 使用率 (SM Util) 在不同的訓練階段之間呈現脈沖式變化 (pulse),而不是持續高水位??紤]到最前沿的稀疏模型可能包含截然不同比例訪(fǎng)存密集型、通信密集型和算力密集型算子,其他硬件資源上也會(huì )出現相似的情況。

實(shí)際生產(chǎn)中的模型往往擁有上百路的 Embedding 特征查詢(xún),每一路的特征查詢(xún)在 Tensorflow 內都會(huì )調用數十個(gè)算子實(shí)例 (operations)。TensorFlow 的引擎在調度上千級別的大量的算子執行需要額外的 CPU 線(xiàn)程開(kāi)銷(xiāo);對于 GPU 設備來(lái)說(shuō), 過(guò)多的 CUDA kernel 提交到流處理器上 (Tensorflow 下每個(gè) GPU 設備只有一個(gè) stream 抽象) 造成了 GPU Stream Multiprocessor (SM) 的冗余調度開(kāi)銷(xiāo),同時(shí)每個(gè)零碎算子處理數據的并發(fā)度又不高,從而很難打滿(mǎn) GPU 的計算單元。類(lèi)似的問(wèn)題在 CV,NLP 等稠密模型的訓練中也有可能發(fā)生,一般采用基于編譯技術(shù)的優(yōu)化手段進(jìn)行算子合并。然而,在 Wide and Deep 模型這樣的稀疏場(chǎng)景下,Embedding 階段的這些算子又往往具有 dynamic shape 的特點(diǎn),在 Tensorflow 靜態(tài)構圖階段無(wú)法獲取準確的算子尺寸進(jìn)行優(yōu)化,導致類(lèi)似 Tensorflow-XLA 等技術(shù)在此類(lèi)場(chǎng)景下沒(méi)有明顯的收益。

工業(yè)級稀疏中,模型復雜度、特征數目、數據總量的日益增加,然而模型交付時(shí)效性要求卻愈加嚴格。為了解決工業(yè)級稀疏訓練的效率問(wèn)題,許多大型企業(yè)根據其業(yè)務(wù)模型的主要瓶頸設計了定制化硬件。然而,對于需要執行大量異構訓練任務(wù)及截然不同結構的稀疏任務(wù)的集群上,部署定制化硬件是很昂貴且效率不高的。對于個(gè)人用戶(hù)或小型機構的用戶(hù)也是如此。所以,在通用硬件設備上進(jìn)行工業(yè)級稀疏訓練優(yōu)化是十分有意義的。

系統架構

圖 2 PICASSO 同步訓練系統架構

圖 3 展示了 PICASSO 同步訓練模式的系統架構,每臺 Worker 各自讀取一部分數據并獨立計算梯度。各臺 Worker 上的梯度會(huì )在聚合后,更新于所有 Worker。相較于 PS(Parameter Server)訓練模式,在 PICASSO 中取消了 PS 角色。Embedding 表被切分放在不同的 Worker 上,切分方法可根據具體 ID 分布情況選擇按 ID 切分、按維度切分、按表劃分;查表操作通過(guò)網(wǎng)絡(luò )傳輸分布式地進(jìn)行。而所有稠密參數是以副本的形式放置在每臺 Worker 上的,考慮到梯度是聚合后更新因而稠密參數的梯度在每臺 Worker 上是完全相同的,所以在訓練過(guò)程中稠密參數在不同 Worker 間能夠保證一致性(諸如 Batch Normalization 的統計量可選擇全局同步或局部維護,已有許多算法論文討論,不在本文的范疇內)。不同 Worker 間通過(guò)集合通訊交換參數及梯度,能夠最大程度利用好網(wǎng)絡(luò )資源(如果配備了 RDMA 網(wǎng)絡(luò ),則性能會(huì )更加出色)。

核心技術(shù)設計

圖 3 PICASSO 核心優(yōu)化策略

數據和算子協(xié)同合并 (Packing)

根據稀疏模型的結構特點(diǎn),大部分細碎的算子來(lái)源于龐大的 Embedding 特征查詢(xún) (lookup) 數量。對于每一路查詢(xún),盡管輸入的訓練數據不同,但使用的算子組合是相同的。對于這種具有數據并行特點(diǎn)的模式,具有相同屬性(維度、初始化器、標定特征組等)的 Embedding 表將被合并為一張新的 Embedding 表,而后后續的訪(fǎng)存查詢(xún)算子也可以被合并為一個(gè)新的大算子。合并算子可以用多線(xiàn)程的方式有序查詢(xún) Embedding,相對于亂序查詢(xún)或分成若干小表查詢(xún),能有顯著(zhù)的性能提升。查詢(xún)完畢后,再依原有代碼需要進(jìn)行反去重和歸位,真正做到了對用戶(hù)透明。此外,根據特征查詢(xún)階段各個(gè)算子在分布式環(huán)境下的語(yǔ)義,我們將部分的計算核進(jìn)行了融合,比如通過(guò)融合 shuffle 和 stitch 算子來(lái)消除冗余的數據拷貝。通過(guò)數據和算子兩個(gè)維度的基于語(yǔ)義的融合,PICASSO 緩解了不同查表操作之間的競爭問(wèn)題,并大大減少了算子數量,對于現如今動(dòng)輒幾百路特征的模型是十分有幫助的;同時(shí),相對于將所有 Embedding 查詢(xún)過(guò)程涉及的算子全部融合在一起的方案,PICASSO 所采用的僅將相鄰的使用相同硬件資源的算子融合的方法,保留了下文敘述的通過(guò)算子間穿插遮掩來(lái)提升硬件利用率的優(yōu)化機會(huì )。

訓練算子間的穿插遮掩 (Interleaving)

為了消除同時(shí)執行相同硬件資源需求的算子而造成的瓶頸, 我們設計了兩種算子穿插遮掩執行 (interleaving) 的優(yōu)化手段。1) D (ata)-interleaving 是通過(guò)對訓練數據 batch 的切分利用 pipeline 的機制來(lái)調度穿插不同資源類(lèi)型的算子,這樣可以在訓練的任何階段緩解某一種資源的瓶頸。比如在大 batch size 的訓練場(chǎng)景下,稀疏模型的 MLP 階段也會(huì )產(chǎn)生很高的 feature map 顯存占用,通過(guò) D-interleaving 就可以有效降低單個(gè) GPU 設備上的峰值顯存占用,從而使得更大的 batch size 訓練成為可能。2)K (ernel)-interleaving 是在 Embedding Layer 內部不同的特征查詢(xún)路數之間做算子的穿插和遮掩,比如將通信密集的 Shuffle 操作和內存訪(fǎng)問(wèn)密集的 Gather 進(jìn)行遮掩,可以有效提升這兩種資源的使用率。

利用數據頻次分布做緩存 (Caching)

PICASSO 利用數據中 ID 明顯分布有偏(即存在大量長(cháng)期或短期熱點(diǎn) ID)的特點(diǎn),將 Embedding 表拆分到 CPU 和 GPU 之上,以最大限度利用內存的大容量和 GPU 的高速內存帶寬。其中,被識別的熱點(diǎn) ID 以及對應的 Embedding 將被放置在顯存上,其余 ID 被放置在內存上,二者之間的 ID 定期進(jìn)行刷新交換。兩張哈希表都是需要正常更新梯度的,以避免梯度更新延遲帶來(lái)的精度損失??紤]到同一 Batch 內的 ID 在訪(fǎng)存查詢(xún)之前已經(jīng)經(jīng)過(guò)去重,這種緩存機制可以大大減少跨不同 Batch 的冗余低速訪(fǎng)存,同時(shí)又沒(méi)有過(guò)多消耗用于計算的顯存。后續這套混合存儲的設計還可以擴展到包含 Intel Persistent Memory, Non-volatile Memory 等更多的硬件設備上。

其他集成技術(shù)

考慮到稀疏訓練任務(wù)的特點(diǎn),PICASSO 集成了近年來(lái)多項業(yè)務(wù)經(jīng)驗,重要設計包括:在 I / O 上開(kāi)啟了預取和列存數據格式的設計,最大限度降低 I / O 對整體訓練的影響;改進(jìn)的集合通訊協(xié)議,減少通信過(guò)程中的冗余傳輸及校驗;算子零拷貝調整,避免不必要的訪(fǎng)存;基于統計值的哈希表最優(yōu)切分放置方案;計算圖算子替換及編譯優(yōu)化;特定算子啟用 GPU 加速庫,如 CUTLASS、CuDNN、CuBLAS 和 Cub 等。除此之外,對于不斷涌現的新技術(shù),用戶(hù)可以通過(guò) Tensorflow 實(shí)現集成到 PICASSO 中。

主要指標表現

Benchmark 效果

我們通過(guò)一組 benchmark 模型的實(shí)驗結果橫向和目前廣泛流行的多種稀疏模型訓練框架進(jìn)行了性能的測試對比。

參與對比的框架有:Tensorflow Parameter Server (TensorFlow-PS) 作為 Baseline,PyTorch 的模型并行版本 (PyTorch),以及采用數據并行模式的 Horovod。參與測試的模型有: 1) MLPerf 采用的 DLRM 模型;2) 工業(yè)界廣泛使用的 DeepFM 模型,3)阿里提出的 DIN / DIEN 模型,擁有更加復雜的特征交叉計算量。我們使用的數據集是公開(kāi)的 Criteo 數據集。

圖 4 多種訓練框架在 Benchmark 上的性能效果

可以看到對比與 TF-PS 的 baseline, PICASSO 在四個(gè) benchmark 模型上都獲得了 1.9x 至 10x 的性能提升,而對比 Facebook 的模型并行版本 PyTorch, 我們也有至少 2 倍以上的性能提升,而且隨著(zhù)模型的計算量復雜度增加 (DIN, DIEN),PICASSO 的優(yōu)勢越明顯。

多種模型性能效果

圖 5 多種模型在 PICASSO 上相對 XDL 的提升

我們隔離了若干訓練節點(diǎn),并分別使用相同的 Worker 數目分別使用內部?jì)?yōu)化的 XDL-PS 模式和 PICASSO 模式實(shí)現若干前沿稀疏模型。我們調整了 Batch size 以記錄極限性能,提升幅度如圖 6 所示??梢钥吹?雖然不同的模型可能處于不同的硬件瓶頸下,但是 PICASSO 都能夠顯著(zhù)地提升平均硬件利用率和整體性能。

內部訓練任務(wù)性能提升

截至目前,PICASSO 已成功在阿里媽媽多條廣告業(yè)務(wù)中落地,本文的實(shí)驗中也介紹了在阿里媽媽 CAN 模型下 PICASSO 相對于上一代的 XDL 訓練框架獲得的性能優(yōu)勢,在下表中可以看到在訓練時(shí)長(cháng)等多個(gè)指標下獲得的顯著(zhù)提升。

圖 6 阿里巴巴內部訓練任務(wù)性能提升

未來(lái)展望

超大規模稀疏模型的訓練效率問(wèn)題已經(jīng)吸引了越來(lái)越多企業(yè)及研究機構的關(guān)注。PICASSO 研發(fā)小隊目前正致力于結合真實(shí)各類(lèi)大小業(yè)務(wù)場(chǎng)景探索通用的稀疏場(chǎng)景的加速優(yōu)化方式,不僅僅是追求指標上的極限,更是追求提供更為易用、經(jīng)濟的解決方案。另一方面,當前 PICASSO 在一定程度上仍然需要用戶(hù)手工調整若干參數以獲取最佳性能,將 PICASSO 變得更加自動(dòng)化也是當前 PICASSO 研發(fā)小隊的重要使命。我們真誠歡迎并邀請業(yè)界研究人員、工程師或感興趣的朋友們與我們交流,共同推進(jìn)稀疏訓練性能乃至所有機器學(xué)習領(lǐng)域模型訓練性能的發(fā)展,支撐深度學(xué)習技術(shù)更好更經(jīng)濟地服務(wù)于更多的領(lǐng)域。

參考文獻

[1] Zhao, Weijie, Jingyuan Zhang, Deping Xie, Yulei Qian, Ronglai Jia, and Ping Li. “Aibox: Ctr prediction model training on a single node.” InProceedings of the 28th ACM International Conference on Information and Knowledge Management, pp. 319-328. 2019.

[2] Zhao, Weijie, Deping Xie, Ronglai Jia, Yulei Qian, Ruiquan Ding, Mingming Sun, and Ping Li. “Distributed hierarchical gpu parameter server for massive scale deep learning ads systems.” Proceedings of Machine Learning and Systems 2 (2020): 412-428.

[3] Mudigere, Dheevatsa, Yuchen Hao, Jianyu Huang, Andrew Tulloch, Srinivas Sridharan, Xing Liu, Mustafa Ozdal et al. “High-performance, distributed training of large-scale deep learning recommendation models.” arXiv e-prints (2021): arXiv-2104.

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

[5] H. Guo, R. TANG, Y. Ye, Z. Li, and X. He,“Deepfm: A factorization-machine based neural network for ctr prediction,”in Proceedings of the Twenty-Sixth International Joint Conference on Artificial Intelligence, IJCAI-17, 2017, pp. 1725–1731.

[6] G. Zhou, X. Zhu, C. Song, Y. Fan, H. Zhu, X. Ma, Y. Yan, J. Jin, H. Li, and K. Gai,“Deep interest network for click-through rate prediction,”in Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, 2018, pp. 1059–1068.

[7] G. Zhou, N. Mou, Y. Fan, Q. Pi, W. Bian, C. Zhou, X. Zhu, and K. Gai,“Deep interest evolution network for click-through rate prediction,”in Proceedings of the AAAI conference on artificial intelligence, vol. 33,2019, pp. 5941–5948.


標簽: