用 DeepSpeed 引擎改造 PyTorch 模型并启动多机训练

DeepSpeed 是面向大规模模型训练的分布式训练框架，核心是一个可以包裹任意 torch.nn.Module 的训练引擎。它把分布式数据并行、混合精度、梯度累积、优化器与学习率调度器管理、检查点保存这些原本需要手写的样板代码收进引擎内部，对外只暴露少量 API。本文要解决的问题是：把一个已经能跑的 PyTorch 训练脚本，改造成由 DeepSpeed 引擎驱动，并用 deepspeed 启动器在单机或多机上拉起任务。适合已经写过 PyTorch 训练循环、需要把训练扩展到多卡多机的读者。如果你只是想先了解这类训练框架的定位，可以在 DeepSpeed 工具页看概览，再回到本文动手改造。

准备工作

安装

安装本身只有一条命令：

pip install deepspeed

DeepSpeed 与 HuggingFace Transformers、PyTorch Lightning 都有直接集成。使用 Transformers 时，可以通过一个 --deepspeed 参数加上配置文件来加速模型训练；使用 PyTorch Lightning 时，则通过 Lightning Trainer 接入 DeepSpeed。这两条路径不需要自己写 deepspeed.initialize，本文主要讲原生 PyTorch 模型的改造方式。

硬件与加速器

DeepSpeed 默认使用 NCCL 后端，这是经过充分测试的组合。除 NVIDIA GPU 外，它还支持 AMD 的 ROCm 环境（有对应的 ROCm 镜像可用）、Intel Xeon CPU、Intel Data Center Max 系列 XPU、Intel Gaudi HPU、华为昇腾 NPU 等加速器，具体安装与配置方式需要参考对应加速器的设置指南。不同加速器支持的算子与特性范围并不一致，选型前应以官方当前文档为准。

代码与配置前置条件

• 一个可运行的 PyTorch 训练脚本，模型是 torch.nn.Module 的子类。
• 一个 DeepSpeed 配置 JSON 文件，通常命名为 ds_config.json。
• 多机训练时，需要一份 hostfile，列出主机名（或 SSH 别名）与每台机器的 GPU 槽位数。
• 多机训练默认要求节点之间可以免密 SSH 登录；如果做不到，可以使用后文的 --no_ssh 模式。

操作步骤

第一步：用 deepspeed.initialize 初始化引擎

DeepSpeed 的训练由引擎完成。引擎包裹模型，并提供训练与检查点相关的最小 API 集合。初始化写法如下：

model_engine, optimizer, _, _ = deepspeed.initialize(
    args=cmd_args,
    model=model,
    model_parameters=params
)

deepspeed.initialize 会在内部完成分布式数据并行或混合精度训练所需的全部准备工作。除了包裹模型，它还会根据传入的参数和 DeepSpeed 配置文件，构造并管理训练优化器、数据加载器和学习率调度器。需要注意，DeepSpeed 会在每个训练步自动执行学习率调度。

第二步：处理已有的分布式初始化代码

如果你的脚本里已经有分布式环境初始化，需要把这一行：

torch.distributed.init_process_group(...)

替换为：

deepspeed.init_distributed()

默认使用 NCCL 后端，也可以覆盖默认值。如果分布式环境在 deepspeed.initialize() 之前并不需要，那么不必调用这个函数，DeepSpeed 会在初始化过程中自动完成分布式环境初始化。无论哪种情况，原来代码里的 torch.distributed.init_process_group 都必须移除。

第三步：改写训练循环

引擎初始化完成后，训练循环只需要三个 API：前向传播（引擎本身可调用）、反向传播（backward）、权重更新（step）。

for step, batch in enumerate(data_loader):
    # 前向传播
    loss = model_engine(batch)
    # 反向传播
    model_engine.backward(loss)
    # 权重更新
    model_engine.step()

在这三个调用背后，DeepSpeed 自动完成了分布式数据并行训练所需的操作，并以混合精度、配合预定义的学习率调度器运行：

• 梯度平均：在分布式数据并行训练中，backward 会在按 train_batch_size 完成训练后，确保梯度在各数据并行进程之间取平均。
• 损失缩放：在 FP16 或混合精度训练中，引擎自动处理 loss 的缩放，避免梯度精度损失。
• 学习率调度：当使用 DeepSpeed 的学习率调度器（在 ds_config.json 中指定）时，每次执行 model_engine.step()，DeepSpeed 都会调用调度器的 step() 方法。

如果不使用 DeepSpeed 的学习率调度器，分两种情况处理：

• 调度需要在每个训练步执行：把调度器传给 deepspeed.initialize，由 DeepSpeed 负责更新以及保存/恢复。
• 调度按其他间隔执行（例如按训练轮次）：不要在初始化时把调度器传给 DeepSpeed，必须自己显式管理。

第四步：编写 DeepSpeed 配置文件

各项特性通过一个 JSON 配置文件开启、关闭或配置，路径由 args.deepspeed_config 指定。一个最小示例：

{
  "train_batch_size": 8,
  "gradient_accumulation_steps": 1,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.00015
    }
  },
  "fp16": {
    "enabled": true
  },
  "zero_optimization": true
}

这里 train_batch_size 是全局训练批大小，gradient_accumulation_steps 控制梯度累积步数，optimizer 段声明优化器类型与参数，fp16.enabled 打开混合精度，zero_optimization 打开 ZeRO 优化。完整特性集合以官方 API 文档为准。

第五步：用 deepspeed 启动器拉起训练

安装 DeepSpeed 会同时安装 deepspeed 入口点，用它来启动分布式训练。假设已经完成模型改造，client_entry.py 是入口脚本，client args 是 argparse 命令行参数，ds_config.json 是配置文件，单机启动命令形如：

deepspeed client_entry.py <client args> \
    --deepspeed --deepspeed_config ds_config.json

第六步：配置多机资源（hostfile）

多机资源通过 hostfile 配置，格式与 OpenMPI、Horovod 兼容。hostfile 每行是一台机器的主机名（或 SSH 别名）加上槽位数，槽位数表示该机器可用的 GPU 数量：

worker-1 slots=4
worker-2 slots=4

这表示 worker-1 与 worker-2 各有 4 块 GPU 参与训练。hostfile 通过 --hostfile 命令行选项指定。如果没有指定 hostfile，DeepSpeed 会查找 /job/hostfile；如果既没有指定也没有找到，DeepSpeed 会查询本机 GPU 数量，以确定本地可用槽位数。

下面这条命令会在 myhostfile 中列出的所有节点和 GPU 上启动 PyTorch 训练任务：

deepspeed --hostfile=myhostfile <client_entry.py> <client args> \
    --deepspeed --deepspeed_config ds_config.json

第七步：限制参与训练的节点与 GPU

可以只使用可用资源的一个子集，通过 --num_nodes 与 --num_gpus 两个参数控制。例如只使用两台节点：

deepspeed --num_nodes=2 \
    <client_entry.py> <client args> \
    --deepspeed --deepspeed_config ds_config.json

也可以用 --include 与 --exclude 精确地包含或排除特定资源。例如使用除 worker-2 的 GPU 0 以及 worker-3 的 GPU 0、1 之外的全部资源：

deepspeed --exclude="worker-2:0@worker-3:0,1" \
    <client_entry.py> <client args> \
    --deepspeed --deepspeed_config ds_config.json

反过来，只使用 worker-2 上的 GPU 0 和 1：

deepspeed --include="worker-2:0,1" \
    <client_entry.py> <client args> \
    --deepspeed --deepspeed_config ds_config.json

第八步：在无免密 SSH 的环境下启动

在 Kubernetes 这类容器编排灵活、搭建 leader-worker 免密 SSH 反而增加复杂度的云环境中，可以使用 no-SSH 模式。该模式下需要在所有节点上分别执行 DeepSpeed 命令，命令结构如下：

deepspeed --hostfile=myhostfile --no_ssh --node_rank=<n> \
    --master_addr=<addr> --master_port=<port> \
    <client_entry.py> <client args> \
    --deepspeed --deepspeed_config ds_config.json

各参数含义：

• --hostfile=myhostfile：指定包含节点与 GPU 信息的 hostfile。
• --no_ssh：启用 no-SSH 模式。
• --node_rank=<n>：节点序号，必须是 0 到 n-1 之间唯一的整数。
• --master_addr=<addr>：leader 节点（rank 0）的地址。
• --master_port=<port>：leader 节点的端口。

这种模式下，hostfile 中的主机名不需要能通过免密 SSH 访问，但 hostfile 仍然必需，因为启动器要靠它收集节点数量、每节点 GPU 数量等环境信息。每个节点必须用唯一的 node_rank 启动，并且所有节点都要拿到 leader 节点的地址与端口。该模式让启动器的行为接近 torchrun。

第九步：传播自定义环境变量

多机训练时，DeepSpeed 默认会传播所有已设置的 NCCL 与 PYTHON 相关环境变量。如果需要额外传播其他变量，可以在执行目录下创建一个名为 .deepspeed_env 的点文件，内容为按行分隔的 VAR=VAL 条目，启动器会读取它。

一个完整示例

下面把前面的步骤串成一个可运行的最小流程。假设入口脚本为 train.py，模型与数据加载器已在脚本内定义。

1. 配置文件 ds_config.json

{
  "train_batch_size": 8,
  "gradient_accumulation_steps": 1,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.00015
    }
  },
  "fp16": {
    "enabled": true
  },
  "zero_optimization": true
}

2. 入口脚本 train.py 的关键片段

import argparse
import torch
import deepspeed

parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1)
parser.add_argument("--deepspeed_config", type=str, required=True)
parser.add_argument("--save_dir", type=str, default="./ckpt")
parser.add_argument("--save_interval", type=int, default=1000)
args = parser.parse_args()

model = MyModel()
params = [p for p in model.parameters() if p.requires_grad]

model_engine, optimizer, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=params
)

client_sd = {}
for step, batch in enumerate(data_loader):
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

    if step % args.save_interval == 0:
        client_sd["step"] = step
        ckpt_id = loss.item()
        model_engine.save_checkpoint(args.save_dir, ckpt_id, client_sd=client_sd)

3. 启动命令

单机多卡，直接使用本机全部 GPU：

deepspeed train.py --deepspeed_config ds_config.json

多机，使用 hostfile 中列出的全部资源：

deepspeed --hostfile=myhostfile train.py \
    --deepspeed_config ds_config.json

4. 从检查点恢复训练

检查点的保存与加载通过 save_checkpoint 与 load_checkpoint 完成，二者都用两个参数唯一标识一个检查点：ckpt_dir 是保存目录，ckpt_id 是目录内唯一标识该检查点的标识符。下面的片段用 loss 值作为检查点标识符，并把训练步数存进客户端状态字典：

# 加载检查点
_, client_sd = model_engine.load_checkpoint(args.load_dir, args.ckpt_id)
step = client_sd["step"]
# 把数据加载器推进到检查点对应的步
dataloader_to_step(data_loader, step + 1)

for step, batch in enumerate(data_loader):
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

    if step % args.save_interval == 0:
        client_sd["step"] = step
        ckpt_id = loss.item()
        model_engine.save_checkpoint(args.save_dir, ckpt_id, client_sd=client_sd)

DeepSpeed 会自动保存和恢复模型、优化器与学习率调度器的状态，把这些细节对用户隐藏。如果训练过程还有模型特有的额外数据需要保存，就通过 save_checkpoint 的 client_sd 参数传入，加载时再从 load_checkpoint 的返回值中取回。上面的例子就是把 step 存进了 client_sd。

注意事项

• 检查点必须所有进程都调用：save_checkpoint 不能只由 rank 0 进程调用。每个进程都需要保存自己的 master weights 以及调度器、优化器状态。如果只让 rank 0 调用，该方法会一直挂起等待与其他进程同步。
• 不要保留旧的分布式初始化：脚本里原有的 torch.distributed.init_process_group 必须删除，否则会与 DeepSpeed 的初始化冲突。
• 学习率调度器的归属要明确：按步执行的调度器交给 DeepSpeed 管理；按轮次等其他间隔执行的调度器不要传给 DeepSpeed，自己显式调用。
• hostfile 的查找顺序：未指定 --hostfile 时先找 /job/hostfile，再退回到查询本机 GPU 数量。多机场景下如果忘了指定 hostfile，任务可能只在本地启动。
• no-SSH 模式仍需 hostfile：主机名不必可免密 SSH 访问，但启动器依赖 hostfile 获取节点数与每节点 GPU 数。
• node_rank 必须唯一：no-SSH 模式下每个节点的 --node_rank 必须是 0 到 n-1 之间不重复的整数，且所有节点都要配置相同的 leader 地址与端口。
• 加速器支持范围：不同硬件后端（ROCm、Intel XPU、Gaudi HPU、昇腾 NPU 等）的可用特性与安装方式不同，具体支持情况、镜像标签与版本对应关系请以官网当前信息为准。
• 配置项与默认值：本文示例中的批大小、学习率、梯度累积步数只是演示值，实际取值需按模型与硬件调整；完整配置项与默认值以官方 API 文档当前版本为准。