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

ingFang SC", "Hiragino Sans GB", "Microsoft YaHei UI", "Microsoft YaHei", Arial, sans-serif;letter-spacing: 1.5px;color: rgb(0, 0, 0);font-size: 14px;text-align: justify;visibility: visible;line-height: 2em;">导读本次分享主要针对京东广告的业务场景，讨论我们在 GPU 吞吐和低延时优化方面的实践工作。

主要内容包括四大部分：

ingFang SC", "Hiragino Sans GB", "Microsoft YaHei UI", "Microsoft YaHei", Arial, sans-serif;letter-spacing: 1.5px;line-height: 2em;visibility: visible;">1.京东广告场景介绍

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

ingFang SC", "Hiragino Sans GB", "Microsoft YaHei UI", "Microsoft YaHei", Arial, sans-serif;letter-spacing: 1.5px;line-height: 2em;">3.京东广告推理场景 GPU 优化实践

ingFang SC", "Hiragino Sans GB", "Microsoft YaHei UI", "Microsoft YaHei", Arial, sans-serif;letter-spacing: 1.5px;line-height: 2em;">4.总结

分享嘉宾｜李健京东算法应用工程师

编辑整理｜王甲君

内容校对｜李瑶

出品社区｜DataFun

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

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

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

在实际实践中，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 显存的额外拷贝。

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

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

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

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

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

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

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

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

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