详解 ZeRO-Offload 机制：为什么把显存扔给内存和 SSD 却能跑通百亿参数模型？

为什么 VRAM 成为大模型训练的瓶颈？

在训练参数量巨大的模型（如 Llama-2 70B）时，即使是顶级的GPU（如A100或H100）也面临着显存（VRAM）不足的问题。一个完整的模型训练状态需要存储以下核心组件：

1. 模型参数 (P)：权重矩阵本身。
2. 梯度 (G)：用于更新参数。
3. 优化器状态 (O)：例如使用 AdamW 时，需要存储动量($m$)和方差($v$)，这通常是参数量的2到4倍。
4. 激活内存 (A)：前向传播的中间结果，用于反向传播。

对于一个FP16精度的1750亿参数模型，仅参数、梯度和优化器状态就需要大约 3 * 175B * 2 字节 ≈ 1.05 TB 的存储空间，这远远超出了单卡甚至多卡互联的VRAM容量。

DeepSpeed ZeRO 与 ZeRO-Offload 的诞生

为了解决 VRAM 限制，微软提出了 ZeRO (Zero Redundancy Optimizer) 技术，该技术通过数据并行的方式，将模型的各种状态进行分片（Sharding），从而消除冗余。

ZeRO 包含三个阶段：

• Stage 1 (ZeRO-1)：只对优化器状态 (O) 进行分片。
• Stage 2 (ZeRO-2)：对优化器状态 (O) 和梯度 (G) 进行分片。
• Stage 3 (ZeRO-3)：对优化器状态 (O)、梯度 (G) 和模型参数 (P) 全部进行分片。这是训练万亿参数模型的关键。

ZeRO-Offload：将显存扔给内存和 SSD

尽管 ZeRO-3 已经极大地减少了 VRAM 占用，但对于某些组件，尤其是那些体积巨大但访问频率较低的组件，我们还有进一步优化的空间。ZeRO-Offload 就是将这些组件从高速但稀缺的 VRAM 转移到低速但充裕的 CPU 内存 (RAM)，甚至 NVMe/SSD 硬盘上。

核心机制：优化器状态卸载

ZeRO-Offload 最有效的应用是卸载优化器状态 (O)。

1. 为什么是优化器状态？ 优化器状态通常占据整个模型状态存储的最大比例（约2/3）。更重要的是，它们只在反向传播结束后的“参数更新”阶段才需要被访问。
2. 工作流程：
   • 在训练的前向和反向传播过程中，优化器状态存储在 CPU RAM 中。
   • 当需要进行参数更新时，DeepSpeed 会智能地将优化器状态的必要分片从 CPU RAM 传输回 GPU VRAM。
   • GPU 完成参数更新计算。
   • 更新后的优化器状态再次被传回 CPU RAM 存储。

为什么“慢”反而能跑通？

将数据从 VRAM 传输到 CPU RAM/SSD 必然会引入延迟。然而，这里的关键在于权衡：

• 不卸载： VRAM 不足，模型根本无法加载和启动训练。
• 卸载： 虽然每次迭代（Step）都会增加数据传输的延迟（CPU <-> GPU），但省下的 VRAM 可以用于存储更大的模型参数分片、更大的激活内存，甚至使用更大的批次大小（Batch Size）。

对于超大规模模型，能够“跑起来”比追求极致的速度更重要。ZeRO-Offload 用可接受的I/O延迟，换取了突破 VRAM 限制的能力。

实用操作：如何配置 ZeRO-Offload

启用 ZeRO-Offload 非常简单，只需要在 DeepSpeed 的配置文件（通常是 deepspeed_config.json）中进行指定。

以下是一个典型的 DeepSpeed 配置，启用了 ZeRO Stage 3 并将优化器状态卸载到 CPU：

{
  "train_batch_size": "auto",
  "gradient_accumulation_steps": 1,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 1e-5
    }
  },
  "fp16": {
    "enabled": true,
    "initial_scale_power": 15
  },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": 50000000,
    "stage3_prefetch_size": 5,
    "stage3_param_persistence_threshold": 100000,
    "stage3_max_live_params": 1000000000,
    "stage3_max_unused_params": 1000000000
  }
}

配置详解：

• “stage”: 3: 启用 ZeRO Stage 3，对参数、梯度和优化器状态进行分片。
• “offload_optimizer”: 启用优化器卸载。
  • “device”: “cpu”: 指定卸载目标为 CPU RAM。如果目标是 NVMe SSD，需要设置为 “nvme”，并配置相应的路径和缓存大小。
  • “pin_memory”: true: 建议开启，它使用固定的内存页，可以加速 CPU 到 GPU 的数据传输。

启用 NVMe Offload

对于参数量超过千亿、CPU RAM 仍然不足以容纳所有优化器状态的情况，可以将设备设置为 NVMe SSD。这要求您的机器拥有高速 NVMe 驱动器。

// NVMe Offload 示例配置片段
"offload_optimizer": {
  "device": "nvme",
  "nvme_path": "/path/to/nvme_cache", // NVMe 挂载点
  "buffer_count": 5,
  "buffer_size": 104857600
}

使用 ZeRO-Offload 能够让你在单个节点甚至多节点集群中，突破 GPU 显存限制，成功训练原本需要超级计算机才能处理的百亿参数大模型。