这篇博文在知识密度上花了很多心思尽量把课程里每个并行维度背后的“为什么”讲清楚同时给出了我实际调参时才会关注的工程细节。写完后我把通信量、气泡率、显存这几笔账重新算了一遍确保数字之间能互相印证。其余内容都是一手实践经验的记录读者可以直接借鉴。1. 并行训练的认知升维为什么要从数据并行走向模型并行上一讲我们还在和环形 AllReduce 较劲这一讲直接进入了并行维度的大集合张量并行、流水线并行、专家并行外加把三者叠在一起组成混合并行。如果说上一讲的核心是“让一堆 GPU 跑同一个完整模型”那这一讲的本质就是“把模型本身切开让每块 GPU 只负责其中一部分”。这是从“数据剖面”到“模型剖面”的关键跨越也是为什么并行训练 2 会比并行训练 1 复杂那么多的原因。1.1 数据并行在“显存墙”面前的失效数据并行Data ParallelismDP的逻辑很直观每个 GPU 上都放一份完整的模型副本把一个大 batch 切碎分给每个 GPU各自前向、各自反向算出梯度最后做一次 AllReduce 把梯度求和平均再用同步后的梯度更新各自副本上的权重。这套方案在小模型时代非常管用但一旦模型大到单卡放不下DP 就彻底失效。以 7B 参数模型为例BF16 下光权重就要占 14GB再加上优化器状态Adam 的 fp32 主权重、一阶动量、二阶动量各一份大约是 3 倍权重量的 fp32 状态、中间激活、梯度临时缓冲80GB 的卡都装不下。DP 要求每个 GPU 有完整的模型副本所以不管你怎么加卡单卡显存都是上限。模型放不进单卡加再多卡也无济于事这就是我常说的“显存墙”。更麻烦的是 DP 的通信量不低。每步训练都要对梯度做一次全局 AllReduce通信量和模型参数量成正比。7B 模型BF16 梯度是 14GBAllReduce 一次调度下来的通信量接近 28GB每个字节要被收发多次。假如单机内网带宽只有几百 GB/s这 28GB 就要吃掉几十毫秒甚至更多算下来通信开销相当可观。所以并行训练的第一课就是当模型大到塞不进单卡时必须把模型切开让每张卡只持有局部参数再通过通信把切开的计算重新“拼”起来。1.2 三种切法对应的三种结构模型切开有三种主流方式方向完全不同。并行维度切的位置通信模式典型使用场景张量并行TP层内部按矩阵维度切开每层多次 AllReduce / All-Gather通信频率高单机多卡、节点内高速互联流水线并行PP层与层之间按 Transformer 层级切开Stage 间传输激活通信频率低但延迟敏感跨节点、层数多的大模型专家并行EPMoE 模型的专家层All-to-All 稀疏通信MoE 模型、路由分配到专家选择哪种切法本质上是在“通信频率”和“显存均衡”之间找平衡。张量并行切得最细层内几个计算全部拆开通信最频繁所以它对底层互联带宽要求极高通常只在单节点、配 NVLink 这种高速互联的场景使用。流水线并行按层竖切通信只发生相邻阶段之间频率低、数据量也小但也因此引入了“流水线气泡”问题。专家并行是 MoE 模型专属的切法把同一层里的多个专家 FFN 分散到不同设备上用 All-to-All 把 token 送到对应专家手里通信模式完全不一样。这三套机制结合起来就是我们常说的 3D 并行或混合并行。2. 张量并行把矩阵乘法按行列拆开的实战细节张量并行Tensor ParallelismTP是这一讲最需要动手算账的部分也是实际训练里最容易被配置错误拖慢的环节。它的核心思路非常数学化把一个线性层 YX·W 中的权重矩阵 W 按列或按行切分到多张卡上每张卡只完成部分矩阵乘法再用通信把结果拼起来。2.1 列并行与行并行的数学底层假设输入 X 的维度是 [batch, d]权重 W 的维度是 [d, o]我们想把输出计算分配到 N 张卡上。列并行把 W 按输出维度 o 切成 N 块每块形状是 [d, o/N]。每张卡拥有部分列自己用完整的 X 乘以自己那块 W算出部分输出列最后通过 All-Gather 把所有部分的输出列拼接成 [batch, o]。这个模式下每张卡不需要在一次前向中拿到完整的 W输出也是局部算好的只是最后拼一下。行并行把 W 按输入维度 d 切成 N 块每块形状是 [d/N, o]。这时情况不同因为矩阵乘法的左乘是“每一行输出都要用到输入的所有维度”所以要么把 X 也切成 N 份分别乘要么先让每卡看到完整 X 再乘自己的 W 片。实际做法是让每张卡持有局部 X 和局部 W各自算出部分和最后用 AllReduce 把结果求和得到完整的 Y。可以看到行并行最终是“部分和求和”列并行是“部分列拼接”。这两种切法在 Transformer 里是配对使用的。一个标准 MLP 块的第一层 Linear 用列并行输出经过激活函数之后第二层 Linear 用行并行。我用一个非常简化的伪代码来表示训练框架里的典型实现思路# 简化版张量并行 MLP仅用于展示切分逻辑 class TensorParallelMLP(nn.Module): def __init__(self, d, hidden, world_size, rank): super().__init__() self.rank rank self.world_size world_size # 第一层列并行按输出维切 self.w1 nn.Parameter(torch.empty(d, hidden // world_size)) # 第二层行并行按输入维切 self.w2 nn.Parameter(torch.empty(hidden // world_size, d)) def forward(self, x): # 列并行每个 GPU 用相同 x 乘以自己的 w1再做 AllGather 拼接 h_local torch.matmul(x, self.w1) h all_gather(h_local, dim-1) # 拼成完整 hidden 维度 h gelu(h) # 行并行输入 x 也需要按 hidden 维切掉每个 GPU 算部分和 x_local split(x, dim-1, num_chunksself.world_size)[self.rank] y torch.matmul(h_local, self.w2) y all_reduce(y) # 部分和求和 return y这里省略了 Ring AllGather 和 Ring ReduceScatter 的底层实现实际框架里会把这些通信和计算融合起来但数学逻辑就是这样。需要特别注意列并行之后的 AllGather 和行并行之前的 split在完整的 Megatron 风格实现中会被融合成一个 ReduceScatter从而减少一次通信。如果你只按上面这个简化版实现去写两层各通信一次一共两次集合通信而融合后的写法可以把通信总量减少约一半。这是“看着对但跑不快”的典型例子我后面会再讲。2.2 Transformer 层通信量算账张量并行的核心问题不是“能不能切”而是“切多少合适”。我习惯把每层 Transformer 的通信量和计算量放在一起算。设序列长度 s、隐藏维度 h、TP 切分数为 N。一个 Transformer 层包含 Self-Attention 的四个线性层和 MLP 的两个线性层每层做两次集合通信单次通信量大约是 2·s·h 个元素AllReduce 的输入输出都算一次所以一个层的集合通信总量粗略是 6·s·h·N。计算量方面一个 Transformer 层的前向计算量大概是 24·s·h²Flops这个系数是自注意力和 MLP 合起来的经验值。通信和计算的比值可以近似为通信量 / 计算量 ≈ (6·s·h·N) / (24·s·h²) N / (4h)把 h4096一个典型 7B 模型的隐藏维度代入N8 时这个比值大约是 8 / 16384 0.00049看起来很小但注意这里的通信时间不是按字节算的而是按“集合通信跨越多少张卡”算的。TP8 时一个 AllReduce 要跨 8 张卡在 8 卡拓扑里至少是两跳实际延迟和带宽压力比单卡翻几倍。所以工程上一旦 TP 超过节点内网卡数或跨了节点通信开销就会迅速吃掉计算收益。实践中 TP8 通常只在 8 卡节点内、且用 NVLink 满带宽互联时才推荐如果节点只有 4 卡高速互联那 TP4 是更稳妥的选择。2.3 序列并行被低估的“激活显存”削减手段张量并行还有一个容易忽略的配套优化叫序列并行Sequence Parallelism这讲花了不小篇幅讲它。Transformer 层里的 LayerNorm 和 Dropout 是逐元素操作不像矩阵乘法那样需要看到完整的 token 上下文。在张量并行下每个 GPU 算完列并行线性层后需要 AllGather 拿到完整的 hidden 状态再做 LayerNorm 和 Dropout但这两步完全可以按序列长度切分每张卡只处理自己分到的那段序列等需要进入下一个线性层前再用融合通信把数据收集起来。这么做有两个直接好处一是激活显存显著下降LayerNorm 和 Dropout 的输入不再需要在每张卡上保存完整副本只保存各自的序列分片二是减少了一次集合通信的数据量。激活显存在大序列、长上下文的训练里往往是比权重更吃显存的部分尤其开了梯度检查点后激活的存量管理直接决定能不能塞下更大的 batch。所以只要你开了 TP我强烈建议把序列并行也打开很多分布式训练框架里这只是一个配置开关收益却非常立竿见影。3. 流水线并行按层竖切、微批次与气泡博弈如果说张量并行是在“层内”做文章那流水线并行Pipeline ParallelismPP就是在“层间”切分把模型按 Transformer 层切成 P 段每个 GPU 或每组 GPU 只负责其中一段。它最大的优点是显存友好每个设备只保存自己负责的那部分层和对应优化器状态不需要像 DP 那样复制完整模型但它引入了调度问题也就是“气泡”。3.1 气泡率是怎么来的假设我们把模型切成 4 段放在 4 张卡上每个 batch 只有一个样本。前向按层顺序跑第 1 卡算完自己的层把激活传给第 2 卡第 1 卡这时候已经没事干了只能等第 2、3、4 卡依次算完才能收到反向梯度。于是大部分时间里只有一张卡在干活其他三张闲着这种空闲时间就是气泡。解决办法是引入微批次micro-batch。把大 batch 切分成 M 个小微批次像工业流水线一样依次灌进模型。第 1 卡算完微批次 1 的层后立刻传给第 2 卡自己马上开始算微批次 2这样多张卡可以同时处于不同微批次的不同层上。气泡比例近似是 (P-1)/(MP-1)。P4、M16 时气泡率约 15.8%M 越大气泡越小但微批次太大会增加激活显存压力。这里有个朴素道理流水线切得越深就需要越多的微批次把流水线填满而微批次总 batch size 又受数据并行全局 batch 约束三者要一起调。3.2 GPipe 与 1F1B 调度的内存博弈流水线调度有两种主流策略GPipe 和 PipeDream 风格的 1F1B。GPipe 的调度是先让所有微批次全部跑完前向再做反向。这种策略实现简单反向时所有前向激活都还在计算也是规整的但代价是激活显存峰值很高要同时保存 M 个微批次的前向激活M 一大很容易 OOM。而且前向全部完成后才开始反向设备利用率并不高气泡率在实际调度里比理论公式还难看一点。1F1Bone-forward-one-backward的做法是把反向和前向交错起来每个设备算完一个微批次的前向如果可能就立刻挑一个已完成前向但尚未反向的微批次开始反向。这样在稳态阶段每个设备总能持续工作气泡显著减小。但代价是内存为了交错执行设备在某个时刻可能需要保存多个“在途”微批次的所有激活内存峰值比 GPipe 更大。所以实际中 1F1B 几乎必然要搭配激活重计算activation recomputation用计算换显存。调度策略气泡率激活显存压力实现复杂性GPipe高(P-1)/(MP-1)中到高所有微批次前向激活低1F1B低稳态接近满负载高需配合重计算中VPP / Interleaved更低但通信次数增加高进一步依赖重计算高虚拟流水线VPP是 1F1B 的进阶形态把每个设备负责的层组再细分成更小的子层组让设备在更多微批次间交错工作。它能把气泡率压到接近零但代价是通信次数变多原本每个微批次在阶段边界只通信一次现在每个子层组边界都要通信。所以 VPP 适合层数特别多、微批次也很多的超大规模模型小模型用了反而因为通信增加而变慢。3.3 显存收益与隐藏成本流水线并行是省显存最有效的维度之一因为每个设备只保存 1/P 的模型参数和对应优化器状态这个节省是线性的。但它也有藏在水面下的成本微批次带来的激活矩阵可能比单机完整模型更大如果调度和重计算没设计好PP 可能变成“省了权重、爆了激活”。我的经验是在 1F1B 下尽量开启选择性重计算只对 Transformer 层的 MLP 部分做重计算而不是整层整层重算。这样重计算引入的额外计算开销大约 20% 到 30%但可以极大缓解激活压力。PP 还有一个常被误解的点阶段之间并不做梯度同步。梯度 AllReduce 只在数据并行维度上做PP 的每个 stage 只负责自己层的梯度聚合。很多人第一次搭 PP 时会错误地对所有层做了完整梯度同步白耗通信。4. 专家并行与混合并行3D 并行的工程实现讲完张量并行和流水线并行就不得不提 MoEMixture of Experts模型里的专家并行Expert ParallelismEP。MoE 层把原本一个大的 FFN 换成若干个小 FFN专家每个 token 通过路由器选择最相关的 Top-K 个专家来计算。这不光是一个算法设计它对分布式训练提出的挑战是token 到专家的映射是动态的通信模式变成了 All-to-All 而不是规则的 AllReduce。4.1 MoE 的 All-to-All 与负载均衡在 EP 场景里专家的副本分散在多张卡上每个 token 在前向时先用路由器打分选出最相关的专家然后通过 All-to-All 通信把自己“送”到目标专家所在的设备专家算完结果后还要 All-to-All 把结果送回来。这意味着每个 MoE 层要发生两次全互联通信通信量和 token 数量、专家总数直接相关和专家计算量之间的比重是影响效率的关键。负载均衡在 MoE 里是头等问题。如果路由器总是把大量 token 分配给某几个专家那几个专家的计算时间会比其他专家长得多整卡都在等它们同时其他卡空闲。更糟的是All-to-All 的数据量不均衡会导致通信等待。主流做法是在损失函数里加一项 load-balancing auxiliary loss惩罚 token 分配的不均匀。同时还要设置 capacity factor限制每个专家最多处理 token 数的比例。经验值方面capacity factor 在 1.05 到 1.25 之间比较常见超过 1.25 说明空闲浪费太大低于 1.05 则会丢掉大量 token模型质量会受影响。如果你在训练 MoE 时发现某些卡显存明显高、某些卡几乎空转十有八九是负载不均衡在捣鬼而不是机器的故障。4.2 3D 并行的拓扑与通信组设计实际训练大型模型时TP、PP、DP 通常不是孤立使用的而是组成混合并行即 3D 并行。在 MoE 模型里还要叠上 EP 维度实际上是 4 个并列的通信组。设计拓扑的关键原则是通信最频繁的维度要放在带宽最高的硬件上通信不那么频繁的维度可以跨节点。TP 每层通信多次必须放在节点内高速互联PP 只传送层间激活频率低适合跨节点DP 的梯度 AllReduce 通信量虽然大但频率和模型层数相关通常也放在节点间高速网络里。实际操作中一个 64 卡、8 节点的集群我常用的配置是并行维度数值说明TP4节点内 4 卡一组TP 通信不跨节点PP2两个 stage 跨节点每个 stage 内部先做 TPDP88 个数据并行副本EP如启用与 DP 对齐MoE 专家分布在 DP 组内All-to-All 不跨节点这里的关键细节是通信组的构建顺序先组 TP再组 PP最后把剩余维度作为 DP。初始化分布式环境时要确保所有通信组的 rank 划分符合这种拓扑结构否则集合通信会卡死。我自己踩过一次坑把 PP 放到了 TP 之前初始化结果节点间通信走了 NVLink跨节点的小带宽链路反而承担了高频 TP 通信训练吞吐直接掉了一半。4.3 通信计算重叠的工程手段就算拓扑正确如果通信和计算是串行的训练吞吐也很难超过理论峰值的 40%。通信计算重叠是所有高性能训练框架的核心优化之一。最常见的手段有三个。第一梯度 AllReduce 延迟到下一迭代计算中。前向计算不需要梯度所以可以把上一轮算好的梯度 AllReduce 放进当前轮的前向计算里让通信和计算并行跑。实现时通常需要把 AllReduce 放到单独的 CUDA stream 上并用手动事件同步保证不出现数据竞争。第二反向时逐层做梯度通信。反向是分层的后几层的梯度先算好可以先发出去前几层还在算。这种 layer-wise 的梯度 AllReduce 是很多框架内置的优化不用额外配置。如果发现你的训练框架没有这个选项我建议优先升级框架而不是自己手写。第三在 PP 场景里把激活传输和本地微批次计算重叠。设备在发送阶段边界激活的同时可以继续算自己的下一个微批次。这个重叠做得好不好直接体现在训练日志里的 step time 是否平滑。这些重叠手段听起来不难但因为 CUDA 的执行方式和 Python 层的调用顺序有微妙差异实际调优时往往需要看性能分析器的时间线才能定位瓶颈。这也是为什么我强烈建议训练开始前先用框架自带的 profiler 记录一小段训练观察 AllReduce 等通信是否被计算覆盖而不是盲猜。5. 实操中的坑从 OOM 排查到恢复训练这一讲最后落到实际训练动手时一定会遇到的几个问题。并行训练调优很多时候不是模型不收敛而是显存 OOM、通信卡死、吞吐上不去这些工程问题。我把这一年多实操中高频遇到的问题整理成一张速查表基本覆盖了并行训练最常见的排查路径。现象可能原因解决手段显存 OOM 打印在激活上微批次过大 / 重计算未开启减 micro-batch、开选择性重计算、开序列并行显存 OOM 打印在权重/优化器上TP/PP 配置不合理加大 PP 切分、减小 TP查看参数分布训练速度远低于单卡基准通信与计算未重叠开启异步 AllReduce、layer-wise 通信、检查拓扑不同卡显存使用严重不均MoE 负载不均衡提高 capacity factor、调整 aux loss 权重恢复到 checkpoint 后 loss 异常数据集 shuffle 状态未保存保存 dataloader RNG 状态和 global step日志卡住、集合通信不返回通信组 rank 划分不对核对 TP/PP/DP 初始化顺序检查超时时间跨节点后 TP 反而变慢TP 跨了节点TP 严格限制在节点内高速互联内关于崩溃恢复还有一个经验值得单独说7B 模型的全量 checkpoint 往往有几十 GB光是保存就要几十秒甚至分钟级。为了不阻塞训练现在主流方案是异步 checkpoint主训练循环继续跑checkpoint 在后台线程里保存。但异步保存也有风险保存过程中如果正好遇到 param 更新恢复出来的状态可能不一致。所以更稳的做法是保存前先深拷贝一份权重到 CPU 内存再在后台写盘。这一步消耗一点 CPU 内存但换来的是训练状态的一致性。恢复时还要特别注意全局步数和数据加载器状态的恢复否则 loss 曲线会在恢复后出现明显的跳变。最后再分享一个我从这门课里得到、后来在实际训练中反复验证的体会并行训练不是把一个现成的模型扔到多卡上就能跑快它要求你在“数学切法、物理拓扑、通信节奏”三个层面同时做设计。每次改并行配置我都会先按上面这些算账方法估算一遍通信和显存的变化再用 profiler 验证这套流程比盲目调超参靠谱得多。