ZeRO：大模型训练的内存优化革命

Zero Redundancy Optimizer（ZeRO）是微软提出的深度学习训练内存优化方案，通过ZeRO-DP（分三阶段优化模型状态内存：优化器状态、梯度、参数分区，最高实现与数据并行度成正比的内存缩减）和ZeRO-R（优化激活内存、临时缓冲区、内存碎片），在保持低通信量和高计算粒度的同时，消除数据并行和模型并行训练中的内存冗余；其实现ZeRO-100B可在 400 块 V100 GPU 上超线性加速训练超 1000 亿参数模型，吞吐量达 15 Petaflops，模型规模较现有技术提升 8 倍、性能提升 10 倍，无需模型并行即可训练 130 亿参数模型，还支持与模型并行结合，理论上可在 1024 块 GPU 上训练万亿参数模型，助力构建了 170 亿参数的 Turing-NLG 这一当时最大且精度创纪录的语言模型。

一、引言

1.1 研究背景

深度学习模型（尤其 NLP 领域的 Transformer 类）参数规模持续增长，从 BERT-large（3 亿参数）、GPT-2（15 亿参数）到 Megatron-LM（83 亿参数）、T5（110 亿参数），更大模型带来显著精度提升，但需向百亿、万亿参数规模突破。
现有训练方案的局限：

1. 数据并行（DP）：模型状态全复制，内存冗余严重，32GB GPU 仅能训练 14 亿参数模型。
2. 模型并行（MP）：垂直拆分模型层，跨节点通信开销大，效率骤降（400 亿参数模型在双 DGX-2 节点上 GPU 效率不足 5%）。
3. 管道并行（PP）：水平拆分模型，存在管道气泡、功能受限（如绑定权重难实现）、收敛受影响等问题。

其他方案：CPU 卸载通信耗时，内存高效优化器可能影响收敛，激活检查点优化不足。
核心痛点：模型训练内存主要消耗于模型状态（优化器状态、梯度、参数，占比极高）和残余状态（激活、临时缓冲区、内存碎片），现有方案无法兼顾内存效率、计算 / 通信效率与易用性。

1.2 论文核心贡献

ZeRO 的核心创新是将内存优化拆分为模型状态优化（ZeRO-DP）和残余状态优化（ZeRO-R），针对性解决不同类型的内存占用问题，同时兼顾通信效率与计算粒度，以下是逐模块深度解析。

优化模块	优化目标	核心技术	内存缩减效果	通信开销变化
ZeRO-DP	模型状态内存（优化器状态 + 梯度 + 参数）	分区策略 + 动态通信调度	1. Pos：4 倍；2. Pos+g：8 倍；3. Pos+g+p：与 DP 度（Nd）成正比	1. Pos/Pos+g：无增加；2. Pos+g+p：增加 50%（达 baseline DP 的 1.5 倍）
ZeRO-R	残余内存（激活 + 临时缓冲区 + 碎片）	1. 分区激活检查点；2. 恒定大小缓冲区；3. 内存碎片整理	1. 激活内存：按 MP 度（Nm）缩减，CPU 卸载后近零；2. 临时缓冲区：不随模型规模增长；3. 碎片：避免无效内存占用	分区激活检查点：通信量增加不足 MP 基线的 10%

---

二、核心概念铺垫

术语名词总结

英文名词	中文释义	简要说明
ZeRO-DP	模型状态零冗余优化	对参数、梯度、优化器状态进行分区，大幅降低模型状态内存
ZeRO-R	残余状态内存优化	优化激活、临时缓冲区、内存碎片
$P_{os}$	优化器状态分区	ZeRO-DP 阶段1，只存 1/Nd 优化器状态，内存降 4 倍
$P_{os+g}$	梯度分区	ZeRO-DP 阶段2，梯度也分区，内存降 8 倍
$P_{os+g+p}$	参数分区	ZeRO-DP 阶段3，参数也分区，内存与 Nd 线性缩减
$P_a$	分区激活检查点	按模型并行度分区存储激活，CPU 卸载后显存占用近 0
$C_B$	恒定大小缓冲区	固定大小的融合缓冲区，内存不随模型变大
$M_D$	内存碎片整理	即时移动张量到连续内存，减少碎片、提升显存利用率
DP degree (Nd)	数据并行度	数据并行的进程/GPU 数量
MP degree (Nm)	模型并行度	模型并行的进程/GPU 数量
All-Gather	全收集操作	从多个进程收集数据拼成完整张量
Reduce-Scatter	归约散射	先归约再分片，减少通信量
Activation Checkpoint	激活检查点	前向只存部分激活，反向重算，节省显存
Offload to CPU	卸载到 CPU	将数据放到 CPU 内存，释放 GPU 显存
Memory Fragmentation	内存碎片	显存空闲但不连续，导致无法分配
Memory Footprint	内存占用量	模型/操作实际占用的显存大小
Non-trivial	显著的、不可忽视的	形容内存开销很大、不能忽略
On-the-fly	即时、动态	训练过程中实时、不暂停地做整理/调度
Democratizing	平民化、普惠化	降低大模型训练门槛，让更多人能做

2.2 传统并行方案的局限性

并行方案	核心逻辑	内存问题	通信问题
数据并行 DP	每块 GPU 都保存完整模型	模型状态冗余严重，显存占用极高	梯度 All-Reduce 通信量大
模型并行 MP	模型按层/张量拆分到多 GPU	激活值多副本冗余，内存浪费	跨设备通信频繁，带宽瓶颈明显
管道并行 PP	模型分层流水线执行	需缓存多批次激活，内存开销大	流水线气泡，通信与计算重叠复杂

---

三、ZeRO核心优化方案

3.1 ZeRO-DP：模型状态的零冗余优化，内存缩减的核心

打破传统 DP「全量复制」、MP「静态分区」的局限，对模型状态（优化器状态 + 梯度 + 参数）进行精细化分区，按数据并行度（Nd​）将模型状态拆分到各 GPU，仅在需要时通过动态通信调度聚合 / 广播，实现「内存效率向 MP 看齐，通信 / 计算效率向 DP 看齐」。

3.1.1 阶段 1：Pos​ 优化器状态分区

核心操作：将优化器状态（如 Adam 的动量、方差）按 Nd​ 拆分，每个数据并行进程仅存储和更新1/Nd​ 的优化器状态，训练步结束后通过 All-Gather 聚合全量更新后的参数。
内存缩减效果：4 倍内存占用降低
通信开销变化：与原生数据并行一致，无额外通信开销

3.1.2 阶段 2：Pos+g​ 梯度分区

核心操作：在 Pos​ 基础上，对梯度也进行分区处理，仅在对应参数分区的进程上执行梯度归约，梯度使用后立即释放内存，无需存储全量梯度。
内存缩减效果：8 倍内存占用降低
通信开销变化：采用 Reduce-Scatter 替代全量 All-Reduce，通信量与原生数据并行一致，无额外开销

3.1.3 阶段 3：Pos+g+p​ 参数分区

核心操作：在 Pos+g​ 基础上，将模型参数也按 Nd​ 拆分，每个进程仅存储自身负责的参数分区，前向 / 反向传播时通过流水线式广播获取所需其他分区参数，使用后立即释放。
内存缩减效果：内存占用与 Nd​ 线性成正比缩减（如 64 路 DP 实现 64 倍缩减，1024 路 DP 实现 1024 倍缩减）
通信开销变化：仅增加50%（为原生数据并行的 1.5 倍），通过流水线调度隐藏部分通信延迟

3.2 ZeRO-R：残余状态内存的精细化优化

3.2.1 ZeRO-R：Pa​ 分区激活检查点

优化目标：解决模型并行中激活值的冗余复制问题，降低激活值内存占用
核心操作：取消激活值在模型并行 GPU 上的全量复制，将激活值按模型并行度（Nm​） 分区存储；前向传播完成后仅保留分区激活检查点，反向传播需要时通过 All-Gather 按需重构全量激活值；超大规模模型场景下，可将分区激活检查点卸载至 CPU，进一步降低 GPU 激活内存占用。
内存缩减效果：激活内存占用按 Nm​ 比例缩减，CPU 卸载后可降至近乎零
通信开销变化：仅增加不足模型并行基线10% 的通信量，可通过模型的高算术强度隐藏

3.2.2 CB​ 恒定大小缓冲区

优化目标：解决临时融合缓冲区随模型规模膨胀的问题，控制临时内存占用
核心操作：传统高性能库（Apex、Megatron）为提升通信效率，会将全量参数 / 梯度融合为单个缓冲区，导致缓冲区大小与模型规模成正比，成为内存负担；ZeRO-R 采用性能高效的恒定大小融合缓冲区，缓冲区大小不随模型规模变化，同时保证缓冲区尺寸足够大以维持通信 / 计算效率。
内存优化效果：临时缓冲区内存占用不再随模型规模增长，彻底避免大模型下的缓冲区内存膨胀

3.2.3 MD​ 内存碎片整理

优化目标：解决训练过程中因张量生命周期交错导致的内存碎片化问题，提升显存利用率
核心操作：摒弃传统被动式内存管理，通过预分配连续内存缓冲区，将激活检查点和梯度主动迁移至预分配的连续内存中；同时根据张量生命周期（短期 / 长期）进行内存规划，执行即时的内存碎片整理，避免因碎片化导致的显存不足（OOM）问题。
内存优化效果：避免无效的碎片化内存占用，提升25% 以上的显存利用率，大幅降低 OOM 风险

3.3 两大模块的协同优化逻辑

ZeRO-DP 与 ZeRO-R 并非独立优化，而是层层递进、协同生效的整体：

先通过 ZeRO-DP 解决占比最高的模型状态内存瓶颈，释放核心显存空间；
再通过 ZeRO-R 优化残余状态内存，消除次要内存瓶颈，进一步提升显存利用率；
两者均兼顾内存优化与通信开销，避免为了内存降低而牺牲训练速度，最终实现内存占用最小化、计算效率最大化、通信开销可控化的三重目标。

---

四、实验结果与性能分析

4.1 实验配置

• 硬件环境：

  • GPU：NVIDIA V100 32GB（核心实验硬件）
  • 节点数 / GPU 数：多节点集群（核心实验采用 64~400 块 V100 GPU，万亿参数模型验证采用 1024 块 V100 GPU）
  • 通信：节点内采用 NVLink 通信（带宽 300GB/s），节点间采用 InfiniBand 通信（带宽 12.5GB/s）

• 模型规模：

  • 模型：Transformer 架构语言模型（含 Turing-NLG 系列模型）
  • 参数规模：7.5B、10B、100B、170B 参数（核心验证 100B 参数模型，理论验证 1T 万亿参数模型）
  • 批次大小：单设备 batch size = 1（确保内存消耗主要来自模型状态，避免批次大小干扰内存优化效果评估）
  • 序列长度：1024（符合 Transformer 主流训练序列长度配置）

• 训练框架与优化策略：

  • 框架：PyTorch + DeepSpeed（ZeRO 集成于 DeepSpeed 框架实现）
  • 优化器：Adam 优化器（采用混合精度训练，优化器状态、梯度以 FP32 格式存储，参数以 FP16 格式存储）
  • 并行策略：ZeRO-DP（三阶段分区优化）、ZeRO-R（残余内存优化），可结合模型并行（MP）协同使用

• 对比基线：

  • 传统数据并行（DP）：全量模型复制，梯度采用 All-Reduce 同步
  • 模型并行（MP）：模型按层拆分，激活值全量复制
  • 激活检查点（Activation Checkpointing）：仅保留部分激活值，反向传播时重算

4.2 核心指标对比

对比项	传统数据并行 DP	模型并行 MP	ZeRO
模型状态内存	高（全量复制）	中（分区存储）	极低（与并行度线性缩减）
激活内存	低	高（多副本冗余）	极低（分区激活 + CPU 卸载）
临时缓冲区	随模型变大	随模型变大	固定大小，不随模型增长
内存碎片	严重	中等	实时整理，利用率高
通信开销	高（All-Reduce）	极高（跨层频繁通信）	低（仅增加约 50%）
可支持最大模型	小（1.4B/32GB V100）	中	极大（可训练万亿参数模型）
计算效率	高	低	高（接近数据并行）

4.3 关键结论

基于上述实验验证，结合 ZeRO 两大优化模块的性能表现，可得出以下核心结论，均来自论文原文验证结果，兼具严谨性与实用性：

1. 内存优化效果显著：ZeRO 可实现模型状态内存与残余状态内存的双重极致优化，相比传统数据并行（DP），内存占用最高可降低 64 倍，32GB V100 GPU 可支持训练 100B 参数模型，1024 块 V100 GPU 可支撑万亿参数模型训练，彻底突破单卡显存与集群内存瓶颈。
2. 通信开销可控，不牺牲计算效率：ZeRO 仅在参数分区阶段增加 50% 通信开销，分区激活检查点仅增加不足 10% 通信开销，且可通过流水线调度、模型高算术强度隐藏通信延迟，计算效率接近原生数据并行，远优于传统模型并行（MP）的低效率问题。
3. 兼容性强，部署成本低：ZeRO 集成于 DeepSpeed 框架，无需对模型结构进行大幅重构，可与现有 Transformer 类模型、Adam 优化器、混合精度训练无缝兼容，同时支持与模型并行（MP）协同使用，适配不同硬件集群规模。
4. 实用性突出，落地价值高：ZeRO 解决了传统并行方案的核心痛点——DP 内存冗余、MP 通信低效、PP 功能受限，真正实现了「内存高效、通信可控、计算高效」的三重目标，为大模型规模化训练提供了可落地的内存优化方案，也为后续万亿参数模型（如 Turing-NLG）的训练奠定了核心基础。

---

五、产业落地与实际价值

5.1 对大模型训练范式的影响

ZeRO 彻底改变了大模型训练的传统范式，打破了「模型规模受单卡显存限制」的核心桎梏，推动大模型训练从「小众化、高门槛」走向「规模化、普惠化」。在 ZeRO 出现之前，万亿参数模型训练仅能由少数拥有顶级硬件集群的机构实现，而 ZeRO 通过内存分区、碎片整理等优化，让普通 GPU 集群也能支撑千亿、万亿参数模型训练，重构了「内存利用-通信效率-计算效率」的平衡逻辑。同时，ZeRO 提出的「模型状态分区+残余内存精细化管控」思路，成为后续大模型内存优化的核心参考范式，影响了后续 Megatron-LM、Colossal-AI 等众多大模型训练框架的优化方向。

5.2 与DeepSpeed的集成应用

ZeRO 作为核心内存优化模块，深度集成于 Microsoft 推出的 DeepSpeed 训练框架，成为 DeepSpeed 最具竞争力的核心功能之一，实现了「技术研发-工程落地」的无缝衔接。集成后，开发者无需单独实现 ZeRO 的优化逻辑，仅需通过简单的配置参数（如指定 ZeRO 优化阶段、并行度、CPU 卸载策略等），即可将 ZeRO 的内存优化能力应用于任意 Transformer 类大模型训练。目前，DeepSpeed+ZeRO 已成为产业界训练大模型的主流组合，广泛应用于自然语言处理、计算机视觉等领域，支撑了 Turing-NLG（170B 参数）、GPT-3 类模型等众多超大规模模型的训练落地，大幅降低了大模型的工程部署成本。

5.3 后续研究方向

基于 ZeRO 的核心优化思路，论文及后续研究主要围绕「进一步降低开销、拓展适用场景」展开，核心方向如下，均来自论文原文及官方后续延伸：

1. 通信开销进一步优化：探索更高效的通信调度策略，减少参数分区、激活重构过程中的通信延迟，实现「内存优化与通信开销」的更优平衡；
2. 异构内存协同利用：深化 CPU、GPU、NVMe 等异构内存的协同调度，进一步将模型状态、激活值等数据卸载至低成本内存设备，降低硬件成本；
3. 多并行策略深度融合：推动 ZeRO 与管道并行（PP）、3D 并行等策略的深度协同，适配更复杂的硬件集群，支撑更大规模模型的高效训练；
4. 适配多任务场景：将 ZeRO 的内存优化思路拓展至大模型微调、推理等场景，解决大模型推理过程中的内存瓶颈，实现「训练-推理」全链路内存优化。

---

六、总结与思考

6.1 核心知识点回顾

1. ZeRO 的核心定位：微软提出的大模型训练内存优化方案，核心目标是消除内存冗余、突破显存瓶颈，支撑万亿参数模型训练，集成于 DeepSpeed 框架；
2. 两大核心模块：ZeRO-DP 针对模型状态内存（参数+梯度+优化器状态），通过三阶段分区实现内存线性缩减；ZeRO-R 针对残余内存（激活+缓冲区+碎片），通过三大策略实现精细化管控；
3. 核心优势：内存优化效果显著（最高64倍缩减）、通信开销可控（仅增加50%）、兼容性强、部署成本低，兼顾内存效率与计算效率；
4. 实验验证：基于 V100 GPU 集群，可支撑 100B 参数模型高效训练，1024 块 V100 可实现万亿参数模型训练，计算效率接近原生数据并行；
5. 产业价值：重构大模型训练范式，推动大模型普惠化，成为 DeepSpeed 核心功能，支撑众多超大规模模型落地。

6.2 个人思考/延伸

** Q1: ZeRO 通过哪些核心技术实现内存优化，不同阶段的内存缩减效果和通信开销如何平衡？**

ZeRO 的内存优化核心是ZeRO-DP（模型状态优化） 和ZeRO-R（残余内存优化） 两大模块。其中 ZeRO-DP 分三阶段：1. 优化器状态分区（Pos），通过将优化器状态按数据并行度（Nd）拆分，实现 4 倍内存缩减，通信量与标准 DP 一致；2. 梯度分区（Pos+g），仅在对应参数分区的进程上归约梯度，内存缩减至 8 倍，通信量仍不变；3. 参数分区（Pos+g+p），参数按 Nd 拆分，需动态广播获取所需参数，内存缩减与 Nd 成正比（如 64 DP 下缩减 64 倍），通信量仅增加 50%。ZeRO-R 通过分区激活检查点（按模型并行度 Nm 缩减激活内存）、恒定大小缓冲区（避免缓冲区随模型膨胀）、内存碎片整理（预分配连续内存）优化残余内存，分区激活仅增加不足 10% 的模型并行通信量，且可通过 CPU 卸载进一步降低内存占用。整体设计通过动态通信调度和分区策略，在极致提升内存效率的同时，最小化通信开销增加。

Q2: ZeRO 在实际应用中表现出哪些关键优势，相比现有技术在模型规模、性能和易用性上有何突破

ZeRO 的核心应用优势体现在三方面：1. 模型规模突破：ZeRO-100B 可在 400 块 V100 GPU 上高效训练 1700 亿参数模型，无需模型并行即可训练 130 亿参数模型（远超现有 DP 方案的 14 亿参数上限），理论上在 1024 块 GPU 上可支持万亿参数模型训练，较当时 SOTA（如 Megatron-LM 的 400 亿参数上限）提升 8 倍；2. 性能飞跃：训练 1000 亿参数模型时吞吐量达 15 Petaflops，单 GPU 性能超 38 TFlops，较基线提升 10 倍，且在 64-400 GPU 区间呈现超线性加速（GPU 数量翻倍，性能提升超翻倍）；3. 易用性革新：无需修改模型代码，仅需通过 DeepSpeed 接口封装 torch.nn.module，兼容标准数据并行使用方式，避免了模型并行（需重构模型、支持算子有限）和管道并行（功能受限）的使用门槛，降低了大模型训练的技术壁垒，同时可兼容现有模型并行方案进一步扩展。

Q3: ZeRO 如何支持万亿参数模型训练，当前面临的主要挑战是什么？

ZeRO 支持万亿参数模型训练的核心逻辑是内存效率的极致提升与并行策略的灵活组合：1. 全量优化（ZeRO-DP 三阶段 + ZeRO-R）下，模型状态内存按数据并行度（Nd）线性缩减，残余内存按模型并行度（Nm）缩减，最大理论内存缩减倍数达 Nd×Nm；2. 结合 16 路模型并行（单 DGX-2 节点内）和 64 路数据并行（跨节点），1024 块 32GB V100 GPU 的聚合内存可满足万亿参数模型（Adam 优化器 + 16 位精度）的 16TB 内存需求（16TB/1024=16GB/GPU）。当前面临的主要挑战是计算能力缺口：万亿参数模型的计算量是 BERT-large（3.3 亿参数）的 3000 倍，即便在 1024 块 V100 GPU 上训练，按现有计算效率仍需 140 天以上（若扩大序列长度和数据量则超 1 年），需 Exa-flop 级计算集群才能将训练时间缩短至合理范围；此外，跨节点通信带宽和长期训练的稳定性也是实际部署中需解决的问题。

---

参考资料

1. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models