阿里开源自研工业级稀疏模型高性能训练框架 PAI-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 变化的硬件资源瓶颈