阿里自研工业级稀疏模型高性能训练框架 PAI-HybridBackend_阿里pai 大模型

HybridBackend 背后的技术框架如何设计？未来有哪些规划？一起来深入了解。

近年来，随着稀疏模型对算力日益增长的需求，CPU 集群必须不断扩大集群规模来满足训练的时效需求，这同时也带来了不断上升的资源成本以及实验的调试成本。

为了解决这一问题，阿里云机器学习 PAI 平台和智能引擎事业部 XDL 训练平台合作开发了稀疏模型高性能同步训练框架 HybridBackend，使得在同成本下 GPU 集群训练吞吐较 CPU 集群提升至 5 倍，大幅降低调试成本，同时 HybridBackend 相关论文《PICASSO: Unleashing the Potential of GPU-centric Training for Wide-and-deep Recommender Systems》也被 ICDE 22’ 所收录。

HybridBackend 是什么

HybridBackend 是阿里云机器学习平台 PAI 自研的、面向稀疏模型训练的高性能同步训练框架，核心能力是大幅提升 GPU 集群单位成本下的训练吞吐性能。目前HybridBackend 已经在阿里巴巴集团内部有多个业务落地，将阿里妈妈智能引擎训练引擎团队的定向广告业务年数据训练任务时间由 1 个月缩短至 2 天，同时 HybridBackend 在公有云多个头部互联网企业中也有成功应用。

项目背景

以搜索、推荐、广告业务为主要应用的稀疏模型训练系统一直是学界和业界研究的热点之一。相比于计算机视觉(CV)和自然语言处理(NLP)为代表的稠密模型训练，稀疏模型针对离散型特征(以 categorical ID 作为训练数据)使用 Embedding 特征表达有着百 GB 至数十 TB 级别的内存占用消耗(比普通的 CV、NLP 模型参数高出一到两个数量级),从而突破了单机的内存容量限制，需要基于分布式系统的训练方案。

早期的此类分布式任务由于模型结构相对简单并且更新迭代缓慢，往往采用定制化的参数服务器(Parameter Server,PS)系统在大规模的 CPU 集群上进行训练。随着 TensorFlow 为代表的通用机器学习编程框架的出现，以及深度神经网络(DNN)在推荐类模型上的流行(deep recommender systems)，业界逐渐转向基于通用机器学习编程框架(TensorFlow、PyTorch 等)进行模型的端到端训练和推理，但是此时依然以参数服务器(PS)和大规模 CPU 集群作为训练的范式和基础设施。

面临挑战

随着稀疏模型对算力日益增长的需求(比如 Attention 等结构的加入)，CPU 集群必须不断扩大集群规模来满足训练的时效需求，这同时也带来了不断上升的资源成本以及实验的调试成本。

以 NVIDIA GPU 为代表的加速器(accelerator)弥补了 CPU 设备单位成本算力低下的劣势，在 CV、NLP 等算力需求大的训练任务上的应用已经成为行业共识。然而实践证明，如只是简单地将 PS 训练范式中的 worker 从 CPU 设备替换为 GPU 设备，并不能有效地提升训练任务的吞吐，通过 profiling GPU 的使用率，发现大量的 GPU 算力资源被闲置浪费。这说明，相比于 CV、NLP 类任务，稀疏模型训练有着自身的模型结构和训练数据的特性，使得传统的PS训练范式不能有效地发挥出 GPU 设备的优势。以深度推荐系统经典的 Wide and Deep 模型结构和 TensorFlow 框架为例,我们分析并总结了在 PS 架构下使用 GPU 设备训练的两个问题。

变化的硬件资源瓶颈

从上图的 Wide and Deep 模型结构可以看出，稀疏训练主要由 Embedding 阶段、特征交叉(feature interation)阶段和多层感知器(MLP)阶段组成，Embedding 阶段在 PS 范式的训练下占据了至少 50% 以上的训练时间。经过分析发现，Embedding 阶段的算子主要以访存密集型(memory access intensive)和通信密集型的算子(communication intensive)为主，主要需要的硬件资源是内存和网络的带宽，而后两个阶段的算子则是计算密集型的算子占主导，需要的资源是算力。这意味着在 PS 的范式训练下，任何一个阶段都有可能存在某一种硬件资源成为瓶颈而其他硬件资源被浪费的现象。以 GPU 的算力资源为例，我们观察 GPU 使用率(SM Util)在不同的训练阶段之间呈现脉冲式变化(pulse)。

算子细碎化(fragmentation)

生产实际中的模型往往拥有上百路的 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 等技术在此类场景下没有明显的收益。

这些问题说明，想要发挥出 GPU 等高性能硬件资源的极致性价比，提高单位成本下的训练吞吐，就必须设计新的训练框架。据我们了解，拥有大型搜索、广告、推荐业务的国内外企业以及硬件厂商都在着手进行新框架的研发，比如 NVIDIA 的 Merlin-HugeCTR [1] 等，然而阿里巴巴集团内云上集群普遍部署的是通用计算节点，且集群上需要执行多种异构的任务，换用专用硬件是很昂贵且不切实际的。

基于这种实际需求，我们推出了 HybridBackend，能够同时适应集团内多元化且不断演进的稀疏模型技术。下文中我们将简要介绍 HybridBackend 的系统架构设计和技术亮点。

HybridBackend 的系统架构

传统的参数服务器(PS)训练范式体现的是通过扩展硬件数量来适应模型训练规模的思路。我们的系统则是同时考虑到了硬件和软件(模型)两个层面的特点，并做到协同设计。高性能 GPU 集群的硬件特性决定了基本的训练范式，而稀疏模型本身的结构特点和数据分布带来的问题则通过更精细的系统优化手段来解决。

利用大 Batch Size 进行同步训练

因为 GPU 设备相对于 CPU 带来的巨大的算力提升，以往需要上百台 CPU 节点的集群可以用几十台机器的 GPU 集群来代替。要保持相同的总训练规模，同时提升单个 GPU 节点上的资源利用率，提升单个 GPU worker 上的 batch size 成为必然的选项。同时，因为集群规模的缩小，可以通过同步训练的方式有效避免过期梯度(staleness)，从而提升模型训练的精度。

相对于 CPU 设备之间通过 PCIe 以及 TCP 进行网络通信，高性能的 GPU 集群在单个节点内的多个 GPU 设备之间往往配备了高速的网络互连(NVLink、NVSwitch),这些高速连接的带宽通常是 TCP 网络带宽的数百倍(第一代 NVLINK 标定达到 300GB/s),而在多个机器节点之间也可以配备基于 RDMA 技术的高速网络设备，达到 100-200Gbps 的带宽。

选择同步训练的第二个好处是，可以使用高性能集合通信算子库(NVIDIA NCCL、阿里自研的 ACCL 等)来有效利用硬件机器的网络拓扑结构，从而提升通信的性能。上述通信库已经在 CV、NLP 之类的基于数据并行的同步训练任务上取得了很好的效果。

使用资源异构而角色同构的训练单元

PS 训练范式在系统的逻辑层面会指定不同的训练角色，比如 server、worker、evaluator 等。server 节点一般分配具有大内存的 CPU 机器，而 worker 节点则会被分配到高主频的计算型 CPU 硬件上。这样形成了训练单元-任务角色-同构资源的耦合，通过增加训练单元数量来水平扩展(scale out)训练的规模。

而在高性能的 GPU 集群上，一个物理的机器节点往往包括多种异构的硬件资源，如 CPU、GPU 处理器、GPU 之间的高速互连、DRAM（动态随机存取内存）、Non-volatile Memory（非易失性内存）等。这样，除了水平扩展节点数量外，还可以通过垂直扩展利用多种异构硬件资源来达到扩大训练规模的目标。

针对这种硬件架构，我们的系统设计中只保留统一的训练执行单元(Executor),每个 Executor 通过内部的异构硬件资源来执行不同的训练任务角色。

• 一方面，Executor 内部任务执行时，可以有效地利用底层硬件资源之间的 locality 加速训练；
• 另一方面，Executor 内部的硬件资源可以同时满足不同的分布式训练范式所需要的硬件资源，以方便我们在模型结构的不同部分进行混合并行训练策略。

深入优化：HybridBackend 的技术亮点

因为稀疏模型结构和训练数据本身的特性，变化的硬件资源瓶颈和算子细碎化，上述的系统架构在实际任务中还是会存在一些影响 GPU 等硬件设备使用率的问题。

举例来说，同步训练范式下，所有 Executor 通过集合通信进行 embedding 的 shuffle 时，网络带宽资源成为瓶颈，而 GPU 的计算资源被闲置。一种解决思路是对硬件资源进行定制化，比如增加网络带宽资源来消除通信瓶颈，但是这样的做法会使得硬件的资源配置和特定的模型结构耦合，是专用推荐系统的老思路。

我们的目标还是希望系统可以架构在云服务上可得的，数量容易水平扩展的通用硬件配置之上(commodity hardware)。某些硬件厂商也尝试通过 Huge kernel 的形式(将 Embedding 层所有的计算手工融合到一个 kernel 内)来解决算子细碎化的问题,这样的做法也很难支持模型结构快速迭代的需求,背离了通用编程架构的设计初衷。

据此，我们从软硬协同的思路出发，设计了如下的几个系统优化手段：

基于数据和算子感知的合并

根据稀疏模型的结构特点，大部分细碎的算子来源于庞大的 Embedding 特征查询(lookup)数量，我们设计了 D-Packing 这一优化技术。

对于每一路查询，尽管输入的训练数据不同，但使用的算子组合是相同的。对于这种具有数据并行特点的模式，具有相同属性（维度、初始化器、标定特征组等）的 Embedding 表将被合并为一张新的 Embedding 表，而后续的访存查询算子也可以被合并为一个新的大算子。合并算子可以用多线程的方式有序查询 Embedding，相对于乱序查询或分成若干小表查询，能有显著的性能提升。查询完毕后，再依原有代码需要进行反去重和归位，真正做到了对用户透明。

此外，通过分析特征查询阶段各个算子在分布式环境下的语义，我们将部分的 kernel 进行融合 K-Packing，比如通过融合 shuffle 和 stitch 算子来消除冗余的数据拷贝。

通过数据和算子两个维度的基于语义的融合，我们既减少了总体的算子数量，降低 fragmentation，同时又避免了所有算子融合在一起而丢失了通过算子间穿插遮掩来提升硬件利用率的优化机会。

基于硬件资源瓶颈感知的交错执行

为了消除同时执行相同硬件资源需求的算子而造成的瓶颈, 我们设计了两种算子穿插遮掩执行(interleaving)的优化手段。

• 其一，D-Interleaving 是通过对训练数据 batch 的切分利用 pipeline 的机制来调度穿插不同资源类型的算子，这样可以在训练的任何阶段缓解某一种资源的瓶颈。比如在大 batch size 的训练场景下，稀疏模型的 MLP 阶段也会产生很高的 feature map 显存占用，通过 D-Interleaving 就可以有效降低单个 GPU 设备上的峰值显存占用，从而使得更大的 batch size 训练成为可能。
• 其二，K-Interleaving 是在 Embedding Layer 内部不同的特征查询路数之间做算子的穿插和遮掩，比如将通信密集的 Shuffle 操作和内存访问密集的 Gather 进行遮掩，可以有效提升这两种资源的使用率。

基于数据频次感知的参数缓存

在解决 Executor 内部多个级别的存储(GPU 显存、DRAM 等)之间的带宽和延迟问题上，我们针对稀疏模型训练数据的分布特点，提出了一种感知数据访问频次分布的 caching 机制。通过统计训练数据的 ID，将最热的访问数据缓存到 GPU 的显存中，而冷数据以及哈希表结构则存放在主内存中，主内存中的数据将根据 ID 的访问频率变化，定期将 top-k 的高频 ID 对应的 embeddings 刷新到 GPU 显存上的缓存中。这样的混合存储可以同时结合 GPU 显存的高带宽和 DRAM 的大容量，后续，这套混合存储的设计还可以扩展到使用 Intel Persistent Memory、Non-volatile Memory 等更多的硬件设备上。

应用场景

HybridBackend 已经成功在阿里妈妈智能引擎训练引擎团队定向广告业务有了落地。在阿里妈妈 CAN 模型下 HybridBackend 相对于上一代的 XDL 训练框架具有明显的性能优势，在下表中可以看到其在训练时长等多个指标下获得的显著提升。

同时，我们还基于阿里妈妈定向广告一年累计的训练数据对模型规模增长下的 HybridBackend 性能表现做了测试，结果如下表所示。可以看到，在使用 128 张 GPU 进行千亿规模参数模型的训练时，同样是消费 1 年的数据量，高性能集群上的 HybridBackend 仅仅需要 2 天的时间就能完成训练任务，而普通集群上的 XDL-PS 模式则需要约 1 个月的时间。

Roadmap

后续我们计划定期发布 Release 版本。近期的 Roadmap 如下：

• v0.6.0 (2022 年 5 月)：支持端到端分布式同步训练与评估。
• v0.7.0 (2022 年 9 月)：优化 GPU 利用率与显存占用。
• v0.8.0 (2023 年 1 月)：进一步优化云上训练性能。

此外，中长期，我们将在训练策略的演进，新硬件的优化，服务化能力的支持等几个探索性方向上持续投入精力，也欢迎各种维度的反馈和改进建议以及技术讨论，同时我们十分欢迎和期待对开源社区建设感兴趣的同行一起参与共建。

开源地址
https://github.com/alibaba/HybridBackend

参考文献
[1] 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.

论文详情
论文标题：PICASSO: Unleashing the Potential of GPU-centric Training for Wide-and-deep Recommender Systems

论文作者：张远行、陈浪石(并列一作)、杨斯然、袁满、易慧民、张杰、王家忙、董建波、许云龙、宋钺、李永、张迪、林伟、曲琳、郑波

论文链接：https://arxiv.org/abs/2204.04903

开源！阿里自研工业级稀疏模型高性能训练框架 PAI-HybridBackend