简介：近年来，随着稠密模型对算力日益增长的需要, CPU集群必须不断扩大集群规模来满足训练的时效需要，这同时也带来了一直回升的资源老本以及试验的调试老本。为了解决这一问题，阿里云机器学习PAI平台开源了稠密模型高性能同步训练框架HybridBackend，使得在同老本下GPU集群训练吞吐较CPU集群晋升至5倍，大幅升高调试老本。那么HybridBackend背地的技术框架如何设计？将来有哪些布局？本文将和大家一起来深刻理解。

作者 | 石浪、满神
起源 | 阿里开发者公众号

近年来，随着稠密模型对算力日益增长的需要, 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绝对于上一代的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/Hy...

参考文献

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

原文链接
本文为阿里云原创内容，未经容许不得转载。