ZeRO-offload内存卸载｜CPU内存不该"躺平"

尽管 GPU 算力持续提升，但显存增量相对缓慢，GPU 的显存容量逐渐成为训练大模型的核心瓶颈。zero-offload异构训练借助 CPU 内存等外部存储资源，对模型状态进行卸载，从而实现单张 GPU 对超大模型的训练支持，同时尽可能减少对训练效率的影响。

1）大模型训练策略横向和纵向扩展方式的分类 2）ZeRO-offload 卸载策略设计方法，以及扩展到多GPU 3）为了不影响训练效率，提出CPU优化器加速方法和一步延迟参数更新策略

1，scale-out and scale-up 模型训练

为解决大模型训练的内存瓶颈问题，有两类核心训练策略：scale-out（横向扩展） 与 scale-up（纵向优化）。二者从不同维度突破显存限制，适用于不同的硬件环境与场景。

1.1，scale-out 横向扩展

scale-out 的核心思想是将模型、数据及训练过程中的中间状态拆分到多个 GPU（或多节点 GPU 集群）上，利用多设备的内存总和承载大模型，同时通过并行计算提升训练速度。其本质是 “用设备数量换内存与算力”，主要包括以下两种并行模式：模型并行和流水线并行。它们都将模型状态和剩余状态切分到不同的GPU设备上，来满足模型训练内存要求。但无法解决单 GPU 训练大模型的问题。

1.2，scale-up 纵向扩展

scale-up 的核心思想是在单设备（如单张 GPU）上，通过算法优化、内存管理等手段，最大化利用硬件资源，减少内存占用。主要包括三种技术路径：

主要有三种方式：

• • 激活值重计算。牺牲计算来节约显存。

训练过程中，反向传播需要依赖前向计算生成的 “激活值”（Activation），若将所有激活值保存在 GPU 显存中，会占用大量空间。激活值重计算策略通过 “牺牲部分计算量换取内存节约”：前向计算时仅保存少量关键激活值，反向传播时重新计算其他激活值。

• • 量化压缩。

传统模型参数与中间数据多采用 32 位浮点数（FP32）存储，量化压缩通过将数据精度降低，如转为 16 位浮点数 FP16、8 位整数 INT8，甚至 4 位整数 INT4，减少内存占用。同时，部分 GPU支持低精度计算指令，可在压缩内存的同时保证算力不下降。

• • 异构训练，使用外部的内存来扩充GPU的内存。

利用 CPU 内存、固态硬盘（SSD）等 GPU 之外的存储资源，将训练过程中暂不使用的数据 “卸载”（Offload）到外部内存，仅在需要时加载回 GPU。这种方式相当于 “用外部内存扩展 GPU 显存”，无需额外 GPU，适用于单 GPU 或小集群场景。

ZeRO-offload 就属于scale-up策略中的第三种方法。相比传统异构训练它不仅利用 CPU 内存扩充 GPU 显存，还充分调度 CPU 算力参与参数更新计算，在减少 GPU 内存占用的同时，最大限度降低了 “卸载” 对训练效率的拖累。

2，ZeRO-offload 策略流程

2.1，内存卸载设计思路

ZeRO-offload 的其设计需围绕三个核心目标：效率优先、高可扩展性、易用性，同时平衡显存节约、通信开销、计算延迟三者的关系。为了达到 SOTA 的模型训练效果，ZeRO-offload的模型训练需要考虑：

• • cpu 的算力限制，如何避免 CPU 算力短板拖累训练速度？

以 FP32 计算为例，CPU 的算力吞吐量仅为GPU 的十分之一到百分之一，若将大量计算任务转移到 CPU，会导致训练效率急剧下降。解决方案：精准划分 CPU 与 GPU 的任务边界。训练过程中的计算任务可分为两类：

高复杂度计算O (MB)：前向计算与反向传播，其计算量与参数数量、批次大小均相关，是训练的核心算力消耗环节。这类任务必须保留在 GPU 上执行，避免效率损失。其中 M 为模型参数数量，B 为批次大小，

低复杂度计算O (M)：如 LayerNorm 参数归一化、Adam动量计算优化器更新等，计算量仅与参数数量相关，复杂度远低于前向计算和反向传播。这类任务可安全转移到 CPU，即使 CPU 算力较低，也不会成为训练的瓶颈。

通过这种 “高复杂度任务 GPU 承担、低复杂度任务 CPU 承担” 的划分，ZeRO-offload 在利用 CPU 内存的同时，规避了 CPU 算力不足的问题。

• • 最小化通信量，如何减少 CPU-GPU 数据交换的开销？

方法：在训练过程中，将模型状态中的必要部分存储在 CPU 内存，并在 CPU 上直接完成参数更新，可将 CPU-GPU 之间的数据传输量压缩至最小。具体来说，每次迭代中仅需传输两类数据：

GPU → CPU：反向传播生成的 FP16 梯度，用于 CPU 端更新参数； CPU → GPU：更新后的 FP16 参数，用于下一轮前向计算。

这种设计使得单次迭代的 CPU-GPU 数据传输量仅为4M（2M FP16 梯度 + 2M FP16 参数，M 为参数数量）。

• • 最大化内存节约，最大化释放内存。

ZeRO-offload 显存优化逻辑：

GPU 显存保留：FP16 模型参数用于前向和反向计算、当前批次的激活值（部分通过重计算优化）、少量临时计算缓冲区。

CPU 内存承载：FP32 模型参数用于高精度更新、FP16 梯度从 GPU 传输而来、优化器状态如 Adam。

通过这一设计，相比不卸载的训练方式，GPU 显存占用可降低一半以上。例如，训练一个 10B 参数的模型，不卸载时 GPU 显存需占用 60GB 以上，而通过 ZeRO-offload，可将 GPU 显存占用降至 20GB 以内，实现单张 GPU的训练支持。下面是不同计算部分在不同内存上进行时，内存节约比例。

2.2 ，内存卸载流程

整个流程中，GPU 始终专注于高复杂度的前向和反向计算，CPU 仅负责低复杂度的参数更新与状态管理，二者通过最小化的通信量协同工作，实现 “显存节约” 与 “效率保障” 的平衡。

• • 前向计算&获得损失（GPU 执行）

GPU 加载保存在本地显存的FP16 模型参数，对输入数据执行前向计算，生成当前批次的预测结果与激活值（部分激活值暂存，部分标记为 “需重计算”）。根据预测结果与标签计算损失函数值，并初始化梯度为 0。

• • 反向传播（GPU 执行）

GPU 基于损失值执行反向传播，计算各层参数的FP16 梯度；对于标记为 “需重计算” 的激活值，实时重新执行前向计算以获取所需数据。

• • 梯度传输（GPU → CPU）

通过 PCIe 总线将 GPU 上的FP16 梯度传输至 CPU 内存，此过程仅需传输 2M 数据（M 为参数数量）。

• • 参数更新（CPU 执行）

CPU 调用优化器，基于保存在本地内存的FP32 模型参数、FP16 梯度及优化器状态，完成参数更新计算，得到新的 FP32 模型参数。

• • 参数转换与回传，迭代循环（CPU → GPU）

CPU 将更新后的FP32 模型参数转换为 FP16 精度，并通过 PCIe 总线传输回 GPU 显存，覆盖旧的 FP16 参数，为下一次迭代做准备。重复步骤，直至模型训练收敛。

图1，模型训练中数据流转图，总的参数量为M。

3，cpu上效率优化技术

当输入的序列长度较短时，cpu上的参数更新相对较慢，影响整个训练的效率，为此zero-offload 提出了针对cpu参数更新的两大优化方案。

3.1，cpu上的Adam优化

Adam 优化器其计算过程包含多个逐元素操作，如平方、开方、指数衰减，在 CPU 上执行时效率较低，成为潜在瓶颈。主要优化措施有三种。

• • 利用 SIMD 指令加速逐元素计算

CPU 支持 SIMD（Single Instruction, Multiple Data）指令集，可单次对 16 个 FP32 数据执行相同操作。ZeRO-offload 将 Adam 的逐元素计算（如乘法、平方）通过 SIMD 指令实现，将计算吞吐量提升 8-16 倍。例如，原本需 16 次操作的 16 个参数，通过一条 AVX-512 指令即可完成。

• • 多线程并行遍历参数

将模型参数按 CPU 核心数划分为多个块，每个核心负责一个块的 Adam 计算，通过多线程并行执行。同时，采用 “无锁编程” 避免线程间竞争，减少同步开销。例如，在 8 核 CPU 上，可将参数分为 8 块，并行更新，理论上可将计算时间缩短至原来的 1/8。

• • 预计算常量减少重复操作

Adam 中的偏差修正项仅与迭代次数t相关，无需逐参数计算。ZeRO-offload 在每次迭代前预计算这些常量，避免在参数遍历过程中重复计算，减少约 20% 的计算量。

经过上述优化，CPU 端 Adam 优化器的计算效率提升5-10 倍。

3.2，一步延迟参数更新

即使对 Adam 进行了优化，CPU 的计算速度仍略慢于 GPU。若严格等待 CPU 完成参数更新后，GPU 才开始下一轮计算，会导致 GPU 出现 “空闲等待”，降低整体训练效率。一步延迟参数更新技术通过 “异步调度” 解决这一问题。

参数缓冲区设计：在 GPU 显存中维护两个 FP16 参数缓冲区 ：“当前缓冲区”和“下一轮缓冲区”，分别用于当前迭代的前向反向计算和接收 CPU 更新后的参数。ZeRO-offload 通过 “双缓冲（Double Buffering）” 机制实现一步延迟更新的流程如下：

• • 迭代 1（初始化）：

GPU 使用 “当前缓冲区” 的初始参数执行前向和反向计算，生成 FP16 梯度；GPU 将 FP16 梯度传输至 CPU 内存，触发 CPU 开始执行参数更新；GPU 不等待 CPU 完成更新，立即开始迭代 2 的前向计算，仍使用 “当前缓冲区” 的初始参数，因 CPU 尚未完成迭代 1 的参数更新；当 CPU 完成迭代 1 的参数更新后，将更新后的 FP32 参数转换为 FP16，传输至 GPU 的 “下一轮缓冲区” 暂存。

• • 迭代 2：

GPU 完成迭代 2 的前向和反向计算，生成新的 FP16 梯度并传输至 CPU；CPU 接收迭代 2 的梯度，开始基于迭代 1 更新后的 FP32 参数执行迭代 2 的参数更新；GPU 切换缓冲区将 “下一轮缓冲区” 中迭代 1 更新后的参数作为 “当前缓冲区” 数据，开始迭代 3 的前向计算；CPU 完成迭代 2 的参数更新后，将新参数传输至 GPU 空闲的 “下一轮缓冲区”。此时原 “当前缓冲区” 已切换为计算用，另一缓冲区变为接收用。

• • 迭代 t（t≥2）：

GPU 始终使用 “当前缓冲区”（存储 CPU 第 t−1轮更新的参数）执行第 t 轮前向和反向计算，生成梯度并发送至 CPU；CPU 接收第 t轮梯度，基于第 t−1轮更新后的参数执行第 t轮参数更新；GPU 完成第 t轮计算后，立即切换缓冲区，用 “下一轮缓冲区”（存储 CPU 第 t 轮更新的参数）开始第 t+1轮计算；CPU 完成第 t轮更新后，将参数传输至 GPU 刚释放的缓冲区（原 “当前缓冲区”），作为下一次切换的 “下一轮缓冲区”。

图2，一步延迟参数更新流程

通过这种 “双缓冲交替 + 计算与更新并行” 的设计，GPU 与 CPU 的任务完全重叠，整个过程中，GPU 无需等待 CPU，彻底消除了 “空闲等待” 时间。

4，如何扩展训练

以上主要介绍了单GPU的内存卸载策略，那如何扩展的多GPU进行训练。

zero-offload 主要结合了zero-2的训练策略，每个GPU上保留一份参数副本，将梯度和优化器分片后，由不同的GPU进行通信和更新。

流程如下：

• • 在不同 GPU 间对梯度和优化器状态进行分区，每个 GPU 将自己拥有的分区卸载到 CPU 内存并在整个训练过程中保持
• • 反向传播时，在 GPU 上通过 reduce-scatter 计算和平均梯度，每个 GPU 仅将其分区中经过平均的梯度卸载到 CPU 内存
• • 梯度在 CPU 上可用后，各数据并行进程直接在 CPU 上并行更新优化器状态分区
• • 更新后，参数分区移回 GPU，随后在 GPU 上执行类似 ZeRO-2 的 all-gather 操作以收集所有参数

图3，多GPU上的zero-offload策略。

总结：

为了“扩充”GPU的显存，训练更大的模型，zero-offload策略将前向和反向计算放到GPU上进行，参数更新放到CPU上进行，并且将优化器状态和FP32的参数副本保留在CPU上，GPU上保留FP16格式的参数，GPU和CPU之间仅进行FP16格式的参数和梯度的通信，需要4M通信量，最大化的减少数据通信。为了不影响整个训练效率，提出cpu上加速Adam优化器和一步参数更新策略。

参考：arXiv:2101.06840