京东广告稀疏大模型训练与推理 GPU 优化实践

本文作者李健针对京东广告的业务场景，分享了在 GPU 吞吐和低延时优化方面的实践工作。

李健：17 年加入京东广告研发部，在算法架构方向深耕 7 年，先后参与了两代京东广告模型系统的架构建设。主导了大规模稀疏模型训练 &推理等能力建设，助力多项算法创新落地，并多次获得京东零售级别奖项。目前作为广告算法架构方向架构师，带领团队建设新一代广告异构算法架构体系。

01 京东广告场景介绍 

京东广告业务场景包括推荐首页的信息流和搜索场景，主要面向 C 端用户，面临百万 QPS 高并发访问压力。为了保障广告展现效率与用户体验，广告系统需要毫秒级响应能力。在算法建模方面，推荐领域模型由浅层大规模离散 DNN 发展到以 Transformer 为主的深层神经网络，参数从百 GB 扩展到 TB 级，算力需求增长了数十倍。目前我们使用 TensorFlow 进行建模训练与推理，其算子级建模和静态图建模方式保证了离在线一致性和算法实验的快速迭代。

面对日益增长的复杂业务需求，我们采用 GPU 专用计算设备提升系统算力，但具体实践中面临许多挑战：

• CTR 模型的高稀疏性特点容易导致 I/O 瓶颈。

• 超大规模稀疏参数模型容易超出显存承载上限。

• 推荐场景的特征计算占用大量 CPU 资源，导致 CPU 与 GPU 资源调度争抢，使得整体利用率释放不充分。

接下来将分享我们在训练和推理场景中是如何解决这些问题的。

02 京东广告训练场景 GPU 优化实践

首先介绍京东广告训练场景 GPU 优化实践。这个场景主要以吞吐为第一优化目标，在这个背景下，GPU 优化面临三个主要挑战：

1. 存储挑战

在推荐领域，模型主要以大规模稀疏参数为主，规模通常会达到千亿级别，存储量可达 TB 级别。即使是 H800 这种高规格 GPU，其显卡显存上限也只有 80G，并且这种 HBM 类型的存储资源价格昂贵。完全将大规模稀疏参数承载在 GPU 上，是一个性价比很低的方案。为解决这个问题，我们采用了多机多卡的训练模式，并外挂了自研的 CPU 稀疏参数服务器，利用价格相对较低且易扩展的 CPU-DRAM 内存来扩展模型的承载能力。

2. 计算挑战

在推荐领域，用户行为特征解析、用户与广告物料的特征交叉等占有很大比重，且计算同样占用大量资源。但由于这些计算重业务逻辑，目前只能由 CPU 进行推理，导致 CPU 与 GPU 资源争抢，算力分配不均。为解决这个问题，我们研发了一种 CPU 与 GPU 异构混合部署的分布式流水线并行训练架构，使特征计算网络和模型训练网络适配不同的硬件类型，能各自发挥其硬件优势。

3. I/O 挑战

在实际实践中，Embedding 层的 I/O 耗时占比超过 30%，导致 GPU 利用率不高。为此，我们在 CPU 稀疏参数服务器基础上，研发了基于 GPU 的 GPU-HBM 参数服务器，以提升整体训练吞吐。

接下来将具体介绍这些方案的实践。

存储方面，我们研发了针对稀疏参数存储的 Embedding 参数服务器，使用 CPU 内存进行分片承载。而对于 MLP 等稠密参数，每个 worker 的 GPU 显存独立复制一份。

参数更新方面，embedding 的 key 通过 AlltoAll 集合通信模式，在卡间和跨机间进行高速通信，对应的 embedding 值通过本机 PCIE 总线访问本地内存。稠密参数则通过卡间的 nvlink 和跨机间的 IB 网络进行 AllReduce 通信。这个方案不仅保证了全同步训练的稳定性，还利用了 CPU 内存的易扩展优势，使参数量扩展到千亿级。然而在这个方案中，每次 embedding 查询与更新都要通过 PCIE 总线，导致参数交换量大、I/O 负载高，占用了整体训练时间。此外，稀疏参数更新集中在 CPU 参数服务器，未能充分利用 GPU 的高算力优势，后面我们会介绍我们如何解决这种 IO 问题。

计算方面，推荐领域模型包含了许多部分，如特征预处理网络和模型训练网络，使其成为 CPU 和 GPU 密集型结合的训练场景。因此，我们研发了计算分图并行训练模式，将完整的算法网络拆分为 CPU 密集型的特征计算网络和 GPU 密集型的模型计算网络，并分别部署于两个不同的集群上。CPU 密集型的特征计算网络独立部署在 CPU 计算集群上；GPU 密集型的模型计算网络则利用全同步训练方案，部署在高性能如 A800、H800 的 GPU 集群上。

原本 CPU 负载高但 GPU 利用率低的场景，通过拆分计算后，CPU 负载降至合理水平，GPU 利用率显著提升。此外，两个集群之间采用分布式流水线方案，使特征预处理和训练并行计算，最大化集群整体吞吐量。由于特征计算网络的独立部署，训练集群中 I/O 负载较重的样本下载过程也在特征计算集群进行，我们可以通过扩容特征计算集群，实现了 I/O 分摊，提升了训练系统的整体 I/O 吞吐能力。

训练方面最后一个挑战是 I/O。在 I/O 方面，我们基于 CPU 参数服务器扩展了稀疏参数的承载能力。为提升吞吐，我们研发了基于 GPU-HBM 的参数服务器作为一级缓存，充分利用 GPU 的高带宽和高并发优势，显著提升训练吞吐。CPU 作为二级参数服务器，利用其高容量和易扩展性，提升参数规模上限。

在模型正向网络中，每次正向训练时，部分 embedding 可能未命中缓存，本地通过 PCIE 总线访问存储服务器。而在梯度更新过程中，由于 GPU-HBM 参数服务器已覆盖当前 step 的所有 embedding key，梯度更新直接在 GPU 显存中进行。同时，我们实现了融合 Adam 优化器，来减少优化器状态多次访问，从而优化 GPU 显存访问次数，降低 IO 压力。

在 I/O 方面，我们分析了训练集群的整体工作流，主要分为五步：样本下载、特征计算、特征拉取、数据从主机到设备的拷贝，以及训练。基于此工作流，我们研发了五级分布式流水线，将 I/O 密集型的特征计算、CPU 密集型计算和 GPU 计算进行了流水线并行优化。

值得注意的是，我们基于 Tensorflow Dataset 机制实现了数据从主机到设备的预拷贝能力，在 TensorFlow dataset 处理过程中，在 map dataset 和 prefetch 之间插入预拷贝 dataset。其可以自动识别特征计算网络输出 tensor 的下游算子的设备依赖，并预拷贝到 GPU 上，同时实现了聚合拷贝来优化传输。拷贝后的 tensor 切片采用零拷贝机制，避免了 GPU 显存的额外拷贝。

03 京东广告推理场景 GPU 优化实践 

推理场景与训练场景不同，训练场景注重吞吐，与 GPU 的高吞吐计算天然契合。而广告推荐场景对低延时和高并发有严格要求，因此需要新的计算框架来适配高吞吐的 GPU 计算设备。具体来说，有几个原因：

• 广告商品排序队列长度不一，较短的请求队列会拉低整体 GPU 利用率。

• 在线低延时高并发场景需要新的方法最大化利用 GPU 计算资源。

• 推荐领域的用户行为序列建模方案多样，针对不同的行为（如点击、加购、购买）分别建模能提升模型效果，但也导致模型结构复杂，包含上千个算子，难以友好调度 GPU。

针对这三个问题，我们研发了各种针对性方案来解决。

首先，我们设计了 TensorBatch 方案，针对在线请求进行 Batch 聚合推理方案。虽然 Batch 方案是 GPU 在线推理常规优化手段，但常规方案主要考虑待推理样本数，以最大化 GPU 吞吐为目标，无法直接用于低延时广告场景。因此，我们设计了具有针对性的 Batch 方案，以最大化吞吐和优化延时为共同目标。在 Batch 聚合过程中，需要考虑 Batch 后的总计算量，由两部分构成：用户序列建模的计算消耗和广告排序队列的计算消耗。这两部分的加权平均为总体计算量，通过动态分配计算量，构造不同的 Batch 聚合结果。这种针对广告推荐场景的 Batch 聚合方案，既提升了整体利用率，又不降低广告系统对外的延迟消耗。

接下来介绍多流计算，这是针对 GPU 高并发请求的优化。在广告推理场景中，我们使用 TensorFlow 作为底层后端，但其 GPU 推理无法处理并发请求。当多个请求同时运行时，底层的计算队列会串行执行，主要原因是 TensorFlow 与 GPU 设备只有一个交互通道，导致多个请求的 GPU 计算串行执行，缺乏并发处理能力。

为此，我们深度定制了 TensorFlow 底层的 device 层，引入了多流计算能力，支持并发请求的并行计算。这包括三个层级：

• 扩展 TensorFlow 底层的 stream group，通过多次 Cuda Stream 实现并发计算。

• 为每个 device 构建多个 Cuda Context，增强 GPU 资源调度的并发能力。

• 基于英伟达的 MPS 工具，减少引入多 Cuda Context 后的上下文切换损耗，最终实现多个请求间算子并行的效果。

最后是关于深度学习编译器。在推荐场景中，用户行为多样，如点击、购买和加购等。我们针对不同行为建模，导致模型结构分支多，算子粒度细。通常一个完整模型包含上千个算子，GPU 调度这些算子的 Kernel Launch 消耗很大。因此，我们的目标是降低整体 GPU 的 Kernel Launch 耗时。

在这方面，业界常用算子融合来优化，常见方案有手动融合和自动融合。手动融合需分析算法结构，提取常见且计算损耗大的模块，手动编写优化算子，如 self-attention 算子融合方案，其可最大化提升性能，但对开发人员要求高，且为了保障离在线一致性，需要实现正向和反向算子，工作量大。自动融合方案则通过采用 TVM、XLA 等深度学习编译器，自动融合算子和优化计算图，这个方案符合我们的优化目标，可以同时提高模型迭代效率和推理速度。

然而，深度学习编译器的特性使其不能直接应用于广告场景，主要原因编译器只支持静态维度，编译结果需要与输入数据维度一一对应。例如，一条请求广告队列为 150，用户行为长度为 1000，针对该请求的编译结果仅支持 150、1000 的特征维度。当请求队列变为 200，用户行为长度为 1100 时，会再次触发编译。直接应用该方案会导致数据维度变化大，产生上万个编译结果。此外，TensorFlow 底层调用 XLA 的运行时编译机制导致编译耗时较长，无法满足低延时需求。因此，我们对深度学习编译器进行了扩展，实现了分图分桶预编译和异步编译能力，解决了深度学习编译器无法直接应用于推荐广告场景的问题。

深度学习编译器的扩展首先是分图分桶预编译。我们将完整模型划分为独立的子图模块，为每个子图分配分桶属性。同时在模型加载的 warmup 阶段，进行分桶属性的预编译，将编译结果限制在几百个量级，避免运行时编译的耗时。

其次是异步编译。在广告流量场景中，偶尔会出现长尾流量，部分流量或特征可能无法匹配预设的分桶，从而触发编译。为了保障在线推理的稳定性，引入异步编译。在原有计算图结构上，同时保留深度学习编译器的结果图和原始图。对于无法命中分桶的情况，会通过 TensorFlow 的原图进行推理。同时，后台会触发异步编译，对未命中分桶结果的 tensor 维度进行编译。编译后，结果会写回分桶缓存，为后续的长尾流量提供服务。

04 总结 

最后进行一下总结。

在京东广告业务中，我们在大规模稀疏模型的建模场景中，主要集中在 I/O 优化技术上，包括 Kernel Launch、I/O 操作，以及主机端到设备端、设备端到主机端和跨网络 TCP 通信等 I/O 优化。未来，随着模型的发展，我们将向稀疏与稠密相结合的复杂场景演变。例如，Transformer 的更大规模应用和生成式模型与推荐系统的结合，要求系统具备更强算力。GPU 推理与训练也将朝着稠密加稀疏结合的方向发展。

未来我们将重点优化计算和 HBM 内的 I/O 操作，并考虑 tensor 并行、流水线并行等模型内部推理并行方案，以增强系统能力。

最后感谢大家的阅读与分享，欢迎大家一起探讨大模型推理与训练的最新动态和挑战。我们相信，通过这样的交流，促进共同进步。同时作为京东广告部的核心方向，算法团队基于京东海量的用户和商家数据，探索最前沿的深度学习等算法技术，创新并应用到业务实践中，赋能千万商家和数亿消费者的消费连接，不断拓展中国乃至全世界的数字经济边界。我们正在寻找充满激情和才华的新成员加入我们的旅程，如果您对未来个性化推荐、DL/RL 领域的前沿技术充满热情，那么欢迎您发送简历到  lijian21@jd.com， 我们期待您的加入。