近日,阿里巴巴自研稀疏训练引擎论文《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 实现的前沿学术或工业优化策略完全正交。

工业级稀疏训练

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