Transformer——Q147 分析ZeRO优化器的内存节省量化模型

墨顿

墨衍会员 · AI 创作全网分发码龄22年

关注

原创最新推荐文章于 2026-11-06 11:13:53 发布 1.2k 阅读 AI 写同款· GEO 优化 ›

大模型训练ZeRO内存优化原理详解本文介绍 Zero Redundancy Optimizer (ZeRO)，如何减少内存使用并加速大规模模型的训练。阅读详情

该问题归类到Transformer架构问题集——训练与优化——分布式训练。请参考LLM数学推导——Transformer架构问题集。

1. 问题背景或来源

在深度学习领域，尤其是大语言模型（LLM）的训练过程中，模型参数规模呈爆炸式增长。从 GPT-3 的 1750 亿参数，到 GPT-4、PaLM 等更为庞大的模型，如此巨量的参数对计算设备的内存提出了极高要求。以 GPT-3 为例，若采用 FP32 格式存储其参数，仅模型参数就需约 700GB 内存，而当前主流的 GPU 设备，如 NVIDIA A100，内存容量仅为 80GB，远远无法满足单卡训练的需求。

传统的训练方式，如数据并行，虽然可以通过多卡并行来处理大规模数据，但随着模型规模的扩大，每卡需要存储完整的模型参数、优化器状态以及中间激活值，导致内存占用急剧增加。当模型参数和中间计算结果超出设备内存容量时，训练过程就会因内存不足而中断。因此，如何在有限的内存资源下训练超大规模模型，成为深度学习领域亟待解决的关键问题。

ZeRO（Zero Redundancy Optimizer）优化器正是为解决这一问题而诞生的。它通过对模型参数、优化器状态和梯度进行分片存储与计算，打破了传统训练方式的内存限制，极大地提升了内存使用效率，使得在普通计算集群上训练超大规模模型成为可能。

2. 技术原理或数学理论的解析

2.1 ZeRO 优化器的核心技术

ZeRO 优化器主要通过三个阶段（ZeRO - Stage1、ZeRO - Stage2、ZeRO - Stage3）来实现内存节省，每个阶段在内存优化和通信开销之间进行不同的权衡。

ZeRO - Stage1：梯度分片

在传统的数据并行训练中，每个计算节点都保存完整的梯度信息。而在 ZeRO - Stage1 中，梯度被均匀地分片存储在不同的节点上。假设共有 N 个计算节点，模型的总梯度为 $\nabla\theta$ ，那么每个节点仅存储 $\frac{1}{N}$ 的梯度信息。在参数更新时，各节点通过 All - Reduce 操作聚合梯度，再进行参数更新。

从内存占用角度来看，设模型参数数量为 P，每个参数占用的字节数为 b（如 FP32 格式下 $b = 4$ 字节），则传统方式下每个节点的梯度内存占用为 $P\times b$ 。而在 ZeRO - Stage1 中，每个节点的梯度内存占用降低为 $\frac{P\times b}{N}$ 。

ZeRO - Stage2：优化器状态分片

除了梯度分片，ZeRO - Stage2 进一步对优化器状态进行分片。优化器状态（如 Adam 优化器中的一阶矩、二阶矩）在传统训练中也是每个节点完整保存。在 ZeRO - Stage2 中，优化器状态同样被划分为 N 份，每个节点仅保存和更新属于自己的那部分优化器状态。

设优化器状态的总大小与模型参数大小相同（实际可能因优化器类型有所差异），在传统方式下，每个节点的优化器状态内存占用为 $P\times b$ 。而在 ZeRO - Stage2 中，每个节点的优化器状态内存占用降低为 $\frac{P\times b}{N}$ 。结合梯度分片，此时每个节点的总内存节省为 $P\times b + P\times b - (\frac{P\times b}{N} + \frac{P\times b}{N}) = 2P\times b(1 - \frac{1}{N})$ 。

ZeRO - Stage3：参数分片

ZeRO - Stage3 是最激进的内存优化阶段，它将模型参数也进行分片存储。每个节点仅保存和更新模型参数的一部分，同时在计算过程中，通过通信机制获取必要的参数来完成前向传播和反向传播。这样，每个节点的模型参数内存占用降低为 $\frac{P\times b}{N}$ 。

此时，每个节点的总内存占用为 $\frac{P\times b}{N} + \frac{P\times b}{N} + \frac{P\times b}{N} = \frac{3P\times b}{N}$ ，相较于传统方式，内存节省比例高达 $1 - \frac{3}{N}$ 。

2.2 内存节省量化模型推导

设 $M_{traditional}$ 为传统数据并行训练方式下每个节点的内存占用， $M_{ZeRO - Stage1}$ 、 $M_{ZeRO - Stage2}$ 、 $M_{ZeRO - Stage3}$ 分别为 ZeRO 优化器三个阶段每个节点的内存占用。

传统数据并行： $M_{traditional} = P\times b + P\times b + P\times b = 3P\times b$ （分别为模型参数、优化器状态、梯度的内存占用）

ZeRO - Stage1： $M_{ZeRO - Stage1} = P\times b + P\times b + \frac{P\times b}{N} = P\times b(2 + \frac{1}{N})$ ，内存节省量 $\Delta M_1 = 3P\times b - P\times b(2 + \frac{1}{N}) = P\times b(1 - \frac{1}{N})$

ZeRO - Stage2： $M_{ZeRO - Stage2} = P\times b + \frac{P\times b}{N} + \frac{P\times b}{N} = P\times b(1 + \frac{2}{N})$ ，内存节省量 $\Delta M_2 = 3P\times b - P\times b(1 + \frac{2}{N}) = 2P\times b(1 - \frac{1}{N})$

ZeRO - Stage3： $M_{ZeRO - Stage3} = \frac{P\times b}{N} + \frac{P\times b}{N} + \frac{P\times b}{N} = \frac{3P\times b}{N}$ ，内存节省量 $\Delta M_3 = 3P\times b - \frac{3P\times b}{N} = 3P\times b(1 - \frac{1}{N})$

从上述公式可以清晰地看出，随着计算节点数量 N 的增加，ZeRO 优化器各阶段的内存节省量逐渐增大，尤其是在 ZeRO - Stage3 阶段，内存节省效果最为显著。

3. 根因分析

3.1 传统训练方式的内存瓶颈

传统数据并行训练方式为了保证各节点模型参数的一致性，每个节点都需要保存完整的模型参数、优化器状态和梯度。这种方式在模型规模较小时可以正常运行，但随着 LLM 模型参数呈指数级增长，内存占用也随之剧增，很快就会超过单个计算节点的内存容量，导致训练无法进行。

3.2 ZeRO 优化器的创新思路

ZeRO 优化器打破了传统方式中每个节点必须保存完整信息的固有模式，基于分布式系统中节点间通信的可行性，将梯度、优化器状态和模型参数进行分片存储。通过节点间的通信和协同计算，在保证训练结果一致性的前提下，大幅减少了每个节点的内存占用，从而突破了内存限制。其核心在于利用分布式系统的特性，在内存节省和通信开销之间找到平衡，实现高效的大规模模型训练。

4. 在 LLM 中的使用示例

4.1 GPT - 3 训练

以 GPT - 3 的训练为例，其拥有 1750 亿个参数。若采用 FP32 格式存储，参数占用内存约 700GB。假设使用 8 个计算节点进行训练：

传统数据并行：每个节点需存储约 700GB 的模型参数，加上优化器状态和梯度，总内存占用远超当前 GPU 设备的内存容量。

ZeRO - Stage1：每个节点的梯度内存占用降低为 $\frac{700}{8} = 87.5$ GB，相较于传统方式，仅梯度部分就节省了大量内存。

ZeRO - Stage2：结合优化器状态分片，每个节点的内存节省量进一步增加，缓解了内存压力。

ZeRO - Stage3：将模型参数也进行分片，每个节点仅需存储 $\frac{700}{8} = 87.5$ GB 的模型参数，使得在普通计算集群上训练 GPT - 3 成为可能。

4.2 BERT - Large 微调

BERT - Large 模型包含 3.4 亿个参数，采用 FP32 格式存储约占用 136GB 内存。在微调过程中，若使用 4 个计算节点：

传统方式：每个节点内存压力较大，可能因内存不足导致训练中断。

ZeRO - Stage3：每个节点仅需存储 $\frac{136}{4} = 34$ GB 的模型参数，同时优化器状态和梯度也进行分片存储，显著降低了内存需求，保证了微调过程的顺利进行。

4.3 OPT 模型训练

OPT（Open Pretrained Transformer）模型同样具有庞大的参数量。在训练过程中，使用 ZeRO 优化器可以有效减少内存占用。例如，当使用 16 个计算节点时，ZeRO - Stage3 可以将每个节点的内存占用降低到原来的 $\frac{1}{16}$ ，使得训练能够在有限的内存资源下高效进行。

5. 优缺点分析

5.1 优点

显著节省内存：通过梯度、优化器状态和模型参数分片，ZeRO 优化器能够大幅降低每个计算节点的内存占用，使得超大规模模型的训练成为可能。

灵活的配置：提供三个不同阶段的内存优化策略，用户可以根据计算节点数量、通信带宽和内存容量等实际情况，选择合适的阶段，在内存节省和通信开销之间找到最佳平衡。

兼容性强：可以与多种深度学习框架和模型结构相结合，无需对模型进行大量修改即可应用，具有广泛的适用性。

5.2 缺点

增加通信开销：由于参数、梯度和优化器状态的分片，节点之间需要频繁进行通信来聚合信息和协同计算，尤其是在 ZeRO - Stage3 阶段，通信开销显著增加。当网络带宽不足时，通信延迟可能会成为训练效率的瓶颈。

训练复杂度提高：分片存储和计算增加了训练过程的复杂性，对分布式系统的管理和调试提出了更高的要求。同时，不同阶段的内存优化策略需要用户根据实际情况进行合理选择和配置，增加了使用难度。

可能影响训练稳定性：在参数分片和通信过程中，可能会引入一些不确定性因素，影响训练的稳定性。例如，通信延迟或数据传输错误可能导致参数更新不一致，从而影响模型的收敛速度和最终性能。

6. 优化策略分析

6.1 通信优化

压缩通信数据：采用梯度压缩技术，如量化、稀疏化等，减少节点间传输的数据量。例如，将梯度从 FP32 格式转换为 FP16 或更低精度格式进行传输，在保证精度损失可接受的前提下，降低通信开销。

重叠计算与通信：利用计算和通信可以重叠进行的特性，在节点进行本地计算的同时，进行数据传输。例如，在计算前向传播的过程中，同时将上一轮的梯度发送出去，提高资源利用率。

6.2 混合并行策略

结合数据并行、模型并行和张量并行等其他并行策略，与 ZeRO 优化器协同使用。例如，在 ZeRO - Stage3 的基础上，再采用模型并行将模型的不同层分布到不同节点，进一步降低每个节点的内存压力，同时平衡计算负载和通信开销。

6.3 动态资源分配

根据训练过程中各节点的内存使用情况、计算负载和通信延迟等动态信息，实时调整 ZeRO 优化器的配置。例如，当某个节点内存占用过高时，自动调整参数分片策略，将部分数据转移到其他节点，实现资源的动态平衡。

7. 代码示例（基于 PyTorch 和 DeepSpeed）

import torch
import deepspeed

# 定义模型
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = torch.nn.Linear(1024, 2048)
        self.fc2 = torch.nn.Linear(2048, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = MyModel()
# ZeRO配置
config = {
    "train_batch_size": 32,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001
        }
    },
    "zero_optimization": {
        "stage": 3,  # 使用ZeRO - Stage3
        "overlap_comm": True,  # 启用计算与通信重叠
        "contiguous_gradients": True,  # 连续梯度以提高通信效率
        "reduce_bucket_size": model.config.hidden_size * model.config.hidden_size,
        "stage3_prefetch_bucket_size": 0.9 * model.config.hidden_size * model.config.hidden_size,
        "stage3_param_persistence_threshold": 0
    }
}

# 初始化ZeRO
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config=config
)

# 训练循环
for epoch in range(10):
    for batch in dataloader:
        inputs, labels = batch
        outputs = model_engine(inputs)
        loss = torch.nn.functional.cross_entropy(outputs, labels)
        model_engine.backward(loss)
        model_engine.step()

8. 代码解读

模型定义：首先定义了一个简单的包含两个全连接层的神经网络模型 MyModel。

ZeRO 配置：通过字典 config 对 ZeRO 优化器进行配置。设置 stage 为 3，表示使用 ZeRO - Stage3 进行内存优化；overlap_comm 设置为 True，启用计算与通信重叠，以提高效率；contiguous_gradients 用于使梯度连续，优化通信过程；同时还设置了与内存和通信相关的其他参数，以适应模型和训练需求。

初始化 ZeRO：使用 deepspeed.initialize 函数对模型、优化器等进行初始化，将普通的 PyTorch 模型转换为支持 ZeRO 优化的模型引擎。

训练循环：在训练循环中，与普通的 PyTorch 训练类似，进行前向传播计算损失，然后通过 model_engine.backward 进行反向传播计算梯度，最后使用 model_engine.step 更新模型参数。在这个过程中，ZeRO 优化器会自动根据配置进行梯度、优化器状态和模型参数的分片处理和通信操作。

9. 总结

ZeRO 优化器通过梯度、优化器状态和模型参数分片的创新方式，为解决大语言模型训练中的内存瓶颈问题提供了有效的解决方案。通过详细的内存节省量化模型推导，我们可以清晰地看到其在不同阶段的内存优化效果。在 LLM 的实际训练中，ZeRO 优化器已经展现出强大的能力，使得超大规模模型的训练在普通计算集群上成为可能。

然而，ZeRO 优化器也存在通信开销大、训练复杂度高和可能影响训练稳定性等问题。通过采用通信优化、混合并行策略和动态资源分配等优化方法，可以在一定程度上缓解这些问题。随着深度学习模型规模的不断扩大，ZeRO 优化器及其相关技术将持续发展和完善，为大语言模型的训练和应用提供更有力的支持。