京东广告模型大规模稀疏训练：从架构演进到高性能实战

1. 项目概述当广告模型遇上“超级稀疏”挑战在广告推荐这个领域干了十几年我见过太多技术方案的迭代但最让我印象深刻的永远是处理“大规模稀疏场景”时的那种如履薄冰。这就像在一片看似平静、实则暗藏无数暗礁的海洋里航行你的模型就是那艘船而稀疏特征就是海面上星星点点的岛屿——数量庞大、彼此孤立但每一个都可能指向一座金矿或者说一个高价值的用户行为信号。京东作为国内电商巨头其广告业务面临的正是这样一个极致场景每天数亿的用户产生千亿级别的曝光和点击背后是百亿甚至千亿维度的稀疏特征比如用户历史点击过的某个冷门商品ID、在某个深夜时段的一次搜索词、或者某个特定地域的偏好品类。传统的模型训练框架在面对这种“超级稀疏”时往往会立刻“趴窝”。想象一下一个拥有千亿维度的嵌入表即使只有极少数维度在每次训练中被激活其存储和通信开销也是天文数字。更棘手的是广告场景的实时性要求极高模型需要快速捕捉热点事件如618大促、新品首发带来的流量和兴趣变化这就要求训练系统不仅要能“装得下”还要“算得快”、“学得准”。京东广告算法架构体系的演进本质上就是一部与“大规模稀疏”持续博弈、并寻求高性能训练最优解的历史。今天我就结合自己的实践和观察来拆解一下这套体系背后的核心思路与方案演变希望能给正在类似泥潭中挣扎的团队一些实在的参考。2. 核心挑战与设计思路演变2.1 大规模稀疏场景的“三座大山”要理解方案为什么这么变首先得搞清楚我们到底在对付什么。大规模稀疏训练尤其是广告场景主要面临三座大山第一座山存储与内存的巨兽。这是最直观的挑战。一个典型的CTR点击率预估模型其稀疏特征部分主要是Embedding层的参数规模轻松超过模型稠密部分的数百甚至上千倍。这些参数无法全部放进单个GPU甚至单个服务器的内存。早期的做法是使用参数服务器Parameter Server, PS架构将庞大的嵌入表分布到多个机器上。但这就引出了第二座山。第二座山通信与同步的瓶颈。在PS架构下训练worker需要从远程PS节点拉取pull它所需的稀疏参数embedding向量计算完梯度后再推送push回去。在广告场景中由于特征的极度稀疏性每次训练batch激活的参数只是总体的极小一部分但为了这点数据worker和PS之间却要发起海量的、细粒度的网络通信。这种“小流量、高并发”的通信模式极易造成网络拥塞和PS节点的性能热点成为训练速度的主要瓶颈。我经历过一个案例模型性能的90%时间都花在了等待网络IO上计算资源利用率低得可怜。第三座山动态性与一致性的权衡。广告系统的特征更新极其频繁。新的商品上线、用户实时行为产生的新特征都需要快速融入模型。这就要求训练系统能支持特征的动态增删Dynamic Vocabulary。同时在分布式训练中如何保证所有worker看到一致的、最新的参数视图一致性模型又是一个难题。严格的一致性如同步SGD会拖慢整体速度而过于宽松的一致性如异步SGD又可能导致模型训练不稳定难以收敛。2.2 架构演进的核心思路从“中心化调度”到“协同计算”面对这些挑战京东广告算法架构的演进脉络非常清晰其核心思路是从传统的“以存储为中心、中心化调度”的PS架构逐步转向“以计算为中心、协同式分布”的新范式。第一阶段参数服务器PS时代。这是解决大规模参数存储的经典方案。它将参数主要是Embedding集中存储在若干PS节点上Worker节点负责计算。这个阶段的优化主要围绕PS本身进行比如开发高性能的PS通信库、优化稀疏参数的拉取合并策略、引入SSD等异构存储来扩展容量。但它的天花板很明显通信瓶颈无法根除PS节点容易成为单点瓶颈系统扩展性受限于PS集群的规模。第二阶段混合并行与模型并行探索。为了缓解PS的压力开始尝试将部分Embedding表“下沉”到Worker本地。例如根据特征的热度访问频率将高频特征放在Worker本地内存极低频特征仍放在远端PS。这相当于在Worker端建立了一层缓存Cache。同时也开始探索纯模型并行即将巨大的Embedding表切分到多个GPU上每个GPU负责一部分特征。但这需要精细的负载均衡和复杂的跨设备通信对算法和工程的要求都很高。第三阶段GPU Embedding Bag与All-to-All通信。这是当前的主流方向也是性能提升的关键。其核心思想是既然通信不可避免那就让它变得高效且规整。具体来说Embedding表完全分布将完整的Embedding表通过某种策略如哈希分区均匀切分到所有训练节点通常是GPU的本地内存中。每个GPU只存储和管理全局表的一部分。规整的All-to-All通信在每一次训练迭代的前向传播中每个GPU根据本batch数据所需的所有特征ID向所有其他GPU发起请求收集自己需要的那些Embedding向量这个过程称为Embedding Lookup。这个收集过程通过高度优化的All-to-All集体通信原语来完成。与PS架构中千军万马挤独木桥式的点对点通信不同All-to-All是GPU间一种规整的、批量化的数据交换能够充分利用高速互联网络如NVLink, InfiniBand的带宽通信效率极高。计算与通信重叠通过CUDA Stream等技术让Embedding Lookup的通信与后续的稠密层计算并行起来进一步隐藏通信延迟。这种方案的本质是将原本网络IO瓶颈的“远程数据访问”问题转化为了一个可利用硬件加速的“高性能计算通信”问题。京东的很多核心业务模型已经基于此架构实现了训练效率的数量级提升。3. 高性能训练的核心技术方案拆解3.1 新一代分布式嵌入训练架构基于上述思路一个现代的高性能稀疏训练架构通常包含以下几个核心组件3.1.1 分层式参数存储并非所有特征都值得用最贵的资源。架构会采用分层的参数存储策略GPU HBM高频带宽内存存放最热门的、访问频率最高的特征Embedding。这是速度最快的一层。CPU内存存放中低频特征的Embedding。通过CPU-GPU的异步拷贝如CUDA Unified Memory或Pinned Memory来参与计算。SSD/持久化内存存放长尾的、极少被访问的冷特征。通过类似缓存换入换出的机制只有当特征在训练中被激活时才将其加载到更快的存储中。 这种分层设计在保证容量的前提下追求整体访问效率的最优化。京东的系统中通常会有一个智能的“特征热度统计与调度器”动态地将特征在各级存储间迁移。3.1.2 基于哈希的分布式嵌入表如何将千亿维度的特征ID映射到相对有限的GPU内存中直接存一个“ID-向量”的字典是不可能的。通用做法是使用哈希分区。假设我们有N个GPU对于一个特征ID通过一个哈希函数hash(id) % N来决定它属于哪个GPU。这个方案简单、负载均衡性好。但存在“哈希冲突”问题两个不同的ID可能被映射到同一个嵌入向量位置。在广告场景中由于特征空间极大而嵌入表相对较小冲突不可避免。业界和学术界普遍认为在规模足够大时适度的哈希冲突对模型最终效果的影响是可控的甚至可以看作一种隐式的正则化。当然也有采用动态哈希表或Cuckoo哈希等更复杂结构来降低冲突概率的方案。3.1.3 高效的All-to-All通信原语这是整个架构的“发动机”。一次Embedding Lookup的通信过程可以分解为本地ID转换与请求打包每个GPU处理本batch的数据得到一组特征ID。根据哈希分区规则它能计算出每个ID应该去哪个目标GPU上获取。然后它将发往同一个目标GPU的ID打包成一个请求包。全局All-to-All通信所有GPU同时将自己打包好的请求包发送给对应的目标GPU。这是一个“全交换”操作。高性能通信库如NCCL对此有极致优化。远程查找与响应目标GPU收到请求包后在自己的本地嵌入表中进行查找取出对应的Embedding向量再打包成响应包。反向All-to-All通信所有GPU将响应包发送回请求方。请求方GPU收到所有需要的向量后组装成本次前向计算所需的完整嵌入输入。这个过程听起来复杂但通过NCCL等库在硬件层面优化后延迟可以做到非常低。关键在于通信的数据量是规整的、批量的而不是零碎的。3.2 训练优化算法与系统协同架构搭好了还需要好的“驾驶技术”才能跑出速度。3.2.1 针对稀疏特征的优化器标准的SGD或Adam优化器在更新Embedding参数时效率不高因为它们需要为每个参数维护动量Momentum等状态这又会消耗大量内存。对于稀疏特征业界广泛采用自适应正则化Adagrad或其变种如FTRL作为优化器。这类优化器为每个参数维护一个累积梯度平方和作为自适应的学习率。其更新是逐元素的element-wise且对于本次训练中未出现的特征其参数和累积状态都不需要更新这天然适合稀疏场景。京东的实践中往往会对稠密部分和稀疏部分采用不同的优化器以达到最佳效果。3.2.2 流水线与计算通信重叠为了进一步压榨硬件性能必须让GPU永远“忙”起来。主要技术是流水线Pipeline数据加载与预处理流水线使用多进程/多线程提前将下一个batch的数据从磁盘加载到内存并进行特征编码、ID化等预处理避免GPU等待数据。计算与通信流水线如前所述利用CUDA Stream让Embedding Lookup的通信发生在Stream A与上一个batch的稠密层反向传播计算发生在Stream B同时进行。当Stream B在计算梯度时Stream A已经在为下一个batch获取Embedding了。梯度更新流水线计算出的稀疏梯度其更新也可以异步进行不阻塞主训练流程。3.2.3 混合精度训练在稀疏训练中混合精度训练Mixed Precision Training同样能带来巨大收益。其主要做法是将Embedding向量、模型权重、激活值等以FP16半精度浮点数格式进行存储和计算这直接减半了内存占用和通信数据量。在计算损失函数和梯度时保留一个FP32的权重副本用于更新以避免梯度下溢带来的精度损失。 对于通信密集型的稀疏训练数据量的减半意味着通信时间近乎减半这是非常可观的性能提升。京东的很多模型训练都已全面启用混合精度。4. 实战构建一个简易高性能稀疏训练流程光讲理论不够我们以一个简化的CTR模型训练为例拆解其关键步骤和实操要点。假设我们使用PyTorch生态并借助一些高性能扩展库。4.1 环境与依赖准备首先你需要一个支持多GPU的环境。深度学习框架选择PyTorch并强烈推荐安装NVIDIA的APEX库用于混合精度训练和Facebook的PyTorch/FBGEMM优化其中包含对稀疏操作的一些优化。对于更接近生产环境的分布式训练可以关注NVIDIA的Merlin HugeCTR或DeepRec等针对推荐系统优化的框架。# 基础环境示例 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install apex -f https://dl.fbaipublicfiles.com/vissl/packaging/apexwheels/py38_cu118_pyt121/download.html4.2 关键代码模块解析4.2.1 分布式嵌入层实现这是核心中的核心。我们不建议从头实现复杂的All-to-All通信可以借助高阶API。以下是一个概念性示例展示了如何使用torch.nn.parallel.DistributedDataParallel(DDP) 配合自定义嵌入层的思想import torch import torch.nn as nn import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP class DistributedEmbeddingBag(nn.Module): 一个简化的分布式嵌入层概念模型。 实际生产环境会使用高度优化的C/CUDA内核。 def __init__(self, num_embeddings, embedding_dim, num_gpus): super().__init__() self.num_gpus num_gpus self.embedding_dim embedding_dim # 假设我们简单地将表哈希分区到各GPU self.embeddings nn.ModuleList([ nn.EmbeddingBag(num_embeddings // num_gpus, embedding_dim, modesum) for _ in range(num_gpus) ]) # 每个EmbeddingBag只负责一部分ID范围 def forward(self, sparse_ids, offsets): sparse_ids: 一个batch中所有样本的特征ID拼接成的长列表。 offsets: 指示每个样本特征ID列表的起始位置。 注意这里省略了复杂的ID路由和All-to-All通信逻辑。 实际实现中需要根据sparse_ids计算出每个ID属于哪个GPU 然后发起通信收集向量最后在本地进行聚合如sum、mean。 # 伪代码实际通信逻辑非常复杂 # 1. 本地计算ID到目标GPU的映射 # 2. 使用 dist.all_to_all_single 或其他集合通信操作交换ID和请求 # 3. 各GPU在本地查找自己负责的ID对应的向量 # 4. 再次通信将查找结果返回给请求方 # 5. 请求方聚合向量得到每个样本的最终嵌入表示 # 此处仅为示意直接返回一个随机值 batch_size len(offsets) - 1 return torch.randn(batch_size, self.embedding_dim) # 初始化分布式进程组 dist.init_process_group(backendnccl) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) # 创建模型 model MyCTRModel(num_sparse_features1e9, embedding_dim128).cuda() model DDP(model, device_ids[local_rank]) # 训练循环中DDP会自动处理梯度同步针对稠密参数 # 稀疏参数的梯度同步需要自定义通常使用上述分布式嵌入层内部的逻辑。注意上述代码是高度简化的概念展示。真实的生产级分布式嵌入层实现极其复杂涉及大量底层通信优化和内核融合。强烈建议基于成熟的工业级框架如Merlin进行开发而非自己从头造轮子。4.2.2 混合精度训练集成使用APEX库可以轻松集成混合精度训练它能自动管理FP16/FP32的转换和损失缩放Loss Scaling。from apex import amp model, optimizer amp.initialize(model, optimizer, opt_levelO2) # opt_levelO2 是常用的优化级别几乎将所有操作转换为FP16同时保持BatchNorm等层为FP32以维持稳定性。 # 在训练循环中 with amp.autocast(): # 自动为前向传播启用混合精度 predictions model(batch_data) loss criterion(predictions, labels) optimizer.zero_grad() # 使用amp进行反向传播它会自动处理损失缩放和梯度转换 scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()4.3 实操心得与避坑指南特征ID哈希冲突的监控至关重要虽然一定程度冲突可接受但必须监控冲突率。可以定期采样数据统计哈希到同一位置的不同特征ID的数量和重要性。如果冲突集中在某些重要特征上可能需要调整哈希函数或扩大嵌入表维度。通信与计算的比例要均衡使用nvprof或PyTorch Profiler工具分析训练过程。理想状态是通信时间被计算时间完全隐藏。如果通信仍然是瓶颈可能需要调整batch size增大batch size通常能让每次通信的数据量更大效率更高或者检查网络硬件和拓扑。嵌入表初始化学习率要区别对待稀疏特征的嵌入向量通常需要更大的初始学习率因为它们在训练初期更新不频繁。在实践中可以为嵌入层和后面的稠密层设置不同的学习率。小心“内存碎片”由于动态特征和哈希GPU显存中嵌入表的访问可能不是连续的这会导致内存碎片影响缓存效率。一些框架提供了“内存整理”的功能需要关注。全量备份与增量更新对于生产系统训练不是一次性的。如何将天级甚至小时级更新的稀疏模型参数高效地同步到线上的推理服务是另一个巨大的挑战。通常会采用“全量模型增量参数”的更新策略。5. 典型问题排查与性能调优实录在实际部署和运行大规模稀疏训练系统时你会遇到各种各样的问题。下面记录几个典型场景和排查思路。5.1 问题一训练速度不稳定时快时慢现象训练迭代时间波动很大有时正常有时突然变长数倍。排查思路检查数据流水线首先怀疑数据加载。查看数据加载进程的CPU/IO使用率是否出现峰值。可能是磁盘IO瓶颈或者数据预处理中有某些样本特别复杂例如某个用户的历史行为序列特别长。检查通信使用NCCL_DEBUGINFO环境变量运行训练观察NCCL通信日志。看是否有某些All-to-All操作耗时异常。可能是网络拥塞或者某个GPU节点响应变慢。检查内存/显存监控GPU显存使用情况。如果显存接近占满会触发昂贵的显存回收和碎片整理操作导致卡顿。可能是batch size过大或某个嵌入表分区不均匀导致单个GPU负载过重。检查系统负载训练服务器上是否混布了其他任务争抢了CPU、内存或网络资源。常见解决措施优化数据存储格式使用TFRecord或更高效的二进制格式并启用多线程预取。调整NCCL通信的缓冲区大小或使用特定的网络拓扑绑定如NCCL_SOCKET_IFNAME指定网卡。实施梯度累积Gradient Accumulation使用更小的物理batch size进行多次前向后向累积梯度后再更新既能维持等效batch size又能降低单次迭代的显存峰值和通信量。5.2 问题二模型效果如AUC下降或不收敛现象切换到新的分布式稀疏训练架构后模型评估指标变差。排查思路验证数据一致性确保分布式训练下每个worker看到的数据是正确且无重复的。检查数据分片Sharding逻辑。检查梯度同步对于自定义的分布式嵌入层确保梯度在所有GPU上正确同步。可以对比小规模数据下单GPU训练和分布式训练几个迭代后的参数值是否一致在允许的误差范围内。检查混合精度训练FP16可能导致梯度下溢特别是对于稀疏特征其梯度本身可能就很小。尝试暂时关闭混合精度opt_levelO0看模型是否收敛。如果问题解决则需要调整损失缩放策略或对嵌入梯度使用动态损失缩放。检查优化器状态像Adagrad这样的优化器其累积状态梯度平方和的初始化很重要。在分布式训练中这些状态也需要正确初始化并同步。监控特征冲突如前所述严重的哈希冲突会破坏特征表示。计算并监控冲突率特别是高频特征之间的冲突。常见解决措施在训练初期使用一个很小的学习率进行“预热”Warm-up让梯度状态稳定下来。对稀疏参数使用更保守的混合精度策略例如保持其优化器状态为FP32。实施更精细的特征哈希策略例如对高频特征进行单独编码或使用双哈希Double Hashing来降低冲突概率。5.3 性能调优检查清单当你的训练系统能跑通后下一步就是让它“跑得快”。下面是一个简单的性能调优检查清单调优维度检查项预期效果与工具数据加载是否启用多进程/多线程数据加载器数据格式是否高效如二进制预处理是否过重减少GPU等待数据时间。使用torch.utils.data.DataLoader的num_workers参数并用torch.profiler分析数据加载耗时。计算图模型前向/反向计算中是否存在大量小算子能否进行算子融合减少内核启动开销。使用PyTorch Profiler的chrome_trace功能可视化计算图寻找优化点。通信All-to-All通信量是否均衡NCCL缓冲区大小是否合适网络带宽是否打满降低通信延迟提高带宽利用率。使用NCCL_DEBUGINFO和nvprof分析通信。内存GPU显存利用率是否健康如80%-90%是否存在内存碎片避免显存溢出和频繁的D2H/H2D拷贝。使用nvidia-smi -l 1监控显存变化。并行策略Batch Size是否最优过小则通信占比高过大则显存可能不足。寻找通信与计算的最佳平衡点。进行简单的缩放实验如128, 256, 512, 1024。框架配置是否使用了最新的CUDA、cuDNN、PyTorch版本它们通常包含性能优化。获得官方的最新性能提升。定期更新基础软件栈。这套从挑战分析、架构演变到实战落地的体系是京东广告算法团队在应对海量数据与极致性能要求下不断打磨出来的工程智慧。它不是一个静态的方案而是一个持续演进的生态系统。随着硬件如新一代GPU、更快的互联技术和软件框架如PyTorch 2.0的编译优化的发展新的优化点又会不断出现。作为从业者理解其背后的核心思想——将稀疏性带来的不规则问题转化为可利用高性能硬件并行处理的规则问题——比记住任何具体的技术细节都更重要。在实际操作中我最大的体会是永远不要忽视监控和 profiling 工具数据驱动的性能分析和问题定位是搞定这类复杂系统的唯一捷径。当你看到训练迭代时间曲线变得平滑稳定AUC指标稳步上升时那种感觉就像一位老船长终于驾驭着他的船平稳地穿过了那片充满暗礁的稀疏之海。