张量并行与流水线并行

所属方向：分布式训练 ｜ 难度：高阶 ｜ 前置知识：数据并行与 ZeRO 显存切分
关键词：张量并行、流水线并行、Megatron-LM、GPipe、通信开销、并行策略

---

01. 为什么需要它

ZeRO-3 能把参数、梯度和优化器状态切到数据并行 rank 上，但每一层计算时仍要聚合这一层的参数。若单个 Transformer 层、词表投影或大矩阵乘本身就放不进一张卡，单纯增加数据并行度解决不了问题；若层能放下，但整网常驻状态或激活放不下，也可考虑沿网络深度切开；能否仅靠 ZeRO-3 和重计算解决，需要先算峰值。

这对应两把不同的刀。张量并行（Tensor Parallelism，TP）在一层内部切矩阵，让多张卡共同完成同一个 GEMM；流水线并行（Pipeline Parallelism，PP）沿层切模型，让不同设备各自保存连续或交错的一段层。前者解决“这一层太宽”，后者解决“模型太深”。数据并行沿 batch 复制模型，三者切的是三个不同维度。

先看一个浪费现场：4 个 pipeline stage 处理 1 个 microbatch。第 0 段前向时其余三段都空闲；数据流到第 3 段后，前几段又空闲。若把一个 global batch 切成 8 个 microbatch，让不同样本在各 stage 重叠，理想 forward sweep 从 $4\times8=32$ 个串行 stage-slot 压到 11 个时间槽，但仍有填充和排空的“气泡”。本文脚本算出 4 段、8 个 microbatch 的理想利用率只有 72.73%，气泡占 27.27%。

再看 TP：若一个 MLP 的第一层权重为 $[d,4d]$，可沿输出维切成 $p$ 份，每卡算 $[d,4d/p]$；第二层 $[4d,d]$ 再沿输入维与中间激活匹配切分。这样两次 GEMM 都不需要在中间把完整 $4d$ 激活拼回每张卡，只在必要边界做集合通信。这正是 Megatron-LM 的 column-parallel 与 row-parallel 配对。

难点不在“切成几份”，而在切后保持数学等价，并让通信落在高带宽拓扑上。TP 频繁同步，通常限制在 NVLink/NVSwitch 节点内；PP 在相邻 stage 前向传激活、反向传对应梯度，更适合跨节点。若反过来布置，同一套卡数会被网络延迟拖垮。

02. 最小可用理解

三句话建立框架：

1. TP 切一层的隐藏维、输出通道或 attention head，每层都会通信，换来单层权重、激活和计算分摊。
2. PP 把层分给不同 stage，再把 batch 切成 microbatch 以重叠各段计算；microbatch 越多，气泡越小，但调度和激活管理更复杂。
3. 生产系统通常组成 DP×TP×PP 的三维网格：TP 放在最快链路内，PP 穿过较慢边界，剩余设备做 DP 与 ZeRO/FSDP。

如果只记一句：TP 是“同一个样本的一层由多卡一起算”，PP 是“同一个样本依次经过多卡上的不同层”。它们都属于模型并行，但通信模式完全不同。

设世界规模 $W$，数据并行度为 $d_p$、张量并行度为 $t_p$、流水线段数为 $p_p$，若没有其他并行轴：

$$W=d_p\,t_p\,p_p$$

这不是越均匀越好。TP 度受矩阵维度、head 数、KV head 数和高速互联限制；PP 度受层数、切分均衡与 microbatch 数限制；最后再用 DP 提升全局吞吐。

03. 数学推导

3.1 Column Parallel：沿输出维切

线性层输入 $X\in\mathbb{R}^{n\times d_{\mathrm{in}}}$，权重 $A\in\mathbb{R}^{d_{\mathrm{in}}\times d_{\mathrm{out}}}$，输出为：

$$Y=XA$$

把 $A$ 沿列切为 $p$ 份：

$$A=[A_0,A_1,\ldots,A_{p-1}],\qquad A_i\in\mathbb{R}^{d_{\mathrm{in}}\times d_{\mathrm{out}}/p}$$

第 $i$ 个 rank 计算：

$$Y_i=XA_i$$

完整结果只是拼接：

$$Y=[Y_0,Y_1,\ldots,Y_{p-1}]$$

前向不必通信，只要输入 $X$ 在 TP 组内复制。若下一层能直接消费分片 $Y_i$，连 all-gather 都可省。反向时，每个 rank 产生一份输入梯度贡献 $dX_i=dY_iA_i^\top$，完整 $dX$ 要求和：

$$dX=\sum_{i=0}^{p-1}dX_i$$

因此 column-parallel 的关键通信落在反向输入梯度 all-reduce。

3.2 Row Parallel：沿输入维切

把输入和权重的匹配维切开：

$$X=[X_0,X_1,\ldots,X_{p-1}],\qquad A=\begin{bmatrix}A_0\\A_1\\ \vdots\\A_{p-1}\end{bmatrix}$$

每卡先算局部部分和 $Z_i=X_iA_i$，完整输出是：

$$Y=XA=\sum_{i=0}^{p-1}X_iA_i=\sum_{i=0}^{p-1}Z_i$$

所以 row-parallel 前向需要 reduce 或 all-reduce，反向则能从复制的 $dY$ 本地得到 $dX_i=dYA_i^\top$。Megatron 把 MLP 的第一层做 column parallel，激活函数逐元素地作用在本地 shard；第二层做 row parallel，最后只求和一次。attention 的 QKV 投影与输出投影也能形成类似配对。

这种设计的精髓不是“每个 Linear 都随便切”，而是让相邻层的分片布局衔接，从而避免在每个算子后都 all-gather。若一个自定义层偷偷需要完整 hidden，通信会突然增加。

3.3 TP 的通信账本

设传输张量有 $N$ 个元素、每元素 $b$ 字节，ring all-reduce 在理想带宽模型中每 rank 通信：

$$V_{\mathrm{AR}}=2\frac{p-1}{p}Nb$$

$p$ 是 TP 组大小。Megatron 风格的一个 Transformer 层在前向与反向关键边界各发生集合通信；真实实现可用 reduce-scatter、all-gather 与 sequence parallel 改写，但总成本仍与 token 数、hidden size 和 TP 频率有关。TP 每层同步，所以延迟不能忽略：小矩阵、多层和大 TP 度会让消息切得太碎，算力利用率反而下降。

TP 对参数显存理想降为 $1/p$，但并非所有东西都跟着切。LayerNorm 参数很小，可能复制；embedding 与 LM head 有词表切分规则；激活是否分片取决于 sequence-parallel 配置；通信临时 buffer 还会增加峰值。因此应从真实 module state 和 profiler 测量，不可把全显存直接除以 TP 度。

3.4 PP 气泡从哪里来

设有 $p$ 个等耗时 stage，global batch 切成 $m$ 个 microbatch，每个 stage 处理一个 microbatch 的单向计算耗时为一格。第一个 microbatch 需要 $p$ 格到达尾部，之后每格完成一个。完成一个 forward sweep 共：

$$T_{\mathrm{sweep}}=m+p-1$$

其中每个 stage 真正工作 $m$ 格，所以理想效率与气泡比例：

$$\eta_{\mathrm{pipe}}=\frac{m}{m+p-1},\qquad f_{\mathrm{bubble}}=\frac{p-1}{m+p-1}$$

当 $p=4,m=8$，效率 $8/11=72.73\%$。要到 90% 以上，需要 $m\ge9(p-1)$。这也解释了为什么 stage 很多却只有几个 microbatch 时 PP 很差。

公式假设各段等耗时、通信完全隐藏、无数据依赖停顿。真实训练的气泡还包括 stage 不均衡、点对点传输、参数同步、数据加载与尾 batch。最后一段若有巨大词表投影，前面切得再平均也会被它卡住。

3.5 GPipe、1F1B 与交错调度

GPipe 先把所有 microbatch 前向跑完，再统一反向，调度简单但要保存更多未反向的激活。1F1B 在 warmup 后交替执行一次 forward 和一次 backward，使稳态内存更接近少量 microbatch；同步更新语义仍要求一个 global batch 的梯度累积完再 step。

交错 1F1B 让每个物理设备承载多个 virtual stage，把一个大气泡拆细。它可改善负载与气泡，却增加更多通信边界、依赖和调度复杂度。Pipeline 并行的优化对象不只是公式里的 bubble，还包括峰值激活、通信重叠和 kernel 连续性。

均衡同步流水线的理想利用率，按 m/(m+p−1) 计算。增加 microbatch 能摊薄填充与排空气泡；图中不含通信、stage 失衡与小 GEMM 效率变化。

04. 代码实现

tp_linear_sim.py 只用标准库，构造 $X:[2,4]$、$W_1:[4,6]$、$W_2:[6,3]$。先完整计算，再把第一层按列、第二层按行切到两个虚拟 rank：

X shape=(2, 4), W1 shape=(4, 6), W2 shape=(6, 3), tp=2
column shards: W1=(4, 3) each, hidden=(2, 3) each
row shards: W2=(3, 3) each, partial output=(2, 3) each
full output=[[52, 42, 72], [132, 114, 180]]
max hidden diff=0, max output diff=0

输出差为 0，验证“按列拼接、按行求和”与完整矩阵乘严格等价。真实浮点并行的归约顺序不同，通常会出现末位误差，不能要求 bitwise 一致，而应设置与 dtype 相称的容差。

pipeline_bubble.py 默认模拟 4 段、8 个 microbatch：

stages=4, microbatches=8, ideal_slots_per_sweep=11
pipeline_efficiency=72.73%, bubble_fraction=27.27%
forward schedule (slot -> active stage:microbatch)
00: S0:M0
01: S0:M1 S1:M0
02: S0:M2 S1:M1 S2:M0
03: S0:M3 S1:M2 S2:M1 S3:M0
04: S0:M4 S1:M3 S2:M2 S3:M1
05: S0:M5 S1:M4 S2:M3 S3:M2
06: S0:M6 S1:M5 S2:M4 S3:M3
07: S0:M7 S1:M6 S2:M5 S3:M4
08: S1:M7 S2:M6 S3:M5
09: S2:M7 S3:M6
10: S3:M7

最前面三格在填充，最后三格在排空，中间四个 stage 才全部忙碌。用参数改变 stages 与 microbatches，就能看到增加 microbatch 如何摊薄固定气泡。

05. 工业级实现对照

以 2026-09 为准，NVIDIA Megatron-LM 的 tensor_parallel/layers.py 仍提供 ColumnParallelLinear 与 RowParallelLinear。工业实现除切权重外，还管理参数初始化、bias、是否 gather 输出、异步梯度 all-reduce、梯度累积融合、sequence parallel、专家并行组和通信 buffer。最小脚本只证明线性代数，没有模拟 autograd 与 NCCL。

ColumnParallelLinear 的 gather_output 决定输出是否立刻汇总；RowParallelLinear 的 input_is_parallel 表示输入是否已经按最后一维切好。错误组合往往不是数值报错，而是多做一次 gather 或得到错位 shard。审计并行模型时，应给每个边界标清 global shape、local shape、分片轴、复制轴和预期 collective。

PP 代码位于 Megatron Core 的 pipeline_parallel 调度模块。配置中的 pipeline_model_parallel_size 切物理 stage，virtual_pipeline_model_parallel_size 启用交错。系统必须传递前向激活和反向梯度，处理 tied embedding、首尾 stage 特殊 loss、不同张量 shape、激活释放与通信重叠。

Megatron 2021 展示了 TP、PP 与 DP 的组合，并提出交错流水调度。工程上的常见布局是：同一节点内组成 TP 组，邻接节点组成 PP 链，跨副本组成 DP 组。原因是 TP 每层通信而 PP 只在 stage 边界传输；把频繁 collective 放在更快互联上更划算。

生产配置还要做 layer partition。按层数平均只在每层耗时相同才合理；MoE 层、cross-attention、视觉模块、embedding 与词表 head 的成本差异很大。应先 profile 单层前后向时间和激活尺寸，再按时间而不是按层数切 stage。

06. 代价与边界

TP 的主要代价是高频 collective。TP 度增加后，每卡 GEMM 变小，计算效率下降，通信延迟占比上升；跨节点 TP 尤其敏感。维度还必须能合理整除：attention head、KV head、MLP intermediate size、词表分片与低精度 tile 对齐都会形成约束。

PP 的主要代价是气泡、激活驻留与调度复杂度。microbatch 增多能降气泡，却让每次 GEMM 的 batch 更小，可能降低算力利用率；还会增加调度次数。梯度累积数也受 global batch、DP 度与 microbatch size 约束：

$$B_{\mathrm{global}}=B_{\mathrm{micro}}\,m\,d_p$$

为了把 $m$ 调大而偷偷改变 global batch，会连学习率与收敛语义一起改变。更稳妥的做法是减小 microbatch size、保持 global batch，再检查小 GEMM 是否仍高效。

PP stage 之间有顺序依赖，单个慢 stage 会拖住全链。设备故障、动态 shape、条件分支和 MoE 路由也比普通 DP 难处理。推理时 batch 与请求长度动态变化，训练得到的均衡切分未必仍均衡。

TP 与 ZeRO/FSDP 并非天然可任意叠加。两者可能切同一参数的不同轴，process group、参数初始化、checkpoint layout 和 optimizer state 都需协调。首先建立二维或三维 rank 映射，再启用框架明确支持的组合，避免自行套两层 wrapper。

什么时候不该用？模型单卡能放下且吞吐受数据不足或 CPU 限制时，TP/PP 只会增加同步；单层很小而层很多时，优先 PP 或 FSDP；单层巨大而层数不多时，TP 更直接；超长序列导致激活爆炸时，还要引入下一篇的序列/上下文并行。

6.1 从模型形状反推 TP 度

假设 decoder hidden size 为 $d$，MLP expansion 为 $4d$，attention head 数为 $H$，每头维度为 $d_h=d/H$。采用 $t_p$ 路 TP 时，至少希望 $H/t_p$ 与 $4d/t_p$ 为整数，且本地矩阵维度满足 Tensor Core tile 对齐。GQA 还要检查 KV head 数 $H_{KV}$；若 Q head 可整除而 KV head 不可整除，框架可能复制 K/V 或根本拒绝配置。

参数容量只是下界。以 MLP 两个权重为例，忽略 bias：

$$P_{\mathrm{MLP}}=d(4d)+(4d)d=8d^2$$

理想 TP 后每卡为 $8d^2/t_p$ 个参数。然而输入 $X$、残差、LayerNorm 与某些输出仍可能复制。是否启用 sequence parallel，会决定 $[B,L,d]$ 激活是每 TP rank 一份还是沿 token 切开。评估 TP 度时要分别列参数、可分片激活、复制激活和通信 buffer，不能只用总显存除 $t_p$。

性能上，本地 GEMM 的算术强度会随 $t_p$ 增大而下降。若 $d/t_p$ 太小，矩阵乘无法占满 SM，即使通信为零也会变慢。实际选型一般从“能放下的最小 TP 度”起步，再尝试少数相邻值；不是卡越多 TP 越大。

6.2 三维 rank 映射为什么决定速度

有 2 个节点、每节点 8 卡，总计 16 卡。假设 TP=8、PP=2、DP=1，合理映射是每个节点内部构成一个 TP 组，两个节点作为相邻 PP stage。这样每层 TP collective 走 NVLink/NVSwitch，只有 stage 边界激活跨节点。

若 rank 编号错误，使每个 TP 组横跨两节点，那么每个 Transformer 层的 all-reduce 都经过网络；PP 反而在节点内。数学结果完全正确，吞吐却可能大幅下降。这类问题从配置数字看不出来，必须导出每个 process group 的物理 GPU、主机名、PCIe/NVLink 路径和 NIC 亲和性。

当 DP>1 时，DP 组应从相同 TP/PP 坐标上取不同副本。例如把 rank 写成坐标 $(r_d,r_p,r_t)$，模型同一 shard 的梯度只在 $r_d$ 维同步；层内 tensor collective 只改变 $r_t$；pipeline 点对点只沿 $r_p$ 邻接。把坐标语义明确下来，checkpoint 分片和故障定位才不会依赖偶然 rank 编号。

多 NIC 节点还要关注通信并发。TP all-reduce、PP send/recv、DP reduce-scatter 可能同时争同一链路。单独 benchmark 每个 collective 很快，不代表组合后仍快；应在真实 schedule 上看每条 stream 和 NIC 的时间线。

6.3 Pipeline 内存不是简单除以段数

PP 将参数按层分段，参数内存近似降为 $1/p_p$，但激活取决于 schedule。GPipe 前向完所有 $m$ 个 microbatch 才反向，每 stage 可能保留 $O(m)$ 份边界与层内激活。1F1B warmup 后尽早反向，可显著减少同时存活的 microbatch 数；不同 stage 的 warmup 长度又不同，峰值不完全一致。

activation checkpoint 将层内保存换成反向重算，但 stage 边界张量通常仍要保留或重新通信。若 PP 与 ZeRO-3 叠加，重算还可能再次触发参数 all-gather；如果 prefetch 与 release 时机不匹配，理论显存节省会被临时完整参数覆盖。

一个更实用的峰值模型是：

$$M_s=M_{\mathrm{params},s}+n_{\mathrm{live},s}M_{\mathrm{act/micro},s}+M_{\mathrm{comm},s}+M_{\mathrm{workspace},s}$$

$s$ 是 stage，$n_{\mathrm{live},s}$ 是该调度下同时存活的 microbatch 数。选切分点的目标应是最小化 $\max_s M_s$ 并平衡每段时间，而不是让每段层数相同。首段 embedding、尾段 vocab projection 和 loss 往往需要单独计量。

6.4 Schedule 的一致性与权重版本

同步 GPipe 或 1F1B 中，一个 global batch 的所有 microbatch 应使用同一版参数，梯度累积后再统一 optimizer step。因此它与非流水同步训练在数学上等价，只改变操作顺序。随机 dropout 若要严格比较，还需保证不同 schedule 消耗一致的随机数流。

异步 pipeline 为减少 flush 可能允许某些 microbatch 使用旧权重，产生 weight staleness。PipeDream 的权重暂存与调度就是为管理这个问题。吞吐更高不等于优化轨迹相同；对需要可复现或大规模预训练的任务，同步 schedule 通常更容易验收。

梯度累积的 loss normalization 也常出错。若每个 microbatch loss 已取 mean，再把 $m$ 份梯度直接相加，最终梯度比全局 mean 大 $m$ 倍；框架可能在 loss、backward 或 optimizer wrapper 某处除 $m$。迁移实现时必须确认缩放发生在哪里，特别是最后一个不完整 microbatch。

6.5 如何系统调优而不是枚举所有组合

第一步做容量约束：根据参数、优化器、激活和临时 buffer，排除会 OOM 的 TP/PP。第二步做整除约束：检查层数、head、KV head、MLP width、词表与 microbatch。第三步按拓扑放组：TP 留在最快域，PP 穿过节点，DP 使用余下副本。第四步才短跑候选方案。

短跑至少记录每个 stage 的 forward/backward 时间、TP collective、PP send/recv、DP collective、bubble、MFU 和峰值显存。若 stage 时间方差大，先重切层；若 collective 裸露，检查异步与 bucket；若 GEMM 利用率低，减少 TP 或增大 microbatch；若 bubble 大，增加 microbatch 或 virtual stage。

还要避免只优化稳态。训练中 checkpoint、评估、数据切换、动态 loss scale 与长短样本混合会改变节拍。视频模型的序列长度可能随分辨率和帧数变化，固定层切分在不同 batch 上会失衡。可以按长度分桶、限制每批 token 数，或使用能处理动态 shape 的 schedule，但都要重新验证 global batch 语义。

6.6 Checkpoint 与并行度变更

TP checkpoint 的一个权重可能沿行或列分片，PP checkpoint 又只在拥有该层的 stage 上出现。保存时应记录全局 shape、分片轴、offset、replica group 与 tied-weight 关系。仅按 rank 存文件却没有布局元数据，会显著增加从 TP=8 改成 TP=4 的恢复难度；若完整掌握原切分约定仍可转换，但必须验证。

成熟 distributed checkpoint 会把逻辑参数名与物理 shard 解耦，加载时重新规划切片。验证转换不能只看“文件读完”，应抽样 all-gather 后与原始全参数比较，并跑一个确定性 forward。optimizer state 也必须按相同参数布局重分片，特别是 Adam 的 m、v 与 FP32 主权重。

PP 的 tied embedding 是典型边界：输入 embedding 在首段，输出投影在尾段，但两者可能共享参数。系统要么在两个 stage 间同步梯度，要么采用明确的复制与更新协议。忽略它会让模型能跑、loss 也下降，却不再是原架构。

最后为每个组合保留机器可读配置：world size、各轴度数、rank 坐标映射、global/micro batch、累积数、schedule、virtual stage、precision 与 checkpoint schema。没有这份清单，性能回归时很难判断是代码变化还是并行布局变化。

6.7 一个具体的 64 卡选型例子

假设 8 个节点、每节点 8 卡，模型单层在 4 卡上能放下，96 层在单节点放不下，目标 global batch 又允许 4 个数据副本。可先试 TP=4、PP=4、DP=4，乘积正好 64。每个节点容纳两个 TP 组，相邻两个节点组成一条四段 pipeline；相同 TP/PP 坐标跨四个副本形成 DP 组。

为什么不直接 TP=8、PP=2、DP=4？它减少 PP 气泡，但本地 GEMM 变小、每层 collective 参与卡数翻倍。为什么不 TP=2、PP=8、DP=4？单层可能放不下，而且 PP 段更多，要求更多 microbatch 才能摊薄气泡。三个方案都满足乘积约束，只有容量、profile 和拓扑能决定赢家。

设 PP=4、microbatch 数 $m=16$，理想气泡为 $3/19=15.79\%$；若为配合 DP=4，global batch 满足 $B_{\mathrm{global}}=B_{\mathrm{micro}}\times16\times4$。当目标 global batch 固定时，microbatch size 可能被压得太小。此时交错 PP 可在不继续增加 $m$ 的情况下缩小气泡，但要多传 stage 边界。

选定后再检查每卡峰值。若尾段词表 head 使 stage 3 明显更慢，可把少量 Transformer 层从尾段移到前段；若 stage 0 embedding 占显存而计算很少，切分目标应同时满足容量与时间，而不是追求参数量绝对相等。

6.8 故障现象与定位路径

若所有 GPU 利用率呈周期性锯齿且空白集中在迭代首尾，优先看 PP bubble；若每层 GEMM 后都有长 NCCL 条带，优先看 TP 通信或 rank 跨节点；若只有某一个 stage 长期满载、其他 stage 等待，是层切分失衡；若迭代末尾集中等待，是 DP 梯度同步没有充分 overlap。

出现 hang 时，先核对各 rank collective 调用序列。条件分支导致某些 rank 少调用一次 all-reduce，或 PP 两端 send/recv shape 不一致，都会永久等待。为通信操作记录 group、sequence number、peer、shape 与 dtype，比只看 Python stack 更有用。设置超时只能让错误更快暴露，不能修复顺序。

数值不一致时，从一个微型模型开始，关闭 dropout 和 fused kernel，分别测试 TP、PP，再测试组合。TP 常见错误是切分轴、bias 重复相加、输出误 gather；PP 常见错误是 loss 缩放、跨段激活 requires-grad、共享参数同步。一次只启用一根并行轴，能把搜索空间从三维降为一维。

OOM 若只发生在第一步，往往是 optimizer state 首次建立或通信 bucket 惰性分配；若在若干步后发生，可能是 graph retention、动态 shape 缓存或碎片；若只在某个 PP stage，先看该段真实 activation 与临时 buffer。不要看到 OOM 就统一减 microbatch，这可能掩盖泄漏却降低所有卡效率。

6.9 上线验收清单

发布配置前确认：世界规模等于各并行轴乘积；每个 rank 坐标唯一；TP group 位于预期高速域；head、KV head、MLP 与词表可整除；各 stage 时间和峰值接近；global batch 算术正确；loss normalization 不依赖 microbatch 数；tied weight 有同步协议；checkpoint 能跨目标并行度恢复。

性能门禁同时保存 tokens/s、MFU、bubble、每类 collective 的裸露时间、峰值显存和最慢 stage。只保存总迭代时间无法解释回归。集群拓扑或通信库升级后，即使模型代码未变，也要重跑基准，因为并行策略本质上是对物理系统的映射。

最后预留降级方案：某节点高速链路异常时能否减少 TP、增加 PP；某个 checkpoint 是否可重分片；global batch 是否仍保持；恢复后数值是否连续。真正稳健的并行配置，不只是峰值最快，也要能在硬件变化和断点恢复时保持语义清楚。

6.10 理论通信量为什么不等于训练时间

同样传 1 GiB，一个大 all-reduce 与数百个小 collective 的耗时不同。常用模型 $T=\alpha n+\beta V$ 中，$n$ 是消息次数，$V$ 是字节量，$\alpha$ 表示每次启动延迟，$\beta$ 表示每字节时间。TP 在每层频繁通信，尤其受 $\alpha$ 影响；PP 消息次数少，但单次边界激活可能很大。

通信是否裸露还取决于依赖。梯度 all-reduce 可与更早层 backward 重叠，PP send 可与下一份本地计算重叠；位于关键路径上的 collective 即使字节少也会直接延长 step。性能报告应区分总 NCCL 时间与 exposed communication time。

bucket 太小会增加消息次数，太大又推迟通信启动，减少 overlap。最优 bucket 与层大小、网络延迟和反向节奏有关。框架默认值是通用折中，不一定适合视频 DiT 的大激活或 MoE 的不均匀参数。

最后还要观察尾延迟。一个 rank 因热降频、ECC、数据抖动或网络拥塞变慢，collective 会让整个组等待。平均 GPU 时间看似健康，最慢 rank 才决定训练。按 rank 记录分位数，并把异常节点与拓扑关联，是大规模并行调优的基本动作。

理论模型最终要由 trace 校准。为每次实验保存并行配置、网络拓扑、逐 stage 时间线与通信矩阵，下一次才能定位回归。若只留一行 tokens/s，既无法判断瓶颈是 GEMM、气泡还是慢 rank，也无法安全迁移到另一代 GPU。并行策略不是模型的附属启动参数，而是计算图和硬件共同组成的一部分，理应像模型结构一样接受版本管理和回归测试。

07. 经典论文脉络

• GPipe（arXiv:1811.06965）：用 microbatch 流水化跨加速器的模型分段，并以同步 mini-batch 语义训练巨型网络。
• PipeDream（arXiv:1806.03377）：探索 1F1B 与异步流水，揭示吞吐、权重版本和一致性之间的权衡。
• Megatron-LM（arXiv:1909.08053）：提出 Transformer 内高效的 tensor model parallel 切法，用少量 collective 支撑数十亿参数。
• Efficient Large-Scale Language Model Training（arXiv:2104.04473）：系统组合 TP、PP、DP，并以交错流水减少气泡，扩展到千卡与万亿参数。

论文演进说明：单一并行轴只能解决一种容量瓶颈，规模继续增长后，真正的问题变成如何把多个轴映射到硬件拓扑并共同调度。

08. 常见误解

误解一：TP 就是把每层平均切开。 切分轴决定通信。Megatron 的 column-row 配对是为了让中间分片直接衔接，不是机械地切一半权重。

误解二：PP 段数越多，显存越省且速度越快。 参数容量会下降，但固定气泡随 $p-1$ 增长；若 microbatch 不够，更多 stage 反而更闲。

误解三：microbatch 越多越好。 它降低理想气泡，却减小 GEMM、增加调度，并可能扩大激活队列。要联动测 MFU 与显存。

误解四：1F1B 是异步优化。 常用同步 1F1B 只是改变前后向顺序，仍在一个 global batch 梯度完成后统一更新；PipeDream 式异步才涉及权重陈旧。

误解五：TP 能把所有显存除以 TP 度。 小参数、复制激活、通信 buffer 与 runtime 不会理想均分。必须看逐项账本。

误解六：三维并行度乘起来等于卡数就配置正确。 整除只是必要条件。拓扑、矩阵 tile、head 数、stage 均衡、global batch 都可能让配置不可用。

09. 动手验证

先运行 tp_linear_sim.py，把 TP 从 2 改成能整除维度的 3，重新切 $W_1$ 的输出与 $W_2$ 的输入。预期完整输出仍一致。然后故意让第二层 shard 顺序交换，结果会静默错误；这说明 distributed tensor layout 必须携带明确语义。

运行 pipeline_bubble.py，固定 stages=8，分别取 microbatches=1、8、32、72。理论效率依次为 $1/8$、$8/15$、$32/39$、$72/79$。脚本当前只实现均衡 stage 的闭式公式；扩展成事件调度模拟后，再给每个 stage 设置不同时间，以最慢 stage 为节拍比较，观察均分层数为何不等于均分时间。

在真实集群做 TP=1/2/4/8 对照，固定 global batch 和模型，记录每层 GEMM 时间、collective 时间、MFU、峰值显存。预期早期 TP 能解除容量或提高总吞吐，超过某点后小 GEMM 和通信使收益反转。用拓扑工具确认 TP rank 是否真的落在 NVLink 域内。

对 PP 做相同实验：固定 PP 度，逐步增加 microbatch 数；再固定 microbatch，增加 virtual stage。记录 bubble、激活峰值与端到端 tokens/s。不要只看 framework 打印的理论 bubble，GPU timeline 中的空白才是真实气泡。

最后做一次等价性验收：单卡、TP、PP、TP+PP 使用相同初始化和样本，关闭 dropout，比较一个 step 的 loss 和若干参数梯度。低精度与归约顺序会造成小差异，但若误差随层爆炸，优先检查 shard 顺序、bias、loss 归一化与 tied weight。

10. 延伸阅读

读完这篇可以继续看：

• 数据并行与 ZeRO 显存切分：DP 沿 batch 复制计算、ZeRO 沿 DP 组切状态，是三维并行的第一根轴（已发布）。
• 序列并行与 Ring Attention：当瓶颈从模型宽度转向超长 token，沿 sequence 切激活与 attention（本文系列下一篇）。
• 混合精度与数值稳定性：低精度改变 TP/PP 通信字节与归约误差，配置并行前要先明确累加精度（本文系列上一篇）。

选型时先问容量瓶颈在哪里：单层太宽选 TP，层总数太深选 PP，副本吞吐选 DP，超长激活选 SP。随后才是把这些轴映射到实际互联。公式给出上限，profile 决定最终配置。

附录：完整代码

09 节用到的脚本全文如下（tp_linear_sim.py、pipeline_bubble.py、make_figures.py）。复制到本地存成同名文件，按各脚本开头的依赖说明准备环境后即可运行。

tp_linear_sim.py

#!/usr/bin/env python3
"""Verify column- and row-parallel linear algebra without any framework."""


def matmul(a, b):
    return [[sum(x * y for x, y in zip(row, col)) for col in zip(*b)] for row in a]


def split_columns(a, parts):
    width = len(a[0]) // parts
    return [[row[i * width:(i + 1) * width] for row in a] for i in range(parts)]


def split_rows(a, parts):
    height = len(a) // parts
    return [a[i * height:(i + 1) * height] for i in range(parts)]


def add(a, b):
    return [[x + y for x, y in zip(ra, rb)] for ra, rb in zip(a, b)]


x = [[1, 2, 3, 4], [5, 6, 7, 8]]       # [tokens=2, hidden=4]
w1 = [[1, 0, 2, 0, 3, 0], [0, 1, 0, 2, 0, 3],
      [1, 1, 1, 1, 1, 1], [2, 1, 0, 1, 2, 1]]  # [4, 6]
w2 = [[1, 0, 1], [0, 1, 1], [1, 1, 0],
      [2, 0, 1], [0, 2, 1], [1, 0, 2]]          # [6, 3]

full_hidden = matmul(x, w1)
full_output = matmul(full_hidden, w2)

# Column parallel W1: every rank computes a slice of output features.
w1_shards = split_columns(w1, 2)
hidden_shards = [matmul(x, shard) for shard in w1_shards]
column_joined = [left + right for left, right in zip(*hidden_shards)]

# Row parallel W2: input features and W2 rows use matching shards, then sum.
w2_shards = split_rows(w2, 2)
partials = [matmul(h, w) for h, w in zip(hidden_shards, w2_shards)]
row_reduced = add(partials[0], partials[1])

max_hidden_diff = max(abs(a - b) for ra, rb in zip(full_hidden, column_joined) for a, b in zip(ra, rb))
max_output_diff = max(abs(a - b) for ra, rb in zip(full_output, row_reduced) for a, b in zip(ra, rb))
print("X shape=(2, 4), W1 shape=(4, 6), W2 shape=(6, 3), tp=2")
print("column shards: W1=(4, 3) each, hidden=(2, 3) each")
print("row shards: W2=(3, 3) each, partial output=(2, 3) each")
print(f"full output={full_output}")
print(f"max hidden diff={max_hidden_diff}, max output diff={max_output_diff}")

pipeline_bubble.py

#!/usr/bin/env python3
"""Idealized GPipe bubble calculator; stdlib only."""

import argparse


parser = argparse.ArgumentParser()
parser.add_argument("--stages", type=int, default=4)
parser.add_argument("--microbatches", type=int, default=8)
args = parser.parse_args()
if args.stages < 1 or args.microbatches < 1:
    parser.error("stages and microbatches must be positive")

p, m = args.stages, args.microbatches
slots = m + p - 1
efficiency = m / slots
print(f"stages={p}, microbatches={m}, ideal_slots_per_sweep={slots}")
print(f"pipeline_efficiency={efficiency:.2%}, bubble_fraction={1-efficiency:.2%}")
print("forward schedule (slot -> active stage:microbatch)")
for t in range(slots):
    active = [f"S{s}:M{t-s}" for s in range(p) if 0 <= t - s < m]
    print(f"{t:02d}: " + " ".join(active))

make_figures.py

"""Regenerate this article's deterministic teaching figure (numpy + matplotlib)."""
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
OUT=Path(__file__).resolve().parents[1]/"figures"
OUT.mkdir(exist_ok=True)
plt.rcParams.update({"font.sans-serif":["PingFang SC","Arial Unicode MS","DejaVu Sans"],"axes.unicode_minus":False})
m=np.arange(1,129)
fig,ax=plt.subplots(figsize=(8,4.8))
for p in [2,4,8,16]:ax.plot(m,m/(m+p-1),label=f"{p} stages")
ax.set(xlabel="Microbatches per global batch",ylabel="Ideal utilization",ylim=(0,1.02),title="Balanced synchronous pipeline: m / (m + p - 1)")
ax.legend();ax.grid(alpha=.25);fig.tight_layout()
fig.savefig(OUT/'pipeline_utilization.png',dpi=170)
plt.close(fig)
print(OUT/'pipeline_utilization.png')

---

---

更多 AIGC 论文解读，关注微信公众号「人工智能炼丹君」

每日更新 · 论文精选 · 深度解读 · 技术脉络

微信搜索 人工智能炼丹君 或扫描下方二维码关注