【DeepSpeed】零冗余优化器（Zero Redundancy Optimizer, ZeRO）

零冗余优化器（Zero Redundancy Optimizer, ZeRO）是 DeepSpeed 提供的一种分布式训练优化技术，旨在通过分片模型参数、梯度和优化器状态，显著降低显存占用，同时支持超大规模模型训练。ZeRO 特别适合显存受限或多 GPU/多节点环境。以下是 ZeRO 的原理、实现方式及使用方法。

---

1. ZeRO 的核心原理

• 背景：在传统数据并行训练中，每个 GPU 持有完整的模型参数、梯度和优化器状态副本，导致显存冗余，尤其在大模型（如 GPT、LLaMA）中显存占用极高。
• 核心思想：ZeRO 将模型参数、梯度和优化器状态分片（partition）到多个设备上，消除冗余存储，仅在需要时通过通信收集所需数据。
• 分片级别：
  • ZeRO Stage 1：分片优化器状态（如 Adam 的动量和方差）。
  • ZeRO Stage 2：分片优化器状态和梯度。
  • ZeRO Stage 3：分片优化器状态、梯度和模型参数。
• 优势：
  • 显存占用随分片数（GPU 数量）线性下降，支持超大模型（如千亿参数）。
  • 通信开销优化，接近数据并行的效率。
  • 可与流水线并行、激活检查点等技术结合。
• 权衡：
  • 增加通信开销（通过高效的 all-gather 和 reduce-scatter 优化）。
  • 配置复杂性稍高，需根据硬件和模型调整。

---

2. ZeRO 的三个阶段

阶段	分片内容	显存节省	通信开销	适用场景
Stage 1	优化器状态	每个 GPU 仅存 1/N 的优化器状态	最低	小规模模型，显存稍紧张
Stage 2	优化器状态 + 梯度	每个 GPU 仅存 1/N 的优化器状态和梯度	中等	中大规模模型，显存受限
Stage 3	优化器状态 + 梯度 + 参数	每个 GPU 仅存 1/N 的参数、梯度和优化器状态	较高	超大规模模型（如 GPT-3、LLaMA）

• 内存节省公式（近似）：
  • 单卡显存占用 ≈ (模型参数 + 梯度 + 优化器状态) / N（N 为 GPU 数量）。
  • 例如，Stage 3 下，一个 10 亿参数模型（FP16 参数约 2GB，优化器状态约 8GB）在 8 个 GPU 上每卡仅需约 1.25GB。

---

3. ZeRO 的实现步骤

步骤 1：安装 DeepSpeed

确保安装 DeepSpeed 和 PyTorch：

pip install deepspeed

步骤 2：配置 DeepSpeed

在 DeepSpeed 的 JSON 配置文件（如 ds_config.json）中启用 ZeRO：

{
  "train_micro_batch_size_per_gpu": 16,
  "gradient_accumulation_steps": 4,
  "zero_optimization": {
    "stage": 2,  // 选择 ZeRO 阶段：1, 2 或 3
    "allgather_partitions": true,
    "reduce_scatter": true,
    "contiguous_gradients": true,
    "overlap_comm": true,
    "offload_optimizer": {
      "device": "cpu",  // 可选：卸载优化器状态到 CPU
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",  // 可选：卸载参数到 CPU（仅 Stage 3）
      "pin_memory": true
    }
  },
  "fp16": {
    "enabled": true,
    "loss_scale": 0
  },
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.001
    }
  }
}

• stage：选择 ZeRO 阶段（1、2 或 3）。
• allgather_partitions：启用参数分片收集。
• reduce_scatter：优化梯度分片聚合。
• contiguous_gradients：优化梯度内存分配，减少碎片。
• overlap_comm：通信与计算重叠，降低通信开销。
• offload_optimizer：将优化器状态卸载到 CPU/NVMe，进一步节省显存。
• offload_param：将参数卸载到 CPU/NVMe（仅 Stage 3）。

步骤 3：代码适配

ZeRO 无需大幅修改模型代码，只需通过 DeepSpeed 初始化模型。例如：

import deepspeed
import torch
from torch import nn

# 示例模型
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(512, 512) for _ in range(8)])

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

# 初始化模型
model = SimpleModel()

# 初始化 DeepSpeed
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config="ds_config.json"
)

# 训练循环
for data, target in dataloader:
    loss = model_engine(data, target)
    model_engine.backward(loss)
    model_engine.step()

• DeepSpeed 自动处理分片和通信，model_engine 封装了 ZeRO 逻辑。

步骤 4：运行训练

使用 DeepSpeed 命令行启动分布式训练：

deepspeed --num_gpus 8 train.py --deepspeed_config ds_config.json

---

4. 性能优化与调优

• 选择合适的 ZeRO 阶段：
  • Stage 1：显存节省有限，适合小模型或 GPU 数量较少。
  • Stage 2：平衡显存和通信，适合中大模型。
  • Stage 3：最大化显存节省，适合超大模型，但通信开销较高。
• 卸载（Offload）：
  • 启用 offload_optimizer 和 offload_param（Stage 3）将数据卸载到 CPU/NVMe，适合低显存设备。
  • 使用 pin_memory 加速 CPU-GPU 数据传输。
  • 若有 NVMe 存储，可配置 device: "nvme"。
• 通信优化：
  • 启用 overlap_comm 重叠通信和计算。
  • 使用高效通信后端（如 NCCL），确保高带宽互联（如 NVLink、InfiniBand）。
• 批次大小：
  • 调整 train_micro_batch_size_per_gpu 和 gradient_accumulation_steps，利用节省的显存增加批次大小，提升吞吐量。
• 结合其他技术：
  • 激活检查点：进一步减少激活内存。
  • 流水线并行：Stage 3 与流水线并行结合，适合千亿参数模型。
  • 混合精度：FP16/BF16 降低参数和梯度内存。

---

5. 与 Hugging Face 集成

ZeRO 与 Hugging Face Transformers 无缝集成。例如：

from transformers import AutoModelForCausalLM, Trainer, TrainingArguments

model = AutoModelForCausalLM.from_pretrained("gpt2")
training_args = TrainingArguments(
    output_dir="./output",
    deepspeed="ds_config.json",  // 引用包含 ZeRO 配置的文件
    per_device_train_batch_size=16
)

trainer = Trainer(
    model=model,
    args=training_args
)
trainer.train()

• 配置 ds_config.json 启用 ZeRO，Hugging Face 会自动应用。

---

6. 性能监控与调试

• 监控：
  • 使用 nvidia-smi 检查显存占用，确保分片生效。
  • 通过 deepspeed.utils.logging 或 torch.profiler 分析通信和计算效率。
• 调试：
  • 单 GPU 测试：设置 --num_gpus 1 验证代码正确性（Stage 1/2 仍有效，Stage 3 需多 GPU）。
  • 日志：启用 deepspeed --log_level DEBUG 查看分片和通信详情。
  • 检查卸载：确保 CPU/NVMe 容量和带宽足够。

---

7. 效果示例

• 场景：训练一个 10 亿参数模型（FP16 参数 2GB，优化器状态 8GB，梯度 2GB），8 个 16GB GPU。
  • 传统数据并行：每卡需 12GB 显存，溢出。
  • ZeRO Stage 2：每卡约 4GB（优化器状态和梯度分片），可运行。
  • ZeRO Stage 3 + Offload：每卡约 1.5GB（参数也分片，优化器卸载到 CPU），支持更大批次。

---

8. 注意事项

• 硬件要求：Stage 3 需要多 GPU 和高效通信（NVLink 或 InfiniBand）。
• 通信开销：Stage 3 通信密集，需优化网络拓扑。
• 卸载性能：CPU/NVMe 卸载可能受限于 PCIe 带宽，需测试实际效果。
• 版本兼容性：确保 DeepSpeed、PyTorch 和 CUDA 版本匹配。

---

9. 资源

• 官方文档：https://www.deepspeed.ai/docs/zero/
• 示例代码：DeepSpeed GitHub 的 examples/ 目录提供 ZeRO 训练示例（如 Megatron-LM）。
• 论文：《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》。