人工智能领域正在进行一场军备竞赛，但武器不是导弹，而是参数。从 BERT (3.4 亿) 到 GPT-2 (15 亿) 再到 T5 (110 亿) ，我们看到了一个清晰的趋势: 更大的模型往往能带来更高的准确率。但这股不懈的增长潮流也伴随着高昂的代价——训练这些庞然大物需要天文数字般的内存，远远超出了单个 GPU 的处理能力。

试想一下: 即便是像 GPT-2 这种参数量为 15 亿的中等规模模型，在使用标准方法训练时，仅训练状态就需要超过 24 GB 的内存。这已经触及了高端 32 GB GPU 的极限——而且这还没算上激活值和所有临时数据。那么，我们究竟如何才能训练拥有数十亿、数百亿甚至万亿参数的模型呢？

标准方法——数据并行——可以扩展计算能力，但无法延展内存容量。它会在每个 GPU 上复制整个模型，所以无论你增加多少 GPU，都会遇到内存瓶颈。模型并行 (将模型切分到不同 GPU 上) 可以绕过这堵墙，但实现起来出了名的困难，并且在跨越单机规模时，由于节点间通信较慢，性能会急剧下降。

这正是微软研究团队着手解决的挑战。他们在开创性的论文《ZeRO: 迈向万亿参数模型训练的内存优化》 (ZeRO: Memory Optimizations Toward Training Trillion Parameter Models) 中，提出了一种新颖且优雅的方案，从根本上改变了训练大规模模型的成本结构。ZeRO，即**零冗余优化器 **(Zero Redundancy Optimizer) 的缩写，是一套内存优化技术，它能让模型大小随着 GPU 数量线性扩展，同时保持高计算效率。

内存消耗大户: GPU 显存都去哪儿了？

在领会 ZeRO 的精妙之前，我们首先需要理解为什么训练模型所消耗的内存远不止存储其参数那么简单。训练过程中的显存占用主要分为两大类:

  • 模型状态 (Model States)
  • 残余状态 (Residual States)

模型状态

对大型模型而言，这是最大的显存消耗者，包含:

  1. 参数 (Parameters) ——训练过程中学习到的权重和偏置。
  2. 梯度 (Gradients) ——在反向传播时计算得到，用于更新参数，大小与参数相同。
  3. 优化器状态 (Optimizer States) ——由优化器 (如 Adam) 维护的额外变量，例如动量 (momentum) 和方差 (variance) ，用于实现自适应学习率。

混合精度训练会加剧这一问题。为了利用现代 GPU 上的 Tensor Core，训练通常采用 FP16 格式。但为了保证稳定性，需要保留 FP32 的主副本以及 FP32 格式的优化器状态。

对于一个使用 Adam 优化器、拥有 \( \Psi \) 个参数的模型:

  • 参数: \( 2\Psi \) 字节 (FP16)
  • 梯度: \( 2\Psi \) 字节 (FP16)
  • 优化器状态:
    • 参数的 FP32 副本: \( 4\Psi \) 字节
    • 动量: \( 4\Psi \) 字节
    • 方差: \( 4\Psi \) 字节

总计:

\[ 2\Psi + 2\Psi + (4\Psi + 4\Psi + 4\Psi) = 16\Psi \ \text{bytes} \]

对于拥有 15 亿参数的 GPT-2: \( 16 \times 1.5 \ \text{B} \approx 24\ \text{GB} \)。仅优化器状态就消耗了 FP16 参数大小的 12 倍。

残余状态

这部分包括其它所有内容:

  • 激活值 (Activations) ——前向传播过程中每一层的输出，用于反向传播。在长序列和大批量的情况下，激活值很容易达到数十 GB。即使使用*激活检查点 *(activation checkpointing，即通过重新计算替代存储) ，开销依然很高。
  • 临时缓冲区 (Temporary Buffers) ——用于梯度聚合等操作；其大小可能与模型本身相当。
  • 内存碎片 (Memory Fragmentation) ——分配和释放显存的循环会造成不连续的小块空闲空间。即便总空闲显存足够，缺乏足够大的连续块也会导致内存不足 (out-of-memory) 错误。

ZeRO 的核心理念: 消除内存冗余

关键洞察是: 在标准数据并行中，大部分显存占用都是冗余的。如果你有 \( N_d \) 个 GPU，就会有 \( N_d \) 份相同的参数、梯度和优化器状态。ZeRO 的使命就是消除这种冗余。

它包括两大类优化:

  1. ZeRO-DP —— 削减模型状态显存占用。
  2. ZeRO-R —— 裁剪残余状态显存占用。

ZeRO-DP: 模型状态分区

ZeRO-DP 通过在 GPU 之间**分区 **(partitioning) 模型状态，而非复制它们，来增强传统数据并行。分为三个递进阶段:

ZeRO-DP 各阶段参数、梯度和优化器状态的内存使用情况示意图。

图 1: ZeRO-DP 的内存节省效果。基线方法会复制所有状态；每个阶段都会对更多组件进行分区 (绿色=优化器状态，橙色=梯度，蓝色=参数) ，显著降低单 GPU 显存占用。

阶段 1 —— 优化器状态分区 (\(P_{os}\))

优化器状态占用了最大比例。ZeRO 将其在 GPU 间分区，使每个 GPU 仅持有 \( 1/N_d \) 的优化器状态。

  • 内存减少: 约 4 倍——从 \( 16\Psi \) 降至约 \( 4\Psi \) 字节。
  • 通信开销: 与标准数据并行相同。

阶段 2 —— 梯度分区 (\(P_{os+g}\))

每个 GPU 只需梯度的相应分区，ZeRO 使用 reduce-scatter 替代 all-reduce，求和并分发分区梯度。

  • 内存减少: 约 8 倍——梯度占用减少 \( N_d \) 倍，每 GPU 总显存约 \( 2\Psi \)。
  • 通信开销: 与基线方法持平 (reduce-scatter + all-gather 的通信量与 all-reduce 相同) 。

阶段 3 —— 参数分区 (\(P_{os+g+p}\))

参数本身也被分区。GPU 在计算每一层时按需获取参数分片，用完后立即释放。

  • 内存减少: 随 \( N_d \) 线性降低；每 GPU 显存占用 = \( 16\Psi / N_d \)。
  • 通信开销: 比基线高约 50%，但能够训练超大模型。

启用全部三阶段后，理论上 ZeRO 可在 1024 块 GPU 上训练一个万亿参数模型，每 GPU 模型状态显存占用仅约 16 GB。

ZeRO-R: 应对残余显存占用

当模型状态冗余被消除后，残余显存会成为新瓶颈。ZeRO-R 通过以下手段缓解:

  • 分区激活检查点 (Partitioned Activation Checkpointing) —— 激活值在 GPU 间分区存储，而非复制；需要时再重建；对于极大型模型，可将其卸载到 CPU 内存。
  • 固定大小缓冲区 (Constant-Size Buffers) —— 限制临时缓冲区的大小，防止随模型增大而膨胀。
  • 内存碎片整理 (Memory Defragmentation) —— 区分存储生命周期短和长的张量，保持连续可用空间。

ZeRO 的实测表现

团队实现了 ZeRO-100B，包括 ZeRO-DP 阶段 2 (\(P_{os+g}\)) 及所有 ZeRO-R 优化。在配备 400 块 NVIDIA V100 GPU 的集群上进行测试，并与当时最先进的大模型框架 Megatron-LM 对比。

前所未有的规模与速度

性能对比图，显示了 ZeRO 的高吞吐量与基线方法在参数超过 400 亿后急剧下降的对比。

图 2: 即使模型超过 1000 亿参数，ZeRO 仍能保持单 GPU 约 38 TFLOPS 的性能；而基线方法在超过 400 亿参数后，由于低效的跨节点模型并行，性能崩溃。

结果显示: ZeRO 可训练高达 1700 亿参数的模型——是 Megatron-LM 极限的 8 倍——并实现最高 10 倍的吞吐量。

超线性扩展性

条形和折线图，显示随着 GPU 数量从 64 扩展到 400，单个 GPU 的吞吐量在增加。

图 3: ZeRO 的超线性扩展性——增加 GPU 不仅提高总性能，还提升单 GPU 吞吐量。

原因是: GPU 数增加，ZeRO-DP 降低了单 GPU 显存占用，使每块 GPU 能处理更大的批次。更大的批次 => 更高的算术密度 => 更好利用率 => 更高 TFLOPS。

大模型训练平民化

图表显示 ZeRO-DP 在没有模型并行的情况下可以训练高达 130 亿参数的模型。

图 4: ZeRO-DP 可在不使用模型并行的情况下训练 130 亿参数的模型。标准数据并行在约 14 亿参数处即失效。

ZeRO 让数据科学家无需复杂的模型重构，便能用普通的数据并行工作流训练巨型模型。

实际影响: Turing-NLG

验证集困惑度图表，显示由 ZeRO 训练的 Turing-NLG (17B) 优于 Megatron-LM (8.3B)。

图 5: ZeRO 让 Turing-NLG (170 亿参数) 的训练成为可能，并取得当时最先进的困惑度表现。

得益于 ZeRO-100B 的高效性，拥有 170 亿参数的语言模型 Turing-NLG 创下了新的准确率纪录。

消融研究: 收益来源分析

表格展示了五种结合了不同 DP 和 R 优化的 ZeRO 配置。

表 3: 用于隔离不同 ZeRO 优化效果的五种配置。

最大模型规模

条形图显示，在启用梯度和激活分区后，可训练的最大模型尺寸从 400 亿跃升至 1400 亿。

图 6: 随着 ZeRO 功能增加，可训练的最大模型规模大幅提升。

内存使用

图表对比了在不同配置下，训练 400 亿和 1000 亿参数模型时的最大缓存内存。

图 7: 随着优化叠加，峰值显存占用降低，为更大的批次释放空间。

吞吐量

图表显示了在不同配置下，训练 600 亿和 1700 亿参数模型时单 GPU 吞吐量。

图 8: 内存节省让批次更大，吞吐量总体上升；CPU 卸载仅在它是唯一可训练方案时才有帮助。

结论: 通向万亿参数之路

ZeRO 标志着大模型训练的范式转移。通过系统性地消除内存冗余，它打破了数据并行的复制瓶颈，并绕开了模型并行的扩展瓶颈。

即便只是部分实现的 ZeRO-100B，如今也已带来 8 倍的模型规模提升和 10 倍的速度提升。完整的 ZeRO-DP 让单设备显存占用与设备数量成反比:

表格展示在三种 ZeRO-DP 阶段下，不同模型规模及数据并行度对应的理论单 GPU 内存占用。

表 1: 在 \(P_{os+g+p}\) 阶段下，一个万亿参数模型在 1024 路数据并行中，每 GPU 仅需 15.6 GB 显存。

是的，训练万亿参数模型仍需巨量计算资源——可能要等到下一代硬件。但 ZeRO 已经解决了显存问题，提供了让计算能力跟上时代时可行的系统技术。

ZeRO 已作为微软 DeepSpeed 库的一部分开源，让大规模 AI 不再高不可攀，为下一代更庞大、更强大、并能改变世界的模型铺平了道路。