Ainexis LogoAinexis.cn

ZeRO优化

Zero Redundancy Optimizer

概述

ZeRO（Zero Redundancy Optimizer）是Microsoft提出的显存优化技术， 通过消除数据并行中的冗余状态，大幅降低训练大模型所需的显存， 使得单GPU可以训练更大的模型。

核心思想：数据并行中每个GPU保存完整的模型、梯度和优化器状态是冗余的，可以切分。

显存分析

训练时显存占用组成

组件大小（FP16）占比（示例）
模型参数2ψ~2GB (7B)
梯度2ψ~2GB (7B)
优化器状态 (Adam)12ψ (FP32参数+momentum+variance)~12GB (7B)
激活值取决于batch和序列长度变化大

ψ为参数量，7B模型使用Adam优化器约需16GB显存（不含激活值）。

ZeRO三个阶段

ZeRO-1: 优化器状态切分

将优化器状态切分到不同GPU：

  • 每个GPU只保存1/N的优化器状态
  • 显存节省：4x（N个GPU时）
  • 通信量：不变

ZeRO-2: + 梯度切分

在ZeRO-1基础上切分梯度：

  • 每个GPU只保存对应优化器状态部分的梯度
  • 显存节省：8x
  • 通信量：不变

ZeRO-3: + 参数切分

在ZeRO-2基础上切分模型参数：

  • 每个GPU只保存1/N的模型参数
  • 前向/反向时通过AllGather获取需要的参数
  • 显存节省：N x（与GPU数量成正比）
  • 通信量：增加50%
策略优化器状态梯度参数节省
标准DP完整复制完整复制完整复制1x
ZeRO-1切分完整复制完整复制4x
ZeRO-2切分切分完整复制8x
ZeRO-3切分切分切分Nx

代码示例

DeepSpeed ZeRO

# deepspeed_config.json
{
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "reduce_scatter": true,
        "overlap_comm": true,
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": 1,
    "train_micro_batch_size_per_gpu": 4
}

# 训练命令
deepspeed train.py --deepspeed_config deepspeed_config.json

Hugging Face + ZeRO

# 使用Hugging Face Accelerate
from accelerate import Accelerator

accelerator = Accelerator()

# 自动使用ZeRO配置
model, optimizer, train_loader = accelerator.prepare(
    model, optimizer, train_loader
)

# 或使用配置文件
# accelerate config -- 选择ZeRO stage

# 训练
for batch in train_loader:
    outputs = model(**batch)
    loss = outputs.loss
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

ZeRO-Offload

将部分状态卸载到CPU内存，进一步减少GPU显存：

  • Optimizer Offload：优化器状态放CPU
  • Param Offload：参数放CPU
  • 利用CPU内存大容量优势
  • 适合显存有限的场景
{
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        }
    }
}

显存节省示例

模型标准DPZeRO-1ZeRO-2ZeRO-3 (64 GPUs)
7B16GB4GB2GB0.25GB
13B32GB8GB4GB0.5GB
30B64GB+16GB8GB1GB
175B无法训练无法训练无法训练2.75GB

ZeRO-Infinity

ZeRO-3的扩展，利用NVMe SSD进一步扩展存储：

  • 参数、梯度、优化器状态可卸载到NVMe
  • 支持训练万亿参数模型
  • CPU和GPU内存统一管理
  • 适合超大模型训练

最佳实践

  • 显存足够时使用ZeRO-1，通信开销最小
  • 显存紧张时使用ZeRO-2，平衡性能和显存
  • 超大模型使用ZeRO-3 + Offload
  • 配合梯度检查点进一步节省显存
  • 使用overlap_comm选项隐藏通信开销
  • ZeRO-3时注意CPU内存使用量

参考资料

  • ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (Rajbhandari et al., 2020)
  • ZeRO-Offload: Democratizing Billion-Scale Model Training (Ren et al., 2021)
  • DeepSpeed Official Documentation
----