分布式训练（DDP / FSDP / ZeRO）

1. 何时需要分布式

• 单卡显存装不下模型；
• 单卡训练时间过长；
• 需要更大 batch size 稳定训练。

2. DDP 基础

DDP（DistributedDataParallel）是 PyTorch 首选数据并行方案。

启动示例：

torchrun --nproc_per_node=8 train.py --config configs/base.yaml

核心要点：

• 每个进程绑定一张 GPU；
• 使用 DistributedSampler 切分数据；
• 只在 rank0 打日志与保存模型。

3. 最小 DDP 代码骨架

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP


def setup():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

local_rank = setup()
model = MyModel().cuda(local_rank)
model = DDP(model, device_ids=[local_rank])

4. FSDP 与 ZeRO

• FSDP：参数/梯度/优化器状态分片，显存收益大；
• ZeRO（DeepSpeed）：分 stage 分片，生态成熟；
• 两者都需要更严格的 checkpoint 与恢复设计。

5. 常见坑

• 忘记设置 sampler.set_epoch(epoch) 导致 shuffle 固定；
• 多卡 batch size 计算错误（全局 batch = 单卡 batch × 卡数 × 累积步数）；
• all-reduce 开销过大导致扩展性差，需要 profiling。