张量并行与序列并行

数据并行和 ZeRO 解决的是「模型状态」的冗余：参数、梯度、优化器状态在每张卡上各存一份太浪费，就把它们切开。但它们有一个共同的前提：每张卡仍然完整地跑一遍每一层的计算，于是每张卡都要为自己那份 micro-batch 存下这一层全部的中间激活。当模型的 hidden size 很大、序列很长时，单卡哪怕只放 1 条序列的激活也装不下，ZeRO 就帮不上忙了。

张量并行（Tensor Parallelism，TP）换了一个思路：不切 batch，而是把每一个矩阵乘法本身拆到多张卡上。这样参数、梯度、优化器状态和（大部分）激活都被切开了，而且不需要在用之前把参数拼回完整形状。代价是：通信从「每步一次」变成了「每层好几次」，而且就在计算的关键路径上。序列并行（Sequence Parallelism，SP）是 TP 的一个小补丁，把 TP 没切开的那部分激活也沿序列维切掉，通信量不变。

预备：记号与集合通信的代价

记号：

• \(b\)：micro-batch 里的序列数；\(s\)：序列长度；\(h\)：hidden size（也写作 \(d_\text{model}\)、\(D\)）；\(F\)：MLP 中间维度（\(d_{ff}\)）；\(V\)：词表大小；\(L\)：层数。
• \(n_h\)：query 头数；\(n_{kv}\)：key/value 头数（GQA 时 \(n_{kv}\lt n_h\)）；\(d_h=h/n_h\)：每个头的维度。
• \(t\)：TP 度，也就是一个 TP 组里有几张卡。
• 激活和权重默认 bf16，每个元素 2 字节。

TP 用到的集合通信只有三种（推导见集合通信与分片矩阵乘法）。设一个张量总大小为 \(M\) 字节，在 \(t\) 张卡上用环形（ring）算法：

• All-Gather：每张卡开始只有 \(M/t\)，结束时每张卡都有完整的 \(M\)。每张卡要发送（也接收）\(\frac{t-1}{t}M\) 字节。
• Reduce-Scatter：每张卡开始有一份完整大小 \(M\) 的「部分和」，结束时每张卡拿到求和结果的 \(1/t\)。每张卡发送 \(\frac{t-1}{t}M\) 字节。
• All-Reduce：每张卡开始有一份部分和，结束时每张卡都有完整的求和结果。环形实现就是 Reduce-Scatter 接 All-Gather，所以每张卡发送 \(2\frac{t-1}{t}M\) 字节。

如果每张卡的出口带宽是 \(W\)（字节/秒），这些操作的耗时近似就是「每张卡发送的字节数 \(/W\)」。特别地，All-Reduce 的代价 = All-Gather + Reduce-Scatter 的代价。

矩阵乘法的两种切法

神经网络里的线性层写成 \(Y = XW\)，其中 \(X\) 是 \([N, k]\) 的输入（\(N=b\cdot s\) 是 token 数），\(W\) 是 \([k, n]\) 的权重。TP 依赖的只是分块矩阵乘法的两个恒等式：

第一个式子说：把 \(W\) 按列切成 \(t\) 块，每张卡拿完整的 \(X\) 和一块 \(W_i\)（形状 \([k, n/t]\)），各自算出输出的 \(n/t\) 列，彼此之间没有任何依赖。这叫列并行线性层（column-parallel linear）。它的输入需要每张卡都有一份完整的 \(X\)，输出是沿最后一维切开的 \(Y_i\)，形状 \([N, n/t]\)。如果下游需要完整的 \(Y\)，就要做一次 All-Gather。

第二个式子说：把 \(W\) 按行切成 \(t\) 块（每块 \([k/t, n]\)），同时 \(X\) 也按最后一维切开（每块 \([N, k/t]\)），每张卡算出一个完整形状 \([N, n]\) 的矩阵 \(X_iW_i\)。但它只是部分和，真正的 \(Y\) 是所有卡的结果加起来，所以要做一次 All-Reduce（或者如果下游只需要一部分，做 Reduce-Scatter）。这叫行并行线性层（row-parallel linear）。

还有一种更紧凑的写法：给张量的维度加下标，表示这一维被切在哪个设备轴上（这里 \(Y\) 就是 TP 所在的那组卡），\(\cdot_k\) 表示沿 \(k\) 维收缩。列并行是 \(X[N, k]\cdot_k W[k, n_Y] \to Y[N, n_Y]\)，没有切收缩维，所以不需要通信；行并行是 \(X[N, k_Y]\cdot_k W[k_Y, n]\to Y[N, n]\{U_Y\}\)，切了收缩维，结果带着一个「未归约」标记 \(\{U_Y\}\)，必须 Reduce 一次。这是判断「这一步要不要通信」的通用规则：只要收缩的那一维被切了，结果就是部分和。

反向传播时情况恰好对调。对 \(Y=XW\)，有 \(\partial L/\partial X = (\partial L/\partial Y)\,W^\top\) 和 \(\partial L/\partial W = X^\top(\partial L/\partial Y)\)。

• 列并行：每张卡有 \(dY_i\)（\([N, n/t]\)）和 \(W_i\)。\(dX = \sum_i dY_i W_i^\top\)：这次收缩的是被切开的 \(n\) 维，所以每张卡算出的是部分和，反向需要 All-Reduce。权重梯度 \(dW_i = X^\top dY_i\) 只用到本地数据，不需要通信。
• 行并行：每张卡有完整的 \(dY\)（如果前向做了 All-Reduce，那么每张卡的 \(dY\) 本来就一样）和 \(W_i\)。\(dX_i = dY\, W_i^\top\) 正好就是 \(X_i\) 对应的那一块梯度，不需要通信；\(dW_i = X_i^\top dY\) 也是本地的。

所以单独看：列并行「前向不通信、反向 All-Reduce 输入梯度」，行并行「前向 All-Reduce 输出、反向不通信」。Megatron 的 f 和 g 算子就来自这对互补关系。

MLP：先列后行，只通信一次

经典 GPT 的 MLP 是 \(Z=\mathrm{GeLU}(XA)\)，\(Y = ZB\)，\(A\) 是 \([h, F]\)，\(B\) 是 \([F, h]\)。怎么切？

如果把 \(A\) 按行切，那么 \(XA\) 是部分和，而 GeLU 是非线性的：\(\mathrm{GeLU}(XA_1+XA_2)\ne \mathrm{GeLU}(XA_1)+\mathrm{GeLU}(XA_2)\)。所以 GeLU 之前必须先 All-Reduce，这是一个额外的同步点。

如果把 \(A\) 按列切，每张卡得到 \(XA_i\)，形状 \([N, F/t]\)，这是输出的一部分列，不是部分和。GeLU 是逐元素的，所以 \(\mathrm{GeLU}(XA)\) 的第 \(i\) 块列就是 \(\mathrm{GeLU}(XA_i)\)，各卡独立算即可。现在每张卡手里正好是 \(Z\) 沿 \(F\) 维切开的一块 \(Z_i\)，这恰好是行并行线性层需要的输入格式。于是 \(B\) 按行切，每张卡算 \(Z_iB_i\)，形状 \([N, h]\)，最后一次 All-Reduce 得到 \(Y=\sum_i Z_iB_i\)。

整个 MLP 前向只有一次 All-Reduce，反向也只有一次（在 MLP 的入口，对 \(dX\)）。两个大矩阵都切成了 \(1/t\)，中间激活 \([N, F]\) 也只存 \(1/t\)。这就是 Megatron-LM 论文的核心构造。

现代 LLM 多用 SwiGLU：\(Y = \big(\mathrm{SiLU}(XW_\text{gate})\odot XW_\text{up}\big)W_\text{down}\)。做法完全一样：gate 和 up 两个矩阵都按列切（同一个 \(F\) 维的同一段，通常把二者拼成一个 \([h, 2F/t]\) 的矩阵一次乘完），逐元素的 SiLU 和乘法在本地做，down 按行切。唯一要小心的是：拼接 gate/up 时要保证每张卡拿到的 gate 列和 up 列对应同一段 \(F\) 下标，否则逐元素乘就乘错了对象。

反过来「先行后列」行不行？行，但更差：第一个矩阵按行切需要输入沿 \(h\) 切开，输出是部分和，必须在 GeLU 前 All-Reduce；第二个矩阵按列切，输出又是切开的，要拼回完整 \(Y\) 还要 All-Gather。两次通信，而且中间那次 All-Reduce 的张量是 \([N, F]\)，比 \([N, h]\) 大好几倍。

注意力：按头切，以及 GQA 的约束

多头注意力天然有一个互相独立的维度：头。第 \(j\) 个头只用 \(Q\)、\(K\)、\(V\) 里属于它的那 \(d_h\) 列，算出自己的 \([N, d_h]\) 输出，和别的头没有关系，直到输出投影 \(W_O\) 把所有头拼起来再乘一个 \([h, h]\) 矩阵。于是：

1. \(W_Q, W_K, W_V\) 按列切，并且切的边界对齐头的边界：每张卡拿 \(n_h/t\) 个完整的 query 头、\(n_{kv}/t\) 个完整的 KV 头。这是列并行，前向不通信。
2. RoPE 是在每个头内部对 \(d_h\) 维做旋转，本地做；\(\mathrm{softmax}(QK^\top/\sqrt{d_h})V\) 对每个头独立，本地做（FlashAttention 照常用）。
3. 拼起来的注意力输出在每张卡上是 \([N, h/t]\)，正好是沿 \(h\) 维切开的一块，所以 \(W_O\) 按行切（\([h/t, h]\)），得到部分和，再 All-Reduce。

结构和 MLP 一模一样：入口一个 f，出口一个 g。所以一层 Transformer（注意力 + MLP）前向共有 2 次 All-Reduce，反向也是 2 次。

几个实现细节：

• 融合 QKV 的权重布局。很多实现把 \(W_Q,W_K,W_V\) 拼成一个 \([h, (n_h+2n_{kv})d_h]\) 的大矩阵一次乘完。如果直接把这个大矩阵按列等分，第一张卡会拿到一堆 Q 头、最后一张卡拿到全部 V，这是错的。正确做法是让每张卡的那一片里依次包含「自己的 Q 头、自己的 K 头、自己的 V 头」，也就是先按 TP rank 分组再拼接。从 Hugging Face 格式的 checkpoint 转换到 Megatron 格式时，这一步重排是最常见的坑。
• 偏置。列并行层的偏置跟着输出列一起切；行并行层的偏置只能加一次，所以在 All-Reduce 之后再加（或者只让一张卡加）。
• TP 度不能超过头数。切的单位是头，所以 \(t\le n_h\) 且 \(t\) 要整除 \(n_h\)。

GQA：KV 头比 query 头少

GQA（Grouped-Query Attention）里 \(n_{kv}\lt n_h\)，每 \(n_h/n_{kv}\) 个 query 头共享一组 K/V。切的时候要让一组 query 头和它们共享的那个 KV 头在同一张卡上：

• \(t\le n_{kv}\) 且 \(t\mid n_{kv}\)：最干净。每张卡拿 \(n_{kv}/t\) 个 KV 头和它们对应的 \(n_h/t\) 个 query 头，没有任何额外开销。
• \(n_{kv}\lt t\le n_h\)：KV 头不够分。常见做法是复制 KV 头：每个 KV 头放到 \(t/n_{kv}\) 张卡上，每张卡只用它服务自己的那几个 query 头。这样 K/V 的权重被复制了，它们的梯度在这几张卡上是不同的部分和，需要额外的 All-Reduce 才能保持同步；KV 的激活也不再按 \(1/t\) 缩小。以 Llama-3 8B 为例：32 个 query 头、8 个 KV 头，TP 最多可以开到 32，但一旦超过 8，每个 KV 头都要复制到多张卡上，并处理好这些副本的梯度同步。推理框架（如 vLLM）在 TP 度大于 KV 头数时也是复制 KV 头。

Llama 3 系列的 8B / 70B / 405B 都是 8 个 KV 头，所以 TP=8 正好每卡一个 KV 头。Llama-3 70B（\(h=8192\)，\(n_h=64\)，\(n_{kv}=8\)，\(d_h=128\)，\(F=28672\)，\(V=128256\)）在 TP=8 下，每张卡上的权重形状如下：

每层每卡的参数从 8.56 亿降到约 1.07 亿。整个模型约 706 亿参数，每卡约 88 亿。

f 和 g：前向与反向到底发生什么

Megatron-LM 把上面的通信封装成两个 autograd 算子：

• f：放在 TP 区域入口（列并行层之前）。前向是恒等（输入在各卡上本来就一样），反向是 All-Reduce（因为 \(dX\) 在各卡上是部分和）。
• g：放在 TP 区域出口（行并行层之后）。前向是 All-Reduce（把部分和加起来），反向是恒等（输出的梯度在各卡上一样，直接往回传）。

它们是一对共轭（conjugate）：一个在前向通信，另一个在反向通信。你在别处也会看到 g 被写成 \(f^*\) 或 \(\bar f\)，指的是同一个算子。用 PyTorch 写出来只有几行：

class CopyToTP(torch.autograd.Function):        # f
    @staticmethod
    def forward(ctx, x):  return x
    @staticmethod
    def backward(ctx, dx):
        dist.all_reduce(dx, group=tp_group); return dx

class ReduceFromTP(torch.autograd.Function):    # g
    @staticmethod
    def forward(ctx, x):
        dist.all_reduce(x, group=tp_group); return x
    @staticmethod
    def backward(ctx, dy): return dy

为什么 f 的反向必须 All-Reduce？f 的输出 \(X\) 被送进各卡的列并行层，每张卡只沿着自己那几列反传，所以每张卡算出的 \(dX\) 只包含「通过自己那部分列」的梯度贡献。\(X\) 真正的梯度是所有路径之和。如果不加起来，各卡的 \(dX\) 不一样，再往前传到 LayerNorm、残差、embedding 就会分叉，各卡上「同一份」复制参数的梯度也会不一样。

复制区域里的 LayerNorm 和 Dropout

在纯 TP 里，g 之后到下一个 f 之前（残差相加、Dropout、LayerNorm/RMSNorm）是「复制区域」：每张卡对完全相同的完整张量做相同的计算。Megatron 的选择是让每张卡都重复算一遍，而不是一张卡算完再广播，因为这些操作很便宜。

• LayerNorm 的参数梯度：每张卡看到的输入一样，反传回来的梯度也一样（f 的反向已经 All-Reduce 过了），所以各卡上 LayerNorm 权重的梯度天然相同，不需要额外同步。
• Dropout 的随机数：复制区域里各卡必须生成相同的 dropout mask，否则各卡的激活会分叉，所以用同一个随机种子；而 TP 区域内部（比如注意力概率上的 dropout，各卡处理不同的头）应该用不同的种子，否则各个头的 mask 都一样，失去随机性。Megatron 为此维护了两套 RNG 状态。现在的 LLM 预训练大多不用 dropout，但这个机制仍然在代码里。

纯 TP 的问题也在这里：复制区域里的激活是完整的 \([b, s, h]\)，每张卡都存一份，没有被 \(t\) 除。序列并行要解决的就是这部分。

词表并行：embedding 与交叉熵

词表很大（Llama 3 是 128256），embedding 表 \([V, h]\) 和输出层（LM head）\([h, V]\) 加起来是一笔不小的参数：Llama-3 8B 的两张表合计约 10.5 亿参数，占全模型的 13%；1B 量级的小模型里这个比例能超过 1/4。而且 logits \([b, s, V]\) 是整个网络里最大的单个激活。所以它们也要切，切在词表维上。

输入 embedding

第 \(i\) 张卡持有词表下标 \([iV/t, (i+1)V/t)\) 这一段行。查表时，每张卡对每个 token 判断：在自己的范围内就查出那一行，不在就输出全零向量。这样每张卡得到一个 \([b, s, h]\) 的张量，每个 token 只有一张卡给出了非零值。All-Reduce（求和）之后，每张卡都拿到完整的 embedding 输出。从矩阵角度看，查表等于 one-hot 矩阵 \([N, V]\) 乘 embedding 表 \([V, h]\)，切的是收缩维 \(V\)，所以它是一个行并行层，出口跟一个 g。

LM head 与交叉熵

输出层 \(W_\text{out}\) \([h, V]\) 按列（词表维）切，是列并行：每张卡算出 logits 的一段 \(z^{(i)}\in\mathbb{R}^{N\times V/t}\)。朴素做法是 All-Gather 出完整 logits 再算 softmax，但这要传 \(b\cdot s\cdot V\) 个元素，非常大。Megatron 的做法是把交叉熵融合进来，只传标量。对一个 token，设目标词是 \(y\)：

每一项都可以分到各卡上算：

1. 每张卡算本地最大值 \(m_i=\max_{v\in\text{本卡}} z_v\)，All-Reduce(MAX) 得到全局 \(m\)。（减最大值只是为了数值稳定，不影响结果。）
2. 每张卡算 \(S_i = \sum_{v\in\text{本卡}} e^{z_v-m}\)，All-Reduce(SUM) 得到 \(S\)。
3. 目标词 \(y\) 只落在一张卡的范围内，那张卡取出 \(z_y\)，其他卡给 0，All-Reduce(SUM) 得到 \(z_y\)。
4. \(\ell = \log S + m - z_y\)，每张卡都能算出同样的 loss。

三次 All-Reduce 的张量都是 \([b, s]\)，相比 \([b, s, V]\) 小了 \(V\) 倍。反向也是本地的：\(\partial\ell/\partial z_v = \mathrm{softmax}(z)_v - \mathbb{1}[v=y]\)，每张卡用全局的 \(m\) 和 \(S\) 就能算出自己那段的 softmax，再在自己负责的目标位置减 1。然后列并行的反向（f）把 \(dX\) All-Reduce。

还有两个小细节。第一，词表大小要能被 \(t\) 整除，而且每卡的那一段最好是 128 的倍数，GEMM 才高效：Megatron-LM 论文把 GPT-2 的 50257 补到了 51200（\(128\times 8 = 1024\) 的倍数）。补出来的行永远不会作为目标出现，相当于浪费一点点算力。第二，如果输入和输出 embedding 共享权重（tied），两边的切法要一致，梯度在同一张卡上相加即可。

序列并行：把剩下的激活也切开

纯 TP 的激活里有多少没被切开？对 GPT 结构的一层逐项记账（bf16，不做重计算），每层要存的激活字节数是：

这里 \(a\) 是头数。\(5as/h\) 那项来自 \(QK^\top\)、softmax、dropout 这些 \(s^2\) 大小的张量（用了 FlashAttention 就不存了）。关键是那个 \(10\)：它来自两个 LayerNorm 的输入（各 \(2sbh\)）、两个 dropout mask（各 \(sbh\)）以及注意力和 MLP 入口线性层的输入（各 \(2sbh\)），这些都在复制区域里，不随 \(t\) 缩小。\(t=8\) 时，\(24/8=3\)，复制部分 10 反而是大头。

观察：复制区域里的操作（LayerNorm、Dropout、残差相加）对每个 token 都是独立的。LayerNorm 要在 \(h\) 维上算均值方差，所以不能沿 \(h\) 切；但沿序列维 \(s\) 切完全没问题：

于是序列并行：在复制区域里，第 \(i\) 张卡只保存和处理 \([b, s/t, h]\) 这一段 token。问题变成两个区域之间怎么转换：

• SP → TP（入口，原来的 f 处）：列并行层需要每张卡都有完整的 \([b, s, h]\)。各卡手里是不同的序列段，所以要 All-Gather（沿 \(s\)）。这个算子叫 \(g\)：前向 All-Gather，反向 Reduce-Scatter。
• TP → SP（出口，原来的 g 处）：行并行层输出完整形状的部分和，下一个区域只需要求和结果的 \(1/t\)（一段序列）。「先求和再切」正好就是 Reduce-Scatter（沿 \(s\)）。这个算子叫 \(\bar g\)：前向 Reduce-Scatter，反向 All-Gather。

注意这里的 \(g\) 是一个新算子，和纯 TP 里那个「前向 All-Reduce」的 g 不是一回事。可以这样记：SP 把纯 TP 入口的 f 换成了 \(g\)（All-Gather），把出口的 g 换成了 \(\bar g\)（Reduce-Scatter）；两者仍然是一对共轭，一个的前向是另一个的反向。

按区域把激活形状列成表（\(h\) 列、\(s\) 列分别表示这一维是完整还是被切开）：

通信量不变

纯 TP 每层前向 2 次 All-Reduce；TP+SP 每层前向 2 次 All-Gather + 2 次 Reduce-Scatter。操作次数翻倍，但环形 All-Reduce 本来就等于一次 Reduce-Scatter 加一次 All-Gather，所以字节数完全一样：每层前向每卡发送 \(4\cdot\frac{t-1}{t}\cdot 2bsh\) 字节。反向是对偶的，同样不变。实际跑起来，SP 还让一层的前向略快一点：在一个 22B GPT 模型上，单层前向从 7.7 ms 降到 7.2 ms，因为 LayerNorm 和 Dropout 只处理 \(1/t\) 的数据。这个收益被另一件事抵消了一部分：同样的数据量，RS + AG 两次调用比一次 All-Reduce 要慢一些。

SP 带来的两个新细节

• LayerNorm 梯度需要同步了。纯 TP 里各卡的 LayerNorm 看到同样的输入，梯度天然一致；SP 里各卡看的是不同的 token，梯度是部分和，需要在 TP 组内 All-Reduce（像数据并行那样）。参数只有 \(h\) 个，代价很小。
• 反向里多一次 All-Gather。列并行层算权重梯度 \(dW_i = Y^\top dZ_i\) 需要完整的输入 \(Y\)（\([b,s,h]\)）。如果前向把 All-Gather 出来的完整 \(Y\) 存着，SP 省下的显存就又吐回去了。所以每张卡只存自己那段 \(Y_i^s\)，反向时再 All-Gather 一次，并且让它和 \(dY\) 的计算重叠。这次额外的 All-Gather 是用通信换显存，严格说来 TP+SP 的总字节数比纯 TP 多了这一项。

显存

有了 SP，每层激活里不再有任何复制项：

每个张量都被 \(t\) 除。再配合选择性重计算（只重算注意力里那些 \(s^2\) 大小、FLOPs 很少的部分；FlashAttention 本身就是这样做的），每层降到 \(34\,sbh/t\)。在 22B 到 1T 参数的 GPT 模型上，序列并行和选择性重计算各自把激活显存减少约一半，合起来约 5 倍。

演示：一层 Transformer 里的 TP / TP+SP

一层 Llama 风格的 Transformer（RMSNorm → 注意力 → 残差 → RMSNorm → SwiGLU MLP → 残差），拆成一步一步的算子。每个小方块代表一张卡上的激活张量：方块的竖边是序列维 \(s\)，横边是当前的特征维（\(h\)、\(F\) 或头）；实心部分是这张卡实际持有的那块，斜线表示「完整形状但只是部分和」。通信步骤用橙色标出，并给出每张卡要发送的字节数（环形算法）。

可以试试：在纯 TP 下数前向的通信，是 2 次 All-Reduce；切到 TP+SP，变成 2 次 All-Gather + 2 次 Reduce-Scatter，走完一层后「累计每卡发送」和纯 TP 一样。再看「当前最大激活」：纯 TP 在复制区域里每卡要放完整的 \([b,s,h]\)，TP+SP 里永远只有 \(1/t\)。反向模式下，TP+SP 多出两步「重新 All-Gather 输入」，那就是拿通信换显存的那次 All-Gather。

精确的通信量

把一层的通信加起来（每卡发送字节数，环形算法，bf16，\(M = 2bsh\) 是一个 \([b,s,h]\) 激活的字节数）：

算一个具体的例子：Llama-3 70B，\(b=1\)，\(s=4096\)，TP=8。\(M=2\times 4096\times 8192\) 字节 = 64 MiB。一次 All-Reduce 每卡发送 \(2\times\frac{7}{8}\times 64 = 112\) MiB，每层前向 + 反向 4 次就是 448 MiB，80 层合计 35 GiB，这只是一条 4K 序列。如果按 H100 NVLink 的 450 GB/s 算，每层约 1.04 ms。

同一层的计算量：每层参数约 8.56 亿，前向 FLOPs 约 \(2\cdot bs\cdot 8.56\times10^8 \approx 7.0\) TFLOP，加上因果注意力的 \(2bs^2h\approx 0.27\) TFLOP，约 7.3 TFLOP；前向 + 反向约 3 倍，再除以 8 张卡，每卡 2.73 TFLOP，按 H100 的 990 TFLOP/s 峰值是 2.76 ms。通信 / 计算 ≈ 0.38。看起来通信不到计算的一半，但要注意三点：一是计算按 100% 峰值算，实际 MFU 在 40%–50% 时计算时间会长一倍，比例反而更好看，但 NVLink 实际也跑不满 450 GB/s：8×H100 上的 All-Reduce 在大消息下约 370 GB/s，58 MB 量级的消息（正好是 70B 一个 MLP 权重分片的大小）只有约 150 GB/s；二是 TP 的通信很难和计算重叠（下面会讲），所以这 38% 大部分是直接加在关键路径上的；三是 \(s\) 和 \(b\) 不影响这个比例（通信和计算都正比于 \(bs\)，注意力项除外），它只取决于 \(h\)、\(F\) 和 \(t\)。

再细看一个线性层：能不能把下一层的 All-Gather（每卡 \(\frac{t-1}{t}\cdot bsh\) 个元素）藏在当前 \([h,h]\) 线性层的计算（\(2bsh^2/t\) FLOPs）后面？两者的时间之比要满足

同样和 \(s\)、\(b\) 无关。对 H100（990 TFLOP/s，450 GB/s，按元素计算时带宽要再除以 2 字节）、\(h=8192\)、\(t=8\)，左边约 \(\frac{7}{16384}\times\frac{990\times10^{12}}{225\times10^{9}}\approx 1.9\)，也就是单个 \([h,h]\) 层的计算不足以盖住通信；MLP 的 \([h,F]\) 层更宽，情况好一些。

显存账

TP 切的东西有两类。

模型状态：除了 RMSNorm 权重（每层 \(2h\) 个，可以忽略），所有参数都被切成 \(1/t\)，梯度和优化器状态跟着参数走。混合精度 + Adam 下每个参数需要 bf16 权重 2 字节、bf16 梯度 2 字节、fp32 主权重 4 字节、Adam 一阶和二阶矩各 4 字节，共 16 字节。Llama-3 70B 是 706 亿参数 × 16 = 1129 GB，TP=8 后每卡 141 GB，仍然超过 80 GB，所以 70B 光靠 TP=8 不够，还要把优化器状态用 ZeRO 在数据并行维度上切（或者加流水线并行）。

激活：按 Llama 结构（RMSNorm、GQA、SwiGLU、FlashAttention、无 dropout）重新记账，每层每条序列存下的张量（bf16）：

• 注意力子层：RMSNorm 输入 \(2sbh\)；QKV 投影输入 \(2sbh\)；FlashAttention 需要的 \(Q\)（\(2sbh\)）、\(K,V\)（\(4sb\,n_{kv}d_h\)）、输出 \(O\)（\(2sbh\)，同时也是 \(W_O\) 的输入）；softmax 的 logsumexp 很小，忽略。
• MLP 子层：RMSNorm 输入 \(2sbh\)；gate/up 投影输入 \(2sbh\)；gate 输出、up 输出、SwiGLU 输出（down 的输入）各 \(2sbF\)。

合计每层 \(12sbh + 4sb\,n_{kv}d_h + 6sbF\) 字节。其中两个 RMSNorm 输入和两个投影输入（\(8sbh\)）在纯 TP 下是复制的，其余被 \(t\) 除；TP+SP 下全部被 \(t\) 除（投影输入只存本卡那段，反向时再 All-Gather）。全重计算时每层只存层输入 \(2sbh\)（SP 下也是 \(/t\)）。最后的 logits 按 fp32 存，\(4sbV/t\)。不同框架具体存哪些张量略有差别（比如有的实现会在 SwiGLU 反向时重算 SiLU），但量级和随 \(t\) 的变化趋势是一样的。

用这个模型，Llama-3 70B、\(s=4096\)、\(b=1\)、TP=8：纯 TP 每层 0.35 GiB，80 层 28 GiB；TP+SP 每层 0.13 GiB，80 层 10.5 GiB。SP 在这里省下了 17 GiB。

为什么 TP 留在节点内

TP 的通信有三个特点，让它对带宽和延迟都极其敏感：

1. 在关键路径上。g 的 All-Reduce 不完成，下一个 LayerNorm 就没法开始；列并行层的输入不 All-Gather 完，矩阵乘法就没法开始。数据并行的梯度 All-Reduce 可以和反向传播的其他层重叠，ZeRO-3 的参数 All-Gather 可以提前预取，但 TP 的通信依赖于刚刚算出来的激活，能重叠的空间很小。
2. 频繁。每层前向 2 次、反向 2 次（SP 下翻倍成 4+4 次更小的操作），一个 80 层的模型每个 micro-batch 有几百次集合通信。
3. 通信量不随 \(t\) 缩小。每卡发送 \(\propto\frac{t-1}{t}M\)，\(t\) 越大越接近 \(M\)；而每卡的计算量是 \(\propto 1/t\)。所以 \(t\) 翻倍，计算减半，通信几乎不变，通信占比翻倍。

第三点可以写成一个干净的上界。我们用一个只有两个矩阵的 MLP 来建模：前向计算 \(T_\text{math} = \frac{4BDF}{Y\cdot C}\)（\(B\) 是 token 数，\(Y\) 是 TP 度，\(C\) 是单卡算力），通信是一次 All-Gather 加一次 Reduce-Scatter，\(T_\text{comms}\approx\frac{2\cdot 2\cdot BD}{W}\)（bf16 两字节，\(W\) 是单卡带宽）。要计算受限，需要

\(B\) 和 \(D\) 都消掉了：TP 能做多大只取决于 MLP 宽度 \(F\) 和硬件的「算力 / 带宽」比 \(C/W\)。对 H100，节点内 \(C/W = 990\times10^{12}/450\times10^{9}\approx 2200\)。Llama-3 70B 的 \(F=28672\)，节点内上界约 \(28672/2200 \approx 13\)，而一个节点只有 8 张卡，所以就是 8。出了节点情况要糟得多：一个 8 卡节点对外总共只有 400 GB/s（8 张 400 Gbps 网卡），由 8 张卡共享，摊到每张卡约 50 GB/s，比 NVLink 慢近一个数量级。即便用分层的集合通信把整个节点的出口带宽都用上，\(C/W\) 也变成 \(990/400\approx 2475\)，TP 在 \(Y \gt F/2475\) 时就会通信受限。所以 TP 通常限制在一个节点内，最多跨两个节点。

在 H100 集群上实测（例如一个 3B 模型、序列长度 4096，扫 TP 度）也是这个形状：单卡吞吐随 TP 增大而下降，从 TP=8 到 TP=16（从 NVLink 变成跨节点的 EFA 网络）掉得最明显，到 TP=32 掉得更厉害；而另一方面，TP 越大，能放下的 batch 越大（激活被切得更碎）。这是一个算得快和放得下之间的取舍。

演示：显存与通信随 TP 度的变化

选一个模型和训练配置，看每张卡的显存和「通信时间 / 计算时间」怎么随 TP 度变化。显存按上一节的记账模型；时间按：每层前向 + 反向每卡发送 \(8\frac{t-1}{t}M\) 字节（SP 时加上重收集的 \(2\frac{t-1}{t}M\)），除以带宽（\(t\) 不超过 NVLink 域时用 NVLink 带宽，否则用每卡跨节点带宽）；计算按 \(3\times\)前向 FLOPs \(/t\) 除以峰值算力。带宽和算力都按标称峰值算，实际能跑到的更低，所以这里的比例用来看趋势，不是性能预测。

几个值得拖一拖的现象：把 \(s\) 拉到 64K 以上，就算 TP=8 加 SP，激活也会撑爆显存，这时需要的是上下文并行，而不是更大的 TP；关掉 SP，TP 增大时激活几乎降不下去（复制项不随 \(t\) 缩小）；8B 模型的 \(F=14336\)，节点内 TP 上界只有约 6.5，所以小模型用 TP 往往得不偿失；切到 NVL72，TP=16、32 的柱子不再撞上跨节点的带宽墙，但比例仍然随 \(t\) 线性增长。

把通信藏到计算后面

TP 的通信难以重叠，但不是完全不能。常见的几种手段：

• 分块流水（collective matmul）。把 All-Gather + 矩阵乘拆成 \(t\) 段：先用本地已有的那段输入算一部分输出，同时去取下一段；Reduce-Scatter 同理，算完一段输出就立刻发出去，同时算下一段。这样通信和计算交错进行。Megatron-LM 和 Nanotron 都实现了 All-Gather 与 MLP 第一个全连接层（FC1）的部分重叠；在 TPU/JAX 的语境里，这种写法就叫 collective matmul。
• 反向里和权重梯度重叠。列并行层反向时，\(dX\) 的 All-Reduce（或 SP 下的 Reduce-Scatter）可以和 \(dW = X^\top dY\) 的计算同时进行，因为后者不依赖前者。SP 下重新 All-Gather 输入 \(Y\) 也可以和 \(dX\) 的计算重叠。这种重叠在实测里看得出来：在一个 22B 模型上，全层重计算的额外开销是 39%，比按 FLOPs 估算的 33% 还高，原因是不重算的基线已经把反向里的 All-Reduce 和权重梯度计算重叠起来，自己先变快了。
• 沿 batch 拆分。把一个 micro-batch 再拆成两半，一半在通信时另一半在计算。Domino 沿着这个方向做得更彻底：把输入沿 batch、权重沿列切成更小的独立块，让 TP 的通信尽量藏到其他块的计算后面。

这些都需要通信和计算在 GPU 上并发执行，会互相抢 SM 和显存带宽，所以实际收益通常小于理论值。

TP 与 FSDP 混合：什么时候各用多少

FSDP 移动权重，TP 移动激活。FSDP 每层要 All-Gather 一次权重（\(\propto DF\)），TP 每层要 All-Gather / Reduce-Scatter 一次激活（\(\propto BD\)）。批次越小，激活越小，TP 越便宜；模型越宽，权重越大，FSDP 越贵。两者可以叠加在不同的设备轴上：\(X\) 个设备做 FSDP，\(Y\) 个做 TP，\(N = XY\)：

好处是相互的：激活沿 batch 切在 \(X\) 上，所以 TP 的 All-Gather 小了 \(X\) 倍；权重沿 \(F\) 切在 \(Y\) 上，所以 FSDP 的 All-Gather 小了 \(Y\) 倍。只看前向，通信时间分别是

\(M_X\)、\(M_Y\) 是两个轴各占用了几个网络维度（TPU 的 3D 环面上两者之和约为 3；GPU 上可以理解为带宽倍数）。\(T_\text{FSDP}\) 随 \(X\) 增大，\(T_\text{TP}\) 随 \(X\) 减小，两者相等时总通信（取 max）最小：

代回去，计算受限的条件变成

对比纯 FSDP 的条件 \(B/N \gt \alpha/M_X\)：混合以后，每卡需要的 token 数从和 \(\alpha\) 成正比变成和 \(\alpha^2/F\) 成正比。代入一组数：TPU v5p 上 \(\alpha\approx 2550\)，取 \(F=32768\)、\(M_XM_Y=2\)，每芯片约 100 个 token 就够；纯 FSDP（三个网络维度全用上）需要约 850，混合后小了约 8 倍。具体例子：在一个 8960 芯片的 TPU v5p pod 上用 400 万 token 的 batch 训练 LLaMA-3 70B，每芯片只分到约 470 个 token，纯 FSDP 是通信受限的；改成约 2048 路 FSDP + 4 路 TP，就能保持计算受限。

这件事在 GPU 上的含义是：当集群很大、全局 batch 被分到每卡只剩几千个 token 时（H100 上跨节点做纯 FSDP，按 \(C/W\approx 2475\) 需要每卡约 2500 个 token），加一点 TP 能把每卡需要的 batch 降下来，这是 TP 除了「省显存」之外的第二个用处。

实践清单

• TP 度整除 \(n_h\)；用 GQA 时最好也整除 \(n_{kv}\)（Llama 3 是 8，TP ≤ 8 正好）。
• TP 不超出一个 NVLink 域（通常 8 卡）。需要更多切分时，用流水线并行或 FSDP 跨节点。Llama 3 论文把并行维度的顺序排为 [TP, CP, PP, DP]，最内层的 TP 要求最高的带宽和最低的延迟，放在同一台服务器内。
• 开 TP 就开 SP：通信量相同，激活显存显著下降。
• 词表补齐到 \(t\times 128\) 的倍数；LM head 和交叉熵用词表并行融合实现，不要 All-Gather logits。
• 按 \(Y\lt F\cdot W/C\) 粗估 TP 上限：\(F\) 小的模型（几 B 参数）在 TP 上收益很低，优先 FSDP。
• 从 HF 格式转换权重时，注意融合 QKV 和 gate/up 的分片顺序。
• 序列一长，TP 区域里的激活仍然是 \([b,s,F/t]\) 这样的完整序列长度，SP 帮不了，这时要上下文并行。

参考与延伸阅读

• Nanotron 团队，The Ultra-Scale Playbook：Tensor Parallelism、Sequence Parallelism 两节，以及附录 A0（Ring AllReduce）与 A3（TP communication analysis，单个线性层的重叠条件）。有 TP 吞吐随 TP 度变化的实测曲线和 picotron 的参考实现，想看实验细节推荐读这两节。
• Austin et al., How to Scale Your Model：Part 5 的 Tensor Parallelism 与 Combining FSDP and Tensor Parallelism；GPU 一章的 H100 带宽实测（All-Reduce 约 370 GB/s）与 TP 上界 \(F/2475\)；Part 6 的 LLaMA-3 70B 例子。
• Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism（f/g 算子、词表并行与融合交叉熵、RNG 处理），2019。
• Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models（序列并行、激活显存公式 \(sbh(34+5as/h)\)、选择性重计算，以及 22B 模型上的单层计时），2022。
• Wang et al., Domino: Eliminating Communication in LLM Training via Generic Tensor Slicing and Overlapping（把 TP 通信藏到独立计算后面），2024。
• Llama Team, The Llama 3 Herd of Models（模型结构参数与并行维度顺序），2024。