【分布式训练④】张量并行与序列并行：把单层参数切碎分到多张卡

【分布式训练④】张量并行与序列并行：把单层参数切碎分到多张卡

本文是「分布式训练」系列的第四篇，带你从零理解张量并行（Tensor Parallelism）和序列并行（Sequence Parallelism）的核心原理。当模型大到单卡连一层都放不下时，数据并行和 ZeRO 已经不够用了------这时候就需要把单层参数"切开"，分到多张卡上一起算。

---

一、开篇引入：当模型大到单卡连一层都放不下时怎么办？

你可能已经听过数据并行（DP）和 ZeRO------它们解决的是"数据太多"的问题，每张卡都持有完整的模型副本，各自处理不同的数据。

但是，当模型大到单卡连一层参数都放不下的时候，数据并行就无能为力了。

举个直观的例子：GPT-3 175B 模型，单是一个 Transformer 层的 MLP 上投影矩阵就有 12288 × 49152 ≈ 6 亿 个参数，FP16 下就是 1.2GB。这还只是一层里的一个矩阵，加上 Attention、LayerNorm、激活值......对于显存只有 80GB 的 A100 来说，模型大到一定程度，单卡真的装不下。

怎么办？把参数本身切开，分到多张卡上。 这就是**模型并行（Model Parallelism）**的核心思想。

而张量并行（Tensor Parallelism, TP）是模型并行中最常用、最高效的一种方式------它不是按层切，而是把单层内的矩阵运算按维度切开，让每张卡只算一部分。

1.1 张量并行 vs 数据并行：一张图分清

维度	数据并行（DP）	张量并行（TP）
切什么	数据（batch）	参数（单层权重）
每张卡有什么	完整模型的副本	模型的一部分参数
通信内容	梯度 AllReduce	激活值 AllGather/AllReduce
通信频率	每 step 一次	每层都要通信
适合场景	模型能放下，数据多	模型太大，单层放不下

打个比方：数据并行就像工厂里多条流水线各做各的产品，最后汇总质检；张量并行就像把一个大零件拆成几块，不同工位各做一块，最后拼起来。

---

二、张量并行原理：矩阵怎么切？

张量并行的本质是矩阵乘法的分块计算。我们从最简单的一维张量并行说起。

2.1 一维张量并行：列切分 vs 行切分

假设有一个线性层 Y=XAY = XAY=XA，其中：

• 输入 XXX 形状：B×S,HinB \\times S, H_{in}B×S,Hin（batch × 序列长度，输入维度）
• 权重 AAA 形状：Hin,HoutH_{in}, H_{out}Hin,Hout（输入维度，输出维度）
• 输出 YYY 形状：B×S,HoutB \\times S, H_{out}B×S,Hout

列切分（Column Parallel）

把权重矩阵 AAA 按列切成两块：

A=A1,A2⇒Y=XA1,XA2 A = A_1, A_2 \quad \Rightarrow \quad Y = XA_1, XA_2 A=A1,A2⇒Y=XA1,XA2

每张卡保存一部分列，各自计算后，把输出在特征维度上拼接起来（AllGather 操作）。

卡0: X @ A1 → Y1  [B*S, H_out/2]
卡1: X @ A2 → Y2  [B*S, H_out/2]
拼接:  [Y1, Y2]   [B*S, H_out]  ← AllGather

特点：输入不需要切分，但输出需要 AllGather。

行切分（Row Parallel）

把权重矩阵 AAA 按行切成两块，同时输入 XXX 也要按列切分：

A=A1A2,X=X1,X2 A = \begin{bmatrix} A_1 \\ A_2 \end{bmatrix}, \quad X = X_1, X_2 A=A1A2,X=X1,X2

Y=X1A1+X2A2 Y = X_1 A_1 + X_2 A_2 Y=X1A1+X2A2

每张卡保存一部分行，各自计算后，把结果加起来（AllReduce 操作）。

卡0: X1 @ A1 → Y1  [B*S, H_out]
卡1: X2 @ A2 → Y2  [B*S, H_out]
求和:  Y1 + Y2     [B*S, H_out]  ← AllReduce

特点：输入需要切分，输出需要 AllReduce（加和）。

小思考 ：为什么列切分用 AllGather，行切分用 AllReduce？

列切分每张卡算出的是输出的"一部分"，需要拼起来；行切分每张卡算出的是"部分和"，需要加起来。

2.2 MLP 中的张量并行

Transformer 的 MLP 层通常是这样的：

Y=GeLU(XWup)⋅Wdown Y = \text{GeLU}(X W_{up}) \cdot W_{down} Y=GeLU(XWup)⋅Wdown

先把 hidden 维度放大 4 倍（up projection），经过激活函数，再缩小回来（down projection）。

聪明的你一定想到了：WupW_{up}Wup 按列切，WdownW_{down}Wdown 按行切，中间就不需要通信了！

X ──列切分W_up──▶ [Y1_up, Y2_up]  (各卡各算各的，不用AllGather)
                      │ GeLU
                      ▼
                 [Y1_up, Y2_up]
                      └─行切分W_down─▶ Y1 + Y2 = Y  (最后一次AllReduce)

整个 MLP 层只需要一次 AllReduce！这就是 Megatron-LM 中 TP 的经典设计------列切分 + 行切分的组合，把中间通信消掉了。

2.3 Attention 中的张量并行

Multi-Head Attention（多头注意力）天生就适合张量并行，因为它本身就是按 head 组织的：

Q=XWQ,K=XWK,V=XWV Q = X W_Q, \quad K = X W_K, \quad V = X W_V Q=XWQ,K=XWK,V=XWV

Attention(Q,K,V)=Concat(head1,...,headh)WO \text{Attention}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W_O Attention(Q,K,V)=Concat(head1,...,headh)WO

TP 的做法很简单：按 head 切分！

• 每张卡只负责一部分 attention head
• WQ,WK,WVW_Q, W_K, W_VWQ,WK,WV 按列切分（每列对应一个 head 的投影）
• 输出投影 WOW_OWO 按行切分，最后 AllReduce

卡0: head_0 ~ head_h/2  的 Q, K, V, attention
卡1: head_h/2 ~ head_h  的 Q, K, V, attention
最后: 各卡的输出拼接后乘 W_O 的行分块 → AllReduce

对于有 32 个 head 的模型，TP=8 的话每张卡只需要算 4 个 head，显存直接降到 1/8。

---

三、为什么 TP 通常限制在节点内？

张量并行虽然强大，但有一个"软肋"------通信太频繁了。

3.1 通信频率有多高？

每一层 Transformer 中，TP 组内至少要进行：

• MLP：1 次 AllReduce（或等价的 ReduceScatter + AllGather）
• Attention：1 次 AllReduce

一个 80 层的模型，每个 step 就要做 160 次集合通信！这还没算反向传播的通信量。

3.2 NVLink vs 跨节点带宽

连接方式	典型带宽	相对速度
NVLink 4 (H100)	~900 GB/s	18x
NVLink 3 (A100)	~600 GB/s	12x
InfiniBand NDR 400G	~50 GB/s	1x
普通万兆以太网	~1.25 GB/s	0.025x

节点内 GPU 之间走 NVLink，带宽是跨节点 IB 网络的十几倍。TP 每层都通信，放在跨节点上，通信时间会远远超过计算时间，训练速度直接崩了。

3.3 通信计算比与 batch/seq_len 的关系

TP 的通信量和计算量之比：

通信量计算量∝1B⋅S \frac{\text{通信量}}{\text{计算量}} \propto \frac{1}{B \cdot S} 计算量通信量∝B⋅S1

• Batch size BBB 越大 → 计算量越大 → 通信占比越小
• Sequence length SSS 越大 → 计算量越大 → 通信占比越小

但即使 batch 和 seq_len 很大，跨节点的带宽差距仍然是数量级的。所以业界的共识是：TP 放在节点内，跨节点用 PP 或 DP。

3.4 实际排布举例

假设有 4 个节点，每个节点 8 张 A100：

节点0 (GPU 0-7):  TP=8  (节点内高速NVLink)
节点1 (GPU 8-15): TP=8
节点2 (GPU 16-23): TP=8
节点3 (GPU 24-31): TP=8

跨节点之间: PP=4 或 DP=4  (走InfiniBand，通信少)

总卡数 32 = TP(8) × PP(4) × DP(1) ------这就是经典的 2D 并行组合。

---

四、序列并行（Sequence Parallelism）

张量并行解决了"参数放不下"的问题，但训练中还有另一座大山------激活值显存。

4.1 激活值显存问题

在长序列训练（比如 8K、32K context）中，激活值的显存占用可能比参数还大。因为：

• 激活值形状：B,S,HB, S, HB,S,H，随序列长度 SSS 线性增长
• 每层都要保存激活值用于反向传播
• 层数越多，激活值总量越大

举个例子：batch=1, seq_len=32K, hidden=8192, 80层，FP16：

• 每层激活值：32K × 8192 × 2 ≈ 512 MB
• 80 层就是 40GB！

TP 把参数切成了 1/8，但激活值在某些算子中并没有减少------比如 LayerNorm、Dropout 这些"逐元素"操作，在 TP 按 head 切分时，输入输出还是完整的序列维度。

4.2 序列并行思想

序列并行（Sequence Parallelism, SP） 的想法很直接：把激活值沿序列维度也切了！

在 MLP 和 LayerNorm 这些不需要跨序列位置通信的算子中，把输入沿序列维度切分到 TP group 内的多张卡上：

输入形状: [B, S, H]
TP group 大小: 8
SP切分后每卡: [B, S/8, H]

这样，LayerNorm、Dropout、激活函数这些算子的中间激活值，显存直接降到原来的 1/TP。

4.3 SP + TP 的收益

SP 不是独立于 TP 存在的，它是 TP 的"增强版"------利用同一个 TP 通信组，在算子之间切换数据的 layout：

• Attention 阶段 ：按 head 切分（标准 TP），数据 layout 是 [B, S, H/TP]
• MLP / LayerNorm 阶段 ：沿序列切分（SP），数据 layout 是 [B, S/TP, H]

两者之间通过一次 AllGather + ReduceScatter（或者等价的 AllToAll）来转换 layout。

收益总结：

1. LayerNorm / Dropout / GeLU 等逐元素算子的激活显存减少为 1/TP1/TP1/TP
2. MLP 的输入输出 AllGather 可以和 SP 的 layout 转换合并，减少通信次数
3. 对于长序列训练（S > 4K），显存收益非常显著

---

五、GQA / MQA 下的 TP 切分策略

现在的大模型大多用 GQA（Grouped Query Attention）或 MQA（Multi-Query Attention）来减少 KV Cache。那它们对 TP 有什么影响呢？

5.1 MHA / GQA / MQA 的区别

类型	Q头数	KV头数	说明
MHA	32	32	每个Q头对应独立的KV头
GQA	32	8	每4个Q头共享1组KV头
MQA	32	1	所有Q头共享1组KV头

5.2 切分挑战

MHA 的 TP 切分很自然：Q/K/V 都按 head 切分，每张卡拿自己的 Q 头、K 头、V 头，各算各的 attention，最后 AllReduce。

GQA 就麻烦了：4 个 Q 头共享 1 组 KV 头。如果 TP=8，而 KV 头只有 8 个（每组 4 个 Q 头），那每张卡只有 1 个 KV 头，但有 4 个 Q 头------这就需要把 KV 头复制到相关的 rank 上。

MQA 更极端 ：全模型只有 1 组 KV 头，TP 切分时 KV 需要广播到所有 TP rank。

5.3 切分策略总结

注意力类型	Q的切分	K/V的切分	额外通信
MHA	按head切分	按head切分	无（自然对齐）
GQA	按group切分	每组内复制广播	少量广播
MQA	按head切分	全量广播	一次广播

结论：GQA/MQA 虽然减少了 KV Cache 显存，但在 TP 训练时增加了 K/V 的通信复杂度。好在 K/V 矩阵比 Q 小很多（KV 头数少），额外通信的开销通常可以接受。

---

六、代码实践：列切分与行切分模拟

光说不练假把式。我们用 NumPy 模拟一下张量并行的列切分和行切分，验证结果是否正确。

"""
张量并行模拟：列切分 vs 行切分
运行环境：只需安装 numpy
"""
import numpy as np


def column_parallel_linear(X, W, world_size):
    """
    列切分线性层：权重按列切分到 world_size 张卡上
    每张卡计算 X @ W_i，最后 AllGather 拼接
    """
    H_out = W.shape[1]
    assert H_out % world_size == 0, "输出维度必须能被world_size整除"
    split_size = H_out // world_size

    # 模拟每张卡的计算
    local_outputs = []
    for rank in range(world_size):
        W_local = W[:, rank * split_size:(rank + 1) * split_size]
        Y_local = X @ W_local  # [B*S, H_out/world_size]
        local_outputs.append(Y_local)

    # AllGather: 拼接所有卡的输出
    Y = np.concatenate(local_outputs, axis=-1)
    return Y


def row_parallel_linear(X, W, world_size):
    """
    行切分线性层：权重按行切分，输入也对应切分
    每张卡计算 X_i @ W_i，最后 AllReduce 求和
    """
    H_in = W.shape[0]
    assert H_in % world_size == 0, "输入维度必须能被world_size整除"
    split_size = H_in // world_size

    # 模拟每张卡的计算
    local_outputs = []
    for rank in range(world_size):
        X_local = X[..., rank * split_size:(rank + 1) * split_size]
        W_local = W[rank * split_size:(rank + 1) * split_size, :]
        Y_local = X_local @ W_local  # [B*S, H_out]
        local_outputs.append(Y_local)

    # AllReduce: 所有卡的输出相加
    Y = sum(local_outputs)
    return Y


def mlp_tensor_parallel(X, W_up, W_down, world_size):
    """
    MLP的张量并行：W_up列切分 + W_down行切分
    中间不需要AllGather，只有最后一次AllReduce
    """
    H_up = W_up.shape[1]
    H_out = W_down.shape[0]
    assert H_up % world_size == 0 and H_out % world_size == 0

    split_up = H_up // world_size
    split_down = H_out // world_size

    # 模拟每张卡的完整MLP计算
    local_outputs = []
    for rank in range(world_size):
        # 第一步：列切分的up projection
        W_up_local = W_up[:, rank * split_up:(rank + 1) * split_up]
        Y_up = X @ W_up_local

        # 中间激活函数（GeLU的简化版，用ReLU代替演示）
        Y_act = np.maximum(Y_up, 0)

        # 第二步：行切分的down projection
        # 注意：这里W_down是按行切分，与W_up的列切分维度对应
        W_down_local = W_down[rank * split_up:(rank + 1) * split_up, :]
        Y_local = Y_act @ W_down_local

        local_outputs.append(Y_local)

    # AllReduce
    Y = sum(local_outputs)
    return Y


def main():
    # 参数设置
    B, S, H_in, H_up, H_out = 2, 8, 16, 64, 16
    world_size = 4  # 模拟4张卡

    np.random.seed(42)
    X = np.random.randn(B * S, H_in).astype(np.float32)
    W_linear = np.random.randn(H_in, H_out * 2).astype(np.float32)
    W_up = np.random.randn(H_in, H_up).astype(np.float32)
    W_down = np.random.randn(H_up, H_out).astype(np.float32)

    # ===== 验证列切分 =====
    Y_ref = X @ W_linear
    Y_col = column_parallel_linear(X, W_linear, world_size)
    diff_col = np.max(np.abs(Y_ref - Y_col))
    print(f"【列切分验证】最大误差: {diff_col:.2e}  {'✓ 正确' if diff_col < 1e-5 else '✗ 有问题'}")

    # ===== 验证行切分 =====
    Y_row = row_parallel_linear(X, W_linear, world_size)
    diff_row = np.max(np.abs(Y_ref - Y_row))
    print(f"【行切分验证】最大误差: {diff_row:.2e}  {'✓ 正确' if diff_row < 1e-5 else '✗ 有问题'}")

    # ===== 验证 MLP TP =====
    Y_mlp_ref = np.maximum(X @ W_up, 0) @ W_down
    Y_mlp_tp = mlp_tensor_parallel(X, W_up, W_down, world_size)
    diff_mlp = np.max(np.abs(Y_mlp_ref - Y_mlp_tp))
    print(f"【MLP TP验证】最大误差: {diff_mlp:.2e}  {'✓ 正确' if diff_mlp < 1e-5 else '✗ 有问题'}")

    # ===== 显存分析 =====
    print("\n【显存分析】(FP16, 单位: MB)")
    param_ref = (W_up.size + W_down.size) * 2 / 1024 / 1024
    param_tp = param_ref / world_size
    print(f"  单卡MLP参数量: {param_ref:.2f} MB → TP后每卡: {param_tp:.2f} MB")
    act_ref = B * S * H_up * 2 / 1024 / 1024  # MLP中间激活
    print(f"  MLP中间激活值: {act_ref:.2f} MB (每层)")


if __name__ == "__main__":
    main()

运行上面的代码，你会看到列切分、行切分和 MLP TP 的计算结果都和原始结果一致（误差在浮点精度范围内）。

自己动手试试：

1. 把 world_size 改成 2 或 8，看看误差是否变化
2. 修改 hidden_size，观察 TP 前后的显存变化
3. 尝试实现 Attention 的 TP 模拟

---

七、常见误区与小结

7.1 常见误区

误区一："张量并行就是数据并行，都是多卡训练嘛"

完全不同！数据并行切的是数据，每张卡有完整模型；张量并行切的是参数，每张卡只有模型的一部分。TP 的通信频率远高于 DP，通常只在节点内使用。

误区二："序列并行是另一种并行方式，和 TP 没关系"

序列并行不是独立的并行维度，而是 TP 的"搭档"------它复用 TP 的通信组，在不同算子阶段切换数据切分维度（head维度 vs 序列维度），进一步降低激活显存。

误区三："GQA 下 TP 和 MHA 一样简单"

GQA/MQA 的 K/V 头数远少于 Q 头数，TP 切分时 K/V 需要在组内广播或复制，增加了通信和实现复杂度。不过 K/V 本身比较小，实际开销可控。

误区四："TP 越大越好，越多卡越快"

TP 受限于矩阵维度（head 数、hidden 大小），而且 TP 越大通信量也越大。通常 TP=8 是比较合理的选择（对应一个 DGX 节点）。

7.2 全文小结

1. 张量并行（TP） 把单层参数按矩阵维度切分到多张卡上，分为列切分和行切分两种基本方式。
2. MLP 的 TP 采用"列切分 + 行切分"的组合，中间不需要通信，整层只需一次 AllReduce。
3. Attention 的 TP 按 head 切分，天生适配多头注意力结构。
4. TP 通信频繁，通常限制在节点内（NVLink 高速互联），跨节点用 PP 或 DP。
5. 序列并行（SP） 沿序列维度切分激活值，进一步降低长序列训练的显存峰值，是 TP 的黄金搭档。
6. GQA/MQA 下的 TP 需要处理 K/V 共享带来的广播问题，但总体开销可控。

---

八、思考问题

1. 为什么 MLP 中通常 WupW_{up}Wup 按列切分、WdownW_{down}Wdown 按行切分？如果反过来（WupW_{up}Wup 行切分、WdownW_{down}Wdown 列切分）会怎样？
2. 张量并行的通信量与 batch size、sequence length 有什么关系？为什么长序列训练中 TP 的通信占比反而更低？
3. 序列并行和流水线并行分别解决什么问题？它们能同时使用吗？如果可以，怎么组合？
4. 在 GQA 的 TP 切分中，如果 TP 大小大于 KV 头数（比如 TP=8，KV 头数=4），K/V 矩阵应该如何处理？
5. 假设你有一个 175B 参数的模型和 256 张 A100（32 节点 × 8 卡），你会设计怎样的 TP/PP/DP 组合？为什么？

---

下一篇预告：我们将学习流水线并行（Pipeline Parallelism）------当模型深到单卡放不下所有层时，如何把模型按层"切"到多张卡上，像流水线一样处理数据。GPipe、1F1B、气泡率......这些名词的含义都会一一揭晓。