【picasso】PICASSO，高效搜索和传播稀疏教育解决方案

作者：张曼、陈朗锡、宋月、元元

资料来源：智能引擎事业部、阿里妈妈-广告技术部、阿里云计算平台事业部。

一、总体情况概述

最近阿里巴巴自行研究稀疏训练引擎论文《PICASSO: Unleashing the Potential of GPU-centric Training for Wide-and-deep Recommender Systems》收录在国际数据工程峰会ICDE 2022(International Conference On Data Engineering)上。介绍了名为Picasso(打包、交互和缓存管理软件系统优化)的培训引擎使用的关键技术。PICASSO是阿里巴巴智能引擎事业部广告培训引擎XDL团队和阿里云计算平台PAI团队深入合作开发的结果。集团内的PICASSO是XDL2、PAI-TensorFlow和PAI-HybridBackend三种产品，负责搜索、推荐和广告业务。

搜索模型教学的性能直接影响迭代开发效率和模型及时性，是目前非常热门的研究主题。作为国际最大的电子商务平台之一，阿里巴巴在搜索稀疏模型培训优化方面有着深厚的经验和积累，近年来，解决类似设计或相关问题的新方案不断出现在国际顶级学术会议和期刊上。(威廉莎士比亚，《Northern Exposure》(美国电视剧)，MESLET)但是，随着模型的复杂性和数据规模的迅速增长，即使密集计算和稀疏处理的优化有很多，通用硬件的培训工作也无法持续显示出高水平的资源利用率。阿里巴巴希望找到进一步提高资源利用效率的方法，在满足商业质量苛刻要求的同时，起到抛砖引玉的作用，吸引更多的研究者，提出更高效、更优雅的水疗解决方案。

二、毕加索技术改造

2.1 Alibaba内部XDL2培训引擎

XDL是Alibaba自主开发的PS培训模式的稀疏培训引擎，提高了混合集群中搜索和宣传任务的培训效率。自2018年开源以来，XDL团队结合行业实践，对XDL框架进行了广泛创新，包括I/O、通信、访问、编译优化、硬件和软件计算加速等。XDL2是XDL多年的业务经验和PICASSO介绍的技术成果，可在高性能集群中充分利用硬件优势，提高培训效率，为阿里巴巴内部多个搜索宣传业务线培训任务保驾护航。

2.2阿里云杂交高性能稀疏教育引擎

HybridBackend是阿里云推出的旨在提高稀疏模型在异构硬件集群中的培训效率的培训框架，迄今为止，它为阿里云的数十名搜索、宣传、商业用户带来了平均2 ~ 3倍的培训性能提升。HybridBackend将PICASSO引入的技术作为Tensorflow的扩展，利用Tensorflow社区的最新加速技术，在异构集群中最大限度地提高教育效率，HybridBackend目前在githu B- Alibaba/hybrid back end : a high-performance framework for training wide-and-deep recommend开放。

三、产业层面的稀疏培训

稀疏训练主要包括嵌入阶段、特征交叉阶段和多层传感器(MLP)阶段。在实际训练中，嵌入阶段在PS范式训练下往往可以占据50%以上的训练时间。分析表明，嵌入阶段的算子主要以访问密集型和通信密集型算子为主，主要需要的硬件资源是内存和网络带宽，接下来两个阶段的算子是计算密集型，即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等技术在此类场景下没有明显的收益。

工业级稀疏中，模型复杂度、特征数目、数据总量的日益增加，然而模型交付时效性要求却愈加严格。为了解决工业级稀疏训练的效率问题，许多大型企业根据其业务模型的主要瓶颈设计了定制化硬件。然而，对于需要执行大量异构训练任务及截然不同结构的稀疏任务的集群上，部署定制化硬件是很昂贵且效率不高的。对于个人用户或小型机构的用户也是如此。所以，在通用硬件设备上进行工业级稀疏训练优化是十分有意义的。

四、系统架构

上图展示了PICASSO同步训练模式的系统架构，每台Worker各自读取一部分数据并独立计算梯度。各台Worker上的梯度会在聚合后，更新于所有Worker。相较于PS(Parameter Server)训练模式，在PICASSO中取消了PS角色。Embedding表被切分放在不同的Worker上，切分方法可根据具体ID分布情况选择按ID切分、按维度切分、按表划分；查表操作通过网络传输分布式地进行。而所有稠密参数是以副本的形式放置在每台Worker上的，考虑到梯度是聚合后更新因而稠密参数的梯度在每台Worker上是完全相同的，所以在训练过程中稠密参数在不同Worker间能够保证一致性（诸如Batch Normalization的统计量可选择全局同步或局部维护，已有许多算法论文讨论，不在本文的范畴内）。不同Worker间通过集合通讯交换参数及梯度，能够最大程度利用好网络资源（如果配备了RDMA网络，则性能会更加出色）。

五、核心技术设计

5.1 数据和算子协同合并(Packing)

根据稀疏模型的结构特点，大部分细碎的算子来源于庞大的Embedding特征查询(lookup)数量。对于每一路查询，尽管输入的训练数据不同，但使用的算子组合是相同的。对于这种具有数据并行特点的模式，具有相同属性（维度、初始化器、标定特征组等）的Embedding表将被合并为一张新的Embedding表，而后续的访存查询算子也可以被合并为一个新的大算子。合并算子可以用多线程的方式有序查询Embedding，相对于乱序查询或分成若干小表查询能有显著的性能提升。查询完毕后，再依原有代码需要进行反去重和归位，真正做到了对用户透明。此外，根据特征查询阶段各个算子在分布式环境下的语义，我们将部分的计算核进行了融合，比如通过融合shuffle和stitch算子来消除冗余的数据拷贝。通过数据和算子两个维度的基于语义的融合，PICASSO缓解了不同查表操作之间的竞争问题，并大大减少了算子数量，对于现如今动辄几百路特征的模型是十分有帮助的；同时，相对于将所有Embedding查询过程涉及的算子全部融合在一起的方案，PICASSO所采用的仅将相邻的使用相同硬件资源的算子融合的方法，保留了下文叙述的通过算子间穿插遮掩来提升硬件利用率的优化机会。

5.2 训练算子间的穿插遮掩(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进行遮掩，可以有效提升这两种资源的使用率。

5.3 利用数据频次分布做缓存(Caching)

PICASSO利用数据中ID明显分布有偏（即存在大量长期或短期热点ID）的特点，将Embedding表拆分到CPU和GPU之上，以最大限度利用内存的大容量和GPU的高速内存带宽。其中，被识别的热点ID以及对应的Embedding将被放置在显存上，其余ID被放置在内存上，二者之间的ID定期进行刷新交换。两张哈希表都是需要正常更新梯度的，以避免梯度更新延迟带来的精度损失。考虑到同一Batch内的ID在访存查询之前已经经过去重，这种缓存机制可以大大减少跨不同Batch的冗余低速访存，同时又没有过多消耗用于计算的显存。后续这套混合存储的设计还可以扩展到包含Intel Persistent Memory, Non-volatile Memory等更多的硬件设备上。

5.4 其他集成技术

考虑到稀疏训练任务的特点，PICASSO已经集成了近年来的多项业务经验，重要设计包括：在I/O上开启了预取和列存数据格式的设计，最大限度降低I/O对整体训练的影响；改进的集合通讯协议，减少通信过程中的冗余传输及校验；算子零拷贝调整，避免不必要的访存；基于统计值的哈希表最优切分放置方案；计算图算子替换及编译优化；特定算子启用GPU加速库，如CUTLASS、CuDNN、CuBLAS和Cub等。除此之外，对于不断涌现的新技术，用户可以通过Tensorflow实现集成到PICASSO中。

六、主要指标表现

6.1 Benchmark效果

我们通过一组benchmark模型的实验结果横向和目前广泛流行的多种稀疏模型训练框架进行了性能的测试对比 参与对比的框架有Tensorflow Parameter Server (TensorFlow-PS)作为Baseline，PyTorch的模型并行版本(PyTorch)，以及采用数据并行模式的Horovod。参与测试的模型有: 1) MLPerf采用的DLRM模型; 2) 工业界广泛使用的DeepFM模型，3) 阿里提出的DIN/DIEN模型，拥有更加复杂的特征交叉计算量。我们使用的数据集是公开的Criteo数据集。

可以看到对比与TF-PS的baseline, PICASSO在四个benchmark模型上都获得了1.9x至10x的性能提升，而对比Facebook的模型并行版本PyTorch, 我们也有至少2倍以上的性能提升，而且随着模型的计算量复杂度增加(DIN, DIEN),PICASSO的优势越明显。

6.2 多种模型性能效果

我们隔离了若干训练节点，并分别使用相同的Worker数目分别使用内部优化的XDL-PS模式和PICASSO模式实现若干前沿稀疏模型。我们调整了Batch size以记录极限性能，提升幅度如上图所示。可以看到，虽然不同的模型可能处于不同的硬件瓶颈下，但是PICASSO都能够显著地提升平均硬件利用率和整体性能。

6.3 内部训练任务性能提升

截至文章被会议录取时，PICASSO已经成功在阿里妈妈多条广告业务中有了落地，本文的实验中也介绍了在阿里妈妈CAN模型下PICASSO相对于上一代的XDL训练框架获得的性能优势，在下表中可以看到在训练时长等多个指标上获得了显著提升。

七、未来规划与展望

超大规模稀疏模型的训练效率问题已经吸引了越来越多企业及研究机构的关注。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." In Proceedings 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.