使用 DeepSpeed 加速大模型训练

Hugging Face 已将 DeepSpeed 的 Zero Redundancy Optimizer（ZeRO）集成到 accelerate 库中，使用户能够通过在 GPU 之间减少内存冗余来更高效地训练大规模模型。此集成使开发者能够避免内存不足（OOM）错误并在几乎不修改代码的情况下增大批量大小。

理解 ZeRO 数据并行

ZeRO（Zero Redundancy Optimizer）通过在可用 GPU 上对训练状态进行分片来优化数据并行期间的内存使用。该框架分为多个阶段，以降低模型训练的内存占用：

  • Stage 1：在数据并行工作节点/GPU 之间分片优化器状态。
  • Stage 2：在工作节点/GPU 之间分片优化器状态和梯度。
  • Stage 3：在工作节点/GPU 之间分片优化器状态、梯度和模型参数。
  • Optimizer Offload：在 Stage 2 基础上，将梯度和优化器状态卸载到 CPU 或磁盘。
  • Param Offload：在 Stage 3 基础上，将模型参数卸载到 CPU 或磁盘。

零代码集成 via Accelerate

用户可以通过 accelerate config 命令和 Accelerate DeepSpeed 插件，在不修改训练代码的情况下使用 DeepSpeed ZeRO Stage‑2。

在使用两块 24 GB NVIDIA Titan RTX GPU 对 DeBERTa‑v2‑xlarge‑mnli（9 亿参数）进行微调的基准测试中，DeepSpeed ZeRO Stage‑2 显著优于 Distributed Data Parallel（DDP）：

方法 最大批量大小 每轮训练时间 (秒) 每轮评估时间 (秒) F1 分数 准确率
DDP 8 103.57 2.04 0.931 0.904
DeepSpeed ZeRO Stage 2 40 28.98 1.79 0.936 0.912

DeepSpeed 实现了 最大批量大小提升 5 倍，以及 训练总时长约提升 3.5 倍，且性能指标没有下降。

使用 DeepSpeed 配置文件进行高级配置

若需要更细粒度的控制，用户可以通过 accelerate config 提供 DeepSpeed 配置 JSON 文件。当配置文件中定义了优化器和调度器时，用户只需将标准的 PyTorch 优化器和调度器替换为 accelerate.utils.DummyOptim 与 accelerate.utils.DummyScheduler，即可完成最小代码改动。

在 MuDoConv 数据集上使用 BlenderBot‑400M‑distill 模型进行测试时，DeepSpeed ZeRO Stage‑2 使最大批量大小达到 200，而 DDP 只能到 100。由此带来了 训练加速 1.44 倍 与 评估加速 1.23 倍。

使用 CPU/磁盘卸载训练超大模型

DeepSpeed ZeRO Stage‑3 配合 CPU 卸载能够训练即使在批量大小为 1 时也放不进 GPU 内存的模型。

在单块 24 GB NVIDIA Titan RTX GPU 上训练 GPT‑XL（15 亿参数）的实验中，DDP 直接触发 OOM 错误。而 DeepSpeed ZeRO Stage‑3 通过对优化器状态、梯度和参数进行 CPU 卸载，成功以 批量大小 16 完成训练，一个 epoch 用时 6608.35 秒。

方法 最大批量大小 每轮训练时间 (秒) 备注
DDP - - OOM 错误
DeepSpeed ZeRO Stage 3 16 6608.35 -

Sources