ZeRO 和分片优化器
ZeRO（零冗余优化器）通过跨 GPU 分片优化器状态、梯度和权重，消除了数据并行性中浪费的内存重复。
概述
It lets you train enormous models with the simplicity of data parallelism but a fraction of the per-GPU memory.
深入探讨
在普通数据并行中，每个 GPU 都会存储优化器状态、梯度和参数的冗余完整副本，这是极大的浪费，特别是对于 Adam 来说，其中优化器状态可能是模型本身大小的几倍。 DeepSpeed 中的 Microsoft 引入的 ZeRO 通过在 GPU 上对这些张量进行分区来消除这种冗余，这样每个设备只拥有一个切片。 ZeRO 分为三个渐进阶段：第 1 阶段对优化器状态进行分片，第 2 阶段添加梯度分片，第 3 阶段对参数本身进行分片。根据需要，GPU 通过通信、计算收集丢失的切片，然后释放它们。其结果是每个 GPU 的内存显着降低，支持数十亿到万亿参数的训练，同时保持数据并行的简单编程模型。
技术洞察
ZeRO 用额外的通信来节省内存。在第 3 阶段，在某个层的前向传递之前，全收集将该层的完整参数收集到每个 GPU 上；之后，非拥有的切片将被丢弃以回收内存。梯度是减少分散的，因此每个 GPU 仅保留与其拥有的参数相匹配的梯度切片。 PyTorch 的 FSDP（完全分片数据并行）原生实现了相同的想法，将模块包装为动态分片和重新分片。
战略影响
成本与预算
多年来，架构决策决定着性能和运营成本。
更清晰的判决
技术教育帮助团队选择正确的堆栈，而不仅仅是最新的堆栈。
质量控制
更好的工程选择可以减少生产中的可靠性事故。
ZeRO 和分片优化器的未来
分片正在成为大规模训练的默认选择，而不是一种奇特的选择。期望与卸载进行更深入的集成（通过 ZeRO-Infinity 将切片推送到 CPU 或 NVMe），通过计算更好地重叠全收集和减少分散以隐藏其成本，以及与张量和管道并行性的组合。随着模型不断增长，内存高效的分片优化器对于将模型适应实际的硬件预算至关重要。
现实世界的实施
使用 DeepSpeed ZeRO Stage 2 微调数十亿参数的语言模型，否则该模型会溢出 GPU 内存。
使用 PyTorch FSDP 进行训练，跨 GPU 分片参数、梯度和优化器状态，并根据需要每层收集它们。
应用 ZeRO-Offload 将优化器状态推送到 CPU 内存，让单个 GPU 训练比其 VRAM 大很多倍的模型。
当 GPU 和 CPU 内存耗尽时，通过从 NVMe 存储流式传输参数分片，使用 ZeRO-Infinity 扩展万亿参数模型。
风险与防护栏
优化一项基准测试可以隐藏更广泛的系统弱点。
基础设施和维护成本常常被低估。
随着系统变得更加复杂，安全性和可观察性差距可能会扩大。
实施路线图
在实施之前定义延迟、质量和成本目标。
在实际负载和数据条件下进行基准测试。
仪器监控错误、漂移和用户影响。
在扩展之前准备回滚和事件响应路径。
不断探索
Free newsletter
Get the daily AI briefing
Three verified AI stories every weekday morning, written in plain English. Free forever, no ads.
One email each weekday. Unsubscribe in one click. We never sell or share your address.
Test yourself
Take the ZeRO and Sharded Optimizers quiz
Instant feedback on every answer, and a shareable certificate with a verifiable ID once you pass a course.
Support free AI education. AI Understanding is a 501(c)(3) nonprofit — no ads, no paywall, ever. Make a donation
常见问题
What is ZeRO and Sharded Optimizers?
ZeRO（零冗余优化器）通过跨 GPU 分片优化器状态、梯度和权重，消除了数据并行性中浪费的内存重复。它可以让您通过简单的数据并行性来训练巨大的模型，但只需要一小部分 GPU 内存。
与普通数据并行相比，ZeRO 消除了哪些冗余？
标准数据并行性在每个 GPU 上存储优化器状态、梯度和权重的完整副本； ZeRO 对这些进行分片，因此每个 GPU 仅保留一个切片。
为什么优化器状态通常是 Adam 占用最大内存的地方？
Adam 维护每个参数的运行估计，例如第一和第二时刻，与 fp32 主权重相结合可以使模型本身的大小相形见绌。
ZeRO 第 3 阶段分片有哪些第 1 阶段和第 2 阶段没有的分片？
第 1 阶段对优化器状态进行分片，第 2 阶段添加梯度，而第 3 阶段则进一步将模型参数分片到 GPU 上。
在 ZeRO Stage 3 中，GPU 如何获取层前向传递所需的完整参数？
在计算某个层之前，全收集会在每个 GPU 上组装其完整参数；一旦完成，非拥有的切片将被释放以回收内存。
哪个 PyTorch 功能本身实现了 ZeRO 风格的分片？
PyTorch 的完全分片数据并行 (FSDP) 对参数、梯度和优化器状态进行分片，动态收集和重新分片，镜像 ZeRO。