DeepSpeed ZeRO-3 全参数分片实践
ZeRO-3 是 DeepSpeed 最激进的显存优化方案，将模型参数、梯度和优化器状态全部分片。本文介绍 ZeRO-3 的生产配置和性能调优实践。
完整配置
python
# ZeRO-3 生产配置
ds_config = {
"zero_optimization": {
"stage": 3,
# 参数分片配置
"overlap_comm": True,
"contiguous_gradients": True,
"reduce_bucket_size": 5e8, # 梯度归约桶大小
"stage3_prefetch_bucket_size": 5e8, # 参数预取桶大小
"stage3_param_persistence_threshold": 1e5, # 小参数不分片
"stage3_max_live_parameters": 1e9, # 同时在 GPU 上的最大参数量
"stage3_max_reuse_distance": 1e9, # 参数复用距离
# CPU Offload（可选）
"offload_optimizer": {
"device": "cpu",
"pin_memory": True,
},
"offload_param": {
"device": "cpu",
"pin_memory": True,
},
# 通信优化
"sub_group_size": 1e8,
"round_robin_gradients": True,
},
"gradient_accumulation_steps": 4,
"train_batch_size": 128,
"fp16": {
"enabled": True,
"loss_scale": 0,
"initial_scale_power": 16,
},
}参数预取优化
ZeRO-3 在前向传播时需要 AllGather 参数，预取策略决定了通信与计算的重叠程度：
python
# 启动训练
deepspeed --num_gpus=8 train.py \
--deepspeed ds_config_zero3.json \
--model_name meta-llama/Llama-2-7b-hf \
--per_device_train_batch_size 4 \
--gradient_accumulation_steps 4关键调优参数
| 参数 | 作用 | 推荐值 |
|---|---|---|
reduce_bucket_size | 梯度归约粒度 | 5e8 - 1e9 |
stage3_prefetch_bucket_size | 参数预取粒度 | 5e8 - 1e9 |
stage3_max_live_parameters | GPU 上活跃参数上限 | 1e9 - 2e9 |
stage3_param_persistence_threshold | 不分片的小参数阈值 | 1e5 |
CPU Offload 实践
ZeRO-3 + CPU Offload 可以在有限 GPU 上训练超大模型：
python
# ZeRO-3 + CPU Offload 训练 70B 模型（4×A100-40GB）
ds_config_offload = {
"zero_optimization": {
"stage": 3,
"offload_optimizer": {"device": "cpu", "pin_memory": True},
"offload_param": {"device": "cpu", "pin_memory": True},
},
}| 配置 | GPU 显存 | CPU 内存 | 训练速度 |
|---|---|---|---|
| ZeRO-3 (8GPU) | 40GB × 8 | 64GB | 100% |
| ZeRO-3 + Offload (4GPU) | 40GB × 4 | 512GB | ~35% |
| ZeRO-3 + Offload (2GPU) | 80GB × 2 | 512GB | ~20% |
Offload 优化
- 使用
pin_memory=True加速 CPU-GPU 数据传输 - NVMe Offload 可将优化器状态卸载到 SSD，进一步降低 CPU 内存需求
- Offload 场景建议使用
gradient_accumulation_steps ≥ 8隐藏 CPU-GPU 传输延迟
Offload 性能代价
CPU Offload 会显著降低训练速度（约 3-5 倍）。仅在 GPU 显存不足时使用。如有条件，优先增加 GPU 数量而非使用 Offload。