入门指南

安装

编写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_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 功能可以通过一个配置 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 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=: 指定节点的排名。这应该是一个从0到n - 1的唯一整数。
  • --master_addr=: 主节点(rank 0)的地址。
  • --master_port=: 主节点的端口。

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

每个节点必须使用唯一的node_rank启动,并且所有节点都必须提供领导者节点(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,我们提供了一个额外的DeepSpeed API调用来支持相同的MPI功能。请将您最初的torch.distributed.init_process_group(..)调用替换为:

deepspeed.init_distributed()

资源配置(单节点)

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

还要注意,CUDA_VISIBLE_DEVICES 不能与 DeepSpeed 一起使用来控制应该使用哪些设备。例如,要仅使用当前节点的 gpu1,请执行以下操作:

deepspeed --include localhost:1 ...

更新: