一、前言

京东广告训练框架随着广告算法业务发展的特点也在快速迭代升级，回顾近几年大致经历了两次大版本的方案架构演变。第一阶段，随着 2016 年 Tensorflow 训练框架的开源，业界开始基于 Tensorflow 开源框架训练更复杂的模型。模型对特征规模和参数规模需求不断提升，大规模稀疏模型具有更强的表征能力，逐渐成为算法的主流趋势。但是 Tensorflow 在大规模稀疏参数的训练机制不完备，因此第一次最大能力升级是通过自研高性能参数服务器，支持超大规模 TB 级稀疏参数模型的建模能力以及基于此架构支持在线学习的能力。 第二阶段，随着用户行为序列建模、多模态建模、多目标等算法技术的发展，模型变得既宽且深，计算算力、通信性能和存储容量逐渐成为瓶颈。基于之前的方案虽然有可能通过扩大训练集群规模满足训练需求，但是在节点拓扑的复杂度、参数通信性能、集群稳定性和模型效果等多个方面都存在较大的问题。随着更先进的 NVIDIA A100 等训练 GPU 硬件资源的出现，基于高性能 GPU 算力构建的新一代软硬深度结合的训练方案成为第二次架构演变的主要方向。接下来，本文将针对大规模稀疏场景，结合模型的发展趋势，详细介绍各阶段的训练方案。

二、持续演进的大规模稀疏场景训练方案

2.1 基于分布式参数服务器的 TB 级大规模稀疏场景训练方案

2.1.1 Tensorflow 在大规模稀疏场景的局限性

随着业务规模和算法能力不断发展，训练样本规模扩展到百亿级，训练参数规模达到千亿级，为了提高模型训练效率和规模，业界通常采用数据并行和模型并行方式来进行分布式训练。由于 Tensorflow 采用静态 Embedding 机制来存储稀疏参数，限制了参数规模，对训练的效率和效果并不友好：

◦静态存储局限性：词表空间过小，hash 冲突加剧；词表空间过大，浪费内存资源，难以支持大规模参数存储。

◦在线学习不支持：针对在线学习场景，无法淘汰不重要的特征，也无法单独释放该特征 Embedding 的内存，参数更新时效性成为瓶颈。

2.1.2 自研高性能参数服务器

为了解决静态 Embedding 的问题，我们通过自研动态 Embedding 的高性能参数服务器，将不同的 Embedding 映射到不同的内存空间，优化存储空间，支持大规模稀疏参数的高效存储，并考虑高并发读写场景，针对稀疏参数设计高性能二级检索方案，优化数据结构，减少并发读写冲突。

图 1 动态 Embedding 参数服务器&训练架构

▪参数二级检索

▪一级参数分片，支撑分布式千亿甚至更大规模参数存储。

▪二级参数分桶，并行无锁读写，支撑数据并行模式下的高并发读写场景。

▪ 稀疏参数存储&训练

▪基于 Map 实现稀疏参数存储，支撑动态 Embedding、增量导出等复杂功能。

▪稀疏参数延迟初始化，支撑特征准入&淘汰等复杂策略。

▪参数&优化器状态同构存储，支撑复杂优化算法（带一二阶动量）的高性能实现。

▪离&在线一体化设计

▪通用化服务接口&模型参数格式，一套核心框架服务离&在线场景，支持在线学习。

图 2 动态 Embedding 高性能参数服务器实现设计

在 2020 年底广告实现了高性能参数服务器架构，2021 年初在搜索广告场景构建了离在线一体实时在线学习闭环系统。性能方面，支持千亿参数规模，提升算法建模空间，对比 TF 原生 PS 训练性能提升 25%，对比阿里 DeepRec 和腾讯 TFRA，内存节省约 15% 至 20%。效果方面，在搜索精排模型场景落地京东零售首个 OnlineLearning 架构，模型时效性从天级提高到分钟级。

2.2 基于高性能算力的全 GPU 训练方案

2.2.1 分布式参数服务器方案的局限性

为了精细刻画用户的行为，捕捉用户的兴趣变化，推荐系统模型进入到更深层次的序列化、多模态建模。推荐领域模型建模逐渐发展到以一定规模 Transformer 结构结合大规模 Embedding 为主，抽象结构则是中等规模稠密参数 + 大规模稀疏参数，对模型训练架构的存储、计算、IO 等方面都提出了更高的要求。虽然可以通过更大的参数服务器规模 + 更大的数据并行规模来提升整体训练集群吞吐，但该方案存在一定局限性：

▪参数服务器规模仅能满足大规模 Embedding 稀疏模型，但是针对 Transformer 等稠密模型无法分片，会导致模型训练性能&效果显著下降；

▪训练节点规模急剧扩大到上百，导致通信拓扑复杂，训练稳定性下降，梯度过期问题加剧；

▪大量参数 IO 使得参数服务器与训练节点之间的参数传输成为通信性能瓶颈。

图 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 上的高性能计算，在业界是一个比较有挑战的问题，在新一代训练架构的落地实践中我们主要面临三个核心挑战：

▪存储挑战：不同于传统 NLP、CV 几十 G 模型规模，推荐领域大规模稀疏模型通常在几百 G 以上，远超单显卡 80G 显存上限，需要设计一种新的参数存储和计算策略；

▪IO 挑战：大规模稀疏参数、稠密参数的拉取与计算为训练集群带来极高的 IO 需求，需要设计一种新的交互式策略，实现卡与卡间，机与机间参数高性能通信；

▪算力挑战：CPU 的算力性能限制了推荐模型的规模与计算复杂度，需要设计一种新的训练并行策略，实现 CPU 与 GPU 的算力均衡。

结合对业界方案的调研以及广告推荐算法的业务场景特点，我们设计了基于高性能算力的全 GPU 训练方案：基于 GPU-HBM+CPU-DRAM 两级交叉缓存来实现全参数的 GPU 存储计算，结合 GPU-RDMA 集合通信技术突破大规模参数训练的通信瓶颈，同时构建 CPU&GPU 异构的分布式流水线并行训练框架，实现了 CPU 和 GPU 计算算力的最大化协同。

图 4 基于高性能算力的全 GPU 训练方案

基于 GPU-HBM 和 CPU-DRAM 的两级交叉参数服务器，实现稀疏参数的跨域存储。GPU 作为一级缓存有效利用 GPU 高带宽、高算力的优势，极大提升训练吞吐，CPU 作为二级缓存有效利用 CPU 内存的高容量、易扩展的优势，提升参数的规模上限。该方案可以进一步扩展 SSD 固态硬盘，形成 HBM-DRAM-SSD 的三级参数服务器方案，支撑更大规模万亿参数级别的高性能训练。

图 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、计算吞吐最大化。

图 6 CPU&GPU 异构的分布式流水线并行训练框架

多机多卡算力水平扩展：基于上述核心能力，借助 IB 网络，进一步实现单机多卡到多机多卡的算力水平扩展，训练加速比达到 1.85+，做到业界领先。

三、总结与展望

综上，新一代基于高性能算力的全 GPU 训练方案，在广告多个业务线进行了落地实践，取得了显著的效果收益：推荐首页 CTR 模型规模从 30G 扩展到 130G，资源 0 增长的情况下，训练性能提升 55%，并通过技术迁移仅一个月时间将技术方案复用至商详、中长尾信息流核心位置，实现模型规模百 G 的突破，算法迭代效率提升 400%，助力推荐大盘整体点击 +8.5%, 消费 +5.77%，搜索精排整体点击 +2%，消费 +5.63%。

算法在不断的迭代发展，硬件也在不断的推陈出新，广告训练框架也在持续的演变，我们正在规划设计下一代算法架构体系，其最显著的特点就是算法、算力、架构的深度融合以及在线、离线一体化的设计方案。算法架构体系建设是一个充满挑战的新的技术领域，需要不断探索、学习和创新。

作者：京东零售 刘凯宁

来源：京东云开发者社区 转载请注明来源