跳转至

ZeRO 优化

💾 ZeRO 的三个阶段分别分片了什么？各能节省多少显存？

ZeRO（Zero Redundancy Optimizer）将训练中的模型状态（优化器状态、梯度、模型参数）在数据并行组内分片，消除冗余。三个阶段逐级深入：

查看内嵌表格

🧮 具体估算（7B 模型，FP16，Adam，假设数据并行度 N=8）：

  • 原生 DP 单卡：权重 14 GB + 梯度 14 GB + 优化器 56 GB + 激活 ≈ 84 GB+。

  • ZeRO-1：优化器分片后 56/8=7 GB，其余不变，总显存 ≈ 14+14+7+激活 ≈ 35 GB+激活。

  • ZeRO-2：梯度假分片 14/8=1.75 GB，总 ≈ 14+1.75+7+激活 ≈ 22.75 GB+激活。

  • ZeRO-3：参数分片 14/8=1.75 GB，优化器和梯度也相应进一步减少，总 ≈ 1.75+1.75+7+激活 ≈ 10.5 GB+激活。

✅ 节省比例：每阶段显存降至原来的约 1/(N)（取决于分片程度），使得千亿参数模型训练成为可能。

🤔 为什么 ZeRO-1 只分片优化器状态？它没有通信代价增加吗？

💡 ZeRO-1 的定位是“用最小的通信代价换取最大的显存节省”，优化器状态占显存大头（~4×权重），且通信开销仅增加一次 Reduce-Scatter，非常划算。

🔍 原因：

  • 优化器状态体量最大：以 Adam 为例，每个参数需要 FP32 的动量和方差，共 8 字节，而 FP16 权重仅 2 字节。优化器状态是权重的 4 倍。因此，首先分片优化器能获得最明显的显存缩减，立竿见影。

  • 通信代价小：ZeRO-1 在反向传播后，每张卡仍拥有完整的梯度，需要将梯度做一次 All-Reduce（DP 本身就有）。ZeRO-1 增加的操作是：在 All-Reduce 后，优化器状态分片仅存在于各卡上，更新参数时，每张卡只更新自己那部分优化器状态和参数，然后需要一次 Reduce-Scatter（或 All-Gather）让各卡得到更新后的完整参数？实际过程：各卡有完整梯度，执行 optimizer.step() 时，每张卡只负责更新自己的参数分片。由于更新后参数不完整，需要在下一轮前向之前通过 All-Gather 收集完整参数。因此新增的通信是一次 All-Gather 参数（或等价操作）。这个通信量等于参数大小，与梯度 All-Reduce 相当，仅增加 50% 的通信量（具体见后文），完全可以接受。

📊 通信成本增加：标准 DP 仅有梯度 All-Reduce（通信量 2×Φ，Φ 为参数量）。ZeRO-1 增加了一次参数 All-Gather（Φ），总通信量变为 3×Φ，增加 1.5 倍。相较于节省的 4 倍显存，收益极高。

✅ 因此，ZeRO-1 的策略是：先吃掉占显存最大的优化器，用少量通信换取巨量空间。

🔄 ZeRO-2 分片梯度后，梯度同步的过程是怎样的？

💡 ZeRO-2 在前向计算时各卡仍拥有完整参数，但在反向计算出梯度后，不再直接 All-Reduce 成完整梯度，而是对梯度进行 Reduce-Scatter，使得每张卡只持有与自己参数分片对应的那一部分完整梯度（即聚合好的部分梯度）。

🔍 具体过程：

  1. 反向传播：每张卡独立计算出完整梯度（针对整个模型），但此时各卡的梯度是基于不同 micro-batch 的，需要平均。

  2. 梯度 Reduce-Scatter：各卡执行 Reduce-Scatter 操作。以参数分片为单位，例如 GPU0 负责第 1/4 的参数，则所有卡上该部分梯度将被求和（Reduce）并分散（Scatter）到 GPU0，最终 GPU0 拥有该部分参数的聚合梯度。其他部分类推。这样，每张卡只得到了自己负责那部分参数的完整梯度，不再存储其他部分的梯度，显存立刻节省。

  3. 优化器更新：每张卡用自己持有的聚合梯度，更新自己负责的优化器状态和参数分片。

  4. 参数全收集：更新后，每张卡只拥有部分更新后的参数，需要通过 All-Gather 将各卡的最新参数收集，组合成完整参数，供下一轮前向使用。

📊 通信变化：ZeRO-2 将原 DP 的 All-Reduce 梯度，拆分成了 Reduce-Scatter 梯度 + All-Gather 参数，通信量仍为约 3×Φ（与 ZeRO-1 相近），但显存进一步节省了梯度冗余。

✅ 因此，ZeRO-2 通过 Reduce-Scatter 实现梯度分片，每卡仅存 1/N 的梯度，进一步降低单卡负载。

🧱 ZeRO-3 分片参数后，前向和反向传播时如何获取参数？增加了什么通信？

💡 ZeRO-3 将模型参数也分片到各卡，前向/反向需要某部分参数时，通过 All-Gather 临时收集完整参数层，用完后立即释放；反向时再重新收集，并伴随 Reduce-Scatter 梯度和分片优化器更新。

🔍 通信模式（以某一层为例）：

  • 前向传播：当计算到某一层时，各卡原本只拥有该层参数的 1/N 分片。通过 All-Gather 将该层的完整参数收集到所有卡上，执行前向计算。计算完毕后，除自己分片外的其他参数片段立即释放，不占用显存。

  • 反向传播：同理，计算该层梯度时，需再次 All-Gather 完整参数（因为之前释放了）。然后计算局部梯度。随后进行 Reduce-Scatter 将梯度聚合到对应的分片所有者，各卡只保留自己负责部分的梯度，其余丢弃。

  • 优化器更新：每卡更新自己的参数分片，无需通信。

  • 上述过程对每一层重复。因此，通信总量大致是：前向每层一次 All-Gather，反向每层一次 All-Gather + 一次 Reduce-Scatter。整体通信量约为标准 DP 的 3 倍（详细见下题）。

📌 显存优势：参数分片后，前向计算时仅临时重建一层完整参数，显存占用从整个模型降为一层大小，极大缓解显存压力。

✅ 所以，ZeRO-3 用额外的 All-Gather/Reduce-Scatter 通信，换取参数完全分片，使得单卡仅需 1/N 的权重+梯度+优化器。

📈 ZeRO-3 的通信量比 DP 大约多多少？

💡 以每层参数 Φ 为例，标准 DP 仅需一次梯度 All-Reduce（2Φ）。ZeRO-3 需要 前向 All-Gather (Φ) + 反向 All-Gather (Φ) + 反向 Reduce-Scatter (Φ) = 3Φ，因此通信量是 DP 的 1.5 倍。但考虑到 ZeRO-3 的优化器更新无需额外通信，总通信体积比 DP 增加约 1.5 倍。

🧮 精确计算：

  • DP 梯度 All-Reduce：每个参数聚合梯度，通信量为 2Φ（每个参与卡发送+接收）。

  • ZeRO-3：前向 All-Gather 参数（Φ），反向 All-Gather 参数（Φ），反向 Reduce-Scatter 梯度（Φ），总共 3Φ。另外参数更新无需通信。

  • 所以 ZeRO-3 的通信量 = 3Φ，DP = 2Φ，倍数 = 1.5×。

🔍 但实际还会有额外开销： 因为 ZeRO-3 是逐层进行，会产生许多小消息，可能降低带宽利用率。通过融合通信和预取技术，实际效率可接近 1.5 倍。因此，ZeRO-3 是以 1.5 倍的通信换来单卡显存降至 1/N，是极为成功的权衡。

🐢 ZeRO-Offload 把优化器状态放到 CPU，会带来多少速度下降？

💡 将优化器状态卸载到 CPU 会导致训练速度下降，但下降幅度取决于 CPU 计算能力、PCIe 带宽以及是否与计算重叠。通常速度下降 20%-50%，极端情况可能更多。

🔍 主要原因：

  • 数据搬运延迟：每一步更新时，需要将梯度从 GPU 搬运到 CPU，在 CPU 上执行优化器步骤，再将更新后的参数搬回 GPU。此过程需要经过 PCIe 总线，带宽远低于 HBM。

  • CPU 计算瓶颈：Adam 等优化器涉及大量逐元素操作，若 CPU 核心不足或向量化差，会成为瓶颈。

  • 重叠优化：DeepSpeed ZeRO-Offload 设计了高效的 CPU 优化器内核，并尽量将数据传输与 GPU 计算重叠。例如在反向传播时，梯度已分批传输到 CPU 并开始优化器步骤，同时 GPU 继续计算后续层。这样可将数据移动隐藏起来，减少速度损失。

📊 实测数据：在 A100 + 双路 Xeon 上，开启 ZeRO-Offload 训练 13B 模型，相比纯 GPU ZeRO-2，吞吐量下降约 20%-30%。如果 CPU 较弱或模型极大导致 PCIe 打满，可能降速 50% 以上。但优势是可以在一张 24GB 显卡上训练数十 B 的模型。

✅ 因此，ZeRO-Offload 用适度的速度换取跨数量级的显存扩展，是“小卡训大模型”的关键技术。

💻 在单卡上使用 ZeRO-Offload 有意义吗？为什么？

💡 非常有意义！单卡使用 ZeRO-Offload 可以将优化器状态和部分参数卸载到 CPU，使得原本装不下的模型能够训练，虽然速度慢，但“从无到有”是质的突破。

🔍 价值体现：

  • 突破显存墙：单卡显存有限，如 24GB。用 ZeRO-Offload 可把优化器（+梯度）放到 CPU，GPU 只保留权重和激活，就能训练原本需要 40-80GB 显存的模型。

  • 单卡实验与研究：对于个人开发者或研究，只有一张消费级显卡，ZeRO-Offload 使得微调 7B-13B 模型成为可能。

  • 成本极低：不需要多卡集群，利用现有 PC 即可。

⚠️ 速度代价：如上题，会有速度下降。但通过减少 batch size 和序列长度，依然可接受。而且 ZeRO-Offload 可与量化训练（QLoRA）结合，进一步降低要求。

📌 例子：一张 RTX 4090 24GB，开启 ZeRO-Offload 和 QLoRA，可微调 33B 模型。若无 Offload，24GB 连 7B FP16 训练都困难。

✅ 所以，单卡 ZeRO-Offload 是让大模型训练平民化的核心功能。

🚀 ZeRO-Infinity 如何利用 NVMe 硬盘进一步扩展显存？

💡 ZeRO-Infinity 在 ZeRO-Offload 基础上，进一步将参数、优化器状态、甚至激活卸载到高速 NVMe SSD，将显存概念扩展到“GPU + CPU + NVMe”的异构存储，使得单卡训练万亿参数模型成为可能。

🛠️ 核心机制：

  1. 层次化存储：将模型状态（参数、梯度、优化器）和激活分区放置在不同设备上：热数据放 GPU，温数据放 CPU，冷数据放 NVMe。通过动态迁移和预取，保证计算设备总有数据可用。

  2. NVMe 作为超大容量内存：现代 NVMe 读写速度可达 3-7 GB/s，虽然远低于 HBM，但容量巨大且比 CPU 内存便宜。ZeRO-Infinity 可以直接从 NVMe 向 GPU 传输参数（通过 GPU Direct Storage），绕过 CPU，减少延迟。

  3. 智能调度与重叠：在计算某一层之前，提前从 NVMe 预取该层参数到 GPU；计算产生的梯度异步写回 NVMe。激活检查点也可以存到 NVMe。通过精心编排，将 I/O 时间隐藏在计算之后。

  4. 支持超长序列：可将注意力中间结果卸载到 NVMe，支持百万 token 上下文训练。

📈 效果：使用单个 GPU 和大量 NVMe 硬盘，可训练高达 1 万亿参数的模型（理论）。实际吞吐量受限于 NVMe 带宽，但突破了显存容量硬限制。

✅ 因此，ZeRO-Infinity 是显存扩展的终极方案，利用 NVMe 实现“无限显存”，使大模型训练资源需求大幅降低。

🎯 在 ZeRO 配置中，如何选择 stage 和 offload 参数？

选择 ZeRO 阶段（stage）和 offload 参数，本质上是在显存节省、通信开销、训练速度三者之间找平衡点。我的决策流程如下：

🗺️ 决策路线图：

  1. 先看单卡能否装下完整模型（权重+优化器+梯度）
  2. 如果可以 → 使用 ZeRO-1（只分片优化器），几乎不牺牲速度，通信开销最小。这也是最保守、最稳定的选择。

  3. 如果单卡装不下 → 进入下一步。

  4. 考虑梯度是否成为瓶颈

  5. 如果显存紧张但梯度+优化器稍大，使用 ZeRO-2，将优化器和梯度都分片。通信量相比 ZeRO-1 增加不多（仍为 1.5 倍 DP），但能进一步省下梯度显存。适用于 7B-13B 模型 + 中长序列。

  6. 权重也装不下或需要更大 batch

  7. 必须上 ZeRO-3。它将参数也分片，单卡显存需求降至 1/N。这适用于 30B 以上模型，或者长序列训练。代价是通信量增加到 1.5 倍 DP，且实现更复杂。

  8. 是否启用 Offload（CPU/NVMe）

  9. CPU Offload（ZeRO-Offload）：当启用 ZeRO-2/3 后显存仍然不够（例如单卡 24GB 想训 13B），可以把优化器状态和（或）梯度卸载到 CPU 内存。设置 offload_optimizer 和 offload_param 等。速度会下降 20%-50%，但能跨数量级扩展。
  10. NVMe Offload（ZeRO-Infinity）：连 CPU 内存都不够时（比如单卡训 175B），启用 NVMe offload，将优化器状态、参数甚至激活写到大容量 SSD 上。速度进一步降低，但使得万亿参数训练成为可能。

📊 经验速查表：

查看内嵌表格

💡 我的原则：能用 ZeRO-2 就不用 ZeRO-3，能不用 offload 就尽量不用，因为 offload 会引入 PCIe 瓶颈。但在资源受限时，offload 就是救命稻草。

🤝 ZeRO 可以和 LoRA 一起使用吗？如何配合？

💡 完全可以，而且这是一种非常高效的大模型微调方案：用 LoRA 冻结原权重、只训练少量低秩矩阵，再用 ZeRO 分片优化器状态和梯度，进一步降低单卡显存。

🔧 配合方式：

  1. ZeRO 负责分片优化器、梯度和主模型参数
  2. 虽然 LoRA 冻结了主模型，但主模型的参数仍然存在于显存中（或需 ZeRO-3 分片以节省空间）。

  3. 优化器状态和梯度仅针对可训练的 LoRA 参数，这部分非常小，所以 ZeRO-1/2 的节省效果对 LoRA 不明显，但 ZeRO-3 可以将冻结的主模型参数分片，极大降低显存占用。

  4. LoRA 权重被视为可训练参数

  5. 配置 DeepSpeed 时，LoRA 注入的线性层是唯一有梯度的参数，ZeRO 会只对这些参数的优化器状态和梯度进行分片（如果 stage>=1）。

  6. 主模型参数被标记为冻结，不产生梯度，也没有对应的优化器状态。在 ZeRO-3 下，这些冻结参数仍会被分片存储，以减少单卡显存。

  7. 实践配置

  8. 使用 PEFT 库（如 Hugging Face PEFT）将 LoRA 适配器挂到模型上。
  9. 启动 DeepSpeed 训练时，指定 ZeRO stage=3，并配置 "stage3_param_persistence_threshold" 等参数（可设大些，因为大部分参数冻结）。
  10. 也可以开启 offload 将冻结参数放到 CPU，进一步释放 GPU 显存。

✅ 组合优势：LoRA 让可训练参数量降至千分之一，ZeRO-3 再将庞大的主模型分片，单卡 24GB 即可微调 70B 模型。这是目前最平民化的大模型微调方案。

🚫 为什么 ZeRO 对激活值显存没有直接影响？

💡 ZeRO 的设计目标是消除模型状态（参数、梯度、优化器）的冗余，而激活值（activations）是前向传播产生的临时张量，与模型状态无关，因此 ZeRO 无法通过分片来减少它。

🔬 原因解析：

  • 激活的本质：每层计算产生的中间结果，依赖于输入数据和模型结构，每张卡的激活通常是不同的（因为数据不同），不存在跨卡冗余。ZeRO 的分片机制是针对多卡间重复存储的模型状态（所有卡存一样的内容）进行切分，而激活本来就不重复，所以无从分片。

  • ZeRO 的节省范围：权重、梯度、优化器是“全局共享但复制”的状态，ZeRO 将它们分成 N 份，每卡存 1/N。但激活是“本地独有”，不能分给其他卡。

  • 激活对显存的影响：在长序列、大 batch 下，激活显存往往成为主要瓶颈。此时需要其他技术：梯度检查点（用计算换显存）、序列并行、FlashAttention 等，这些才直接针对激活。ZeRO 不解决激活问题，所以即使开启了 ZeRO-3，激活依然可能撑爆显存。

📌 实际中：我们常看到“ZeRO-3 + 梯度检查点 + FlashAttention”的组合，分别应对参数/优化器、激活、注意力矩阵的显存压力，它们各司其职。

⚙️ 在 DeepSpeed 中，overlap_comm 和 contiguous_gradients 对显存有何影响？

💡 这两个参数主要用于优化通信与计算的并行度，以及梯度内存布局，对显存的影响通常是增加少量临时缓冲区，但不会显著改变总占用。

🔍 具体影响：

  • overlap_comm（重叠通信） 开启后，DeepSpeed 会在反向传播计算当前层梯度的同时，异步执行上一层的梯度 Reduce-Scatter（或 All-Reduce）。这需要额外的通信缓冲区来暂存异步处理的梯度，通常为梯度的 1-2 倍大小。由于梯度只是模型参数量大小，这些缓冲区会占用一定显存（例如 7B FP16 梯度 14 GB，缓冲区可能额外 1-2 GB）。因此，开启此参数会略微增加显存占用，但能显著提升训练速度。若显存紧张，可关闭它以节省空间。

  • contiguous_gradients（连续梯度） 将梯度张量在内存中重新排列成连续块，以减少通信时的碎片开销，提高带宽利用率。这需要拷贝梯度到连续缓冲区，同样需要额外临时缓冲。显存开销与 overlap_comm 类似，增加一个梯度大小的临时空间。在显存充裕时建议开启，能加速通信；若显存不足，关闭可省下这份缓冲。

✅ 总结：两者都是以较小的显存代价（额外 10%-20% 梯度大小）换取通信性能。在显存临界情况下，可以关闭它们来释放几百 MB 到几 GB 的空间。

🚀 ZeRO++ 对 ZeRO-3 的通信做了什么优化？

💡 ZeRO++ 针对 ZeRO-3 的大量 All-Gather 和 Reduce-Scatter 通信，通过量化、分层通信和细粒度重排等策略，大幅削减通信量，同时保持显存节省效果。

🔧 核心优化技术：

  1. qZDP（量化 ZeRO 数据并行） 在 All-Gather 参数时，先对参数分片进行 INT8 量化再传输，接收后反量化。因为 ZeRO-3 需要频繁收集参数，此方法可将前向参数收集的通信量减少近一半。由于量化开销小，几乎不影响精度。

  2. hpZ（分层分割 ZeRO） 将模型参数根据通信需求分层处理：对于大参数层（如 FFN 的权重），使用标准的 All-Gather/Reduce-Scatter；对于小参数层（如 bias/LayerNorm），直接复制而不分片（因为复制开销小于通信开销），从而避免大量小消息通信，提升效率。

  3. 通信合并与重排 优化通信调度，将多个小张量的 All-Gather 合并成一个大消息，减少延迟。同时使用更优的梯度 Reduce-Scatter 算法。

📊 效果：ZeRO++ 在百亿参数模型上实测通信量降低 30%-50%，训练吞吐提升 1.3-1.5 倍，同时保持了 ZeRO-3 的显存节省能力。这使得跨机训练的带宽压力大大减轻。

✅ 因此，ZeRO++ 是 ZeRO-3 的通信增强版，适合带宽有限的集群。

🔄 FSDP 与 ZeRO-3 有何异同？

💡 两者理念高度相似：都通过分片参数、梯度和优化器状态来降低单卡显存。主要区别在于实现、API 和生态，FSDP 是 PyTorch 原生支持，ZeRO-3 是 DeepSpeed 的招牌。

📊 详细对比：

查看内嵌表格

🔍 关键差异：

  • 粒度：FSDP 默认对整个模块的参数进行 All-Gather（可配置为逐层），而 ZeRO-3 逐层收集，后者显存峰值更低，因为可以更快释放参数。

  • 通信库：FSDP 使用 NCCL，与 PyTorch 分布式一致；DeepSpeed 用自定义通信后端，有更多优化空间。

  • 社区与维护：FSDP 由 PyTorch 团队维护，紧跟框架演进；DeepSpeed 针对大模型特化，创新多。

✅ 选择建议：如果你希望轻量原生且与 PyTorch 深度融合，选 FSDP；如果你追求极致显存和训练性能，且需要 offload、自动调优等高级功能，选 ZeRO-3。两者也可互换，许多框架已提供适配。

💾 使用 ZeRO-3 时，如何保存和加载模型 checkpoint？

💡 ZeRO-3 将模型参数分片到所有卡上，因此保存和加载时需要聚合或分别处理分片。DeepSpeed 提供了合并保存、分片保存以及加载的专用 API，确保模型可恢复并可用于推理或再训练。

🔧 保存 Checkpoint：

  • 合并保存（推荐用于最终模型） 调用 model.save_checkpoint(save_dir, tag="global_step") 时，ZeRO-3 会自动将所有卡的参数分片 All-Gather 到 rank0，然后由 rank0 保存一份完整的 Hugging Face 格式的模型（含权重和配置文件）。这种方式占用一定 CPU 内存，但生成的 checkpoint 可直接用于 Hugging Face 推理。 也支持 save_16bit_model 保存 FP16 合并权重。

  • 分片保存（用于恢复训练） 默认情况下，DeepSpeed 保存的是 ZeRO 分片格式：每个 rank 保存自己那部分参数和优化器状态。这要求保存时所有卡都写磁盘，产生多个文件。恢复训练时，必须用相同规模（卡数）加载。这种方式速度快，不占用额外显存。

📥 加载 Checkpoint：

  • 分片加载（恢复训练） 使用 model.load_checkpoint(load_dir, tag="global_step") 自动读取各卡对应的分片文件，恢复到之前训练的完整状态（包括优化器、学习率调度器等），继续训练。

  • 单卡加载（用于推理或再训练） 如果保存的是合并的 Hugging Face 模型，可直接用 AutoModel.from_pretrained(load_dir) 加载，不依赖 DeepSpeed。然后再根据需要包装成 DeepSpeed 引擎（如果需继续训练，需重新初始化优化器，但可以从 checkpoint 加载权重）。

⚠️ 注意事项：

  • ZeRO-3 保存合并模型时，需要所有 GPU 协同工作，如果卡数很多，聚合过程会消耗大量 CPU 内存，甚至 OOM。此时可以用 --save-interval 配合分片保存，或在单卡上加载分片后手动合并。

  • 跨版本兼容性：DeepSpeed 尽量保持 checkpoint 兼容，但不同 ZeRO stage 或 offload 配置可能不通用，加载时需匹配原始配置。

✅ 最佳实践：训练期间用分片保存（快速、低开销），训练结束后，用 ds_to_hf.py 等工具将分片 checkpoint 转换为 Hugging Face 格式，方便部署。