使用 DeepSpeed 增强多模态训练和内存效率

概述

本篇博客介绍了 DeepSpeed 的两项重要更新：(1) 与 PyTorch 完全一致的 backward API，支持多模态、多组件模型的高效训练（包括非标量 backward 调用）；(2) 低精度模型训练，可显著降低峰值内存占用。

对于多模态工作负载（例如视觉编码器与大语言模型结合），训练循环往往变得复杂且包含多个组件。首项更新引入了与 PyTorch 完全一致的 backward API，使得编写此类循环变得简单直接，通过简洁的代码实现复杂的并行方案，同时由 DeepSpeed 透明地管理各种性能优化。作为一个示例，该 API 的灵活性实现了解耦混合并行 (disaggregated hybrid parallelism)，在多模态 AI 模型训练中实现了 30% 的加速，并使 DeepSpeed 模型开发更接近“原生 PyTorch”。

同时，针对大模型微调，我们引入了将所有模型状态（参数、梯度和优化器状态）保持在低精度（如 BF16 或 FP16）的新选项。这大幅降低了内存占用，使研究人员能在受限硬件上训练更大的模型。低精度训练在监督微调 (SFT)、强化学习 (RL) 和多模态训练等广泛应用中极具价值。我们的实验显示，在保持数值稳定性的前提下，峰值内存降低了 40%（基准测试脚本）。这种数值稳定性通过与 torch.autocast 的集成实现，确保了模型质量不受影响。

本篇博客的其余部分将详细阐述这些更新如何直接促进前沿训练工作负载的开发。

1. 与 PyTorch 完全一致的 backward API

DeepSpeed 现在支持 PyTorch 原生的 backward() 语法，同时保留了其所有优化功能。传统上，DeepSpeed 的训练循环依赖于引擎的 backward API。

loss = model_engine(batch)
model_engine.backward(loss)
model_engine.step()

引擎的 backward API 对于传统的预训练和微调管道已足够。然而，近期复杂的训练管道需要更高的灵活性。过去主要存在两个限制：

它仅接受标量损失函数（scalar loss）。
必须调用 model_engine.backward(loss)，而不是使用通常的 PyTorch loss.backward() 风格。

由于这些限制，用户无法简单地实现原生 PyTorch 允许的模式。以下是一些示例：

# 1. Combine multiple models and losses
output1 = model1(batch1)
output2 = model2(batch2)
loss = criterion(output1, output2)
loss.backward()

# 2. Define a loss function separately from the main model
output = model(batch)
loss = loss_fn(output)
loss.backward()

# 3. Call backward through non-scalar tensors with custom gradients
output = model(batch)
output.backward(grad)

DeepSpeed 引擎过去可以通过内部 API 处理这些用例，但这需要大量的代码修改，且容易引入错误。随着与 PyTorch 一致的 backward API 的加入，现在可以使用与原生 PyTorch 相同的代码，同时保留 DeepSpeed 强大的优化功能，包括 ZeRO 和卸载（offloading）。

与 PyTorch 一致的 backward API 的一个应用场景是使用 Ray 对多模态模型进行解耦混合并行。在此训练管道中，两个 Ray Actor 组分别处理视觉编码器和 LLM。在反向传播过程中，LLM 将梯度传递给视觉编码器，视觉编码器使用该梯度调用 backward 函数。由于该梯度是一个非标量张量，DeepSpeed 的旧版 API 并不支持此场景。解耦混合并行证明，backward API 的灵活性结合 DeepSpeed 的优化和 DeepSpeed-Ulysses（高效序列并行），可实现 30% 的训练加速。

以下是在不同 Actor 上运行的两个模型的伪代码。由于它们在不同进程中运行，我们通过 Ray Actor 通信传递梯度。如代码所示，视觉嵌入的梯度是非标量张量。尽管此代码与 PyTorch API 完全相同，但它将根据您的配置激活各种 DeepSpeed 优化。

# Runs on LLM actors
def text_backward_step(self):
# ...
  self.loss.backward()
  return self.vision_embeddings.grad.detach().clone()

# Runs on Vision actors
def vision_backward_step(self, vision_embedding_grad):
  self.vision_output.backward(gradient=vision_embedding_grad)

请查阅仓库获取完整的训练管道。

2. 内存高效的低精度模型状态

现在，您可以将所有模型状态（参数、梯度和优化器状态）保持在 BF16 或 FP16 中，从而显著减少内存消耗。

传统上，DeepSpeed 的混合精度会保留 FP32 的主参数、梯度和优化器状态，这在技术上更安全但非常消耗内存。虽然 DeepSpeed 此前已通过配置支持 torch.autocast（参考 API 文档），但由于缺乏跳过创建 FP32 状态的选项，限制了在受限硬件上训练大模型的能力。在实践中，许多训练工作负载无需 FP32 状态即可稳定收敛。

借助低精度模型状态选项，您可以轻松跳过创建 FP32 状态，并将低精度选项与 torch.autocast 支持结合使用（具体配置参考文档和示例）。这种组合在不牺牲收敛性的前提下，极大地提高了内存效率。

{
...
  "zero_optimization": {
    "stage": 3,
    ...
  },
  "bf16": {
    "enabled": true,
    "bf16_master_weights_and_grads": true,
    "bf16_optimizer_states": true
  },
  "torch_autocast": {
    "enabled": true,
    "dtype": "bfloat16"
  }
}

我们的示例脚本展示了显著的内存节省效果：

配置	分配内存	峰值内存	平均步时
基准（fp32 master）	25.74 GB	31.38 GB	0.6016s
BF16 低精度（master + opt 状态）	16.17 GB	18.93 GB	0.6427s

实验（7B 模型，ZeRO3，4 GPU）表明峰值内存降低了 40%。为了验证 BF16 低精度训练的数值稳定性，我们在 Wikitext-103 数据集上训练了 1000 步：

损失曲线对比

配置	最终损失	平均损失
基准（fp32 master）	3.09	2.78
BF16 低精度	3.12	2.90

总结

此 DeepSpeed 更新带来了关键进展：

支持复杂的多模态工作负载：全新的 PyTorch 一致 backward API 实现了复杂的多组件训练循环（如多模态模型所需），代码简单清晰。作为一个示例，该 API 助力解耦混合并行实现了 30% 的速度提升。
扩展至更大模型： 低精度模型状态与 torch.autocast 的结合可在不牺牲收敛性的情况下将峰值内存降低高达 40%，使您能够在相同硬件上训练更大的模型。

我们非常期待看到您在自己的训练方案中使用这些新 API 和功能，并在尝试过程中欢迎在 GitHub 上提出反馈和问题。

通过 DeepSpeed 增强多模态训练与内存效率

概述

1. 与 PyTorch 完全一致的 backward API

2. 内存高效的低精度模型状态

相关测试

总结

文档

教程

资源

保持联系以获取更新、活动信息和最新新闻