起因#
上一篇把 ZeRO 的三个 Stage 走了一遍之后，我以为关于参数分片这件事已经想得差不多了——ZeRO-3 / FSDP 把参数、梯度和 Optimizer States 全部分片存储到多张 GPU 上，单卡 Model States 从 112 GB 降到约 28 GB，问题解决。
但想了一会儿发现有个地方不太对。
ZeRO-3 的 forward 发生了什么？每张 GPU 先通过 AllGather 把当前层的完整参数临时收集回来，然后用完整的权重矩阵对自己的 mini-batch 做完整的矩阵乘法。换句话说，参数被分片存储了，但计算量并没有被分片执行。 7B 模型 4 卡 ZeRO-3，单卡 Model States 确实降了，但每个 Linear Layer 做 \(Y = XW\) 时，那个完整的矩阵乘法仍然在每张 GPU 上各做一遍——只是不同 GPU 算的是不同 mini-batch。
继续看 MIT 6.S191 Lecture 9 关于 Tensor Parallelism 的部分，发现 TP 解决的是另一个层面的问题：不止分片参数的存储，而是把单个算子的计算拆开，让多张 GPU 分担。每张 GPU 只持有权重矩阵的一个分片，只负责输出的一部分，局部结果合起来数学上等价于单卡的完整计算。
一个让我觉得需要分清楚的对比：ZeRO-3 的典型数据并行组里，各 GPU 处理不同样本，用临时收集来的完整参数各自独立计算。而 Column TP 中，多张 GPU 处理同一批输入数据，但计算不同的输出特征。
两者不是简单的升级替代关系——FSDP 和 TP 解决不同的瓶颈，通信代价的形式也不一样，实际训练中往往组合使用。
这篇想把 TP 的两种基本拆分方式以及它们的 backward 通信需求搞清楚，最后看看两层连续 Linear 之间的通信是怎么被优化掉的。
一个 Linear，两种拆法#
全文统一使用下面的记号。一个 Linear Layer（不带 bias）：
$$ Y = XW $$$$ X \in \mathbb{R}^{B \times 4}, \quad W \in \mathbb{R}^{4 \times 8} $$使用两张 GPU。
先提一下记号的一个容易混淆的地方。上面的 \(Y = XW\) 用的是数学中常见的写法，\(W\) 的 shape 是 \((\text{in\_features}, \text{out\_features})\)，也就是 \((4, 8)\)。但 PyTorch 的 nn.Linear 实际存储的 weight 是转置的：shape 是 \((\text{out\_features}, \text{in\_features}) = (8, 4)\)，计算时执行的是 \(Y = X \cdot \text{weight}^T\)。全文按数学记号讨论。了解这一点就不会在看框架源码时被搞糊涂。
Column Parallel#
把 \(W\) 沿列方向（输出特征维度）切开：
$$ W = [W_0 \mid W_1], \quad W_0, W_1 \in \mathbb{R}^{4 \times 4} $$两张 GPU 都持有完整的输入 \(X\)，分别计算：
$$ Y_0 = X W_0 \in \mathbb{R}^{B \times 4}, \quad Y_1 = X W_1 \in \mathbb{R}^{B \times 4} $$完整输出逻辑上是两个分片拼接的结果：
$$ Y = [Y_0 \mid Y_1] \in \mathbb{R}^{B \times 8} $$几点需要记住：
- \(X\) 在两张 GPU 上是复制的，数值相同。
- 不同 GPU 负责的是不同的输出列——GPU 0 计算前 4 列，GPU 1 计算后 4 列，彼此不重叠。
- 所以合并方式是拼接，不是求和。
- 完整 \(Y\) 是逻辑拼接的结果，不一定需要立刻执行 AllGather 通信。是否需要物理拼接，取决于下一个算子需要什么样的输入布局。
写成一般形式：\(X [W_0 \mid W_1] = [X W_0 \mid X W_1]\)。矩阵乘法对列维度的线性性，让按列拆分后的结果可以直接拼接还原。
Row Parallel#
换一种拆法。这次沿 \(W\) 的行方向（输入特征维度）切开：
$$ W = \begin{bmatrix} W_0 \\ W_1 \end{bmatrix}, \quad W_0, W_1 \in \mathbb{R}^{2 \times 8} $$既然 \(W\) 的行被拆了，输入 \(X\) 也需要沿对应的列维度分片：
$$ X = [X_0 \mid X_1], \quad X_0, X_1 \in \mathbb{R}^{B \times 2} $$两张 GPU 分别计算：
$$ Y_0 = X_0 W_0 \in \mathbb{R}^{B \times 8}, \quad Y_1 = X_1 W_1 \in \mathbb{R}^{B \times 8} $$\(Y_0\) 和 \(Y_1\) 的 shape 都是 \((B, 8)\)——和最终输出一样大。但它们不是完整输出的不同列。
把某个输出元素展开就能看清楚。完整计算中，\(Y\) 的第 \(j\) 列是输入向量和 \(W\) 第 \(j\) 列的内积。GPU 0 用 \(X\) 的前两个特征和 \(W\) 的前两行算出了内积的一部分，GPU 1 用后两个特征和后两行算出了另一部分。两个 GPU 算的是同一个输出元素的不同部分和，要得到完整值，必须逐元素求和：
$$ Y = X_0 W_0 + X_1 W_1 $$一般形式：\([X_0 \mid X_1] \begin{bmatrix} W_0 \\ W_1 \end{bmatrix} = X_0 W_0 + X_1 W_1\)。如果每张 GPU 需要完整的 \(Y\)，可以使用 AllReduce(SUM)——但和 Column Parallel 一样，是否需要立刻通信取决于下游算子的布局需求。
对比一下：
| Column Parallel | Row Parallel | |
|---|---|---|
| 权重拆分 | 沿输出特征（列）：\(W = [W_0 \mid W_1]\) | 沿输入特征（行）：\(W = [W_0; W_1]\) |
| 输入要求 | 每卡持有完整 \(X\) | 输入也需分片：\(X = [X_0 \mid X_1]\) |
| 每卡输出含义 | 完整输出的一个列分片 | 完整输出的部分和 |
| 合并方式（数学） | 拼接 | 求和 |
Backward：为什么有些梯度需要求和#
Forward 的拆分理清之后，Column 和 Row 其实挺直观——无非是矩阵乘法对不同维度的可分解性。但 backward 的通信分析并不平凡，让我停下来想了比较久。
两种梯度#
先回顾一下 \(Y = XW\) 的 backward 需要算什么。假设已经从上游拿到了损失函数对 \(Y\) 的梯度 \(dY\)，backward 需要计算：
$$ dW = X^T dY \quad \text{（Weight Gradient）} $$$$ dX = dY \, W^T \quad \text{（Input Gradient）} $$- Weight Gradient \(dW\)：用于更新当前层的权重。
- Input Gradient \(dX\)：把梯度继续向前传给上一层。
这里的"输入"指当前算子 \(Y = XW\) 的输入 \(X\)，不特指模型的原始输入。如果 \(X\) 前面还有需要梯度的计算层，\(dX\) 就是那一层 backward 所需的上游梯度。如果 \(X\) 前面没有可训练的计算（比如 \(X\) 就是原始数据），\(dX\) 可以省略。
在 TP 中，weight gradient 和 input gradient 的通信需求不同。
Column Parallel 的 backward#
回到 Column Parallel 的 forward：\(Y_0 = XW_0\)，\(Y_1 = XW_1\)。
Backward 时，每张 GPU 从下游拿到自己那一部分的梯度 \(dY_0\) 和 \(dY_1\)。
Weight Gradient。 GPU 0 计算 \(dW_0 = X^T dY_0\)，GPU 1 计算 \(dW_1 = X^T dY_1\)。每张 GPU 本地就有完整的 \(X\)（forward 时复制的）和本地的 \(dY_i\)，所以各自独立计算即可，不需要跨 GPU 通信。
Input Gradient。 这里让我想了比较久。
GPU 0 能算出：
$$ G_X^{(0)} = dY_0 \, W_0^T $$GPU 1 能算出：
$$ G_X^{(1)} = dY_1 \, W_1^T $$这两个结果的 shape 都是 \((B, 4)\)，和 \(X\) 一样大。但它们各自都不是完整的 \(dX\)。完整的 input gradient 是：
$$ dX = G_X^{(0)} + G_X^{(1)} $$我一开始看到这个求和有点困惑：两个 shape 都是 \((B, 4)\) 的东西为什么要加起来？Forward 中 Column Parallel 的输出明明是拼接的。
从多变量链式法则看为什么要求和#
后来从计算图的角度想通了。
\(X\) 在计算图中被两条路径共同使用：它既参与了 \(Y_0 = XW_0\)，也参与了 \(Y_1 = XW_1\)。Loss \(L\) 通过这两条路径都依赖于 \(X\)。
多变量链式法则告诉我们：当一个变量通过多条路径影响最终结果时，总导数是各条路径贡献的加和。
$$ \frac{dL}{dX} = \frac{\partial L}{\partial Y_0} \frac{\partial Y_0}{\partial X} + \frac{\partial L}{\partial Y_1} \frac{\partial Y_1}{\partial X} = dY_0 \, W_0^T + dY_1 \, W_1^T $$用一个标量的例子可以把这件事看得更清楚。假设 \(x\) 是标量，\(y_1 = x \cdot w_1\)，\(y_2 = x \cdot w_2\)，损失函数是 \(L = y_1 \cdot y_2\)（乘法，不是加法）。直接展开：
$$ L = x^2 w_1 w_2, \quad \frac{dL}{dx} = 2x w_1 w_2 $$用链式法则从 \(y_1\) 和 \(y_2\) 分别走：
$$ \frac{dL}{dx} = \frac{\partial L}{\partial y_1} \frac{dy_1}{dx} + \frac{\partial L}{\partial y_2} \frac{dy_2}{dx} = y_2 \cdot w_1 + y_1 \cdot w_2 = (xw_2) w_1 + (xw_1) w_2 = 2x w_1 w_2 $$结果一致。即使 \(y_1\) 和 \(y_2\) 在 forward 中不是通过加法组合的，\(x\) 的梯度仍然是两条路径贡献的加和。 求和来自计算图的分支结构——\(x\) 通过多条路径影响 \(L\)——和 forward 中 \(y_1\)、\(y_2\) 具体的合并方式无关。
回到矩阵的 Column Parallel：\(G_X^{(0)}\) 和 \(G_X^{(1)}\) 是同一个 \(X\) 沿两条不同路径产生的梯度贡献，不是 \(X\) 的不同分片。它们的 shape 都是 \((B, 4)\)，和 \(X\) 完全一样，因为它们度量的是"Loss 通过 \(Y_0\) 路径对 \(X\) 每个元素的偏导"和"通过 \(Y_1\) 路径的偏导"。两者必须相加才是 \(X\) 的完整梯度。
实现上，如果上游需要完整的 \(dX\)，需要一次 AllReduce(SUM) 来聚合两张 GPU 各自计算的贡献。
Row Parallel 的 backward#
Row Parallel 的 forward：\(Y = X_0 W_0 + X_1 W_1\)。
GPU 0 持有 \(X_0, W_0\)，GPU 1 持有 \(X_1, W_1\)。经典布局下 forward 通过 AllReduce 使两张 GPU 都拿到完整的 \(Y\)，因此 backward 时两张 GPU 拿到的 \(dY\) 也相同。
Weight Gradient。 GPU 0 计算 \(dW_0 = X_0^T dY\)，GPU 1 计算 \(dW_1 = X_1^T dY\)。各自本地有 \(X_i\) 和 \(dY\)，独立计算即可。
Input Gradient。
$$ dX_0 = dY \, W_0^T, \quad dX_1 = dY \, W_1^T $$GPU 0 算出的 \(dX_0\) 是 \(X_0\) 的完整梯度，GPU 1 算出的 \(dX_1\) 是 \(X_1\) 的完整梯度。
和 Column Parallel 的关键区别：\(dX_0\) 和 \(dX_1\) 对应的是不同的输入变量。\(X_0\) 和 \(X_1\) 是 \(X\) 沿特征维度的不同分片，它们在计算图中不共享——\(X_0\) 只参与了 \(X_0 W_0\) 这一条路径，\(X_1\) 只参与了 \(X_1 W_1\) 这一条路径。没有多条路径汇合的问题，不需要对它们求和。
如果上游需要完整的 \(dX = [dX_0 \mid dX_1]\)，可以通过 AllGather 收集。如果上游本身也接受分片输入，各 GPU 保持自己的 \(dX_i\) 就够了。
Weight Gradient 不需要 TP 组内通信#
两种方式下，每张 GPU 都可以用本地数据独立计算自己那一片权重的 weight gradient：
- Column：\(dW_i = X^T dY_i\)，GPU 有完整 \(X\) 和本地 \(dY_i\)。
- Row：\(dW_i = X_i^T dY\)，GPU 有本地 \(X_i\) 和完整 \(dY\)。
在当前 TP Group 内部，weight gradient 不需要 AllReduce。 每张 GPU 独立持有自己分片的权重和对应的完整梯度，直接用于本地的参数更新。
如果 TP 与 Data Parallelism 组合——比如节点内 TP 拆分算子，节点间 DP 同步不同 mini-batch 的梯度——那 DP 组内仍然需要额外的梯度同步。但这属于 DP 的通信需求，不是 TP 本身的。
两层之间的通信能不能省掉#
搞清楚单层的 forward 和 backward 之后，自然会想：真实模型中矩阵乘法不是孤立的。如果两层 Linear 前后相连，中间的通信能不能避免？
考虑一个简化的两层 MLP：
$$ X \xrightarrow{\text{Linear}_1} \xrightarrow{\text{ReLU}} \xrightarrow{\text{Linear}_2} Y $$$$ W_1 \in \mathbb{R}^{4 \times 8}, \quad W_2 \in \mathbb{R}^{8 \times 4} $$Forward：Column + Row#
第一层 Column Parallel。把 \(W_1\) 沿列切成两半：
$$ H_i = \text{ReLU}(X W_{1,i}), \quad W_{1,i} \in \mathbb{R}^{4 \times 4}, \quad H_i \in \mathbb{R}^{B \times 4} $$两张 GPU 各自算出 activation 的一个列分片。ReLU 是逐元素算子——每个输出元素只依赖对应位置的输入，不依赖其他位置——所以可以在各 GPU 本地独立执行。
第二层 Row Parallel。Row Parallel 的输入本来就需要沿特征维度分片，而 Column Parallel 的输出恰好就是按特征维度分片的。\(H_0, H_1\) 直接作为第二层的分片输入：
$$ Z_i = H_i W_{2,i}, \quad W_{2,i} \in \mathbb{R}^{4 \times 4}, \quad Z_i \in \mathbb{R}^{B \times 4} $$$$ Y = Z_0 + Z_1 $$最后一步求和通过 AllReduce 完成。
两层之间不需要 AllGather 完整的 \(H\)。 Column 的分片输出直接流入 Row 的分片输入，中间的 ReLU 在各 GPU 本地完成。整个 MLP 的 forward 只需要一次 AllReduce。
为什么不反过来 Row + Column#
如果第一层用 Row Parallel，第二层用 Column Parallel，会怎样？
Row Parallel 的输出是部分和——每张 GPU 拿到的 \(Z_i\) 需要相加才是完整输出。如果不归约直接对 \(Z_i\) 执行 ReLU：
$$ \text{ReLU}(Z_0) + \text{ReLU}(Z_1) \neq \text{ReLU}(Z_0 + Z_1) $$取 \(Z_0 = 1, Z_1 = -2\) 就能看到：\(\text{ReLU}(1) + \text{ReLU}(-2) = 1\)，但 \(\text{ReLU}(-1) = 0\)。
非线性激活函数对加法不满足分配律。 部分和必须先归约为完整输出，才能正确执行 ReLU。这意味着 Row → ReLU 之间会被迫插入一次通信。
Column → ReLU → Row 之所以能省掉中间通信，是因为 Column 的分片输出不需要归约就可以独立通过逐元素算子——各分片之间彼此独立，没有"先合并才能继续算"的约束。
Backward 的通信#
从输出端倒推。
第二层 Row Parallel 的 backward。 两张 GPU 都有完整的 \(dY\)（forward 的 AllReduce 使 \(Y\) 在两卡上一致，所以上游传回来的 \(dY\) 也一致）。
Row Backward 的 input gradient：
$$ dH_i = dY \, W_{2,i}^T $$GPU 0 算出 \(dH_0\)，GPU 1 算出 \(dH_1\)。这是对不同输入变量的梯度——和前面 Row Backward 分析的一样——各自独立，不需要互相求和。
ReLU backward。 ReLU 的梯度是逐元素 mask（输入大于 0 的位置保留，否则归零）。每张 GPU 本地的 \(H_i\) 决定了 mask，本地的 \(dH_i\) 乘以 mask 得到 ReLU 之前的梯度。完全在本地完成。
第一层 Column Parallel 的 backward。
$$ G_X^{(i)} = (\text{ReLU backward 输出}_i) \, W_{1,i}^T $$$$ dX = G_X^{(0)} + G_X^{(1)} $$如果上游需要完整的 \(dX\)，需要一次 AllReduce(SUM)。
把 forward 和 backward 放在一起看这个 MLP 的通信：
- Forward：一次 AllReduce（Row Parallel 的输出求和）。
- Backward：一次 AllReduce（Column Backward 的 input gradient 求和）。
- 中间的 activation 和梯度始终保持分片状态，不需要额外 AllGather。
这是在"输入复制、输出复制"这一特定布局下的通信路径。不同框架的实现可能选择不同的布局约定，通信操作会相应变化，不能机械套用。
关于 Megatron-LM，我觉得 Column 后接 Row 这个组合本身并不是什么难以想到的数学技巧——两层矩阵乘法共享中间维度 \(M\)，Column 沿 \(M\) 拆分输出，Row 沿 \(M\) 拆分输入，只要分片规则一致就能直连。Megatron-LM 的贡献在于把这种布局优化系统化地应用到了 Transformer 的 MLP 和 Multi-Head Attention 上，并在千卡规模验证了可扩展性。
工程权衡#
最后说几个工程上的考虑。
TP degree 增大，单 GPU 上的权重存储和局部矩阵乘法的规模通常减小。但通信开销、同步延迟、小 GEMM 的 kernel 利用率都会受到影响。TP 不保证训练更快。
TP 对 GPU 间的互连带宽比较敏感。TP 的通信发生在算子的 forward 和 backward 内部，粒度比 FSDP 的 AllGather / ReduceScatter 更细，延迟直接影响 step time。这也是为什么 TP 通常部署在节点内的 NVLink 互连上——跨节点走较慢的网络时，收益会被通信延迟吃掉。
如果模型的单层计算能很好地在单卡上完成，FSDP 可能更简单、更合适。TP 的价值更多体现在单层参数量大到单卡放不下、或单卡计算效率低的场景。真实训练中，TP 和 FSDP / DP 往往组合使用：节点内跑 TP，节点间跑 FSDP。
回到这篇的主线。从 FSDP 的"参数分片存储"到 TP 的"计算分片执行"，看起来只多走了一步。但仅仅拆分单个算子的计算还不够。Column + Row MLP 的例子说明了一件更重要的事：设计相邻算子的分片布局，使上一个算子的输出能直接成为下一个算子的输入，减少中间不必要的跨 GPU 数据搬运。
分布式训练的核心不仅是把数据或计算拆成几份，更是设计这些分片在不同算子之间如何流动。
参考资料#
- 从 DDP 到 ZeRO：为什么分布式训练要对参数、梯度和优化器状态分片？
- 从 AllReduce 到 Ring：分布式训练中的通信原语
- 深入理解 PyTorch DDP：为什么更少的通信，不一定意味着更快的训练？
- 7B 模型只有 14 GB 权重，一张 80 GB GPU 为什么仍可能 OOM？
- MIT 6.S191: Introduction to Deep Learning — Lecture 9: Secrets to Massively Parallel Training. https://introtodeeplearning.com/
- Shoeybi et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2019.