5.3 分布式训练

训练大模型需要海量算力。GPT-3 拥有 1750 亿参数，训练一次消耗约 355 个 GPU 年的计算量，仅凭单卡根本无法完成。分布式训练是将训练任务拆分到多台机器、多张 GPU 上并行执行的技术体系，是大模型时代的必备工程能力。

5.3.1 分布式训练的三大范式

数据并行（Data Parallelism, DP）

数据并行是最直观的并行方式：每张 GPU 持有完整的模型副本，将训练数据的不同部分（mini-batch）分配到各 GPU 上独立计算梯度，然后聚合（AllReduce）梯度并同步更新参数。

GPU 0: 完整模型 + batch 0 → grad 0
GPU 1: 完整模型 + batch 1 → grad 1
GPU 2: 完整模型 + batch 2 → grad 2
GPU 3: 完整模型 + batch 3 → grad 3
                ↓ AllReduce（梯度同步）
            平均梯度 → 同步更新参数
​

优点：实现简单，代码改动最小。缺点：每张 GPU 都存储完整模型，显存成为瓶颈。当模型大到单卡放不下时，数据并行无能为力。

模型并行（Model Parallelism, MP）

模型并行将模型的不同部分分布到不同 GPU 上：

张量并行（Tensor Parallelism）：将模型层的参数矩阵按列或行切分，分配到不同 GPU。例如，一个 4096→4096 的线性层可以切分为两个 4096→2048 的子层，分别在不同 GPU 上计算。

流水线并行（Pipeline Parallelism）：将 Transformer 的不同层分配到不同 GPU。GPU 0 处理第 1-6 层，GPU 1 处理第 7-12 层，通过流水线方式执行，提高 GPU 利用率。

混合并行

现代大模型训练通常组合使用多种并行策略。Megatron-LM（NVIDIA 开源框架）将张量并行与流水线并行结合，DeepSpeed 则在此基础上加入了数据并行的优化。实际生产中的大模型训练方案通常是混合并行。

5.3.2 PyTorch 分布式数据并行（DDP）

基本原理

PyTorch 的 DistributedDataParallel（DDP）是对数据并行的优化实现。相比早期的 DataParallel（DP），DDP 的核心改进是：每个进程独立进行反向传播，然后在后台异步同步梯度，避免了 DP 中因梯度收集导致的单点瓶颈。

DDP 的梯度同步使用 Ring AllReduce 算法，时间复杂度为 O(n)，其中 n 是梯度大小，与 GPU 数量无关。

DDP 代码示例

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def setup(rank, world_size):
    """初始化分布式环境"""
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

def train_ddp(rank, world_size):
    setup(rank, world_size)
    
    # 创建模型并移到当前 GPU
    model = nn.Sequential(
        nn.Linear(784, 4096),
        nn.ReLU(),
        nn.Linear(4096, 4096),
        nn.ReLU(),
        nn.Linear(4096, 10),
    ).to(rank)
    
    # 用 DDP 包装模型
    ddp_model = DDP(model, device_ids=[rank])
    
    # 分布式采样器（确保不同 GPU 处理不同数据）
    sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True,
    )
    dataloader = DataLoader(dataset, batch_size=64, sampler=sampler)
    
    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=1e-3)
    
    for epoch in range(10):
        sampler.set_epoch(epoch)  # 确保 shuffle 在每轮不同
        
        for batch in dataloader:
            inputs, labels = batch
            inputs, labels = inputs.to(rank), labels.to(rank)
            
            optimizer.zero_grad()
            outputs = ddp_model(inputs)
            loss = nn.CrossEntropyLoss()(outputs, labels)
            loss.backward()  # 梯度在 backward 中自动同步
            optimizer.step()
        
        if rank == 0:
            print(f"Epoch {epoch}, Loss: {loss.item():.4f}")
    
    cleanup()

# 启动方式：torchrun --nproc_per_node=4 train.py
​

DDP 性能调优

梯度累积：当显存不足以使用大 batch size 时，通过多步前向传播累积梯度，等效增大 batch size：

gradient_accumulation_steps = 4
optimizer.zero_grad()

for i, batch in enumerate(dataloader):
    loss = model(batch) / gradient_accumulation_steps
    loss.backward()
    
    if (i + 1) % gradient_accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()
​

混合精度训练（AMP）：使用 FP16/BF16 进行前向和反向计算，FP32 维护主权重，兼顾速度和精度：

scaler = torch.cuda.amp.GradScaler()

for batch in dataloader:
    optimizer.zero_grad()
    
    with torch.cuda.amp.autocast():
        outputs = model(batch)
        loss = criterion(outputs, labels)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
​

5.3.3 DeepSpeed ZeRO

ZeRO 三阶段

DeepSpeed 的核心贡献是 ZeRO（Zero Redundancy Optimizer），通过消除冗余来打破数据并行的显存瓶颈：

ZeRO-Stage 1：优化器状态分片。将 AdamW 的动量（m）和方差（v）按参数切分到不同 GPU，每张 GPU 只存储 1/N 的优化器状态。显存减少约 4x。

ZeRO-Stage 2：梯度分片。在 Stage 1 基础上，梯度也按参数切分存储，每张 GPU 只存储 1/N 的梯度。显存进一步减少。

ZeRO-Stage 3：参数分片。模型参数本身也按切分存储。任何一层在计算时，从各 GPU 收集该层的完整参数，计算完后再释放。显存减少约 N 倍（N 为 GPU 数），可以在有限 GPU 上训练超大模型。

标准 DP:  GPU 0 = 完整模型 + 完整优化器 + 完整梯度
ZeRO-1:   GPU 0 = 完整模型 + 完整梯度 + 1/N 优化器
ZeRO-2:   GPU 0 = 完整模型 + 1/N 梯度 + 1/N 优化器
ZeRO-3:   GPU 0 = 1/N 模型 + 1/N 梯度 + 1/N 优化器
​

DeepSpeed 配置

# deepspeed_config.json
{
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 3e-5,
            "betas": [0.9, 0.95],
            "eps": 1e-8
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 3e-5,
            "warmup_num_steps": 100
        }
    },
    "fp16": {
        "enabled": "auto"
    },
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": "auto",
        "stage3_max_reuse_distance": "auto",
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "gradient_clipping": 1.0,
    "prescale_gradients": false,
    "wall_clock_breakdown": false
}
​

DeepSpeed 与 HuggingFace 集成

from transformers import AutoModelForCausalLM, TrainingArguments, Trainer

model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m")

training_args = TrainingArguments(
    output_dir="./ds_output",
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    num_train_epochs=3,
    learning_rate=3e-5,
    fp16=True,
    
    # DeepSpeed 配置
    deepspeed="./deepspeed_config.json",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)

trainer.train()
​

只需一个配置文件和一行参数，就能让 HuggingFace Trainer 使用 DeepSpeed 进行分布式训练。

5.3.4 FSDP（Fully Sharded Data Parallel）

FSDP 概述

FSDP 是 PyTorch 2.0 原生引入的分布式训练方案，与 ZeRO-3 类似，将模型的参数、梯度和优化器状态分片到所有 GPU 上，但有以下优势：

• PyTorch 原生集成：无需额外安装第三方库
• 更灵活的单元：以模块为单位进行分片，可以精细控制分片粒度
• 与 PyTorch 编译器兼容：可与 torch.compile 配合使用

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from torch.distributed.fsdp import ShardingStrategy

# 分片策略
sharding_strategy = ShardingStrategy.FULL_SHARD  # 等同于 ZeRO-3

# 定义需要分片的 Transformer 层
wrap_policy = transformer_auto_wrap_policy(
    transformer_layer_cls={transformers.BertEncoderLayer},
)

# 用 FSDP 包装模型
fsdp_model = FSDP(
    model,
    sharding_strategy=sharding_strategy,
    auto_wrap_policy=wrap_policy,
    device_id=torch.cuda.current_device(),
)

# 训练循环与普通 PyTorch 完全一致
for batch in dataloader:
    loss = fsdp_model(batch)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
​

5.3.5 工程实践建议

选择合适的并行策略

常见问题与解决方案

OOM（显存不足）：优先使用梯度累积而非增大 batch size；开启梯度检查点（gradient checkpointing），用计算换显存；升级到 ZeRO-3/FSDP。

通信开销过大：增大 per_device_train_batch_size 减少通信频率；使用 NCCL 环境变量优化带宽；选择 ZeRO-2 而非 ZeRO-3（如果显存允许）。

训练不稳定：使用 BF16 替代 FP16（更大的数值范围）；降低学习率；增加 warmup 步数；开启梯度裁剪。

GPU 利用率低：检查数据加载是否成为瓶颈（使用 DataLoader 的 num_workers 和 pin_memory）；确保 batch size 足够大；使用 torch.compile 编译模型。

分布式训练是大模型工程的基石。从 DDP 到 DeepSpeed 再到 FSDP，工具链日趋成熟。掌握这些技术，你就能驾驭任意规模的 Transformer 模型训练任务。