Deep Optimizer States: Towards Scalable Training of Transformer Models Using Interleaved Offloading

该论文提出了名为“Deep Optimizer States"的新型技术，通过利用训练过程中 GPU 内存利用率的波动动态地在主机与 GPU 内存间交错卸载优化器状态，并结合性能模型智能调度更新阶段，从而在大规模 Transformer 模型训练中实现了比现有最先进方法快 2.5 倍的迭代速度。

要点

这篇论文介绍了一种名为 "Deep Optimizer States"（深度优化器状态） 的新技术，旨在解决训练超大型人工智能模型（如大语言模型 LLM）时遇到的“内存墙”难题。

为了让你轻松理解，我们可以把训练 AI 模型的过程想象成在一个拥挤的厨房里做一道极其复杂的菜。

1. 遇到的难题：厨房太小，食材太多

• 背景：现在的 AI 模型（比如 GPT-3、LLaMA）非常巨大，拥有数百亿甚至数千亿个“参数”（你可以把它们想象成食谱里的调料和步骤）。
• 问题：
  • GPU（显卡） 是厨房里的超级大厨，手速极快，但操作台（显存）非常小。
  • CPU（处理器） 是助手，虽然操作台很大（内存大），但手速很慢。
  • 现状：为了训练模型，我们需要把大量的“调料”（优化器状态，包括动量、方差等数据）放在操作台上。但是，超级大厨的操作台太小了，根本放不下所有调料。
  • 传统做法：把多余的调料搬到助手（CPU）那边的操作台上。
  • 瓶颈：
    1. 搬运太慢：大厨和助手之间只有一条狭窄的走廊（PCIe 总线）。每次大厨做完一步，都要等助手把新调料搬过来，或者把旧调料搬走。这条走廊经常堵车。
    2. 助手太慢：助手处理调料的速度比大厨慢几十倍。一旦轮到助手干活，整个厨房就停滞了。
    3. 资源浪费：在大厨忙着炒菜（前向/反向传播）的时候，助手那边的操作台其实空了一大半；而在助手忙着搬运时，大厨却在发呆。

2. 核心洞察：发现“时间差”

作者观察到一个有趣的现象：

• 当大厨在炒菜（前向/反向传播）时，操作台上很挤，但助手那边的走廊其实很空。
• 当大厨炒完菜准备更新配方（更新阶段）时，操作台上的空间突然腾出了一大块（因为之前的临时数据被清理了），但此时助手还在慢吞吞地搬运。

这就好比：大厨在炒菜时，助手在搬砖；大厨炒完菜休息时，助手还在搬砖。其实，大厨休息的那一小会儿，完全可以自己顺手把一部分砖搬了，或者让助手在搬砖的同时，大厨也能处理一部分。

3. 解决方案：Deep Optimizer States（ interleaved Offloading）

这项新技术的核心思想是**“ interleaved"（交错/穿插），就像接力赛**一样，而不是传统的“等一个人做完，另一个人再做”。

比喻：智能的“双人舞”

想象大厨（GPU）和助手（CPU）在跳一支双人舞，他们不再按部就班地轮流干活，而是同时配合：

1. 把大任务切碎：
   把那一堆巨大的“调料”（优化器状态）切成很多小块（子组）。
2. 动态分配：
   • 有些小块，因为大厨操作台刚好有空位，就直接在大厨手里处理（在 GPU 上更新）。
   • 有些小块，因为操作台满了，就交给助手处理（在 CPU 上更新）。
3. 无缝衔接（重叠）：
   • 当助手在搬运下一块调料时，大厨已经在处理当前这块调料了。
   • 当大厨在炒菜时，助手在搬运之前的废料。
   • 关键点：他们不再互相等待。走廊（PCIe）被充分利用，大厨的手（GPU 算力）和助手的手（CPU 算力）都在同时工作。

具体的“魔法”技巧：

• 精准计算：作者设计了一个“数学模型”，像是一个智能调度员。它会根据当前厨房的拥挤程度、走廊的宽度、大厨和助手的手速，精确计算出：“每让助手搬 2 块砖，就应该让大厨自己搬 1 块砖”。这样能确保没有人闲着，也没有人堵车。
• 直接转换：以前，调料从助手搬到大厨时，需要换个包装（数据格式转换），这很浪费时间。现在，他们直接在搬运过程中就换好包装，省去了额外的步骤。

4. 效果如何？

通过这种“穿插式”的协作：

• 速度提升：训练一轮的速度比目前最先进的技术（如 DeepSpeed TwinFlow）快了 2.5 倍。
• 资源利用：原本闲置的走廊和大厨的休息时间都被充分利用了。
• 成本降低：这意味着用户可以用更少、更便宜的显卡（比如只有 4 张显卡的单机）来训练以前需要巨大集群才能训练的模型。

总结

这篇论文就像是在说：

“以前我们训练 AI，就像让一个超级大厨和一个慢吞吞的助手在狭窄的厨房里干活，两人总是互相等对方，导致效率极低。

现在，我们发明了一种**‘智能交错舞步’。通过把任务切得细碎，并让大厨和助手在搬运和处理的间隙无缝配合、同时工作**，我们不仅消除了等待时间，还让那条狭窄的走廊跑出了高速公路的速度。结果就是，训练大模型变得更快、更便宜、更可行。”

这项技术对于让 AI 模型在资源有限的设备上（比如单台服务器）也能高效运行，具有非常重要的意义。

技术摘要

1. 研究背景与问题 (Problem)

随着大语言模型（LLM）和 Transformer 架构的规模迅速膨胀（达到数百亿甚至万亿参数），训练这些模型面临着严峻的**“内存墙”（Memory Wall）**挑战：

• 显存不足：即使使用 3D 并行（流水线、张量、数据并行）和聚合多块 GPU 的显存，往往仍不足以容纳模型参数、梯度、激活值以及庞大的优化器状态（Optimizer State）。
• 现有方案的瓶颈：为了应对显存限制，现有的最先进方法（如 DeepSpeed Offload, ZeRO-Offload）通常将优化器状态（包含 FP32 参数、动量、方差）卸载到主机（CPU）内存中。然而，这种混合 CPU-GPU 计算存在两个主要瓶颈：
  1. I/O 带宽瓶颈：PCIe 链路带宽有限（通常 25-50 GB/s），且在与节点间通信竞争时进一步受限。在反向传播和更新阶段，数据在 GPU 和 CPU 之间的传输往往未被充分利用。
  2. 计算能力瓶颈：CPU 的计算能力远低于 GPU（在测试环境中，CPU 更新速度比 GPU 慢约 20 倍）。
  3. 资源利用不均：在现有的静态卸载方案（如 DeepSpeed TwinFlow）中，GPU 显存在前向/反向传播阶段被激活值占用，而在更新阶段，GPU 显存大部分空闲，但优化器更新却主要在慢速的 CPU 上进行，导致 PCIe 链路和 GPU 计算资源在更新阶段处于低效或闲置状态。

2. 核心方法论 (Methodology)

作者提出了 Deep Optimizer States，一种基于**交错卸载（Interleaved Offloading）**的新颖技术，旨在动态利用 GPU 显存波动和 PCIe 带宽，加速优化器更新过程。

2.1 关键观察

• 显存波动：在训练的前向（Forward）、反向（Backward）和更新（Update）阶段，GPU 显存利用率存在显著波动。特别是在更新阶段，激活值被释放，GPU 显存有大量空闲空间可用于暂存部分优化器状态。
• PCIe 利用率低：在反向传播和更新阶段，PCIe 链路的实际吞吐量远低于峰值，存在巨大的优化空间。
• 子组并行性：利用 DeepSpeed ZeRO-3 的**子组（Subgroup）**划分机制，优化器状态可以被细分为多个独立的子组，这些子组的更新操作是“尴尬并行”（embarrassingly parallel）的，互不依赖。

2.2 核心设计原则

1. 交错更新（Interleaved Updates）：
   • 不再静态地将优化器划分为“全在 GPU"或“全在 CPU"。
   • 在更新阶段，动态调度一部分子组在 GPU 上更新，另一部分在 CPU 上更新。
   • 利用 GPU 空闲显存暂存待更新的子组状态（参数、动量、方差）。
2. 重叠执行（Overlapping Execution）：
   • 计算与传输重叠：当 CPU 计算一组子组的更新时，GPU 同时计算另一组子组的更新，同时 PCIe 链路异步地进行数据搬运（H2D 和 D2H）。
   • 流水线设计：在 CPU 计算当前子组时，异步预取（Prefetch）下一个要在 GPU 上更新的子组；在 GPU 计算时，异步刷新（Flush）CPU 更新好的参数。
3. 高效梯度管理：
   • 利用前向传播释放的显存空间来存储需要在 GPU 上更新的子组对应的梯度。
   • 如果显存不足，则动态卸载到主机。
4. 高精度传输优化：
   • 传统方法在传输梯度时通常使用 FP16 以减少带宽消耗，但在主机端进行 FP16 到 FP32 的转换和内存分配非常耗时。
   • Deep Optimizer States 在 GPU 端直接进行FP16 到 FP32 的片上转换（利用 GPU 极高的计算吞吐量），然后将 FP32 数据直接传输到主机已分配的 pinned 内存中。这避免了主机端的昂贵转换和内存分配开销，显著提升了传输效率。

2.3 性能模型与调度算法

• 性能模型：作者建立了一个数学模型（公式 1），用于计算最优的**"GPU 更新步长”（Update Stride, $k$）**。该模型平衡了 CPU 更新速度、GPU 更新速度、FP32 到 FP16 的下采样速度以及 PCIe 的传输带宽，以确定在多少个 CPU 更新后插入一个 GPU 更新，从而实现计算与传输的最大重叠。
• 调度算法：基于计算出的 $k$ 值，算法动态决定每个子组是在 CPU 还是 GPU 上执行更新，并管理异步的数据流（使用 CUDA Stream 确保依赖关系正确）。

3. 主要贡献 (Key Contributions)

1. 深入分析：详细研究了在优化器卸载场景下，训练迭代中 GPU 显存利用率和 PCIe 链路利用率的动态变化特征，揭示了现有静态方案的资源浪费问题。
2. 新架构设计：提出了 Deep Optimizer States，实现了基于子组的动态交错卸载，能够根据性能模型动态调整 GPU/CPU 的更新比例。
3. 系统实现：将该技术集成到广泛使用的 DeepSpeed 和 Megatron-LM 框架中，作为中间件支持 ZeRO-3 阶段。
4. 性能优化：通过异步传输、片上精度转换和流水线重叠，显著减少了更新阶段的延迟。

4. 实验结果 (Results)

作者在 4×H100 (80GB) GPU 和 192 核 CPU 的测试平台上，对高达 20B 参数的模型进行了广泛评估：

• 迭代速度提升：
  • 在优化器完全卸载到 CPU 的场景下，Deep Optimizer States 比 DeepSpeed ZeRO-3 快 2.5 倍。
  • 在部分卸载（类似 TwinFlow）场景下，即使 GPU 仅静态驻留 20% 的优化器状态，其迭代速度仍比 TwinFlow 快 1.7 到 2.3 倍。
  • 甚至在 0% 静态 GPU 驻留（完全动态调度）的情况下，其性能优于 TwinFlow 在 50% 静态驻留下的表现，且节省了约 45% 的 GPU 显存。
• 更新吞吐量：
  • 优化器更新吞吐量平均提升了 70%。
  • 在 20B 参数模型上，实现了约 15.4 Billion 参数/秒 的更新速度（ZeRO-3 仅为 8.8）。
• 端到端训练时间：
  • 端到端训练时间缩短了 2.5 倍。
  • 证明了重叠的数据传输不会阻塞后续迭代，性能提升具有累积效应。
• 资源利用率：
  • 实现了接近 100% 的 GPU 利用率 和 40% 的 PCIe 峰值带宽利用率（相比之下，ZeRO-3 的 PCIe 利用率极低）。
  • 在 50% 更新调度在 GPU 上时，TFLOPs 达到 75，是 ZeRO-3 的 2.5 倍。
• 可扩展性：
  • 在增加数据并行度（Data Parallelism）和微批次大小（Micro-batch size）时，该方法均表现出良好的扩展性。
  • 性能模型在不同硬件平台（如 V100 和 H100）上均被验证为有效。

5. 意义与影响 (Significance)

• 突破内存墙限制：Deep Optimizer States 提供了一种在有限显存下高效训练大模型的新范式，使得在单节点或少量 GPU 上训练更大规模的模型成为可能。
• 硬件资源最大化：通过动态调度，它解决了 CPU 和 GPU 在混合训练中的“木桶效应”，充分利用了 PCIe 带宽和 GPU 的计算能力，避免了传统方案中 CPU 慢速更新拖累整体速度的问题。
• 面向未来架构：随着 CPU-GPU 互联技术（如 NVIDIA Grace Hopper 的 200 GB/s C2C 互联）的发展，这种动态交错卸载策略将变得更加重要，能够进一步释放下一代异构系统的性能潜力。
• 开源与实用性：作为开源中间件集成到 DeepSpeed 中，该方法可直接被现有 LLM 训练工作流采用，无需修改模型代码，具有极高的实用价值。

总结：Deep Optimizer States 通过精细化的资源管理和动态调度策略，成功将大模型训练中的优化器更新瓶颈从“内存受限”和"CPU 计算慢”转变为“计算与传输的高效重叠”，显著提升了大规模 Transformer 模型的训练效率。