v4：Tensor Parallel 与 Sequence Parallel
ZeRO 切的只是“存”，算的时候每卡用的还是完整的 weight，activation 也没切，模型到 70B 以上就不够了。这一课把每一层的矩阵乘法本身切成 t 份，t 张卡合算同一层（TP），再把 TP 漏掉的那部分 activation 沿 sequence 切开（SP）。这是 Megatron 的核心，也是 NVIDIA 考过的原题。
ZeRO 到 70B 以上会撞上三堵墙
就是第 4 课末尾的三个问题，这里补上大模型的数：
- activation 不切：8B 模型一条 8K 样本就有约 41 GB（每层约 1.14 GB），用多少张卡都不变
- DP 的卡数有上限：Llama 3 405B 一个 batch 约 16M tokens，按 8K 一条只有约 2000 条，它却用了 16K 张卡
- all-gather 太慢：ZeRO-3 每个 step 每卡要发约 6Ψ bytes，70B 就是 420 GB，走 IB 约 8.4 s。而每卡算一条 8K 样本只要 6 × 70B × 8192 ÷ 400 TFLOPS ≈ 8.6 s，通信和计算差不多长，很难藏住
办法是把每一层的计算本身切开，activation 跟着变小，weight 也不用拼回来。
TP 的数学：分块矩阵乘法
一个 linear layer 就是 \(Y = XA\)，\(X\) 是 [token 数, h]。把 \(A\) 切成两块放到两张卡上，有两种切法。先看一个能心算的例子：
\[ X = \begin{bmatrix}1 & 2\end{bmatrix},\quad A = \begin{bmatrix}1 & 2\\ 3 & 4\end{bmatrix},\quad XA = \begin{bmatrix}7 & 10\end{bmatrix} \]- 按列切（column parallel）：GPU0 拿第 1 列，算 1·1 + 2·3 = 7；GPU1 拿第 2 列，算 1·2 + 2·4 = 10。拼起来就是 [7, 10]
- 按行切（row parallel）：GPU0 拿第 1 行和 \(X\) 的第 1 个数，算 1 × [1, 2] = [1, 2]；GPU1 拿第 2 行和第 2 个数，算 2 × [3, 4] = [6, 8]。逐个相加才是 [7, 10]。每卡算出的是 partial sum，形状完整、数值只是一部分，要 all-reduce 求和
写成一般形式：
\[ \text{column：}\ X[A_1, A_2] = [XA_1,\ XA_2] \qquad \text{row：}\ [X_1, X_2]\begin{bmatrix}A_1\\ A_2\end{bmatrix} = X_1A_1 + X_2A_2 \]你的背景线性代数里的分块矩阵乘法就是 TP 的全部数学。column parallel 是把右边的矩阵按列分块，row parallel 是把两个矩阵沿相乘的那一维（inner dimension）对应分块。面试时先写出上面两个等式，再讲通信，最不容易乱。
Column 和 Row 各在哪里通信
| Column parallel | Row parallel | |
|---|---|---|
| weight [k, n] | 按列切，每卡 [k, n/t] | 按行切，每卡 [k/t, n] |
| 每卡的输入 | 完整的 \(X\) | \(X\) 的一块 \(X_i\) |
| 每卡的输出 | 输出的几列，拼起来是完整结果 | 完整形状的 partial sum |
| forward | 不通信 | all-reduce |
| backward | all-reduce（\(dX\) 是 partial sum） | 不通信 |
backward 为什么正好反过来：column 的 \(dX = dY_1A_1^\top + dY_2A_2^\top\)，每卡只能算出其中一项，要 all-reduce 相加；row 的输出已经 all-reduce 成完整结果，每卡都有完整的 \(dY\)，自己那块 \(dX_i = dY\,A_i^\top\) 在本卡就能算。
完整流程：column parallel，t = 3
\(X = [1,\ 2]\)，\(W = [c_1 \mid c_2 \mid c_3] = \begin{bmatrix}1 & 2 & 3\\ 4 & 5 & 6\end{bmatrix}\)，GPU i 拿第 i 列。backward 时假设从后面传回的 \(dY = [1,\ 2,\ 3]\)。
| 步骤 | GPU0 | GPU1 | GPU2 | 通信 |
|---|---|---|---|---|
| forward：\(Y_i = X c_i\) | 9 | 12 | 15 | 无 |
| backward 收到 \(dY_i\) | 1 | 2 | 3 | 无 |
| \(dW_i = X^\top dY_i\) | [1, 2] | [2, 4] | [3, 6] | 无，就是自己那列 weight 的 gradient |
| dX 的一份：\(dY_i\, c_i^\top\) | [1, 4] | [4, 10] | [9, 18] | — |
| all-reduce 相加（f 的 backward） | [14, 32] | [14, 32] | [14, 32] | 每卡都拿到完整的 dX |
X 在 forward 被三张卡都用了，所以它的 gradient 是三份之和。W 的每一列只在一张卡上用过，它的 gradient 在本卡就算完了。
Megatron 把这两次通信包成一对算子 f 和 g，插在每个 TP block 的入口和出口：
- f（入口）：forward 原样传，backward all-reduce
- g（出口）：forward all-reduce，backward 原样传，因为每份 \(Y_i\) 的 gradient 就是 \(dY\)，每卡本来就有
面试要写 pseudocode 时写到下面这个程度就够了。PyTorch 的 weight 存成 [out, in]，所以 column parallel 切的是 out 那一维（省略了 @staticmethod 和初始化）：
class f(torch.autograd.Function): # 入口
def forward(ctx, x): return x
def backward(ctx, dx): dist.all_reduce(dx, group=tp_group); return dx
class g(torch.autograd.Function): # 出口
def forward(ctx, y): dist.all_reduce(y, group=tp_group); return y
def backward(ctx, dy): return dy
class ColumnParallelLinear(nn.Module): # weight [out/t, in]，bias 也切成 out/t
def forward(self, x): # x：完整输入
return F.linear(f.apply(x), self.weight, self.bias) # 输出：自己那几列
class RowParallelLinear(nn.Module): # weight [out, in/t]，bias 不切
def forward(self, x_i): # x_i：本卡那一块输入
return g.apply(F.linear(x_i, self.weight)) + self.bias # 求和后只加一次 biasMLP：先 column，再 row
\[ Z = \mathrm{GeLU}\big(X[A_1, A_2]\big)\begin{bmatrix}B_1\\ B_2\end{bmatrix} = \mathrm{GeLU}(XA_1)\,B_1 + \mathrm{GeLU}(XA_2)\,B_2 \]- 第一层 \(A\)（h → 4h）按列切，每卡得到中间结果的几列
- GeLU 是逐元素的，每卡对自己那几列直接算，算完正好是 row parallel 要的那块输入
- 第二层 \(B\)（4h → h）按行切，输出 partial sum，最后由 g 做 1 次 all-reduce
- Qwen3 用的是 SwiGLU：gate 和 up 两个矩阵都按列切，逐元素相乘也在本卡做。t = 8 时每卡拿 12288 ÷ 8 = 1536 列
反过来先 row 后 column 为什么不行：第一层的输出成了 partial sum，而 GeLU 不能分别作用在 partial sum 上。拿 ReLU 举例，两卡的 partial sum 是 3 和 −2：ReLU(3 + (−2)) = 1，ReLU(3) + ReLU(−2) = 3。所以 GeLU 之前就得 all-reduce 12288 宽的中间结果（h 的 3 倍），通信更多，每卡还要存完整的中间 activation。
Attention：按 head 切
每个 head 各算各的。Qwen3 有 32 个 query head、8 个 KV head（GQA），t = 8 时每卡拿 4 个 query head 和 1 个 KV head：
- \(W_Q, W_K, W_V\) 按 head 做 column parallel，每卡的 attention 在本卡算完，不通信
- \(W_O\) 做 row parallel，最后 1 次 all-reduce
- 所以 head 数要能被 t 整除
一层通信几次，为什么出不了机器
attention 和 MLP 各是一个“f → column → 本地计算 → row → g”的 block。forward 两个 g、backward 两个 f 各 all-reduce 一次，每层 2 + 2 = 4 次，36 层就是每个 micro-batch 144 次。每次 all-reduce 的是这一层的输出 [b, s, h]，也就是 64 MiB。t = 8 的 ring all-reduce 每卡要发 2 × 7/8 × 64 = 112 MiB：
- 机内 NVLink（450 GB/s）：一次约 0.26 ms，一层 4 次约 1.0 ms，36 层约 38 ms
- 跨机 IB（50 GB/s）：一次约 2.3 ms，一层约 9.4 ms，36 层约 338 ms
- 对比计算：按 6ND 粗估，t = 8 时每卡每层的 forward + backward 约 6 × 0.19B × 8192 ÷ 8 ÷ 400 TFLOPS ≈ 3 ms。走 NVLink 通信约是计算的 1/3，走 IB 约是 3 倍
- 而且藏不住：all-reduce 的结果就是下一个算子的输入，GPU 只能干等，这叫在 critical path 上。DP 的 gradient 可以边 backward 边发，TP 不行
所以 TP 留在一台机器的 NVLink 里，t ≤ 8。TP 是几个厨房合做同一道菜，每一步都要互相递半成品，只能用同一栋楼里的传菜电梯。换成货车跨楼送，厨师大部分时间都在等车。
t 开多大：通信 1× / 1.5× / 1.75×
| t | 每卡计算 | 每卡每次 all-reduce 发多少 | 通信 ÷ 计算（t = 2 记为 1） | 每卡 FFN 列数 / head 数 |
|---|---|---|---|---|
| 2 | 1/2 | 64 MiB（1×） | 1× | 6144 / 16 |
| 4 | 1/4 | 96 MiB（1.5×） | 3× | 3072 / 8 |
| 8 | 1/8 | 112 MiB（1.75×） | 7× | 1536 / 4 |
t 每翻一倍，每卡的计算减半，通信却几乎不降（ring 的 \(2(t-1)/t\) 趋近 2），GEMM 也越切越窄、越难跑满 Tensor Core。所以 memory 放得下就用小一点的 t。
SP：把 TP 漏掉的 activation 也切开
TP 只切了 attention 和 MLP 里的 GEMM。夹在 block 之间的 LayerNorm、dropout、residual add，每卡都对完整的 [s, h] 算一遍、存一份。按 Korthikanti 的公式（开 FlashAttention），TP 之后每层 activation 是 \(sbh(10 + 24/t)\) bytes，那个不随 t 变小的 10 就是这些没切的部分。
这些操作每个 token 各算各的，所以可以沿 sequence 切：每卡只管 8192 ÷ 8 = 1024 个 token。通信只是换个拆法：
- row parallel 之后的 all-reduce 换成 reduce-scatter：求完和，每卡只拿自己那 1024 个 token
- 下一个 column parallel 之前加一个 all-gather，把完整序列拼回来
- all-reduce = reduce-scatter + all-gather：每卡 56 + 56 = 112 MiB，通信量不变
| activation（8K tokens，FlashAttention） | 系数 × sbh | 每层 | 36 层 |
|---|---|---|---|
| 不切 | 34 | 1.14 GB | 41 GB |
| TP，t = 8 | 10 + 24/8 = 13 | 0.44 GB | 15.7 GB |
| TP + SP，t = 8 | 34/8 = 4.25 | 0.14 GB | 5.1 GB |
Embedding 和 LM head：按 vocab 切
Qwen3 的 embedding 是 151936 × 4096 ≈ 0.62B 个 parameter，算上 Adam 约 10 GB。Megatron 把它按 vocab（行）切成 8 份，每卡管 18992 个 token id。查表时不归自己管的 token 输出 0，all-reduce 一加，每个 token 正好拿到自己那一行。LM head 也按 vocab 切，就是 column parallel，一条 8K 样本的 fp32 logits 从约 5 GB 降到每卡 0.62 GB。cross-entropy 不用把 logits 拼回来，每个 token 只 all-reduce 三个数：max、\(\sum e^{z}\) 和目标 logit。
你的经验你在 TikTok 用的 DeepSpeed ZeRO 和 TP 都“切 weight”，区别在算的时候：ZeRO-3 先把完整 weight all-gather 回来再算，每卡算的是自己那几条样本的完整一层；TP 从不拼回 weight，每卡只算每一层的 1/t，同组的卡处理同一批样本。去查一下你用 vLLM 时有没有设过 tensor_parallel_size？vLLM 的 TP 用的也是这套 column / row 切法。
30 秒英文回答
Tensor parallelism splits each weight matrix across t GPUs. Column-parallel splits the output columns: every GPU takes the full input and produces a slice of the output, so the forward needs no communication, but the input gradient is a partial sum, all-reduced in backward. Row-parallel splits the input dimension: each GPU produces a full-size partial sum, so the forward needs an all-reduce and the backward none. Megatron's MLP is column then row, so GeLU stays local; attention splits by heads. That's two all-reduces forward and two backward per layer. Weights, optimizer state and activations shrink by t, but the all-reduces block the next op, so TP stays on NVLink.
自测
Column 和 Row Parallel Linear 怎么实现？pros 和 cons？（NVIDIA 框架岗 2023 原题）
column 的 forward 不通信，backward 用 f 对 \(dX\) 做 all-reduce，row 正好相反（写法见上面的 pseudocode）。pros：weight、optimizer state、block 内的 activation 都缩成 1/t，global batch 不变。cons：all-reduce 在 critical path 上，只能放在 NVLink 内，t 大了 GEMM 变窄、通信占比上升，LayerNorm、dropout 的 activation 仍是整份，要靠 SP。
TP 开得越大越快吗？
不是。通信 ÷ 计算从 t = 2 到 t = 8 涨了 7 倍，GEMM 也越切越窄。t 超过 8 还要走 IB，一次 all-reduce 从约 0.26 ms 变成约 2.3 ms。
bias 怎么处理？
Column 的 bias 跟着输出列一起切。Row 的 bias 不切，all-reduce 之后加一次，否则每卡各加一遍，bias 就被加了 t 次。
怎么验证你的实现是对的？
固定 seed，用同一份完整 weight 切给各卡，和单卡 nn.Linear 对比 output 和 gradient（\(dX\)、\(dW\)、bias 的 gradient）。t = 1、2、4 都测，有无 bias 都测。没有多卡就用 gloo backend 在 CPU 上起多个进程。
想看细一点 6.1 TP Linear（NumPy 模拟 t = 2，含反例）· 6.2 TP communication（华为的 80 MB 计算题）· 6.3 Sequence Parallel（SP 下 MLP 每一步的公式）· 6.4 Vocab parallelism（vocab-parallel cross-entropy 的代码）· 6.12 all-reduce