AI 棱镜站

主题

张量并行：Megatron-LM 方案 ​

Megatron-LM 张量并行是训练超大模型的核心技术，通过将权重矩阵按维度切分到多个 GPU 上，突破单 GPU 显存限制并利用并行计算加速。

Megatron-LM张量并行

核心思想 ​

将矩阵乘法 Y = XA 按列切分：

A = [A1, A2]  (列切分到2个GPU)
Y = X[A1, A2] = [XA1, XA2]  (各GPU独立计算，无需通信)

对于 MLP 的两层结构，巧妙组合列切分和行切分使得只需一次 AllReduce：

Y = GeLU(XA)B
  = GeLU(X[A1, A2])[B1; B2]  (A列切分, B行切分)
  = [GeLU(XA1), GeLU(XA2)][B1; B2]
  = GeLU(XA1)B1 + GeLU(XA2)B2  (AllReduce求和)

Attention 张量并行 ​

python
# Megatron-LM Attention 张量并行伪代码
class TensorParallelAttention:
    def __init__(self, hidden_dim, num_heads, tp_rank, tp_size):
        self.num_heads_per_partition = num_heads // tp_size
        # QKV 投影：列切分
        self.qkv = ColumnParallelLinear(hidden_dim, 3 * hidden_dim // tp_size)
        # Output 投影：行切分
        self.o_proj = RowParallelLinear(hidden_dim // tp_size, hidden_dim)

    def forward(self, x):
        # 本地 QKV 计算（无通信）
        qkv = self.qkv(x)  # [batch, seq, 3*hidden/tp]
        q, k, v = split_qkv(qkv)

        # 本地多头注意力
        attn_out = scaled_dot_product_attention(q, k, v)

        # Output 投影 + AllReduce
        output = self.o_proj(attn_out)  # 内部执行 AllReduce
        return output

通信分析 ​

每个 Transformer 层需要 2 次 AllReduce：

组件切分方式通信操作数据量
Attention QKV列切分无0
Attention O行切分AllReduceb×s×h
MLP FC1列切分无0
MLP FC2行切分AllReduceb×s×h

通信优化

Megatron-LM 将两次 AllReduce 合并为一次，利用 reduce-scatter + all-gather 分解，与计算重叠执行。这称为 Overlap 通信优化，可减少 30-50% 的通信开销。

显存节省 ​

张量并行的显存节省：

TP 度权重显存梯度显存优化器状态总节省
1100%100%100%0%
250%50%50%50%
425%25%25%75%
812.5%12.5%12.5%87.5%

TP 度选择

TP 度必须整除注意力头数。TP 增加时通信开销线性增长，建议 TP 限制在单节点内（利用 NVLink 高带宽）。跨节点应使用 PP 或 DP。

相关资源 ​

最近更新