快速入门
安装
- 安装非常简单,只需使用
pip install deepspeed
,查看更多详细信息。 - 要开始在 AzureML 上使用 DeepSpeed,请参阅 AzureML 示例 GitHub
- DeepSpeed 与 HuggingFace Transformers 和 PyTorch Lightning 有直接集成。HuggingFace Transformers 用户现在可以通过简单的
--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 管理其更新或保存/恢复。 - 如果调度程序应该在任何其他间隔(例如训练 epoch)执行,则用户不应在初始化期间将调度程序传递给 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
的一部分。
**重要事项**:所有进程都必须调用此方法,而不仅仅是等级为 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 的示例用法
- 您已将 DeepSpeed 集成到您的模型中
client_entry.py
是模型的入口脚本client args
是argparse
命令行参数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 和 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 ...