快速入门

安装

编写 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 管理其更新或保存/恢复。
    • 如果调度程序应该在任何其他间隔(例如训练 epoch)执行,则用户不应在初始化期间将调度程序传递给 DeepSpeed,并且必须显式管理它。

模型检查点

保存和加载训练状态通过 DeepSpeed 中的 save_checkpointload_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 的一部分。

**重要事项**:所有进程都必须调用此方法,而不仅仅是等级为 0 的进程。这是因为每个进程都需要保存其主权重和调度程序+优化器状态。如果仅对等级为 0 的进程调用此方法,则此方法将挂起等待与其他进程同步。

DeepSpeed 配置

DeepSpeed 功能可以使用配置文件启用、禁用或配置,该文件应指定为 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 argsargparse 命令行参数
  4. ds_config.json 是 DeepSpeed 的配置文件

资源配置(多节点)

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

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

指定名为worker-1worker-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 和 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>: 主节点(排名0)的地址。
  • --master_port=<port>: 主节点的端口。

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

每个节点都必须使用唯一的node_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分布式NCCL后端,而不是MPI后端。

要使用mpirun + DeepSpeed或AzureML(使用mpirun作为启动器后端)启动训练作业,您只需安装mpi4py python包。DeepSpeed将使用它来发现MPI环境并将必要的状态(例如,世界大小、排名)传递给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一起使用来控制应使用哪些设备。例如,要仅使用当前节点的gpu1,请执行以下操作:

deepspeed --include localhost:1 ...

更新: