自PyTorch 1.x发布迭代后，使用PyTorch原生自带的DDP进行分布式训练逐渐形成了主流。本文为您介绍如何使用AIACC-Training，对基于PyTorch框架搭建的模型进行分布式训练加速的方法，以及可能遇到的问题和解决办法。

适配PyTorch DDP API（推荐）

背景信息

关于PyTorch DDP的更多信息，请参见PyTorch官网。

代码适配与运行

  1. 适配代码。
    您仅需在训练代码main函数对应的文件上增加一行导入aiacc代码即可，但请务必确保导入Perseus的位置在导入torch之前。命令行示例如下所示：
    import perseus    # aiacc的导入一定要放在torch的导入之前
import torch
import torch.nn as nn
import torch.distributed as dist
......
  2. 启动脚本。
    使用标准的torch.distributed.launch方式运行DDP的分布式训练，以下示例展示使用双机进行PyTorch DDP方式启动脚本，每台机器挂8张卡，命令行示例如下所示：
    ### 机器一上的命令
python -m torch.distributed.launch --nproc_per_node=8 --nnodes=2 --node_rank=0 --master_addr=${machine1_ip} --master_port=6007 --use_env ${TRAIN_SCRIPT}

### 机器二上的命令
python -m torch.distributed.launch --nproc_per_node=8 --nnodes=2 --node_rank=1 --master_addr=${machine1_ip} --master_port=6007 --use_env ${TRAIN_SCRIPT}
    部分参数说明如下：
    • ${machine1_ip}为机器一的内网ip，如192.168.2.211。您可以通过执行ifconfig命令进行查看。
    • ${TRAIN_SCRIPT}为实际训练的脚本。

示范用例

AIACC-Training软件包路径中为您提供了已适配DDP的示例代码。您可以通过以下操作体验训练过程。

  1. 进入示例代码目录。
    cd `echo $(python -c "import perseus; print(perseus)") | cut -d\' -f 4 | sed "s/\_\_init\_\_\.py//"`examples/
  2. 启动训练。
    采用原生DDP的启动方式运行PyTorch_ddp_benchmark.py脚本，以下示例为单机8卡训练的启动命令。
    NP=8
ADDR=localhost
PORT=6006
python -m torch.distributed.launch --nproc_per_node=$NP --nnodes=1 --node_rank=0 --master_addr=$ADDR --master_port=$PORT \
  pytorch_ddp_benchmark.py

适配Horovod API

前提条件

已将Perseus更新至1.3.2或以上版本，否则会影响您使用PyTorch的相关加速功能。

操作步骤

AIACC-Training for PyTorch支持Horovod API。如果您之前是使用Horovod进行分布式训练，只需替换import模块即可。替换内容如下：
import perseus.torch.horovod as hvd

如果您的训练代码是非分布式代码，可以参考以下操作步骤将训练代码升级为Horovod接口的分布式训练代码。

  1. 在main函数的开头部分，执行如下命令，初始化Perseus Horovod模块。
    说明 请务必在使用其他Perseus API之前调用hvd.init()。 
    hvd.init()
  2. 将当前process绑定对应的GPU卡。
    torch.cuda.set_device(hvd.local_rank())
  3. 通常情况下，对于训练的step数和warmup step参数，需要除以总的进程数hvd.size()，学习率需要对应增大hvd.size()倍。
    说明 部分模型不需要增大学习率，如BERT模型，具体请根据训练收敛情况作判断。 
    step = step // hvd.size()
learning_rate = learning_rate * hvd.size()
  4. 重载Optimizer。
    optimizer = hvd.DistributedOptimizer(
                          optimizer,named_parameters=model.named_parameters())
    如果是多个模型，有多个named_parameters，则需要进行合并处理，例如：
    all_named_parameters = []
for name, value in model1.named_parameters():
    all_named_parameters.append((name, value))
for name, value in model2.named_parameters():
    all_named_parameters.append((name, value))
optimizer = hvd.DistributedOptimizer(
                           optimizer, named_parameters=all_named_parameters)
  5. 广播全局变量参数到所有机器。
    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)
    如果是多个模型，有多个state_dict，则需要进行合并处理，例如：
    all_state_dict={}
all_state_dict.update(model1.state_dict())
all_state_dict.update(model2.state_dict())
hvd.broadcast_parameters(all_state_dict, root_rank=0)
  6. 将数据集划分为子数据集。
    train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset, num_replicas=hvd.size(), rank=hvd.rank())
loader = torch.utils.DataLoader(
    train_dataset, batch_size=batch_size, sampler=train_sampler, **kwargs)
  7. 模型单机单卡化。

    由于Perseus运行的是单机单卡程序，然后下发给各个GPU进行单机多卡或分布式训练，因此程序的运行方式需要为单机单卡。

    原程序为：
    model = nn.DataParallel(model.cuda())
    使用以下任一种方法修改：
    • 方式1
      model = nn.DataParallel(model.cuda(), device_ids=[hvd.local_rank()])
    • 方式2
      model = model.cuda()

      此方式中，cuda()默认会调用步骤2中设置的当前process绑定的GPU卡。

  8. 保存checkpoint。
    save_checkpoint = True if hvd.rank() == 0 else False
verbose = 1 if hvd.rank() == 0 else 0
log_writer = tensorboardX.SummaryWritter(log_dir) if hvd.rank() == 0 else None
  9. 加载checkpoint。
    if hvd.rank() == 0:
    checkpoint = torch.load(filepath)
    model.load_state_dict(checkpoint['model'])
    optimizer.load_state_dict(checkpoint['optimizer'])
  10. 启动分布式训练。
    • 单机8卡运行训练
      mpirun -allow-run-as-root -np 8 -npernode 8 -x NCCL_DEBUG=INFO ./train.sh
    • 4机8卡运行训练
      mpirun -allow-run-as-root –bind-to none -np 32 -npernode 8 \
                   -x NCCL_DEBUG=INFO -x PATH -x LD_LIBRARY \
                   -x PERSEUS_ALLREDUCE_STREAMS=8
                   -hostfile mpi_host.txt ./train.sh

常见问题

模型数据存放位置冲突，导致报错Input type (CUDAFloatTensor) and weight type (CPUFloatTensor) should be the same

通常情况下，是因为存在模型的参数不在GPU中，而输入数据在GPU中，遗漏了model的cuda操作导致。您可以通过增加model.cuda()将模型转移到GPU上来解决该问题。

出现报错RuntimeError: all tensors must be on devices[0]

请排查：

  • 是否在程序初始化时候指定了gpu_ids，如torch.cuda.set_device(hvd.local_rank())。
  • 是否采用DataParallel的时候没有统一设置。请设置一致的device_ids，如nn.DataParallel(model.cuda(), device_ids=str(hvd.local_rank()))。

加载模型显存OOM

一般是由于模型数据太大导致，可以通过将模型导入为CPU内存来解决此问题，如torch.load(pretrain, map_location='cpu')。

系统提示退出异常

首先请确保您已经正确完成了代码的适配操作。具体操作，请参见代码适配与运行。

发生该异常可能出自多个模型的组合问题。您需要将多个模型的named_parameters进行合并，然后再传递给Perseus进行Optimizer的封装。另外，多个模型组合时的广播参数需要设置为广播所有参数，因此，也需要进行state_dict的参数合并。

系统提示端口被占用

请先确认已有的端口是否被其他进程占用，如果有，可以使用pkill python命令进行进程杀除。

在Torch 1.9版本，默认使用的启动方式为torch.distributed.run，该方式使用默认的rdzv_backend作为训练的launcher，但是在启动训练之前就已经创建了服务，而代码里如果有dist.init_process_group()则会继续创建该服务，因此导致冲突，单机多卡情况下可以增加配置选项--standalone解决该问题。