作者：石浪、满神

近日，阿里巴巴投稿的论文《PICASSO: Unleashing the Potential of GPU-centric Training for Wide-and-deep Recommender Systems》被数据库和数据挖掘方向CCF-A类国际三大顶会之一的ICDE 2022 (International Conference on Data Engineering) 接收。ICDE 2022 将于今年5月9日至5月13日在马来西亚吉隆坡线上召开，届时将汇集工业级和学术界的顶级学者共同探讨数据密集型的系统和应用相关的热点问题。此次论文入选，意味着阿里云机器学习平台PAI自研的面向工业级稀疏模型的高性能训练框架(PAI-HybridBackend)达到了世界先进水平，得到了业界的广泛认可。

为了回馈社区并推动技术进步，阿里云机器学习平台PAI已经将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)的优化手段。1) D-Interleaving是通过对训练数据batch的切分利用pipeline的机制来调度穿插不同资源类型的算子，这样可以在训练的任何阶段缓解某一种资源的瓶颈。比如在大batch size的训练场景下，稀疏模型的MLP阶段也会产生很高的feature map显存占用，通过D-Interleaving就可以有效降低单个GPU设备上的峰值显存占用，从而使得更大的batch size训练成为可能。2）K-Interleaving是在Embedding Layer内部不同的特征查询路数之间做算子的穿插和遮掩，比如将通信密集的Shuffle操作和内存访问密集的Gather进行遮掩，可以有效提升这两种资源的使用率。

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

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

业务落地

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

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

论文详情

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

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