零冗余优化器 ZeRO：分布式训练内存革命

分布式训练的内存困境

在深度学习领域，模型参数规模正以指数级增长。从 BERT 到 GPT-4，参数数量已从数亿跃升至数万亿。训练这样庞大的模型无法依靠单块 GPU，必须借助分布式训练。

传统分布式训练中最经典的方式是模型并行与数据并行。模型并行将网络切分到不同设备，但跨设备通信巨大，开发与调试非常复杂；数据并行则复制整个模型到每个 GPU，各 GPU 处理不同 batch，然后对梯度做 All-Reduce 来保持参数同步。数据并行实现简单，扩展性好，但有一个致命伤——显存冗余。

在标准数据并行中，每个 GPU 除了保存一份完整的模型参数外，还必须保存一份完整的优化器状态（例如 Adam 中的动量和方差）。如果模型参数占用 10 GB，那么优化器状态可能额外占用 20 GB 甚至更多，而且每个 GPU 上的内容完全相同。这种惊人的冗余使得可用显存成为限制模型规模的最大瓶颈。

一个 100 亿参数的模型，使用半精度训练（参数约 20 GB），加上 Adam 优化器状态（约 80 GB），再加梯度（20 GB），单卡至少需要 120 GB 显存——远超主流 GPU 容量。而使用 32 张卡做数据并行，每张卡依然需要 120 GB，因为冗余部分没有减少。

零冗余优化器 ZeRO 总览

ZeRO（Zero Redundancy Optimizer） 正是为了解决上述显存冗余而诞生。它由微软 DeepSpeed 团队提出，核心思想可以用一句话概括：

不再在每个 GPU 上保存整个优化器状态、梯度和模型参数的完整副本，而是将它们分片（shard）到所有数据并行的 GPU 上，需要时再动态通信重组。

ZeRO 并非一种全新的并行范式，而是对数据并行的内存管理革命。它让数据并行不再是显存的桎梏，而是可以利用聚合起来的分布式显存来训练超大模型。ZeRO 本身由三个阶段组成，每个阶段都消除了一类冗余，你可以根据硬件条件灵活选择。

ZeRO 的三个阶段

阶段 1：优化器状态分片 (ZeRO-1)

数据并行下最大的冗余来自于优化器状态。以 Adam 为例，每个参数都需要维护一阶动量 m 和二阶动量 v，两者大小均与参数本身相同，总共占用是参数的 2 倍（半精度训练下参数为 FP16，而优化器状态通常为 FP32，实际倍数可能更大）。所有这些状态在每张卡上完全一样。

ZeRO-1 将优化器状态平均切分到所有数据并行进程。 每个 GPU 只存储并负责更新与自己分片对应的那一部分参数所属的优化器状态。前向传播和反向传播保持不变，所有 GPU 依然保持完整模型参数和梯度。反向传播后，各 GPU 拥有自己 batch 计算出的完整梯度，在更新参数前，先执行一次 Reduce-Scatter 操作，将梯度按优化器分片方式进行聚合和分发，保证每个 GPU 只获得自己负责的那部分梯度。然后每个 GPU 独立地使用自己的优化器状态和聚合后的梯度来更新对应的参数片段。最后再通过 All-Gather 将更新后的参数片段收集到所有 GPU，从而获得新一轮的完整模型参数。

• 显存节约：优化器状态占用降为原来的 1/N（N = 数据并行度）。例如使用 64 张卡，该部分显存仅为原来的 1/64。
• 额外通信：仅增加一次 Reduce-Scatter 和 All-Gather，通信量与数据并行原有的 All-Reduce 相当。
• 适用场景：优化器状态是显存的首要瓶颈，模型及梯度尚可容纳时，推荐使用 ZeRO-1。

阶段 2：梯度分片 (ZeRO-2)

除了优化器状态，梯度在数据并行中也存在完全冗余。每个 GPU 上的完整梯度经过 All-Reduce 后才会用于更新，因此在更新前所有梯度的副本都是多余的。

ZeRO-2 在 ZeRO-1 的基础上进一步将梯度也进行分片。 反向传播计算出每一层的梯度后，立刻对这部分梯度执行 Reduce-Scatter，将属于某个 GPU 分片的梯度聚合过去，同时释放在其他 GPU 上的梯度副本。这样，每个 GPU 只需要存储与自己负责的参数相对应的梯度即可。随后的优化器更新流程与 ZeRO-1 相同。参数更新完成后，同样使用 All-Gather 分发完整参数。

• 显存节约：优化器状态 + 梯度占用均降为 1/N。梯度原本占用与模型参数等量的显存，现在也变成近乎无冗余。
• 额外通信：由于梯度在产出时即进行 Reduce-Scatter，通信与计算可以很好地重叠，最终通信量与标准数据并行几乎一致。
• 适用场景：梯度占用开始成为瓶颈，或需要进一步压榨显存时选择 ZeRO-2。它已成为多数大模型训练的基础配置。

阶段 3：参数分片 (ZeRO-3)

数据并行中最后的冗余是模型参数本身。在正向、反向传播过程中，所有 GPU 需要在每一层持有完整的参数才能进行计算，但分层来看，其实并非每一时刻都需要全部参数。

ZeRO-3 将模型参数也进行分片。 每个 GPU 持久存储的只是参数的 1/N。当计算某一层的前向传播时，负责该层分片的 GPU 广播（或 All-Gather）其持有的参数片段到所有需要计算的 GPU；计算结束后，其他 GPU 即可释放这部分参数。反向传播同理，按照从后向前的顺序，需要计算某层梯度时，先集合完整参数，计算后丢弃。这样，在整个训练过程中，任何一个 GPU 上驻留的参数始终保持为分片大小，而不是完整模型。

• 显存节约：优化器状态 + 梯度 + 参数，三者均降为 1/N。理论上，只需 N 张卡的聚合显存 ≥ 总需求，即可训练。ZeRO-3 能够训练万亿参数级别模型。
• 额外通信：每次前向和反向传播都需要为每一层执行一次短促的参数集合通信，通信量显著增加，但都经过精心设计以与计算重叠。
• 适用场景：模型参数本身单卡无法存下，或需要将批大小调至远大于原有内存限制时，必须用 ZeRO-3。

突破物理边界：ZeRO-Offload 与 ZeRO-Infinity

ZeRO 的三个阶段将 GPU 显存的使用效率推向极致，但有时 GPU 内存总和仍然不足，或者需要极致降低 GPU 占用以增加批大小。

ZeRO-Offload 将一部分数据和计算卸载到 CPU 内存或 NVMe 上。主要卸载目标为优化器状态和优化步骤的更新计算。由于优化器更新本身计算密度较低，卸载到 CPU 并不会显著拖慢训练，而 GPU 专注于高吞吐的矩阵运算。ZeRO-Offload 可以在 ZeRO-2 的基础上将优化器状态全部放在 CPU，GPU 仅持有参数和梯度，从而极大地释放 GPU 内存。

ZeRO-Infinity 则是对 Offload 的进一步延伸，它可以将参数、梯度、优化器状态全部卸载到 CPU/NVMe 存储，并在需要时以流式方式异步预取到 GPU。它打破了“GPU 显存墙”，让训练数千亿甚至万亿参数的模型在少量 GPU 上成为可能。

实践指南：启用 ZeRO

ZeRO 主要通过微软开源的 DeepSpeed 库实现，需配合 PyTorch 使用。基本实例如下：

import deepspeed
import torch

# 定义模型、数据加载器等
model = MyLargeModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# DeepSpeed 配置字典，重点在于 ZeRO 阶段
ds_config = {
    "train_batch_size": 64,
    "gradient_accumulation_steps": 1,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 1e-4
        }
    },
    "zero_optimization": {
        "stage": 2,               # 可选 1、2、3
        "offload_optimizer": {
            "device": "cpu"       # 启用 ZeRO-Offload (需 stage=2 或 3)
        },
        "offload_param": {
            "device": "cpu"       # 参数卸载 (需 stage=3)
        },
        "allgather_partitions": True,
        "allgather_bucket_size": 5e8,
        "overlap_comm": True,
        "reduce_scatter": True,
        "contiguous_gradients": True
    },
    "fp16": {
        "enabled": True
    }
}

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    optimizer=optimizer,
    config=ds_config
)

# 训练循环
for batch in data_loader:
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

阶段选择建议

ZeRO-1 到 ZeRO-3 的通信量依次递增，但即使 ZeRO-3 也远小于模型并行中的全对全通信，且可以借助高速网卡和计算重叠大幅隐藏延迟。一般来说，只要能容纳模型单卡，优先使用 ZeRO-2；否则使用 ZeRO-3 再配合 Offload。

总结

零冗余优化器 ZeRO 从根本上改变了分布式训练的显存面貌。 它用灵活的分片思想消除了数据并行中优化器状态、梯度、参数这三个级别的冗余，使得 GPU 显存从孤岛变为连通的海。开发者几乎无需修改模型代码，只需通过 DeepSpeed 的简单配置，就能训练以前只能仰望的模型规模。

从 ZeRO-1 到 ZeRO-3，再到 Offload 与 Infinity，ZeRO 形成了一套完整的显存优化方案谱系，已经成为大模型训练领域的基石技术之一。掌握 ZeRO，就是掌握了打开大模型训练之门的钥匙。