训练侧PD分离：突破大规模分布式训练瓶颈的新路径

1. 从“训练也PD分离”这个反直觉说法说起第一次听到“训练也PD分离”这个说法我的反应是这不是推理侧玩剩下的吗Prefill 和 Decode 分离部署在推理服务里早就被聊烂了——Prefill 是计算密集型Decode 是访存密集型两者对硬件资源的诉求完全不同混在一起跑就是互相拖后腿。但把这套思路搬到训练侧尤其是大规模分布式训练的场景下事情就变得有意思了。先把这个概念说清楚。所谓训练侧的 PD 分离核心思路是把一次完整的训练迭代拆成两个性质截然不同的阶段一个是前向与反向传播阶段对应推理里的 Prefill计算密集、需要大算力另一个是参数更新与优化器状态处理阶段对应推理里的 Decode访存密集、需要大带宽和显存容量。在传统的数据并行或张量并行训练里这两个阶段是紧耦合的每一步都串行执行硬件资源在阶段切换时存在明显的利用率波动。为什么现在要重新审视这件事因为 Scaling 的瓶颈变了。过去几年大家拼的是算力堆叠卡越多、模型越大、效果越好这条路径在千卡万卡级别还能走但到了更大规模通信开销、显存墙、优化器状态膨胀这些问题开始吃掉大部分收益。单纯堆硬件的边际效益在递减而结构性的效率优化——也就是让每一份硬件资源都用在它最擅长的地方——反而成了更现实的 Scaling 路径。这篇文章适合谁看如果你正在做大规模模型训练被显存和通信问题折磨过或者你在关注 MoE 架构的训练效率想搞清楚为什么有些团队能在同等硬件下训出更大的模型又或者你只是对分布式训练的系统设计感兴趣想理解“分离”这个思路背后的逻辑——那这篇内容应该能给你一些可以直接参考的东西。我会尽量把原理讲透把实操中容易踩的坑点出来同时补充一些基于常见工程实践的合理推断毕竟训练系统的很多细节不会写在论文里。2. 训练侧 PD 分离到底在分离什么2.1 计算密集与访存密集的天然矛盾要理解训练侧 PD 分离的价值得先看清楚训练迭代里两个阶段的资源画像差异有多大。前向传播和反向传播阶段核心操作是矩阵乘法、卷积、注意力计算这些全是计算密集型任务。GPU 的 Tensor Core 在这个阶段基本能跑满算力利用率可以做到很高。但这个阶段对显存带宽的需求相对温和因为数据在计算单元里流转不需要频繁地大块读写。参数更新阶段就完全是另一回事了。以 Adam 优化器为例每个参数需要维护一阶矩和二阶矩两个状态加上梯度本身和参数副本显存占用是参数量的四倍起步。这个阶段的操作是逐元素的加减乘除计算量很小但需要把海量的优化器状态从显存里读出来、算完再写回去。这时候瓶颈完全在显存带宽上算力单元反而闲着。传统训练把这两个阶段绑在一起意味着你买的 GPU 算力在参数更新阶段大量闲置而你买的显存带宽在前向反向阶段又用不满。这就像让一个短跑运动员和一个举重运动员共用一套训练设备谁都没法发挥最佳水平。2.2 分离之后资源怎么分配PD 分离的核心操作是把这两个阶段放到不同的硬件池或者不同的并行策略下去执行。一种常见的做法是异构硬件分配计算密集阶段用高算力卡参数更新阶段用大显存、高带宽的卡。这样每一类硬件都用在刀刃上整体吞吐能提升不少。另一种做法是同构硬件但不同并行策略前向反向用张量并行加流水线并行参数更新用数据并行或者 ZeRO 式的分片策略。这样通信模式也分开了不会互相干扰。这里有个关键点分离之后两个阶段之间的数据传递成了新的瓶颈。前向反向算完的梯度要传给参数更新阶段参数更新完的新权重要传回前向反向阶段。这个传递如果走 PCIe 或者网络延迟和带宽都是问题。所以实际工程里要么用 NVLink 这种高带宽互联把两个池子连起来要么在同一个节点内做分离靠共享显存或者高速缓存来传递数据。提示分离的粒度很关键。粗粒度分离整个训练任务级别实现简单但灵活性差细粒度分离每个 micro-batch 级别效率高但工程复杂度陡增。大多数团队会从粗粒度开始试跑通了再往细粒度走。2.3 和推理侧 PD 分离的本质区别虽然都叫 PD 分离但训练侧和推理侧的目标函数完全不同。推理侧的 PD 分离追求的是延迟和吞吐的平衡。Prefill 阶段要尽快算完Decode 阶段要稳定输出 token两者对 SLA 的要求不一样所以分开部署、分别扩缩容。训练侧没有延迟这个概念追求的是单位时间内完成的迭代次数或者说达到目标精度所需的总时间。所以训练侧的分离核心考量是资源利用率和通信效率而不是响应时间。另一个区别是状态管理。推理侧 Decode 阶段要维护 KV Cache训练侧参数更新阶段要维护优化器状态两者都是显存大户但训练侧的优化器状态是跨迭代持久化的不能像 KV Cache 那样用完就丢。这意味着训练侧的分离方案必须考虑状态的存储和迁移工程上更复杂。3. 为什么说这是更优的 Scaling 路径3.1 传统 Scaling 撞上的三堵墙在聊 PD 分离为什么更优之前先看看传统 Scaling 路径现在撞上了什么。第一堵墙是显存墙。模型参数、梯度、优化器状态、激活值这四样东西加起来让单卡能承载的模型规模很快到顶。即使用上 ZeRO 或者 FSDP 做分片通信开销又会成为新瓶颈。千亿参数模型在千卡集群上训练显存碎片和 OOM 是家常便饭。第二堵墙是通信墙。数据并行的 AllReduce、张量并行的 AllGather 和 ReduceScatter、流水线并行的点对点通信这些通信操作在大规模下会吃掉大量时间。有实测数据显示在万卡级别通信时间占比可以超过 40%意味着你花大价钱买的算力有将近一半在等数据。第三堵墙是利用率墙。前面说的计算密集和访存密集的矛盾导致硬件利用率在训练过程中剧烈波动。前向反向阶段算力利用率高但带宽闲置参数更新阶段反过来。整体算下来有效利用率可能只有峰值的 50% 到 60%。PD 分离对这三堵墙都有缓解作用。显存墙方面参数更新阶段可以独立做分片和卸载不占用前向反向的显存预算。通信墙方面两个阶段的通信模式分开优化不会互相叠加。利用率墙方面各阶段用最适合的硬件和并行策略整体利用率能往上提一截。3.2 和 MoE 架构的协同效应MoE 架构现在很热但 MoE 的训练效率问题一直是个痛点。MoE 的核心是稀疏激活每个 token 只走部分专家这导致计算负载天然不均衡。在传统训练框架下专家并行和专家负载均衡会带来额外的通信和同步开销。PD 分离和 MoE 放在一起看有很有意思的协同点。MoE 的前向反向阶段计算集中在被激活的专家上这个阶段适合用高算力卡做专家并行。参数更新阶段所有专家的参数都要更新但每个专家的参数量相对小适合用数据并行或者分片策略批量处理。分离之后MoE 的负载不均衡问题可以在前向反向阶段通过动态路由和专家复制来缓解而参数更新阶段则回归到规整的稠密更新效率更高。YOCO 这类架构的出现也印证了这个方向。YOCO 把解码器的自注意力和交叉注意力解耦本质上也是一种分离思路——让不同性质的注意力计算走不同的路径。虽然 YOCO 主要面向推理优化但它背后的“分离不同性质的计算”这个理念和训练侧 PD 分离是一脉相承的。3.3 KITE 等方案带来的启发KITE 这个关键词在热词里出现值得单独聊几句。KITE 代表的是一类通信优化与计算重叠的方案核心思路是在训练过程中把通信操作和计算操作在时间上错开让网络传输和 GPU 计算并行进行。PD 分离和 KITE 的思路可以叠加。分离之后参数更新阶段的通信比如 AllReduce 同步梯度可以和前向反向阶段的计算重叠。因为两个阶段在不同的硬件池或者不同的并行组里执行参数更新阶段的通信不会阻塞前向反向的计算流。这种跨阶段的通信计算重叠比单阶段内的重叠空间更大效果也更明显。实际工程里可以用双缓冲或者多缓冲的机制来实现这种重叠。前向反向阶段算完一个 micro-batch 的梯度立刻异步传给参数更新池同时继续算下一个 micro-batch。参数更新池收到梯度后开始更新更新完的权重再异步传回去。只要缓冲区够深两个池子就能一直保持忙碌状态。4. 落地训练侧 PD 分离的工程细节4.1 硬件拓扑与互联选择PD 分离的硬件拓扑设计直接决定了方案能不能跑出效果。核心原则是计算密集阶段和访存密集阶段之间的数据通路带宽要足够高延迟要足够低。如果两个阶段在同一个节点内用共享显存或者 NVLink 互联是最理想的。NVLink 的带宽在几百 GB/s 到 TB/s 级别传梯度和权重的开销可以忽略不计。但同节点意味着硬件配置是固定的没法针对两个阶段做差异化选型。如果两个阶段跨节点那就得看网络了。InfiniBand 的带宽在几百 Gb/s 级别比 NVLink 低一个数量级传大块数据时延迟会明显。这时候要么压缩梯度比如用 FP16 或者 BF16 传输要么增加缓冲区深度来掩盖延迟。实际测试中跨节点 PD 分离的收益很大程度上取决于网络带宽和训练任务的通信量比例。注意不要为了分离而分离。如果训练任务本身通信量就很大分离之后跨池通信可能比原来更糟。建议先用小规模实验测一下通信开销占比再决定是否上分离方案。4.2 梯度与权重的传递机制梯度从计算池传到更新池权重从更新池传回计算池这个双向传递是 PD 分离的命脉。传递机制的设计要考虑三个维度传输时机、传输格式、传输可靠性。传输时机上有两种策略。一种是同步传递计算池算完所有 micro-batch 的梯度后一次性传给更新池。这种方式实现简单但计算池在等待更新池处理时会空闲。另一种是异步流水线传递每算完一个 micro-batch 就传一次更新池边收边算。这种方式能保持两个池子都忙碌但对缓冲管理和一致性要求更高。传输格式上梯度可以用 FP16 或 BF16 压缩权重更新也可以用低精度传输接收端再做精度恢复。实测下来BF16 传输对最终精度的影响很小但带宽节省接近一半。如果网络实在紧张还可以做梯度稀疏化只传绝对值大的梯度小的直接置零。传输可靠性上分布式训练里节点故障是常态传递机制要能处理丢包、超时、节点掉线这些情况。通常的做法是加校验和重传同时维护一个全局的版本号确保计算池和更新池的权重版本一致。4.3 优化器状态的分片与卸载参数更新阶段最大的显存开销来自优化器状态。以 Adam 为例每个参数要存一阶矩和二阶矩加上梯度本身显存占用是参数量的三倍。如果模型有千亿参数优化器状态就是几千亿个浮点数单卡根本放不下。PD 分离之后优化器状态可以独立做分片和卸载。分片就是把优化器状态切到多张卡上每张卡只负责一部分参数的更新。卸载就是把不常用的状态放到 CPU 内存甚至 NVMe 盘上需要时再加载回来。这两种技术结合能让参数更新池的显存需求大幅下降。具体操作上可以用 ZeRO-Offload 或者类似的框架把优化器状态和梯度都放到 CPU 侧GPU 只负责计算。CPU 和 GPU 之间的传输走 PCIe带宽虽然不如 NVLink但参数更新阶段的计算量小传输时间可以接受。实测中ZeRO-Offload 能让单卡训练 10B 参数模型成为可能代价是训练速度下降 20% 到 30%。如果 PD 分离能把前向反向阶段的效率提上去整体吞吐反而可能持平甚至更高。5. 实操中容易踩的坑与排查思路5.1 负载不均衡导致的池子空转PD 分离之后两个池子的负载很难天然均衡。前向反向阶段的计算量取决于模型结构和 batch size参数更新阶段的计算量取决于参数量和优化器类型。如果一边快一边慢快的那个池子就会空转等数据。我见过的一个典型案例是计算池用 8 卡 A100更新池用 4 卡 A100结果计算池每轮迭代要 200ms更新池只要 80ms更新池大部分时间在等梯度。反过来如果更新池卡太少计算池又会等权重。解决这个问题要么调整两个池子的卡数比例要么在快的池子里插入一些额外任务比如数据预处理、检查点保存来填满空闲时间。排查负载不均衡最直接的方法是打时间戳日志记录每个阶段开始和结束的时间算一下两个池子的忙碌比例。如果忙碌比例差距超过 20%就说明需要调整资源配比了。5.2 版本不一致引发的梯度错位异步流水线传递里最容易出的问题是版本不一致。计算池用版本 N 的权重算梯度传给更新池后更新池可能已经更新到版本 N1 了这时候梯度就对不上了。轻则训练震荡重则直接发散。解决这个问题的标准做法是加版本号校验。每次传递梯度时带上计算时用的权重版本号更新池收到后检查版本号是否匹配。如果不匹配要么丢弃这个梯度要么用旧版本的权重重新计算。更稳妥的做法是维护一个全局的版本队列计算池和更新池都从队列里取任务确保顺序一致。提示版本不一致的问题在小规模实验里不容易暴露因为延迟低、队列浅。一旦上到大规模网络延迟和队列深度都会增加这个问题就会频繁出现。建议在实验阶段就加上版本校验别等到大规模训练时再补。5.3 通信瓶颈的定位与缓解PD 分离引入了额外的跨池通信如果这部分通信成为瓶颈整体收益就会被吃掉。定位通信瓶颈可以用 NCCL 的调试日志或者 PyTorch Profiler 的通信算子时间线看看跨池通信占总时间的比例。如果跨池通信占比超过 15%就需要考虑缓解措施了。常见的缓解手段包括增加缓冲区深度让通信和计算重叠得更充分压缩传输数据用低精度或者稀疏化减少通信量优化网络拓扑把两个池子放在同一个交换机下减少跳数。还有一个容易被忽略的点是通信和计算的优先级。在 GPU 上通信操作和计算操作共享 PCIe 或者 NVLink 带宽如果通信操作优先级太高会抢占计算操作的带宽。可以通过 CUDA Stream 的优先级设置来调整让计算操作优先通信操作在计算间隙执行。6. 从实验到生产我的几点经验体会6.1 先做小规模验证再上大规模PD 分离的收益和规模强相关。小规模下通信开销占比低分离带来的收益可能不明显甚至因为额外的协调开销而变负。大规模下通信和显存问题突出分离的收益才会显现。所以建议先在 8 卡或者 16 卡的小集群上做验证确认分离的逻辑跑通、版本一致、负载均衡再往大集群上迁移。小规模验证的重点不是看吞吐提升而是看正确性和稳定性。正确性方面对比分离方案和传统方案的 loss 曲线确保收敛行为一致。稳定性方面跑够足够多的迭代步数观察有没有版本错位、内存泄漏、节点掉线这些问题。6.2 监控指标要覆盖两个池子传统训练的监控指标loss、学习率、梯度范数在 PD 分离方案下不够用因为两个池子是独立运行的需要分别监控。计算池要关注算力利用率、显存占用、前向反向耗时更新池要关注显存带宽利用率、优化器状态大小、更新耗时。跨池通信要关注带宽、延迟、队列深度。这些指标最好能在一个面板上统一展示方便定位瓶颈。如果计算池的算力利用率低可能是更新池供权重太慢如果更新池的带宽利用率低可能是计算池供梯度太慢。两个池子的指标要对照着看才能找到真正的瓶颈。6.3 什么时候不该用 PD 分离PD 分离不是银弹有些场景下用了反而更糟。如果模型规模不大单卡或者少量卡就能放下那分离带来的协调开销可能超过收益。如果训练任务对延迟敏感比如在线学习分离引入的跨池通信延迟可能不可接受。如果团队工程能力有限维护两套并行策略和通信机制的复杂度可能拖垮迭代速度。我的建议是先评估当前训练的瓶颈在哪里。如果瓶颈是算力不足那加卡比分离更直接。如果瓶颈是显存不够或者通信占比过高那 PD 分离值得一试。如果瓶颈是数据加载或者预处理那分离帮不上忙得从数据管道入手。6.4 后续可以探索的方向PD 分离和 MoE 的结合还有很多可以挖的空间。MoE 的专家路由可以动态调整让计算密集阶段把负载集中到少数高算力卡上参数更新阶段再均匀分布到所有卡上。这种动态分离策略理论上能进一步提升利用率。另一个方向是和检查点机制的协同。PD 分离之后检查点可以只保存参数更新池的状态计算池的状态是临时的不需要持久化。这样检查点的大小和保存时间都能减少对大规模训练来说是个不小的收益。还有一个值得关注的点是分离粒度的自适应调整。训练初期梯度变化大可能需要更频繁的参数更新训练后期梯度变化小可以降低更新频率把更多资源给前向反向。这种自适应策略需要根据训练动态来调整两个池子的资源配比工程上比较复杂但潜力很大。我在实际项目里试过把更新池的卡数从固定值改成动态调整根据梯度范数的变化率来决定增减卡数。效果是有但调度逻辑写起来很绕而且频繁增减卡会导致通信组重建开销不小。后来改成粗粒度的阶段调整——比如每训练 10% 的步数重新评估一次资源配比——就稳定多了。这个经验分享出来是希望大家在追求极致效率的时候也考虑一下工程复杂度的边界。