当前位置：首页 > news >正文

大模型训练显存优化全方案：ZeRO、Offload与重计算技术对比

news 2025/10/21 9:12:42

点击 “AladdinEdu，同学们用得起的【H卡】算力平台”，注册即送-H卡级别算力，一站式沉浸式云原生集成开发环境，80G大显存多卡并行，按量弹性计费，教育用户更享超低价。

当大语言模型的参数量从亿级迈向万亿级，我们面临的不再仅仅是算力瓶颈，更是显存墙的严峻挑战。训练一个1750亿参数的模型，仅存储FP32格式的参数、梯度和优化器状态就需要超过2TB的显存，这远远超过了当前最强单卡GPU的容量。如何在有限的硬件条件下突破显存限制，让普通研究机构和开发者也能参与大模型训练，成为了AI领域亟待解决的核心问题。

本文将从显存占用的本质出发，系统解析ZeRO、Offload和重计算三大核心技术的原理与实现，并提供在混合精度训练场景下的综合优化策略，为读者构建完整的大模型训练显存优化知识体系。

一、大模型训练的显存占用分析

在深入优化技术之前，我们必须精确理解训练过程中显存都被哪些组件消耗。

1.1 显存组成模型

训练过程中的显存占用主要由以下几部分组成：

模型参数：模型中所有可训练权重的存储。对于Transformer架构，这包括Attention层的QKV投影矩阵、输出投影矩阵、FFN层的两个线性变换矩阵等。
梯度：反向传播过程中计算得到的损失函数对每个参数的偏导数。
优化器状态：优化算法所需的中间变量。不同的优化器状态大小不同：
- SGD：无额外状态（或动量），与参数量相同。
- Adam：动量（momentum）和方差（variance），每个都是FP32，因此优化器状态是参数量的2倍（FP32模型）或8倍（混合精度训练，后续详述）。
激活值：前向传播过程中每一层的中间输出，在反向传播时用于计算梯度。这是除了参数相关存储外最大的显存占用项。
临时缓冲区：诸如矩阵乘法、All-Reduce通信等操作的中间工作区。

1.2 显存占用的数学建模

以一个参数量为 Φ 的模型为例，我们可以建立如下显存占用模型：

仅存储模型（推理）：
- FP32: 4Φ 字节
- FP16/BF16: 2Φ 字节

完整训练（朴素数据并行）：

FP32训练:

参数: 4Φ
梯度: 4Φ  
优化器状态(Adam): 8Φ (动量4Φ + 方差4Φ)
激活值: A
总显存 ≈ 16Φ + A

混合精度训练（最常见）：

参数(FP16): 2Φ
梯度(FP16): 2Φ
优化器状态(FP32): 8Φ  # 主参数副本、动量、方差均为FP32
激活值: A
总显存 ≈ 12Φ + A

从这个模型可以看出，优化器状态是训练时显存占用的最大头。对于一个70亿参数的模型（Φ=7×10⁹），混合精度训练仅参数、梯度、优化器状态就需要约 12 × 7×10⁹ ≈ 84GB 显存，这已经超过了单张A100（80GB）的容量，更不用说激活值了。

二、 Zero Redundancy Optimizer (ZeRO) 深度解析

ZeRO是微软DeepSpeed库的核心技术，其核心思想是通过分区消除数据并行中的冗余存储。

2.1 数据并行的显存冗余

在传统数据并行（DP）中，每个GPU都复制完整的模型参数、梯度和优化器状态。当模型大到单个GPU无法容纳时，这种冗余就成为瓶颈。ZeRO通过将状态分区到多个GPU上来解决这个问题。

2.2 ZeRO的三个阶段

ZeRO的优化分为三个渐进的阶段，其优化过程与内存节省效果如下图所示：

a) ZeRO-Stage 1 (Os Partitioning)：优化器状态分区

原理：将优化器状态（FP32的主参数、动量、方差）均匀分区到所有数据并行进程（GPU）上。每个GPU只负责存储和更新其中一部分优化器状态。
实现机制：
1. 前向和反向传播与传统DP相同，每个GPU都有完整的FP16参数和计算完整的梯度。
2. 在优化器步骤前，通过All-Gather操作从所有GPU收集完整的梯度（每个GPU已经通过平均得到完整梯度）。
3. 每个GPU只更新自己负责的那部分优化器状态和参数。
4. 通过All-Gather或广播将更新后的参数同步到所有GPU。
显存节省：优化器状态从8Φ减少到8Φ/N（N为GPU数量），总显存从12Φ + A降至(4 + 8/N)Φ + A。当N较大时，趋近于4Φ + A。

b) ZeRO-Stage 2 (Gradient Partitioning)：梯度分区

原理：在Stage 1的基础上，进一步将梯度分区。每个GPU只存储与其负责的优化器状态对应的那部分梯度。
实现机制：
1. 前向传播不变。
2. 反向传播计算得到完整梯度后，每个GPU只保留自己负责的那部分梯度，其余丢弃。
3. 通过Reduce-Scatter操作高效地完成梯度分区和平均：每个GPU只得到自己那部分梯度的平均值。
4. 优化器步骤与Stage 1类似，但每个GPU已经有了自己需要的梯度部分。
显存节省：梯度从2Φ减少到2Φ/N，总显存从(4 + 8/N)Φ + A降至(2 + 8/N)Φ + A。

c) ZeRO-Stage 3 (Parameter Partitioning)：参数分区

原理：在Stage 2的基础上，进一步将模型参数本身分区。每个GPU只存储与其负责的优化器状态和梯度对应的那部分参数。
实现机制：
1. 前向传播时，当需要当前GPU不持有的参数时，通过All-Gather从其他GPU收集。
2. 计算完成后，立即释放不属于自己分区的参数。
3. 反向传播类似，也需要通过All-Gather获取参数，计算得到梯度后，通过Reduce-Scatter进行梯度分区。
显存节省：参数从2Φ减少到2Φ/N，总显存从(2 + 8/N)Φ + A降至(0 + 10/N)Φ + A（因为需要一些临时缓冲区）。当N足够大时，显存占用主要由激活值A决定。

2.3 ZeRO的通信分析

ZeRO通过增加通信量来换取显存节省，其通信量约为传统数据并行的1.5倍。但在现代高速互联（如InfiniBand）的集群上，这个开销通常是可接受的，因为它使得训练超大模型成为可能。

三、 Offload技术：利用主机内存

当GPU显存仍然不足时，Offload技术将部分数据卸载到CPU主机内存甚至NVMe SSD。

3.1 CPU Offload

ZeRO-Offload：将优化器状态和梯度卸载到CPU内存，在CPU上执行优化器步骤。
- 工作流程：
  1. GPU进行前向和反向传播。
  2. 梯度被复制到CPU。
  3. CPU更新优化器状态和参数。
  4. 更新后的参数复制回GPU。
- 优势：将优化器状态的显存占用降为0，仅需2Φ + A显存（参数+激活值）。
- 瓶颈：GPU-CPU之间的数据传输可能成为瓶颈，需要仔细的重叠计算与通信。
ZeRO-Infinity：ZeRO-Offload的增强版，支持将参数、梯度和优化器状态全部卸载到CPU/NVMe，并结合更精细的并行策略。
- 技术特点：
  - 带宽优化的分层卸载：频繁访问的数据放在CPU，不频繁的放在NVMe。
  - 张量切片：将大张量切片，仅按需加载所需切片。
  - 异步预取：提前预测并加载下一步计算需要的数据。

3.2 NVMe Offload

对于远超CPU内存容量的大模型，可以使用NVMe SSD作为扩展存储。

适用场景：当模型状态（参数+梯度+优化器）超过CPU内存容量时。
性能考量：NVMe的带宽远低于GPU显存和CPU内存，需要极其精细的数据调度来隐藏延迟。

四、激活重计算（Gradient Checkpointing）

激活值占据了除模型状态外的大部分显存，尤其是在Transformer模型中，激活值与序列长度、隐藏维度的平方成正比。

4.1 基本原理

激活重计算，也称为梯度检查点，其核心思想是以计算时间换取显存空间。它只保存网络中的部分子图（检查点）的激活值，在反向传播时根据保存的检查点重新计算被丢弃的激活值。

4.2 实现策略

均匀策略：每N层设置一个检查点。例如，一个100层的网络，每10层保存一个检查点，则只需要保存10层的激活值，而不是100层。
动态规划策略：通过动态规划算法选择最优的检查点位置，在给定显存预算下最小化重计算开销。

4.3 显存-计算权衡

假设一个网络有L层，激活值总大小为A：

不启用重计算：显存占用包含所有激活值A，无额外计算。
均匀策略（每K层检查点）：显存占用降为A/K，但需要额外重计算大约(K-1)/K比例的前向传播计算量。
极限情况（只保存输入）：显存占用最小，但需要几乎重新运行整个前向传播。

在实践中，对于Transformer模型，激活重计算通常可以节省3-5倍的激活显存，而额外增加20-30%的训练时间。

五、混合精度训练的显存-计算平衡

混合精度训练是另一个关键的优化技术，但它也带来了独特的显存-精度平衡挑战。

5.1 混合精度的工作流程

FP16前向：使用FP16权重计算前向传播，得到FP16激活值。
Loss Scaling：将损失值放大（如1024倍）以避免梯度下溢。
FP16反向：计算FP16梯度。
FP32优化：将梯度反缩放后，在FP32主副本上更新优化器状态和参数。
权重更新：将更新后的FP32参数截断为FP16，用于下一轮前向。

5.2 显存-精度平衡策略

BF16 vs FP16：
- FP16：表示范围小（~5.96×10⁻⁸ to 65504），易出现梯度下溢/上溢。
- BF16：具有与FP32相同的指数范围，但精度较低。更适合大模型训练，减少需要Loss Scaling的调参。
- 选择策略：如果硬件支持，优先使用BF16以获得更稳定的训练。
精度配置调优：
- 参数精度：核心参数使用FP16/BF16，优化器状态使用FP32。
- 激活精度：激活值通常使用与参数相同的精度，但对于敏感操作（如LayerNorm、Softmax）可保持FP32。
- 梯度精度：梯度通常与参数精度一致，但通信时可使用FP16以减少带宽。
动态Loss Scaling：
- 自动检测梯度溢出，动态调整缩放因子。
- 避免手动调参，同时保证训练稳定性。

六、综合方案对比与实践配置

6.1 技术对比表

技术	显存节省	计算/通信开销	适用场景
ZeRO-Stage 1	中等 (~4Φ)	低	优化器状态占主导的中等模型
ZeRO-Stage 2	高 (~6Φ)	中	大多数大模型训练场景
ZeRO-Stage 3	极高 (~10Φ)	高	极端大模型，GPU数量充足
CPU Offload	极高 (优化器状态→0)	非常高	单卡或少量卡训练大模型
激活重计算	节省3-5x激活值	增加20-30%计算	激活值占主导，序列长的模型
混合精度	节省50%参数显存	可忽略	所有训练场景，硬件支持时

6.2 实战配置示例

以下是一个使用DeepSpeed配置综合优化方案的示例，用于训练一个130亿参数的模型：

{"train_batch_size": 1024,"gradient_accumulation_steps": 1,"fp16": {"enabled": true,"loss_scale": 0,"loss_scale_window": 1000,"initial_scale_power": 16},"zero_optimization": {"stage": 2,"allgather_partitions": true,"allgather_bucket_size": 2e8,"reduce_scatter": true,"reduce_bucket_size": 2e8,"overlap_comm": true,"contiguous_gradients": true},"activation_checkpointing": {"partition_activations": true,"contiguous_memory_optimization": true,"cpu_checkpointing": false},"offload_optimizer": {"device": "cpu","pin_memory": true}
}