ZeRO技术原理详解：如何消除大模型训练中的内存冗余

大型深度学习模型，尤其是大语言模型（LLM），在带来显著性能提升的同时，也伴随着巨大的训练成本。训练数十亿到数万亿参数的模型面临诸多挑战，其中显存占用是最核心的瓶颈之一。传统的并行训练策略在内存效率与计算效率之间往往难以兼顾。

以 bf16 混合精度预训练一个 LLaMA2-7B 模型为例，其显存占用通常高达近120GB，远超出单张A100/H100（80GB）显卡的显存上限，因此多卡并行训练成为必然选择。显存占用主要由以下几部分构成：

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

ZeRO（Zero Redundancy Optimizer）是由微软在2020年前后提出的一项革命性内存优化技术。其核心目标是消除传统数据并行和模型并行中的内存冗余，同时保持低通信开销和高计算粒度。它使得模型规模能够与设备数量近似成比例地扩展，理论上支持训练超过1万亿参数的模型。

在实际评估中，ZeRO在 400 个 GPU 上成功训练了超过 1000 亿参数的模型，实现了超线性加速，吞吐量高达15 Petaflops。相比之前的方案，ZeRO 将可训练的模型大小提升了8倍，训练速度提升了10倍，而且无需引入复杂的模型并行策略。对于使用PyTorch等框架进行AI开发的工程师来说，ZeRO提供了一套覆盖训练全流程的端到端内存优化解决方案，其底层主要通过ZeRO-DP和ZeRO-R两大组件协同实现。

一、 ZeRO-DP：模型状态分片

ZeRO-DP是内存优化的核心。它不再让每张GPU都保存一份完整的模型、梯度和优化器状态副本，而是将这三者统称为 “模型状态” ，并对其进行智能分片，分散存储到集群的所有GPU上。这极大降低了单卡显存占用，使得用有限显存的GPU训练超大模型成为可能。

以使用 Adam优化器 和 64张GPU 为例，ZeRO-DP分为三个渐进的优化阶段：

1. 优化器状态分片 (P_os)：将优化器状态（动量、方差、参数备份）分散存储到所有GPU上，每张卡只保存一部分。这能实现约 4倍 的显存压缩，且通信开销与传统数据并行相同。
2. 增加梯度分片 (P_os+g)：在上一阶段基础上，进一步将梯度也进行分片存储。每张卡只负责计算和存储自己那部分参数对应的梯度。这能实现约 8倍 的显存压缩，通信开销仍保持不变。
3. 增加模型参数分片 (P_os+g+p)：在前两个阶段基础上，将模型参数本身也进行分片存储。每张GPU只负责维护和更新自己分配到的参数。此时显存减少与GPU数量成线性关系。在64张GPU上，理论最大可实现 64倍 的显存压缩。代价是通信量会增加大约50%。

二、 ZeRO-R：激活与临时内存优化

如果说ZeRO-DP优化了“模型状态”的内存，那么ZeRO-R（R 代表 Residual）的目标就是优化训练过程中除此以外的其他主要内存开销，实现对显存的全方位极致压榨。

1. 分区激活检查点

• 问题：在模型并行中，前向传播产生的中间结果（激活值）为了用于反向传播，通常会在多个GPU间完整复制，造成巨大的冗余存储。这是大模型训练的主要内存负担之一。
• 解决方案：ZeRO-R采用“分区激活检查点”技术。前向传播后，不保存完整的激活值副本，而是将其分割成分区，分散存储在不同GPU上。当反向传播需要时，再通过一次 all-gather 通信操作临时重建。结合激活重计算技术，甚至可以将这些分区卸载到CPU内存，将GPU上的激活内存占用降至近乎为零（以增加CPU-GPU通信为代价）。

2. 恒定大小缓冲区

• 问题：为了提升计算效率，框架通常会将所有参数融合到一个巨大的临时缓冲区中再调用核函数。对于万亿参数模型，这个缓冲区本身就可能大得令GPU无法承受。
• 解决方案：ZeRO设计并使用一个固定大小的融合缓冲区。无论模型多大，这个缓冲区都保持恒定。计算时，数据被分块送入这个固定缓冲区进行处理，从而打破了临时缓冲区内存与模型规模的线性增长关系。

3. 内存碎片整理

• 问题：大模型训练中，由于长期存活的激活检查点和短期存活的临时张量频繁分配释放，GPU显存会产生严重的内存碎片。这会导致即使总空闲内存足够，也因找不到连续大内存而报错，同时内存分配器性能下降。
• 解决方案：ZeRO在训练开始前，就为激活检查点和梯度等张量预先分配好连续的、长期持有的内存块。训练过程中产生的数据被实时复制到这些预分配的“内存池”中，而非动态申请。这极大地减少了内存碎片，提高了内存利用率和训练效率，这也是解决大模型训练通用挑战的关键工程优化之一。