入门

安装

  • 安装就像 pip install deepspeed 一样简单，查看更多详情。
  • 要在 AzureML 上开始使用 DeepSpeed，请参阅 AzureML 示例 GitHub
  • DeepSpeed 已与 HuggingFace Transformers 和 PyTorch Lightning 直接集成。HuggingFace Transformers 用户现在可以通过一个简单的 --deepspeed 标志和配置文件轻松使用 DeepSpeed 加速他们的模型 查看更多详情。PyTorch Lightning 通过 Lightning Trainer 提供了对 DeepSpeed 的便捷访问 查看更多详情。
  • DeepSpeed 可以在 AMD 上通过我们的 ROCm 镜像 使用，例如，docker pull deepspeed/rocm501:ds060_pytorch110。
  • DeepSpeed 还支持英特尔至强 CPU、英特尔数据中心 Max 系列 XPU、英特尔 Gaudi HPU、华为昇腾 NPU 等，请参阅加速器设置指南

编写 DeepSpeed 模型

DeepSpeed 模型训练通过 DeepSpeed 引擎完成。该引擎可以封装任意 torch.nn.module 类型的模型，并提供最少的 API 集用于训练和检查点保存模型。请参阅教程获取详细示例。

初始化 DeepSpeed 引擎

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

deepspeed.initialize 确保分布式数据并行或混合精度训练所需的所有必要设置都在后台妥善完成。除了封装模型之外，DeepSpeed 还可以根据传递给 deepspeed.initialize 的参数和 DeepSpeed 配置文件 来构建和管理训练优化器、数据加载器和学习率调度器。请注意，DeepSpeed 会在每个训练步骤自动执行学习率调度。

如果您已经设置了分布式环境，则需要替换

torch.distributed.init_process_group(...)

为

deepspeed.init_distributed()

默认使用 NCCL 后端，DeepSpeed 已对其进行彻底测试，但您也可以覆盖默认设置。

但如果您在调用 deepspeed.initialize() 之后才需要分布式环境设置，则无需使用此函数，因为 DeepSpeed 将在其 initialize 期间自动初始化分布式环境。无论如何，如果您已经使用了 torch.distributed.init_process_group，则需要将其移除。

训练

DeepSpeed 引擎初始化后，可以使用三个简单的 API 来训练模型：前向传播（可调用对象）、反向传播（backward）和权重更新（step）。

for step, batch in enumerate(data_loader):
    #forward() method
    loss = model_engine(batch)

    #runs backpropagation
    model_engine.backward(loss)

    #weight update
    model_engine.step()

在底层，DeepSpeed 自动执行分布式数据并行训练、混合精度训练以及预定义学习率调度器所需的必要操作

  • 梯度平均：在分布式数据并行训练中，backward 确保在 train_batch_size 上的训练后，梯度在数据并行进程之间进行平均。

  • 损失缩放：在 FP16/混合精度训练中，DeepSpeed 引擎自动处理损失缩放以避免梯度中的精度损失。

  • 学习率调度器：当使用 DeepSpeed 的学习率调度器（在 ds_config.json 文件中指定）时，DeepSpeed 会在每个训练步骤（当 model_engine.step() 执行时）调用调度器的 step() 方法。当不使用 DeepSpeed 的学习率调度器时

    • 如果调度器应该在每个训练步骤执行，那么用户可以在初始化 DeepSpeed 引擎时将调度器传递给 deepspeed.initialize，并让 DeepSpeed 管理其更新或保存/恢复。
    • 如果调度器应该在任何其他间隔（例如，训练周期）执行，那么用户在初始化时不应将调度器传递给 DeepSpeed，并且必须显式管理它。

模型检查点

训练状态的保存和加载通过 DeepSpeed 中的 save_checkpoint 和 load_checkpoint API 处理，这两个 API 接受两个参数来唯一标识一个检查点

  • ckpt_dir: 检查点将保存到的目录。
  • ckpt_id: 唯一标识目录中检查点的标识符。在以下代码片段中，我们使用损失值作为检查点标识符。
#load checkpoint
_, client_sd = model_engine.load_checkpoint(args.load_dir, args.ckpt_id)
step = client_sd['step']

#advance data loader to ckpt step
dataloader_to_step(data_loader, step + 1)

for step, batch in enumerate(data_loader):

    #forward() method
    loss = model_engine(batch)

    #runs backpropagation
    model_engine.backward(loss)

    #weight update
    model_engine.step()

    #save checkpoint
    if step % args.save_interval:
        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 的一部分存储。

重要提示：所有进程都必须调用此方法，而不仅仅是 rank 0 的进程。这是因为每个进程都需要保存其主权重和调度器+优化器状态。如果仅对 rank 0 的进程调用此方法，它将挂起等待与其他进程同步。

DeepSpeed 配置

DeepSpeed 功能可以通过配置文件 JSON 文件启用、禁用或配置，该文件应指定为 args.deepspeed_config。下面显示了一个示例配置文件。有关完整功能集，请参阅API 文档。

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

启动 DeepSpeed 训练

DeepSpeed 安装了入口点 deepspeed 来启动分布式训练。我们以下列假设为例说明 DeepSpeed 的用法

  1. 您已将 DeepSpeed 集成到您的模型中
  2. client_entry.py 是您模型的入口脚本
  3. client args 是 argparse 命令行参数
  4. ds_config.json 是 DeepSpeed 的配置文件

资源配置（多节点）

DeepSpeed 使用与 OpenMPI 和 Horovod 兼容的主机文件配置多节点计算资源。主机文件是主机名（或 SSH 别名）的列表，这些机器可通过无密码 SSH 访问，以及插槽计数，指定系统上可用的 GPU 数量。例如，

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

指定名为 worker-1 和 worker-2 的两台机器各拥有四个 GPU 用于训练。

主机文件通过 --hostfile 命令行选项指定。如果未指定主机文件，DeepSpeed 将搜索 /job/hostfile。如果未指定或找到主机文件，DeepSpeed 将查询本地机器上的 GPU 数量以发现可用的本地插槽数量。

以下命令在 myhostfile 中指定的所有可用节点和 GPU 上启动一个 PyTorch 训练作业

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

或者，DeepSpeed 允许您将模型的分布式训练限制为可用节点和 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 和 GPU 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 启动

DeepSpeed 现在支持无需无密码 SSH 即可启动训练作业。此模式在 Kubernetes 等云环境中特别有用，在这些环境中可以进行灵活的容器编排，而设置带有无密码 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 信息的宿主机文件。
  • --no_ssh: 启用无 SSH 模式。
  • --node_rank=<n>: 指定节点的等级。这应该是一个从 0 到 n - 1 的唯一整数。
  • --master_addr=<addr>: Leader 节点（rank 0）的地址。
  • --master_port=<port>: Leader 节点的端口。

在此设置中，主机文件中的主机名无需通过无密码 SSH 即可访问。但是，启动器仍然需要主机文件来收集有关环境的信息，例如节点数量和每个节点的 GPU 数量。

每个节点都必须使用唯一的 node_rank 启动，并且所有节点都必须提供 Leader 节点（rank 0）的地址和端口。此模式使启动器类似于 torchrun 启动器，如 PyTorch 文档 中所述。

多节点环境变量

在跨多个节点进行训练时，我们发现支持传播用户定义的环境变量很有用。默认情况下，DeepSpeed 将传播所有已设置的 NCCL 和 PYTHON 相关环境变量。如果您想传播其他变量，可以在名为 .deepspeed_env 的点文件中指定它们，该文件包含以换行符分隔的 VAR=VAL 条目列表。DeepSpeed 启动器将在您执行的本地路径和您的主目录 (~/) 中查找。如果您想用自己的文件或路径和名称覆盖此文件的默认名称，可以使用环境变量 DS_ENV_FILE 来指定。这在您启动多个需要不同变量的作业时最有用。

举一个具体的例子，一些集群在训练前需要设置特殊的 NCCL 变量。用户只需将这些变量添加到其主目录中名为 .deepspeed_env 的文件中，其内容如下

NCCL_IB_DISABLE=1
NCCL_SOCKET_IFNAME=eth0

DeepSpeed 将确保在整个训练作业中，在每个节点上启动每个进程时，这些环境变量都已设置。

MPI 和 AzureML 兼容性

如上所述，DeepSpeed 提供自己的并行启动器来帮助启动多节点/多 GPU 训练作业。如果您更喜欢使用 MPI（例如 mpirun）启动训练作业，我们也支持这样做。需要注意的是，DeepSpeed 仍将使用 torch distributed NCCL 后端，而不是 MPI 后端。

要使用 mpirun + DeepSpeed 或 AzureML（使用 mpirun 作为启动器后端）启动训练作业，您只需安装 mpi4py Python 包。DeepSpeed 将使用它来发现 MPI 环境，并将必要的状态（例如，世界大小、rank）传递给 torch 分布式后端。

如果您正在使用模型并行、流水线并行，或者在调用 deepspeed.initialize(..) 之前需要进行 torch.distributed 调用，我们提供相同的 MPI 支持，并额外提供一个 DeepSpeed API 调用。将您的初始 torch.distributed.init_process_group(..) 调用替换为

deepspeed.init_distributed()

资源配置（单节点）

如果仅在单个节点（带有一个或多个 GPU）上运行，DeepSpeed 不需要如上所述的主机文件。如果未检测到或未传入主机文件，DeepSpeed 将查询本地机器上的 GPU 数量以发现可用的插槽数量。--include 和 --exclude 参数正常工作，但用户应将 'localhost' 指定为主机名。

另请注意，CUDA_VISIBLE_DEVICES 可以与 deepspeed 一起使用，以控制在单个节点上应使用哪些设备。因此，以下任一命令都可以在当前节点的设备 0 和 1 上启动

deepspeed --include localhost:0,1 ...

CUDA_VISIBLE_DEVICES=0,1 deepspeed ...

更新日期：