## 1.构建DeepSpeed模型
DeepSpeed 模型训练使用 DeepSpeed 引擎完成。该引擎可以包装任何类型为 torch.nn.module 的任意模型,并具有用于训练和检查点模型的最小API集。
要初始化 DeepSpeed 引擎:
```python
model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args,
model=model,
model_parameters=params)
```
`deepspeed.initialize` 确保在幕后适当地完成了分布式数据并行或混合精度训练所需的所有必要设置。除了包装模型外,DeepSpeed 还可以根据传递给 `deepspeed.initialize` 和 **DeepSpeed 配置文件**的参数构建和管理`training optimizer`, `data loader`以及 `the learning rate scheduler`。请注意,DeepSpeed 自动在每个训练步骤中执行学习率调度。
如果您已经有了分布式环境设置,您需要将:
```python
torch.distributed.init_process_group(...)
```
替换为:
```python
deepspeed.init_distributed()
```
默认情况下使用NCCL后端,DeepSpeed已经对其进行了彻底测试,但您也可以覆盖默认设置。
但是,如果您直到在`deepspeed.initialize()`之后才需要分布式环境设置,您就不必使用此函数,因为DeepSpeed将在其初始化期间自动初始化分布式环境。无论如何,您都需要删除`torch.distributed.init_process_group`,如果您已经安装了它。
## 2.训练
DeepSpeed**引擎初始化**后,可以使用三个简单的API进行模型训练:前向传播(callable object) ,反向传播(backward)和权重更新(step)。
```python
model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args,
model=model,
model_parameters=params)
```
在底层,DeepSpeed使用预定义的学习率调度器,自动执行分布式数据并行训练和混合精度训练所需的操作:
- **Gradient Averaging**:在分布式数据并行训练中,反向传播`backward`确保在一次`train_batch_size`的训练之后,在各数据并行进程之间平均梯度。
- **Loss Scaling**:在FP16 /混合精度训练中,DeepSpeed引擎自动处理缩放损失以避免梯度中的精度损失。
- **Learning Rate Scheduler**:当使用DeepSpeed的学习率调度器(在ds_config.json文件中指定)时,DeepSpeed在每个训练步骤(执行`model_engine.step()`时)调用调度器的`step()`方法。如果不使用DeepSpeed的学习率调度器:
(1) 如果计划在每个训练步骤执行调度,则用户可以在初始化DeepSpeed引擎时将调度器传递给deepspeed.initialize,并让DeepSpeed管理它以进行更新或保存/恢复。
(2) 如果计划在任何其他间隔(例如,训练周期)执行调度,则用户不应在初始化期间将调度程序传递给DeepSpeed,并且必须显式管理它。
## 3. Model Checkpointing
在DeepSpeed中,通过`save_checkpoint`和`load_checkpoint` API处理训练状态的保存和加载,它们需要两个参数来唯一标识一个检查点:
- ckpt_dir:检查点保存的目录。
- ckpt_id:一个标识符,用于唯一标识目录中的检查点。在下面的代码片段中,我们使用损失值作为检查点标识符。
```python
#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`中检索。在上面的示例中,步骤值存储为`client_sd`的一部分。
`重要提示:所有进程都必须调用此方法,而不仅仅是rank=0的进程。这是因为每个进程都需要保存其主权重和scheduler+optimizer states。如果仅为rank=0的进程调用此方法,它将挂起等待与其他进程同步。`
## 4. DeepSpeed配置文件
DeepSpeed的功能可以通过配置JSON文件进行启用、禁用或配置,该文件应作为`args.deepspeed_config`指定。以下是一个示例配置文件。
```json
{
"train_batch_size": 8,
"gradient_accumulation_steps": 1,
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.00015
}
},
"fp16": {
"enabled": true
},
"zero_optimization": true
}
```
## 5. 启动DeepSpeed训练
DeepSpeed安装了入口点`deepspeed`以启动分布式训练。我们通过以下假设说明了DeepSpeed的示例用法:
- 您已经将DeepSpeed集成到了您的模型中。
- `client_entry.py`是您的模型的入口脚本。
- `client args`是`argparse`命令行参数。
- `ds_config.json`是DeepSpeed的配置文件。
## 6. 资源配置
DeepSpeed配置了与OpenMPI和Horovod兼容的主机文件来配置多节点计算资源。主机文件是主机名(或SSH别名)和槽计数的列表,它们是通过无密码SSH可访问的机器和指定系统上可用的GPU数量。例如,
```python
worker-1 slots=4
worker-2 slots=4
```
指定了两台名为worker-1和worker-2的机器,每台机器都有四个GPU可用于训练。
主机文件使用`--hostfile`命令行选项指定。如果没有指定主机文件,DeepSpeed会搜索`/job/hostfile`。如果没有指定或找到主机文件,则DeepSpeed会查询本地机器上的GPU数量,以发现可用的本地槽数。
以下命令在`myhostfile`中指定的所有可用节点和GPU上启动一个PyTorch训练作业:
```shell
deepspeed --hostfile=myhostfile <client_entry.py> <client args> \
--deepspeed --deepspeed_config ds_config.json
```
另外,DeepSpeed允许您将您的模型的分布式训练限制在可用节点和GPU的子集上。这个特性通过两个命令行参数--num_nodes和--num_gpus启用。例如,可以使用以下命令将分布式训练限制为仅使用两个节点:
```shell
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之外的所有可用资源:
```shell
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:
```shell
deepspeed --include="worker-2:0,1" \
<client_entry.py> <client args> \
--deepspeed --deepspeed_config ds_config.json
```
## 7.多节点环境变量
在跨多个节点进行训练时,我们发现支持传播**用户定义的环境变量**非常有用。默认情况下,DeepSpeed将传播所有设置的NCCL和PYTHON相关环境变量。如果您想传播其他变量,可以在名为`.deepspeed_env`的点文件中指定它们,该文件包含一个新行分隔的`VAR = VAL`条目列表。 DeepSpeed启动器将查找您执行的本地路径以及您的主目录`(〜/)`。
具体而言,一些集群在训练之前需要设置特殊的NCCL变量。用户只需将这些变量添加到其主目录中的.deepspeed_env文件中,该文件如下所示:
```shell
NCCL_IB_DISABLE=1
NCCL_SOCKET_IFNAME=eth0
```
然后,DeepSpeed将确保在启动每个节点上的每个进程时设置这些环境变量,以进行训练作业。
## 8.MPI和AzureML兼容性
如上所述,DeepSpeed提供了自己的并行启动器来帮助启动多节点/多GPU训练作业。如果您希望使用MPI(例如mpirun)启动训练作业,我们提供对此的支持。需要注意的是,DeepSpeed仍将使用torch分布式NCCL后端,而不是MPI后端。
要使用mpirun + DeepSpeed或AzureML(使用mpirun作为启动器后端)启动训练作业,您只需要安装mpi4py Python包即可。 DeepSpeed将使用此来发现MPI环境并将必要的状态(例如world size、rank)传递给torch分布式后端。
如果您正在使用模型并行性、管道并行性或者在调用deepspeed.initialize(..)之前需要torch.distributed调用,则我们提供了相同的MPI支持,需要额外的DeepSpeed API调用。将初始的`torch.distributed.init_process_group(..)`调用替换为:
```python
deepspeed.init_distributed()
```
## 9. 资源配置-单节点
在我们仅在单个节点上运行(具有一个或多个GPU)的情况下,DeepSpeed不需要像上面描述的主机文件。如果未检测到或传递主机文件,则DeepSpeed将查询本地计算机上的GPU数量以发现可用的插槽数。--include和--exclude参数按照正常工作,但用户应将“localhost”指定为主机名。
还要注意,CUDA_VISIBLE_DEVICES不能用于DeepSpeed来控制应使用哪些设备。例如,要仅使用当前节点的gpu1,请执行:
```shell
deepspeed --include localhost:1 ...
```