【大模型】理论知识：ZeRO-3、FSDP 和 分页优化器的介绍与对比

前言

先梳理下模型训练过程中，显存中存储的参数。

---

1 训练过程

【前向传播】→ 【loss】→【梯度】→【优化器(m/v)】→【Δθ】→【权重 −= Δθ】

1. 前向传播：用 旧参数 θ 算 loss
2. 反向传播 → 得到 原始梯度 g
3. 优化器内部： （以Adam / AdamW 为例）
   m = β₁m + (1−β₁)g
   v = β₂v + (1−β₂)g²
   Δθ = α · m̂ / (√v̂ + ε)  ← 这一步 只产出增量，g 本身没被改
4. 参数更新：θ  ← θ − Δθ
5. 梯度 g 使命完成，内存直接释放（或被覆盖）

---

2 显存存储参数-静态：

显存 ≈ 模型参数 + 梯度 + 优化器状态

• 梯度参数量：等于 模型参数量

• 优化器状态参数量：当前训练大模型几乎清一色用 AdamW + FP32 Master Param。AdamW 优化器包含 动量m方差v。所以优化器状态 ≈ 3 × 模型参数。

  • 【1】 混合精度下为了数值稳定，还要再存一份 FP32 精度的“主参数”，长度 = 参数量。
  • 【2】 模型计算流程如下：
    前向(16) → loss(16) → (×scale) → 梯度(16) → (cast+÷scale)→FP32 → 优化器 mv/Δθ(32) → 更新主参数(32)
  • 【3】 其中 scale 为什么要放大？
    FP16 的正数下界是 ≈ 6.0×10⁻⁸。深度网络后期很多激活梯度会小到 10⁻⁸～10⁻⁹，直接变成 0（underflow）。
    先 ×32 768 就能把这类小值拉到 > 2⁻¹⁴ 的安全区，避免被 flush to zero。
     

  三者加起来就是 3 × 模型参数量，而且全是 FP32，所以字节数再 ×4。
  （梯度虽然也是 FP16，但只在反传后短暂存在，生命周期短；优化器状态却常驻显存。）
   

优化器	每参数额外状态	状态总量 / 模型参数量
SGD（无动量）	0	0 × params
SGD+Momentum	1 动量向量	1 × params
Adam / AdamW	1 动量向量 + 1 方差向量	2 × params
Adam + FP32 Master（混合精度常见）	2 状态 + 1 FP32 主拷贝	3 × params
Shampoo/AdaFactor	矩阵统计量	1–6 × params 不等

---

3 显存存储参数-动态：

显存 ≈ 动态激活 + 临时梯度缓存

• 激活（activations）——batch 线性放大
  • 每一层都要保存前向输出，用于反向做链式法则。
  • 显存 ∝ [sequence_length] × [hidden_size] × [batch_size] × [layer_num] × [精度字节数]。
  • 例：7B 模型 40 层、hidden=4096、seq=2048、FP16，单样本激活 ≈ 0.5 GB；batch=8 就瞬间 +4 GB，batch=64 直接 +30 GB 以上。
• 梯度临时缓存——与 batch 成正比
  • 虽然最终梯度 tensor 形状跟参数一样（14 GB for 7B），但 大 batch 下每张卡先局部累加，CUDA kernel 会申请同样大小的 workspace，峰值阶段=2×梯度体积（28 GB）。
  • 用 梯度累加（micro-batch） 可以把峰值压回“单 micro-batch” 水平，但 总激活省不了。
     

其它-杂项

• CUDA kernel workspace、NCCL buffer、CUDA context 一般 < 1 GB，可忽略。

---

4 总览

一张直观加法表（7B 模型，FP16）
静态 70 GB 是“地板”；batch 越大，激活这块天花板越高，最终显存瓶颈往往由它决定。
因此大模型训练先算“静态”够不够，再看“最大能放多少 batch”就取决于激活体积。

batch size	静态权重+状态	激活（估算）	峰值梯度缓存	总显存
1	70 GB	~0.5 GB	~14 GB	~85 GB
8	70 GB	~4 GB	~14 GB	~88 GB
64	70 GB	~32 GB	~14 GB	~116 GB

一、 ZeRO-3

ZeRO 的全称是 Zero Redundancy Optimizer，是由微软开发的一种内存优化技术，旨在解决大规模深度学习模型训练中的数据并行内存冗余问题。

---

1. 核心思想

• 在传统的数据并行中，每个GPU都保存着一份完整的模型参数、优化器状态和梯度的副本。当模型非常大时，这些副本会占用大量显存，成为训练瓶颈。ZeRO的核心思想是将模型训练过程中产生的状态（优化器状态、梯度、参数）在多个GPU之间进行分区保存，而不是每个GPU都存一份完整的副本，从而消除内存冗余。

---

2. 三大优化阶段

ZeRO通过三个阶段，逐步优化内存，其优化能力依次增强：

• ZeRO-1（优化器状态分区）：
  • 只对优化器状态进行分区。例如，Adam优化器需要保存参数、动量、方差，这些状态会被分到多个GPU上。
  • 当一个GPU需要更新它负责的那部分参数时，其他GPU会将所需的优化器状态通过集合通信传递给它。
  • 效果：内存显著降低，通信量略有增加。
• ZeRO-2（梯度分区）：
  • 在ZeRO-1的基础上，进一步对梯度进行分区。
  • 在反向传播计算完梯度后，每个GPU只保留它负责的那部分梯度，用于后续的优化器更新。
  • 效果：内存进一步降低，通信量与ZeRO-1类似。
• ZeRO-3（参数分区）：
  • 在ZeRO-2的基础上，进一步对模型参数本身进行分区。
  • 每个GPU只保存模型参数的一个子集。在前向和反向传播过程中，当需要其他分区的参数时，通过集合通信实时获取，使用完后立即释放。
  • 效果：内存降低达到极致，可以实现几乎线性的内存减少，但通信开销也是最大的。

---

3. 优点

• 极高的内存效率：使得在有限显存的GPU上训练超大模型成为可能。
• 良好的扩展性：几乎保持了数据并行的简单性，同时能高效地利用大量GPU。
• 无缝集成：通过 deepspeed 库，可以相对方便地集成到PyTorch训练脚本中。

---

4. 缺点

• 通信开销：随着分区粒度变细（从ZeRO-1到ZeRO-3），通信量会增加，对集群的网络带宽要求更高。

二、 FSDP

FSDP 的全称是 Fully Sharded Data Parallel，是PyTorch官方在 torch.distributed 模块中实现的类ZeRO-3的数据并行策略。可以将其理解为PyTorch官方实现的、与社区深度集成的ZeRO-3。

---

1. 核心思想

• FSDP的核心思想与ZeRO-3完全一致：将模型参数、梯度和优化器状态全部分片 across 所有数据并行进程。

---

2. 工作流程

1. 包装模型：使用 FullyShardedDataParallel 包装模型。它会将模型的每个子模块（如Transformer的一个层）视为一个“分片单元”。
2. 前向传播：
   • 对于当前处理的子模块，将所有必需的参数分片收集到当前GPU上。
   • 计算完成后，立即释放这些收集来的参数。
3. 反向传播：
   • 同样，在计算某个子模块的反向传播前，再次收集其参数。
   • 计算得到的梯度会立即被平均并分片回各自负责的GPU上。
4. 优化器步骤：
   • 每个GPU只更新它负责的那部分参数和优化器状态。

---

3. 与ZeRO-3的关系

FSDP 本质上是ZeRO-3的一种实现。但它做了很多工程上的优化和与PyTorch生态的深度集成：

• 灵活性：允许更灵活的分片策略，如可以按层分片，平衡内存和通信效率。
• CPU Offload：原生支持将参数、梯度或优化器状态卸载到CPU内存，进一步节省GPU显存。
• 与PyTorch生态无缝衔接：作为PyTorch原生组件，与DDP、torch.compile 等特性结合更好，调试也更方便。

---

4. 优点

• 极致的显存节省：与ZeRO-3相同。
• 官方支持与易用性：作为PyTorch的一部分，有更好的长期支持和文档，集成更简单。
• 性能优化：持续受益于PyTorch团队的性能优化。

---

5. 缺点

• 通信开销：与ZeRO-3相同，对网络要求高。
• 配置复杂性：虽然易用，但要达到最佳性能，仍需仔细调优分片策略、混合精度等参数。

三、 分页优化器

分页优化器是一种针对优化器状态的内存优化技术，与ZeRO和FSDP解决的是不同维度的问题。

---

1. 核心思想

• 传统优化器在创建状态（如动量、方差）时，会分配一块固定的、与模型参数同样大小的GPU显存。分页优化器的核心思想是利用CPU和GPU的统一内存管理，将优化器状态存储在CPU的“分页”内存中。当GPU需要更新某一部分参数时，相应的优化器状态页面被按需交换到GPU显存中，更新后再交换回CPU。

---

2. 工作方式

1. 统一虚拟地址空间：利用NVIDIA的Unified Memory技术，在CPU和GPU之间创建一个统一的地址空间。
2. 按需页面迁移：优化器状态分配在CPU内存上。当GPU内核试图访问某个优化器状态时，如果该数据不在GPU显存中，会触发一个“页面错误”。
3. 自动数据传输：驱动程序捕获到这个错误，自动将所需的“页面”（一块内存）从CPU内存传输到GPU显存，然后继续执行计算。
4. 更新后写回：计算完成后，被修改的页面可能会在后台被迁移回CPU内存。

---

3. 优点

• 透明易用：通常只需一行代码（如 torch.optim.Optimizer 的 nesterov 参数在较新版本中已支持，或使用 nvSmith 等第三方实现），无需重构训练代码。
• 有效利用CPU内存：将庞大的优化器状态转移到廉价的、大容量的CPU内存中，极大地缓解了GPU显存压力。
• 与其它技术互补：它可以与ZeRO或FSDP结合使用。例如，在使用FSDP分片优化器状态的基础上，再使用分页优化器将分片后的状态卸载到CPU。

---

4. 缺点

• 性能开销：在CPU和GPU之间频繁交换数据会带来额外的延迟，可能成为训练速度的瓶颈，尤其是在PCIe带宽不足的情况下。
• 不是分布式解决方案：它主要解决单个GPU的显存不足问题，本身不具备跨GPU分片模型状态的能力。

四、对比总结

特性	ZeRO	FSDP	分页优化器
核心目标	通过分布式分片消除数据并行中的内存冗余	同ZeRO-3，是PyTorch官方的全分片实现	通过CPU-GPU统一内存将优化器状态卸载到CPU
解决的问题	模型参数、梯度、优化器状态的跨GPU冗余	同ZeRO	单个设备上优化器状态的显存占用
技术范畴	分布式训练算法/系统	分布式训练算法/系统	单机内存优化技术
内存节省级别	极高（ZeRO-3 > ZeRO-2 > ZeRO-1）	极高（等同于ZeRO-3）	中等（主要节省优化器状态）
通信开销	中到高（取决于阶段）	高（等同于ZeRO-3）	无（单机内部数据迁移，非集合通信）
易用性	通过 deepspeed 配置，有一定学习成本	PyTorch原生，API相对统一，易用性更好	非常简单，通常一行代码或一个参数
最佳适用场景	使用Deepspeed库的超大模型分布式训练	使用PyTorch进行的大模型训练（尤其是新项目）	1. 显存主要被优化器状态占用的场景 2. 作为FSDP/ZeRO的补充，进一步节省内存
相互关系	FSDP的理论基础和实践先驱	可看作是ZeRO-3的PyTorch官方实现和演进	与ZeRO/FSDP是正交且互补的技术，可以结合使用

如何选择？

1. 训练超大模型（数十亿到万亿参数）：
   • 首选 FSDP，特别是如果技术栈以PyTorch为主。它是当前PyTorch生态下的标准和未来。
   • 如果已经深度使用Deepspeed库，或者需要其独有的功能（如推理引擎、压缩等），ZeRO 也是一个非常强大和成熟的选择。
2. 需要极致节省内存：
   • 将 FSDP（或ZeRO-3） 与 分页优化器 结合使用。FSDP负责分片参数、梯度和优化器状态，而分页优化器进一步将已经分片的优化器状态卸载到CPU，实现“双重”内存节省。
3. 中等规模模型，优化器状态是瓶颈：
   • 如果通信开销是无法承受的，或者模型还没有大到必须使用FSDP/ZeRO，可以单独尝试分页优化器，它通常能以很小的代码改动带来显著的内存收益。
      

总而言之，FSDP/Zero-3 和 分页优化器 是现代大模型训练工具箱中不同层级的利器，前者是解决分布式内存问题的“重型武器”，而后者是轻巧灵活的“内存减压器”，在实践中根据需求组合使用它们已成为常态。