ZeRO 和 FSDP 是什么？它们如何降低训练显存？

简化版

经典 DDP 在每个 rank 复制参数、梯度和优化器状态。ZeRO 逐级在数据并行组内分片这些状态：Stage 1 分优化器状态，Stage 2 再分梯度，Stage 3 再分参数；需要计算时通过 collective 恢复当前单元所需数据。PyTorch FSDP 与 ZeRO-3 的核心思想相近，都会完全分片参数、梯度和优化器状态并按单元 All-Gather，但包装粒度、状态字典、预取与运行时语义不应简单视为完全相同实现。

详细版

问题：数据并行的显存冗余

用 Adam + 混合精度训练时，每个参数要占的「模型状态」显存约为（常见估算）：

FP16 参数 2B + FP16 梯度 2B + FP32 参数副本 4B + FP32 动量 m 4B + FP32 方差 v 4B ≈ 16 字节/参数

所以一个 7.5B 模型光「模型状态」就约 120GB。在普通数据并行下，N 张卡各存一份这 120GB，冗余 N 倍，卡再多也放不下更大的模型。

ZeRO 三级分片

级别	分片内容	显存节省（DP 度 = N）	通信开销
ZeRO-1	优化器相关状态	DP 很大时、特定 16B 账本上限趋近约 4×	与 DDP 同量级，模式不同
ZeRO-2	+ 梯度	同一账本上限趋近约 8×	常用 Reduce-Scatter/All-Gather
ZeRO-3	+ 参数	长期状态理想上接近按 DP 度分摊	增加按需参数 All-Gather

ZeRO-1：在本文传统账本里把 m、v 与可能的 FP32 master 等优化器相关状态分片，每卡只保留其负责部分；具体字节归属依实现。
ZeRO-2：梯度也分片，进一步省。
ZeRO-3：连参数本身也分片，每卡平时只存 1/N 的参数，用到某层时临时 all-gather 出完整参数、算完立即释放。显存随卡数近似线性下降，能训最大的模型。

FSDP

PyTorch 原生的全分片数据并行，与 ZeRO-3 在“全状态分片 + 按需聚合”上相近，但并非所有配置与接口逐项等价。
前向到某层时 all-gather 聚合该层完整参数 → 计算 → 释放；反向同理，再 reduce-scatter 分发梯度。
它可与 PyTorch 训练栈集成，是否采用取决于模型结构、并行组合、checkpoint 和性能验证。

完整版教学

一、先看清冗余到底浪费在哪

标准数据并行的最大问题是重复存储。N 张卡组成数据并行组，每张卡都持有：完整参数、完整梯度、完整优化器状态。这三样里，优化器状态最占地方——用 Adam 时，每个参数要额外存 FP32 的一阶动量 m、二阶动量 v，加上 FP32 参数副本，混合精度下「模型状态」高达约 16 字节/参数，其中优化器相关就占约 12 字节。

关键洞察：数据并行里，每张卡其实只负责更新”一部分”参数就够了，没必要每张卡都存全套。 ZeRO 就是把「本可以分工的东西」真正分给各卡，消除冗余。

二、ZeRO 的核心思想：分片 + 按需聚合

ZeRO（Zero Redundancy Optimizer）不改变数据并行「各算各的数据、同步梯度」的本质，而是把模型状态沿卡分片存储，需要完整数据时临时用集合通信（all-gather / reduce-scatter）拼出来，用完就丢。

平时：每张卡只存自己负责的那 1/N 份状态，显存大降。
用时：需要完整参数/梯度时，临时从各卡聚合，算完立即释放，不长期占显存。

这是“以 collective 通信和调度复杂度换长期状态显存”。是否划算取决于分片单元、网络拓扑、计算通信重叠和峰值 All-Gather 缓冲，网络较慢时可能显著拖累吞吐。

三、三级分片：省得越多，通信越多

ZeRO 分三级，逐级把更多东西分片：

ZeRO-1（分片优化器相关状态）：把 Adam m/v 以及实现中归入该状态的 master 参数分片。在 16B/参数的简化账本里，DP 度趋大时每卡从 16B 趋近参数+梯度的 4B，因此理论上限约 4 倍；有限卡数、缓冲和实现会改变结果。

ZeRO-2（再分片梯度）：每卡只长期保留负责参数的梯度。在同一账本且 DP 度很大时理论上限趋近 8 倍；通信量可与 DDP 同量级，但 collective 形态、重叠和临时缓冲不同。

ZeRO-3（连参数也分片）：最激进，参数本身也分片，每卡平时只存 1/N 的参数。前向/反向到某层时，all-gather 临时拼出该层完整参数 → 计算 → 立即释放。显存随卡数近似线性下降，理论上卡越多能训越大的模型。代价是多了参数聚合的通信（每层前向反向都要 all-gather 参数）。

记忆钩子：Stage 1 分优化器状态，Stage 2 加梯度，Stage 3 再加参数；分得越彻底，长期状态越少，对通信、预取和分片粒度的要求越高。

四、FSDP 与 ZeRO-3 的共同点和差异

FSDP 是 PyTorch 的全分片数据并行实现，在核心机制上与 ZeRO-3 相近：

模型按「分片单元」（如每个 Transformer 层）拆开，参数、梯度、优化器状态全部分片；
前向：进入某层时 all-gather 聚合出该层完整参数 → 计算 → 用完立即释放显存；
反向：同样 all-gather 参数算梯度，再用 reduce-scatter 把梯度分发回各自负责的卡；
各卡用本地那份梯度更新本地那份参数。

二者都可配合混合精度与激活重计算，但 auto-wrap/分片单元、预取、offload、checkpoint 格式和与其他并行的组合方式不同。迁移时应比较峰值显存、通信 Trace、吞吐和状态保存恢复，而不是只替换类名。

五、ZeRO 和三种并行是什么关系

初学者容易混淆 ZeRO 和张量/流水线并行。关键区分：

ZeRO / FSDP 本质是”去冗余的数据并行”：它没有把单层计算拆到多卡（那是张量并行），也没按层分段（那是流水线并行）。每张卡仍然算完整的前向/反向，只是参数/梯度/优化器状态的存储被分片了。
因此 ZeRO 可以和 TP、PP 叠加：3D 并行解决”计算怎么拆”，ZeRO 解决”数据并行维度上的存储冗余”。

可用一句边界来区分：TP/PP 改变单个样本的计算如何跨设备分工；ZeRO/FSDP 主要消除数据并行副本的状态冗余，但计算前仍需按需聚合参数。它们可以组合，却会共享通信带宽，并非没有耦合成本。

六、按 8 张卡估算分片收益

仍以 7B、16 字节/参数的混合精度 Adam 粗算，普通数据并行每卡模型状态约 112GB。8 卡 ZeRO-1 只分优化器相关状态，参数与梯度仍复制；ZeRO-3/FSDP 若忽略临时聚合，三类状态都分片，理论每卡约 112/8=14GB。

方案	分片内容	每卡长期状态趋势	主要代价
DDP	无	约 112GB	冗余最大
ZeRO-1	优化器状态	明显下降	少量额外管理
ZeRO-2	再加梯度	继续下降	reduce-scatter
ZeRO-3/FSDP	再加参数	理论接近 1/N	频繁 all-gather

DDP states per GPU ≈ 112 GB
ZeRO-3 ideal shard ≈ 112 / 8 = 14 GB
actual peak = shard + activations + gathered layer + buffers

理论值不含激活、临时完整参数、通信 buffer 和碎片，所以不能直接当显卡最低容量。实际选分片单元时要平衡峰值显存与 collective 次数，并验证 checkpoint 在不同并行度下能否正确保存和恢复。

七、常见误区与追问

误区：ZeRO 是一种模型并行（拆计算）。 不是，它是分片存储的数据并行，计算仍在每卡完整进行。
误区：优化器状态不占多少显存。 FP32 Adam m/v 已约 8B/参数，若另有 FP32 master 再增 4B；具体布局依实现，但通常是重要分片对象。
追问：为什么每参数约 16 字节？ 混合精度 Adam：FP16 参数 2 + FP16 梯度 2 + FP32 参数副本 4 + FP32 动量 4 + FP32 方差 4 = 16。
追问：ZeRO-3 为什么能近线性省显存？ 参数/梯度/优化器状态全按卡数分片，每卡只存 1/N，用时临时聚合即释放。
追问：FSDP 和 ZeRO-3 有什么关系？ 全状态分片思想相近，但包装、预取、状态字典、offload 和生态接口不同，不能声称逐项等价。
追问：ZeRO-3 的代价？ 每层前向/反向都要 all-gather 参数，通信量增加，需高带宽。
追问：ZeRO 和张量并行冲突吗？ 不冲突，正交——一个分存储、一个分计算，常一起用。

八、加强记忆

把 ZeRO 记成逐级去冗余：Stage 1 分优化器状态，Stage 2 加梯度，Stage 3 再加参数；计算当前单元前 All-Gather，梯度再 Reduce-Scatter。4×、8×与按 DP 度分摊都是传统 16B 账本下的理想上限，不是每个集群的实测保证。FSDP 与 Stage 3 思想相近而实现语义不同。它们主要分片数据并行状态，TP/PP 则改变计算分工；组合时仍要共同核算峰值聚合缓冲、通信带宽、吞吐与 checkpoint 可恢复性。

简化版

详细版