模型并行训练

Model Parallelism for Large Language Models

概述

当模型参数量超过单个GPU的显存容量时，需要使用模型并行技术将模型分布到多个GPU上。 模型并行主要分为张量并行和流水线并行两种方式。

并行策略对比

张量并行 (Tensor Parallelism)

原理

将单个算子（如矩阵乘法）的计算分布到多个GPU上：

• 列并行：按列切分权重矩阵，每个GPU计算部分输出
• 行并行：按行切分权重矩阵，每个GPU计算部分输入

MLP层的张量并行

注意力层的张量并行

将注意力头分布到不同GPU上：

• 每个GPU负责部分注意力头
• Q、K、V投影按列切分
• 输出投影按行切分
• 最后AllReduce合并输出

流水线并行 (Pipeline Parallelism)

原理

将模型的层按顺序分配到不同GPU上，形成流水线：

流水线气泡问题

简单流水线存在GPU空闲时间（气泡）：

• 前向传播时，后面的GPU空闲等待
• 反向传播时，前面的GPU空闲等待
• 气泡比例与流水线深度相关

解决方案：微批次 (Micro-batching)

GPipe和PipeDream等算法通过微批次减少气泡：

• 将大批次分成多个微批次
• 流水执行多个微批次
• 最后同步梯度

3D并行

实际训练超大模型时，通常组合使用多种并行策略：

示例：GPT-3 175B训练

总GPU数：8 × 8 × 128 = 8192 GPUs

代码示例

Megatron-LM张量并行

from megatron.core import tensor_parallel

# 初始化张量并行组
tensor_parallel.initialize_model_parallel(
    tensor_model_parallel_size=4  # 4路张量并行
)

# 并行线性层
class ColumnParallelLinear(nn.Module):
    """按列切分的并行线性层"""
    def __init__(self, in_features, out_features):
        super().__init__()
        self.out_features = out_features
        self.in_features = in_features
        
        # 每个GPU只有部分输出
        self.weight = nn.Parameter(
            torch.empty(out_features // tp_size, in_features)
        )
    
    def forward(self, x):
        # 每个GPU计算部分输出
        output = torch.matmul(x, self.weight.t())
        return output

PyTorch流水线并行

import torch.distributed as dist
from torch.distributed.pipeline.sync import Pipe

# 定义模型分段
model = nn.Sequential(
    nn.Linear(1024, 4096),
    nn.ReLU(),
    nn.Linear(4096, 4096),
    nn.ReLU(),
    nn.Linear(4096, 1024),
)

# 创建流水线并行
pipe = Pipe(model, chunks=8)  # 8个微批次

# 分布式训练
dist.init_process_group(backend='nccl')
pipe.to(device)

显存分析

最佳实践

• 张量并行用于单节点内GPU间通信（高带宽）
• 流水线并行用于跨节点通信（较低带宽）
• 张量并行度通常设为2-8（与注意力头数匹配）
• 使用ZeRO优化器状态减少显存占用
• 考虑激活重计算进一步节省显存