为什么需要张量并行？

当模型参数量超过单GPU显存容量时，模型并行成为必需。张量并行（Tensor Parallelism, TP）是最常用的模型并行方式——将单个层的权重矩阵切分到多个GPU上，每个GPU计算一部分结果，然后通过通信合并。

NVIDIA的Megatron-LM提出了张量并行的标准方案，成为2026年大模型训练的基础设施。

张量并行的核心原理

线性层的切分

一个标准线性层 Y = XW，其中 W ∈ ℝ^{d×h}。将W切分为两个子矩阵有两种方式：

列切分（Column Parallel）

将W按列切分：W = [W₁, W₂]，每个GPU持有 W_i ∈ ℝ^{d×h/N}

GPU 1: Y₁ = XW₁  ∈ ℝ^{n×h/N}
GPU 2: Y₂ = XW₂  ∈ ℝ^{n×h/N}
Y = [Y₁, Y₂]     ∈ ℝ^{n×h}

列切分后，每个GPU独立计算Y的一部分列，最终通过All-Gather合并。但更聪明的方式是——不需要显式合并，因为后续操作可以利用分块结构。

行切分（Row Parallel）

将W按行切分：W = [W₁; W₂]，每个GPU持有 W_i ∈ ℝ^{d/N×h}

GPU 1: Y₁ = X₁W₁  (X₁是X的前d/N列)
GPU 2: Y₂ = X₂W₂  (X₂是X的后d/N列)
Y = Y₁ + Y₂        (All-Reduce)

行切分后需要All-Reduce求和。

Megatron的FFN并行

Megatron-LM将FFN的两个线性层巧妙地组合列切分和行切分，使得整个FFN只需要一次All-Reduce：

class MegatronFFN(nn.Module):
    def __init__(self, d_model, d_ff, n_gpus, rank):
        super().__init__()
        # 第一个线性层：列切分
        self.w1 = nn.Linear(d_model, d_ff // n_gpus)  # 每GPU持有d_ff/n_gpus列
        # 第二个线性层：行切分
        self.w2 = nn.Linear(d_ff // n_gpus, d_model)  # 每GPU持有对应行
        self.rank = rank
        self.n_gpus = n_gpus
    
    def forward(self, x):
        # x在所有GPU上相同
        h = F.silu(self.w1(x))  # 列切分：各GPU独立计算，无需通信
        # h的形状: [batch, seq, d_ff/n_gpus]
        
        out = self.w2(h)  # 行切分：各GPU独立计算
        # out的形状: [batch, seq, d_model]，但只是部分和
        
        out = all_reduce(out, op=SUM)  # All-Reduce求和
        return out

关键洞察：第一个线性层的列切分输出直接作为第二个线性层行切分的输入，中间不需要任何通信。只有最终结果需要All-Reduce。

注意力的并行

多头注意力天然适合张量并行——将不同的注意力头分配到不同GPU：

class MegatronAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_gpus, rank):
        super().__init__()
        self.local_n_heads = n_heads // n_gpus
        self.head_dim = d_model // n_heads
        
        # QKV投影：列切分（按头切分）
        self.wq = nn.Linear(d_model, self.local_n_heads * self.head_dim)
        self.wk = nn.Linear(d_model, self.local_n_heads * self.head_dim)
        self.wv = nn.Linear(d_model, self.local_n_heads * self.head_dim)
        
        # 输出投影：行切分
        self.wo = nn.Linear(self.local_n_heads * self.head_dim, d_model)
    
    def forward(self, x):
        # 各GPU独立计算自己负责的注意力头
        Q = self.wq(x)  # [batch, seq, local_n_heads * head_dim]
        K = self.wk(x)
        V = self.wv(x)
        
        # 重排为多头格式
        Q = Q.view(batch, seq, self.local_n_heads, self.head_dim).transpose(1, 2)
        K = K.view(batch, seq, self.local_n_heads, self.head_dim).transpose(1, 2)
        V = V.view(batch, seq, self.local_n_heads, self.head_dim).transpose(1, 2)
        
        # 本地注意力计算（无需通信）
        attn = flash_attention(Q, K, V)
        attn = attn.transpose(1, 2).contiguous().view(batch, seq, -1)
        
        # 输出投影 + All-Reduce
        out = self.wo(attn)
        out = all_reduce(out, op=SUM)
        return out

每个GPU独立计算自己负责的注意力头，最后通过All-Reduce合并输出。整个注意力层只需要一次All-Reduce。

通信开销分析

每层通信次数

一个Transformer层包含一个注意力和一个FFN，每个都需要一次All-Reduce。因此每层2次All-Reduce。

通信量

每次All-Reduce的通信量为 2 × batch × seq × d_model × (N-1)/N × 4 bytes（FP32），其中N是GPU数量。通信量与序列长度和batch size线性相关。

通信-计算比

计算量为 O(batch × seq × d_model² / N)，通信量为 O(batch × seq × d_model)。当 d_model » N 时（大模型的典型情况），通信-计算比较低，张量并行效率高。

与流水线并行的配合

张量并行在层内切分，流水线并行在层间切分。两者正交，可以组合使用：

GPU拓扑（8卡，TP=4, PP=2）:
  Pipeline Stage 1: GPU 0, 1, 2, 3 (张量并行)
  Pipeline Stage 2: GPU 4, 5, 6, 7 (张量并行)

数据流:
  输入 → [GPU 0-3 TP] → 中间输出 → [GPU 4-7 TP] → 输出

典型的3D并行配置：

  • 数据并行（DP）：不同GPU组处理不同batch
  • 张量并行（TP）：GPU组内切分层权重
  • 流水线并行（PP）：GPU组间切分模型层
# Megatron-LM 3D并行配置示例
parallel_config = {
    "tensor_model_parallel_size": 4,   # TP=4
    "pipeline_model_parallel_size": 2, # PP=2
    "data_parallel_size": 8,          # DP=8 (总共4*2*8=64 GPU)
}

工程实现细节

通信优化：All-Reduce融合

将同一层的注意力All-Reduce和FFN的All-Reduce合并为一次操作，减少通信延迟：

class TransformerLayerWithFusedAR(nn.Module):
    def forward(self, x):
        # 注意力和FFN的前向可以部分并行
        attn_out = self.attention_partial(x)  # 不做All-Reduce
        ffn_out = self.ffn_partial(attn_out)  # 不做All-Reduce
        
        # 一次All-Reduce同时合并注意力+FFN的结果
        combined = all_reduce(attn_out + ffn_out, op=SUM)
        return combined

KV Cache的张量并行

推理时，KV Cache也需要按头切分存储。每个GPU只存储自己负责的头的KV Cache，注意力计算完全本地完成。

序列并行

Megatron-LM的序列并行将LayerNorm和Dropout的激活值按序列维度切分，进一步减少显存：

class SequenceParallelLayerNorm(nn.Module):
    def forward(self, x):
        # x按序列维度切分
        # All-Gather恢复完整序列
        x_full = all_gather(x, dim=1)  # 沿seq维度
        x_norm = layer_norm(x_full)
        # 重新切分
        return scatter(x_norm, dim=1)

序列并行将LayerNorm的激活值从 O(batch × seq × d) 减少到 O(batch × seq/N × d)，代价是额外的All-Gather和Reduce-Scatter通信。

2026年的张量并行发展

TP+SP融合

张量并行和序列并行不再作为独立选项，而是融合为统一的并行策略。在注意力计算前做All-Gather（恢复完整序列），计算后做Reduce-Scatter（切分序列）。这减少了LayerNorm和Dropout的显存占用，同时保持了注意力的计算效率。

异构张量并行

不同层使用不同的TP度——注意力层用TP=8（计算密集），LayerNorm用TP=1（计算轻量，不值得切分）。这需要精细的通信调度。

与MoE的集成

在MoE模型中，张量并行同时处理专家内的权重切分和专家间的路由通信。All-to-All通信（专家路由）和All-Reduce（张量并行）需要精心调度以避免通信拥塞。

实践建议

  1. TP度选择：通常为2/4/8，受限于节点内GPU数。跨节点TP延迟太高
  2. TP与PP的平衡：优先增加TP（层内通信延迟低），TP达上限后增加PP
  3. 梯度同步：TP的反向传播中，梯度也需要All-Reduce
  4. Checkpoint兼容：梯度检查点与张量并行完全兼容，但需要注意checkpoint的输入需要是完整副本

结语

张量并行是大模型分布式训练的基石。Megatron-LM的列-行切分组合将每层通信减少到最少，与流水线并行和数据并行组合后可以高效训练千亿参数模型。随着序列并行和异构并行的发展，张量并行的效率仍在持续提升。

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞，认知差异创造无限可能。