12 · 张量并行

⏱ 26 分钟 ｜ ⭐ 把一层切开放到多张卡上

🎯 一句话

前面的并行都是"每卡一份完整的层"。 张量并行把【一层内部的矩阵】横着或竖着切开，分给多张卡。 它是唯一能救"单层都装不下"的办法 —— 代价是每层都要通信。

✂️ 一、两种切法

设一层是 $Y = \text{GeLU}(XA)$，权重 $A$ 要切开。

① 列并行（按列切 A）

② 行并行（按行切 A）

🧩 二、Megatron 的经典组合（⭐ 这一节是核心）

MLP 块：$Y = \text{GeLU}(XA) \cdot B$

🔑 为什么这样组合是最优的： 列并行的输出天然是"按列切开的"， 而行并行需要的输入恰好就是"按列切开的" —— 严丝合缝。 如果顺序反过来（先行后列），中间就要多一次通信。

Attention 块：更自然

📉 三、通信代价：为什么它必须在机内

🔑 这个数字直接给出了铁律： 张量并行度不要超过单机的 GPU 数（通常是 8）。 这不是建议，是带宽算出来的硬约束。

⭐ 另外注意：通信量和层数、hidden、序列长度都成正比，但和张量并行度无关 —— 也就是说，TP=2 和 TP=8 的通信量一样，但 TP=8 分摊的计算更少 → 并行度越高，通信占比越大，效率越低。

🔧 四、代码长什么样

# 概念示意：列并行的线性层
import torch
import torch.nn as nn
import torch.nn.functional as F
class ColumnParallelLinear(nn.Module):
    def __init__(self, in_f, out_f, tp_size, tp_rank):
        super().__init__()
        assert out_f % tp_size == 0
        self.weight = nn.Parameter(torch.empty(out_f // tp_size, in_f))
        #                                      ↑ 只存自己那一片

    def forward(self, x):
        return F.linear(x, self.weight)      # ⭐ 不需要通信

# 行并行
class RowParallelLinear(nn.Module):
    def __init__(self, in_f, out_f, tp_size, tp_rank):
        super().__init__()
        assert in_f % tp_size == 0
        self.weight = nn.Parameter(torch.empty(out_f, in_f // tp_size))

    def forward(self, x):
        out = F.linear(x, self.weight)
        dist.all_reduce(out, group=TP_GROUP)  # ⭐ 这里必须通信
        return out

实际使用（别自己写）：

# 🧩 骨架：`model` 来自你自己的代码，这一段只看写法
# ⭐ PyTorch 原生（2.x）
from torch.distributed.tensor.parallel import (
    parallelize_module, ColwiseParallel, RowwiseParallel)

parallelize_module(model, tp_mesh, {
    "attn.qkv_proj": ColwiseParallel(),
    "attn.out_proj": RowwiseParallel(),      # ⭐ 列→行的配对
    "mlp.gate_proj": ColwiseParallel(),
    "mlp.up_proj":   ColwiseParallel(),
    "mlp.down_proj": RowwiseParallel(),
})

💡 别自己实现张量并行。用 Megatron-LM、PyTorch 原生 TP、 或推理侧的 vLLM/TensorRT-LLM（tensor_parallel_size=8）。 自己写极容易在权重初始化、随机种子、梯度缩放上出微妙的错。

🧵 五、序列并行：一个重要的补充

💡 又是那个恒等式（第 5 章、第 11 章）： AllReduce = ReduceScatter + AllGather。 这已经是它第三次出场了 —— 它是分布式训练里最有用的一个等式。

# 🧩 骨架：`sequence` 来自你自己的代码，这一段只看写法
# Megatron 里开启
--sequence-parallel     # ⭐ 几乎总是该开，和 TP 配套

⚠️ 六、四个坑

💥 随机种子那条特别隐蔽：如果 TP 组内所有卡的 dropout mask 相同， 那么被切开的那些神经元会被同步地丢弃，等效的 dropout 率完全不对 —— 训练能跑，loss 也在降，但正则效果没了。

🔗 和站内其他章的关系

✅ 检查点

1. 列并行和行并行的区别？哪个需要通信？
2. Megatron 的 MLP 为什么用"先列后行"？反过来会怎样？
3. Attention 怎么切？为什么它比 MLP 更自然？
4. 一个 Transformer 层前向反向共几次 AllReduce？
5. 算一下：80 层、hidden=8192、batch=8、seq=2048 时每步通信多少？在 NVLink 和跨机各要多久？
6. 张量并行度为什么不该超过 8？
7. 通信量和 TP 度的关系是什么？这有什么含义？
8. 序列并行解决什么问题？它靠哪个恒等式？
9. TP 组内的 dropout 种子为什么必须不同？

🛑 可以停在这里

⚡ 走神救援

张量并行把一层内部的矩阵切开分给多卡——⭐ 唯一能救「单层都装不下」的办法，代价是每层都要通信。

两种切法：一种切完各算各的、不用通信；另一种切完必须求和。⭐⭐ Megatron 的精妙之处是把两者按顺序组合：前一种的输出天然就是后一种需要的输入形状，严丝合缝——⭐ 于是整个前馈块只要一次通信，反过来排就要多一次。 注意力那边更自然：多个头本来就独立，直接分给不同卡。

⭐⭐ 通信代价直接推出一条铁律：把每层前向反向的同步量乘上层数，在机内高速互联上是可接受的，走 PCIe 就慢一个量级，跨机则完全不可用——⭐ 所以张量并行度不要超过单机的卡数。这不是经验，是带宽算出来的硬约束。

⭐ 另一条也是算出来的：通信量和并行度无关，而计算被分摊了——⭐ 所以并行度越大，通信占比越高，到某个点之后再加收益已经很小。

序列并行补的是它的盲区（归一化和随机丢弃没被切），⭐ 激活显存再降一档而通信总量不变——又是「一次同步等于两次半同步」那个等式，这已经是它第三次出场，是分布式里最有用的一个。

⚠️ 四个坑里最阴的一个：⭐⭐ 同一组内的随机种子必须不同——相同的话，被切开的神经元会被同步丢弃，训练照跑、loss 也降，但正则效果没了。

⭐ 最后一句：别自己实现，用现成的。

下一节 👉 13-流水线并行.md