京东广告算法架构体系建设--大规模稀疏场景高性能训练方案演变
一、前言
京东广告训练框架随着广告算法业务发展的特点也在快速迭代升级,回顾近几年大致经历了两次大版本的方案架构演变。第一阶段,随着2016年Tensorflow训练框架的开源,业界开始基于Tensorflow开源框架训练更复杂的模型。模型对特征规模和参数规模需求不断提升,大规模稀疏模型具有更强的表征能力,逐渐成为算法的主流趋势。但是Tensorflow在大规模稀疏参数的训练机制不完备,因此第一次最大能力升级是通过自研高性能参数服务器,支持超大规模TB级稀疏参数模型的建模能力以及基于此架构支持在线学习的能力。 第二阶段,随着用户行为序列建模、多模态建模、多目标等算法技术的发展,模型变得既宽且深,计算算力、通信性能和存储容量逐渐成为瓶颈。基于之前的方案虽然有可能通过扩大训练集群规模满足训练需求,但是在节点拓扑的复杂度、参数通信性能、集群稳定性和模型效果等多个方面都存在较大的问题。随着更先进的NVIDIA A100 等训练GPU硬件资源的出现,基于高性能GPU算力构建的新一代软硬深度结合的训练方案成为第二次架构演变的主要方向。接下来,本文将针对大规模稀疏场景,结合模型的发展趋势,详细介绍各阶段的训练方案。
二、持续演进的大规模稀疏场景训练方案
2.1 基于分布式参数服务器的TB级大规模稀疏场景训练方案
2.1.1 Tensorflow在大规模稀疏场景的局限性
随着业务规模和算法能力不断发展,训练样本规模扩展到百亿级,训练参数规模达到千亿级,为了提高模型训练效率和规模,业界通常采用数据并行和模型并行方式来进行分布式训练。由于Tensorflow采用静态Embedding机制来存储稀疏参数,限制了参数规模,对训练的效率和效果并不友好:
2.1.2 自研高性能参数服务器
为了解决静态Embedding的问题,我们通过自研动态Embedding的高性能参数服务器,将不同的Embedding映射到不同的内存空间,优化存储空间,支持大规模稀疏参数的高效存储,并考虑高并发读写场景,针对稀疏参数设计高性能二级检索方案,优化数据结构,减少并发读写冲突。
![](https://s3.cn-north-1.jdcloud-oss.com/shendengbucket1/2023-12-29-15-47S36dnbsRSn47pjYD.png)
图1 动态Embedding参数服务器&训练架构
![](https://s3.cn-north-1.jdcloud-oss.com/shendengbucket1/2023-12-29-14-0929VWfCq2EYhXyBem.png)
图2 动态Embedding高性能参数服务器实现设计
在2020年底广告实现了高性能参数服务器架构,2021年初在搜索广告场景构建了离在线一体实时在线学习闭环系统。性能方面,支持千亿参数规模,提升算法建模空间,对比TF原生PS训练性能提升25%,对比阿里DeepRec和腾讯TFRA,内存节省约15%至20%。效果方面,在搜索精排模型场景落地京东零售首个OnlineLearning架构,模型时效性从天级提高到分钟级。
2.2 基于高性能算力的全GPU训练方案
2.2.1 分布式参数服务器方案的局限性
为了精细刻画用户的行为,捕捉用户的兴趣变化,推荐系统模型进入到更深层次的序列化、多模态建模。推荐领域模型建模逐渐发展到以一定规模Transformer结构结合大规模Embedding为主,抽象结构则是中等规模稠密参数 + 大规模稀疏参数,对模型训练架构的存储、计算、IO等方面都提出了更高的要求。虽然可以通过更大的参数服务器规模 + 更大的数据并行规模来提升整体训练集群吞吐,但该方案存在一定局限性:
![](https://s3.cn-north-1.jdcloud-oss.com/shendengbucket1/2023-12-29-14-09iDA0cyvNS2kXDDv.png)
图3 分布式参数服务器存在的问题
随着Nvidia推出新一代NVIDIA A100 80GB SXM GPU服务器,伴随NvLink、IB网络等GPU-RDMA硬件加持,相较P40,在存储容量(显存640GB、内存1~2TB、SSD10+TB级别)、通信性能 (由PCIe总线提供的32GB/s升级到NVLink卡间互联的600GB/s的高速带宽,多机之间通过IB网络采用RDMA进行通信,通信带宽相比于TCP而言,从1G/s 提升到50G/s)以及计算能力(从12TFLOPS升级到156 TFLOPS) 都有质的突破。更高算力的硬件资源为模型迭代和优化带来了更大的机遇,同时也给训练架构设计带来了更大的挑战,基于高性能GPU算力构建的新一代软硬深度结合的训练方案成为第二次架构演变的主要方向。
2.2.2 基于高性能算力的全GPU训练方案
如何把分布式训练中上百个节点的复杂拓扑结构,融入到高算力GPU的一体化方案中,实现千亿甚至更大规模参数在GPU上的高性能计算,在业界是一个比较有挑战的问题,在新一代训练架构的落地实践中我们主要面临三个核心挑战:
结合对业界方案的调研以及广告推荐算法的业务场景特点,我们设计了基于高性能算力的全GPU训练方案:基于GPU-HBM+CPU-DRAM两级交叉缓存来实现全参数的GPU存储计算,结合GPU-RDMA集合通信技术突破大规模参数训练的通信瓶颈,同时构建CPU&GPU异构的分布式流水线并行训练框架,实现了CPU和GPU计算算力的最大化协同。
![](https://s3.cn-north-1.jdcloud-oss.com/shendengbucket1/2023-12-29-14-097c12Goh12C2VmZkEB.png)
图4 基于高性能算力的全GPU训练方案
基于GPU-HBM和CPU-DRAM的两级交叉参数服务器,实现稀疏参数的跨域存储。GPU作为一级缓存有效利用GPU高带宽、高算力的优势,极大提升训练吞吐,CPU作为二级缓存有效利用CPU内存的高容量、易扩展的优势,提升参数的规模上限。该方案可以进一步扩展SSD固态硬盘,形成HBM-DRAM-SSD的三级参数服务器方案,支撑更大规模万亿参数级别的高性能训练。
![](https://s3.cn-north-1.jdcloud-oss.com/shendengbucket1/2023-12-29-14-09zJOSo0O27OYlRJc.png)
图5 GPU-HBM+CPU-DRAM 两级交叉存储方案
基于GPU-RDMA集合通信的全参数同步训练范式。在参数通信方式上:参数的拉取方式相比于传统的分布式参数服务器架构,由CPU-to-CPU的TCP通信升级为GPU-to-GPU的RDMA通信,带宽吞吐量实现了数量级的显著提升(1GB/s->600GB/s)。在参数更新方面:通过AllReduce、AllToAll等集合通信技术,实现全参数同步训练范式,保证模型效果。基于上述能力,框架一体化设计实现了参数交互&参数训练的最佳配速。
CPU&GPU异构的分布式流水线并行训练框架,将模型按CPU密集型子图(重逻辑)与GPU密集型子图(重计算)进行分图,部署于CPU&GPU异构集群进行分布式分图训练,在发挥各自硬件算力优势的同时,解决CPU与GPU算力不匹配的问题,达到算力均衡。同时构建多级流水线并行训练模式,使训练集群整体IO、计算吞吐最大化。
![](https://s3.cn-north-1.jdcloud-oss.com/shendengbucket1/2023-12-29-14-09KIJy2VCRdadLx9r.png)
图6 CPU&GPU异构的分布式流水线并行训练框架
多机多卡算力水平扩展:基于上述核心能力,借助 IB 网络,进一步实现单机多卡到多机多卡的算力水平扩展,训练加速比达到1.85+,做到业界领先。
三、总结与展望
综上,新一代基于高性能算力的全GPU训练方案,在广告多个业务线进行了落地实践,取得了显著的效果收益:推荐首页CTR模型规模从30G扩展到130G,资源0增长的情况下,训练性能提升55%,并通过技术迁移仅一个月时间将技术方案复用至商详、中长尾信息流核心位置,实现模型规模百G的突破,算法迭代效率提升400%,助力推荐大盘整体点击+8.5%, 消费 +5.77%,搜索精排整体点击+2%,消费+5.63%。
算法在不断的迭代发展,硬件也在不断的推陈出新,广告训练框架也在持续的演变,我们正在规划设计下一代算法架构体系,其最显著的特点就是算法、算力、架构的深度融合以及在线、离线一体化的设计方案。算法架构体系建设是一个充满挑战的新的技术领域,需要不断探索、学习和创新。
作者:京东零售 刘凯宁
来源:京东云开发者社区 转载请注明来源