3、张量并行（Tensor Parallelism）：原理与实现，以Transformer为例的列切分与行切分

张量并行，说白了就是把一个巨大的矩阵运算拆成几块，让多张卡一起干活。你想想看，当模型大到单卡放不下的时候，除了流水线并行那种按层切的方式，还有一种更细粒度的切法——把每一层的参数矩阵都切成片，分到不同设备上。

我个人习惯把张量并行叫做“真正的并行”。为什么？因为流水线并行其实还是串行的，只不过把串行任务分到了不同设备上。但张量并行不一样，它是真正让多张卡同时计算同一个矩阵乘法的一部分。嗯，这里要注意，它和流水线并行的核心区别就在这里。

3.1 为什么需要张量并行？

我在项目中遇到过这样一个场景：训练一个 175B 的 GPT 模型，单卡显存 80GB，但光一个参数矩阵就 200GB 了。流水线并行虽然能解决，但你会发现一个问题——GPU 利用率上不去。因为流水线里总有设备在空等。

张量并行能解决两个核心痛点：

• 单卡放不下大矩阵：把一个大矩阵切成几块，每块放一张卡上
• 减少通信瓶颈：相比数据并行那种全量同步，张量并行只在矩阵乘法前后做一次 all-reduce

3.2 列切分（Column-wise Partition）

我们先看列切分。以 Transformer 的 FFN 层为例，它有两个线性变换：

# 标准实现
hidden = gelu(x @ W1 + b1)  # W1: [d_model, d_ff]
output = hidden @ W2 + b2    # W2: [d_ff, d_model]

列切分怎么做？把 W1 按列切成两块：

# 假设 2 张卡
# 卡0 持有 W1[:, :d_ff//2]
# 卡1 持有 W1[:, d_ff//2:]

# 每张卡独立计算
hidden_local = gelu(x @ W1_local + b1_local)
# 此时 hidden_local 是 [batch, d_ff//2]

这里有个关键点：列切分后，每张卡只需要计算自己那部分。但注意，列切分不需要通信，因为输入 x 是完整的，每张卡拿到的 W1 列不同，计算出的 hidden 也不同维度。

3.3 行切分（Row-wise Partition）

行切分就有点意思了。还是那个 FFN 层，但这次我们切 W2：

# 标准实现
output = hidden @ W2 + b2  # W2: [d_ff, d_model]

# 行切分：把 W2 按行切成两块
# 卡0 持有 W2[:d_ff//2, :]
# 卡1 持有 W2[d_ff//2:, :]

# 每张卡计算
output_local = hidden_local @ W2_local + b2_local
# 此时 output_local 是 [batch, d_model]

等等，这里有个坑。hidden_local 是列切分后的结果，每张卡只有一半的维度。但 W2_local 是完整行数的一半。所以每张卡算出来的 output_local 都是完整的 [batch, d_model] 维度，但数值不同。

为什么会这样？因为行切分后，每张卡拿到的 W2 行不同，但列相同。所以每张卡算出的 output 是完整维度的，但只包含了部分行的贡献。最后需要做一次 all-reduce 求和。

3.4 Transformer 中的完整切分方案

我们来看一个完整的 Transformer 层怎么切。以 Megatron-LM 的方案为例：

# 伪代码：2 卡张量并行
class TransformerLayerTP:
    def __init__(self, tp_size=2):
        # 列切分：Q, K, V 的权重矩阵
        self.qkv_w = [W_qkv[:, :d//tp_size], W_qkv[:, d//tp_size:]]
        # 行切分：输出投影
        self.out_w = [W_out[:d//tp_size, :], W_out[d//tp_size:, :]]
        # FFN 第一层：列切分
        self.ffn1_w = [W1[:, :d_ff//tp_size], W1[:, d_ff//tp_size:]]
        # FFN 第二层：行切分
        self.ffn2_w = [W2[:d_ff//tp_size, :], W2[d_ff//tp_size:, :]]

    def forward(self, x, rank):
        # 1. 自注意力：列切分 QKV
        qkv_local = x @ self.qkv_w[rank]
        q, k, v = split_qkv(qkv_local)

        # 2. 注意力计算（每张卡独立算自己的头）
        attn_out = attention(q, k, v)

        # 3. 输出投影：行切分，需要 all-reduce
        out_local = attn_out @ self.out_w[rank]
        out = all_reduce(out_local)

        # 4. FFN 第一层：列切分
        hidden_local = gelu(out @ self.ffn1_w[rank])

        # 5. FFN 第二层：行切分，需要 all-reduce
        ffn_out_local = hidden_local @ self.ffn2_w[rank]
        ffn_out = all_reduce(ffn_out_local)

        return ffn_out

3.5 张量并行的通信开销分析

我建议你用这个公式估算通信量：

切分方式	通信操作	通信量	适用场景
列切分	无	0	第一层线性变换
行切分	all-reduce	2 × 输出大小	第二层线性变换
混合切分	2 次 all-reduce	4 × 隐藏层大小	完整 Transformer 层

你想想看，如果隐藏层大小是 4096，batch size 是 4，那一次 all-reduce 也就传 64KB 数据。相比模型参数动辄几百 GB，这点通信开销完全可以接受。

3.6 实现注意事项

嗯，这里有几个实际工程中容易忽略的点：

• 负载均衡：确保每张卡分到的参数数量尽量相等。我见过有人把 d_ff 设成质数，结果切分不均匀，慢卡拖慢整体。
• 数值精度：all-reduce 求和时，不同卡的浮点数精度可能不同。建议用 FP16 时做一次精度对齐。
• 显存碎片：张量并行会创建很多小张量，容易产生显存碎片。我习惯用缓存分配器来管理。

3.7 知识体系图

下面这张图展示了张量并行的核心逻辑：

这张图把整个流程串起来了。从输入开始，经过列切分（无通信），中间激活，再到行切分（需要 all-reduce），最后输出。说白了，张量并行就是通过巧妙的切分策略，把通信次数降到最低。

---