作者：玄慈、石浪、松岳、滿神

概述

近日，阿里巴巴自研稀疏訓練引擎論文《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）是阿里巴巴智能引擎事業部廣告訓練引擎XDL團隊和阿里云計算平臺PAI團隊深度合作研發的成果，在集團內部PICASSO作為XDL2、PAI-Tensorflow和PAI-HybridBackend三個產品服務于搜索、推薦、廣告業務。本文將對PICASSO訓練引擎的設計方案及思路展開介紹，歡迎閱讀交流。

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

PICASSO技術轉化

阿里巴巴XDL2訓練引擎

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

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

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

GitHub地址：https://github.com/alibaba/HybridBackend

工業級稀疏訓練

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

稀疏訓練主要由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等技術在此類場景下沒有明顯的收益。

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

系統架構

圖2 PICASSO同步訓練系統架構

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

核心技術設計

圖3 PICASSO核心優化策略

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

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

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

為了消除同時執行相同硬件資源需求的算子而造成的瓶頸, 我們設計了兩種算子穿插遮掩執行(interleaving)的優化手段。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的冗余低速訪存，同時又沒有過多消耗用于計算的顯存。后續這套混合存儲的設計還可以擴展到包含Intel Persistent Memory, Non-volatile Memory等更多的硬件設備上。

其他集成技術

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

主要指標表現

Benchmark效果

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

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

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

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

多種模型性能效果

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

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

內部訓練任務性能提升

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

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

未來展望

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

參考文獻

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