文档

使用AIACC-Training PyTorch版

更新时间:

自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解决该问题。

  • 本页导读 (1)
文档反馈