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

概述

近日，阿里巴巴自研稀疏训练引擎论文《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.