近日,阿里巴巴自研稀疏訓練引擎論文《PICASSO: Unleashing the Potential of GPU-centric Training for Wide-and-deep Recommender Systems》被國際數據工程頂級會議 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 三個產品服務于搜索、推薦、廣告業(yè)務。本文將對 PICASSO 訓練引擎的設計方案及思路展開介紹,歡迎閱讀交流。

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

PICASSO 技術轉化

阿里巴巴 XDL2 訓練引擎

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

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

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

工業(yè)級稀疏訓練

圖 1 典型推薦系統(tǒng)模型結構

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

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

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

系統(tǒng)架構

圖 2 PICASSO 同步訓練系統(tǒng)架構

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

核心技術設計

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

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

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

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

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

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

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

其他集成技術

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

主要指標表現

Benchmark 效果

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

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

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

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

多種模型性能效果

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

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

內部訓練任務性能提升

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

圖 6 阿里巴巴內部訓練任務性能提升

未來展望

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

參考文獻

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