ZeRO：大模型训练的显存救星

一句话解释

ZeRO是微软提出的一套显存优化策略，它让模型训练时不再在每个GPU上完整复制优化器状态、梯度和参数，而是将这三部分拆分并均匀分布到所有参与训练的GPU上。这样每个GPU只存储自己的那一份，需要使用时通过通信动态获取，从而将显存占用从原来的O(N)降低到O(1/N)。

为什么会被关注

随着GPT、LLaMA等大模型参数突破千亿，单卡显存（80GB）完全装不下完整模型。传统的分布式策略如数据并行会浪费大量显存放在冗余的优化器状态上；模型并行则增加通信开销且实现复杂。ZeRO几乎在不牺牲计算效率的前提下，让用户用更少的GPU就能训练更大的模型，极大降低了大模型研发的硬件门槛。

另外，ZeRO被集成在DeepSpeed库中，与PyTorch无缝衔接，用户只需修改几行代码即可获得数倍显存节省。微软还开源了ZeRO-Infinity，进一步将显存卸载到CPU或NVMe上，实现单GPU训练万亿参数模型的愿景。

核心逻辑

ZeRO的核心思想是“分而治之”——消除数据并行中的冗余。它定义了三个阶段：ZeRO-1只切分优化器状态（如Adam的动量和方差），每个GPU只维护部分；ZeRO-2进一步切分梯度，每个GPU只存储自己负责的梯度和优化器状态；ZeRO-3则将参数也切分到各GPU。每次前向/反向传播时，通过all-gather通信收集所需参数或梯度，计算完成后即丢弃非本地的部分。

这种“以通信换显存”的策略，使得显存占用随GPU数量线性下降。在1024张GPU上，ZeRO-3可将单卡显存需求降低至原版的1/1024，同时通信开销可通过重叠和带宽优化控制在可接受范围内。

常见场景

大模型预训练：训练GPT-3、LLaMA等百亿到千亿参数模型时，ZeRO是标配方案。结合混合精度训练，通常使用ZeRO-2或ZeRO-3。

大模型微调（LoRA/QLoRA）：ZeRO-3可将基座模型参数分片存储，只更新少量适配器参数，显著降低微调所需的机器数。

单卡跑大模型推理：ZeRO-Offload或ZeRO-Infinity可将模型卸载到CPU内存，让消费级显卡也能运行70B模型进行推理演示。

容易混淆的点

ZeRO ≠ 模型并行（也叫张量并行）。模型并行是将单个算子（如Transformer层）切分到多卡计算，需要改动模型结构；而ZeRO只改变参数存储和通信方式，不改变计算图，用户代码无需修改。

ZeRO也不同于数据并行。数据并行中每张卡都存完整模型副本，ZeRO通过分片去除了冗余，因此常被称为“零冗余数据并行”。

注意ZeRO和DeepSpeed的关系：DeepSpeed是一个包含ZeRO、混合精度训练、梯度检查点等多项技术的框架，ZeRO是其中最具代表性的优化组件。