deepspeed

DeepSpeed 分布式训练优化

补充说明

DeepSpeed 提供 ZeRO 显存优化、混合精度、Offload 和推理加速。通过 JSON 配置启用不同 ZeRO stage，适合单机多卡和多机大模型训练。

常用示例

安装：

pip install deepspeed
ds_report

ZeRO-2 配置：

cat > ds_config.json <<'EOF'
{
  "train_batch_size": 32,
  "gradient_accumulation_steps": 4,
  "fp16": {"enabled": true},
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "overlap_comm": true
  }
}
EOF

启动训练：

deepspeed --num_gpus=4 train.py --deepspeed ds_config.json

ZeRO-3 + CPU Offload：

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {"device": "cpu"},
    "offload_param": {"device": "cpu"}
  }
}

多机启动：

deepspeed --hostfile hosts --master_addr 10.0.0.10 \
  --master_port 29500 train.py --deepspeed ds_config.json

易错点与排错

• train_batch_size 与实际 batch、梯度累积、世界大小不匹配，启动时报错。
• ZeRO-3 下没有使用 deepspeed.zero.Init 或配置正确，模型加载即 OOM。
• CPU Offload 需要足够内存和快速 NVMe，否则训练极慢。
• 混合精度配置与模型 dtype 冲突，出现 NaN。
• 多机 hostfile 的槽位数和 GPU 数不一致。
• 只在一张卡上保存 checkpoint，恢复时缺少优化器状态。
• DeepSpeed 版本与 PyTorch/CUDA 不兼容。
• 没有开启通信 overlap，扩展效率低。