DeepSpeed分布式训练：大模型训练的工程利器

训练大模型不是加显存就能解决的事。当模型参数量超过单卡容量，分布式训练就是唯一出路。DeepSpeed是微软开源的分布式训练框架。

三大核心技术

ZeRO（零冗余优化器）把模型状态（优化器状态、梯度、参数）分片到多张GPU上。ZeRO-1分片优化器状态，ZeRO-2加梯度分片，ZeRO-3加参数分片。显存节省依次递增。

梯度累积在多步前向传播后才做一次参数更新，等效于更大的batch size，但不需要同时在显存中保存所有样本的梯度。

混合精度训练用FP16或BF16做前向和反向传播，用FP32维护主权重副本。显存占用减半，速度提升2-3倍。

配置文件

{
    "train_batch_size": 32,
    "gradient_accumulation_steps": 4,
    "optimizer": {"type": "AdamW", "params": {"lr": 2e-5}},
    "fp16": {"enabled": true},
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {"device": "cpu"}
    }
}

stage 2是性价比最高的选择。stage 3显存更省但通信开销更大。offload可以把优化器状态卸载到CPU内存，进一步节省显存。

实际训练经验

学习率从2e-5开始，warmup比例设5%。gradient_accumulation_steps根据显存调整，目标是等效batch size在16-64之间。loss曲线出现震荡就降低学习率。

DeepSpeed把大模型训练从"只有大公司能做的事"变成了"有几张好显卡就能做的事"。