共计 6023 个字符，预计需要花费 16 分钟才能阅读完成。
近年来，随着稠密模型对算力日益增长的需要, CPU 集群必须不断扩大集群规模来满足训练的时效需要，这同时也带来了一直回升的资源老本以及试验的调试老本。
为了解决这一问题，阿里云机器学习 PAI 平台和阿里妈妈智能引擎训练引擎团队合作开发了稠密模型高性能同步训练框架 HybridBackend，使得在同老本下 GPU 集群训练吞吐较 CPU 集群晋升至 5 倍，大幅升高调试老本，同时 HybridBackend 相干论文《PICASSO: Unleashing the Potential of GPU-centric Training for Wide-and-deep Recommender Systems》也被 ICDE 22′ 所收录。HybridBackend 背地的技术框架如何设计？将来有哪些布局？明天一起来深刻理解。
一 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 设施训练的两个问题。
1 变动的硬件资源瓶颈
从上图的 Wide and Deep 模型构造能够看出，稠密训练次要由 Embedding 阶段、特色穿插 (feature interation) 阶段和多层感知器 (MLP) 阶段组成，Embedding 阶段在 PS 范式的训练下占据了至多 50% 以上的训练工夫。通过剖析发现，Embedding 阶段的算子次要以访存密集型 (memory access intensive) 和通信密集型的算子 (communication intensive) 为主，次要须要的硬件资源是内存和网络的带宽，而后两个阶段的算子则是计算密集型的算子占主导, 须要的资源是算力。这意味着在 PS 的范式训练下，任何一个阶段都有可能存在某一种硬件资源成为瓶颈而其余硬件资源被节约的景象。以 GPU 的算力资源为例，咱们察看 GPU 使用率 (SM Util) 在不同的训练阶段之间出现脉冲式变动(pulse)。
2 算子细碎化(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 集群的硬件个性决定了根本的训练范式，而稠密模型自身的构造特点和数据分布带来的问题则通过更精密的系统优化伎俩来解决。
1 利用大 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 之类的基于数据并行的同步训练任务上获得了很好的成果。
2 应用资源异构而角色同构的训练单元
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 内) 来解决算子细碎化的问题, 这样的做法也很难反对模型构造疾速迭代的需要, 背离了通用编程架构的设计初衷。
据此，咱们从软硬协同的思路登程，设计了如下的几个系统优化伎俩:
1 基于数据和算子感知的合并
依据稠密模型的构造特点，大部分细碎的算子来源于宏大的 Embedding 特色查问 (lookup) 数量，咱们设计了 D -Packing 这一优化技术。
对于每一路查问，只管输出的训练数据不同，但应用的算子组合是雷同的。对于这种具备数据并行特点的模式，具备雷同属性（维度、初始化器、标定特色组等）的 Embedding 表将被合并为一张新的 Embedding 表，而后续的访存查问算子也能够被合并为一个新的大算子。合并算子能够用多线程的形式有序查问 Embedding，绝对于乱序查问或分成若干小表查问，能有显著的性能晋升。查问结束后，再依原有代码须要进行反去重和归位，真正做到了对用户通明。
此外，通过剖析特色查问阶段各个算子在分布式环境下的语义，咱们将局部的 kernel 进行交融 K -Packing, 比方通过交融 shuffle 和 stitch 算子来打消冗余的数据拷贝。
通过数据和算子两个维度的基于语义的交融，咱们既缩小了总体的算子数量，升高 fragmentation, 同时又防止了所有算子交融在一起而失落了通过算子间交叉遮掩来晋升硬件利用率的优化机会。
2 基于硬件资源瓶颈感知的交织执行
为了打消同时执行雷同硬件资源需要的算子而造成的瓶颈, 咱们设计了两种算子交叉遮掩执行 (interleaving) 的优化伎俩。
其一，D-Interleaving 是通过对训练数据 batch 的切分利用 pipeline 的机制来调度交叉不同资源类型的算子，这样能够在训练的任何阶段缓解某一种资源的瓶颈。比方在大 batch size 的训练场景下，稠密模型的 MLP 阶段也会产生很高的 feature map 显存占用，通过 D -Interleaving 就能够无效升高单个 GPU 设施上的峰值显存占用，从而使得更大的 batch size 训练成为可能。
其二，K-Interleaving 是在 Embedding Layer 外部不同的特色查问路数之间做算子的交叉和遮掩，比方将通信密集的 Shuffle 操作和内存拜访密集的 Gather 进行遮掩，能够无效晋升这两种资源的使用率。
3 基于数据频次感知的参数缓存
在解决 Executor 外部多个级别的存储 (GPU 显存、DRAM 等) 之间的带宽和提早问题上，咱们针对稠密模型训练数据的散布特点，提出了一种感知数据拜访频次散布的 caching 机制。通过统计训练数据的 ID，将最热的拜访数据缓存到 GPU 的显存中，而冷数据以及哈希表构造则寄存在主内存中，主内存中的数据将依据 ID 的拜访频率变动，定期将 top- k 的高频 ID 对应的 embeddings 刷新到 GPU 显存上的缓存中。这样的混合存储能够同时联合 GPU 显存的高带宽和 DRAM 的大容量，后续，这套混合存储的设计还能够扩大到应用 Intel Persistent Memory、Non-volatile Memory 等更多的硬件设施上。
六 利用场景
HybridBackend 曾经胜利在阿里妈妈智能引擎训练引擎团队定向广告业务有了落地。在阿里妈妈 CAN 模型下 HybridBackend 绝对于上一代的训练框架 (应用 Parameter Server 模式) 具备显著的性能劣势，在下表中能够看到其在训练时长等多个指标下取得的显著晋升。
同时，咱们还基于阿里妈妈定向广告一年累计的训练数据对模型规模增长下的 HybridBackend 性能体现做了测试，后果如下表所示。能够看到，在应用 128 张 GPU 进行千亿规模参数模型的训练时，同样是生产 1 年的数据量，高性能集群上的 HybridBackend 仅仅须要 2 天的工夫就能实现训练任务，而一般集群上的 PS 模式则须要约 1 个月的工夫。
七 Roadmap
后续咱们打算定期公布 Release 版本。近期的 Roadmap 如下：
- v0.6.0 (2022 年 5 月)：反对端到端分布式同步训练与评估。
- v0.7.0 (2022 年 9 月)：优化 GPU 利用率与显存占用。
- v0.8.0 (2023 年 1 月)：进一步优化云上训练性能。
此外，中长期，咱们将在训练策略的演进，新硬件的优化，服务化能力的反对等几个探索性方向上继续投入精力，也欢送各种维度的反馈和改良倡议以及技术探讨，同时咱们非常欢送和期待对开源社区建设感兴趣的同行一起参加共建。
参考文献
[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.
