本小节将会以深度学习框架PyTorch为例，讲解如何在OpenPAI集群上提交一个分布式（多机多卡）任务作业。本文假设读者已经掌握了PyTorch分布式训练的相关知识，如未掌握如何在普通机器上进行分布式训练，请查阅官方文档和官方示例。

提交分布式任务

提交分布式任务与普通任务的步骤非常类似，但是对比普通任务，分布式任务还需要进行两个额外的设置：1）指定节点的数量，2）在高级设置中指定一个预留的端口。

从OpenPAI环境变量获取初始化参数

除了在提交任务页面进行额外的设置之外，分布式训练的代码也需要进行一些改动才能够在OpenPAI集群中运行。下面将会介绍如何使用PyTorch原生分布式框架和horovod分布式框架在OpenPAI集群中运行时的代码改动。

PyTorch原生分布式框架

使用PyTorch自带的分布式训练启动器（torch.distributed.launch）一般会从命令行参数或者环境变量中获得参数用以初始化分布式框架（torch.distributed.init_process_group），一般会包括分布式后端（BACKEND），通信IP地址（MASTER_ADDR），通信的端口号（MASTER_PORT），节点数量（WORLD_SIZE），当前节点的编号（RANK）等参数。然而在OpenPAI集群中进行分布式训练时，我们无法直接指定如通信IP地址和通信的端口号这些参数。为了解决这一问题，OpenPAI在运行我们提交的任务时预设了一些环境变量，我们可以从这些环境变量获得分布式初始化参数。具体代码如下：

  1.  # 这个是我们设置的port label，假设我们设置的label为"distributed_sync_port"，也可以设置其他名称
  2. DISTRIBUTED_SYNC_PORT = "distributed_sync_port"
  3. # 当前任务名称，OpenPAI预设的环境变量都与这个名称有关
  4. taskrole = os.environ["PAI_CURRENT_TASK_ROLE_NAME"] 
  6. # 以下环境变量是分布式启动所需要的参数
  7. os.environ['WORLD_SIZE'] = os.environ[f"PAI_TASK_ROLE_TASK_COUNT_{taskrole}"]
  8. os.environ['MASTER_ADDR'] = os.environ[f'PAI_HOST_IP_{taskrole}_0'] # 0号节点的IP
  9. os.environ['MASTER_PORT'] = os.environ[f'PAI_{taskrole}_0_{DISTRIBUTED_SYNC_PORT}_PORT']
  10. os.environ["RANK"] = os.environ["PAI_CURRENT_TASK_ROLE_CURRENT_TASK_INDEX"]
  11. os.environ["LOCAL_RANK"] = str(0) # 由于每个节点只使用一块GPU，因此将LOCAL_RANK设置为“0”

只要在分布式初始化之前将上述环境变量设置即可，对代码的改动非常小。
OpenPAI的部分预设环境变量会与taskRole名称有关，因此我们首先从 PAI_CURRENT_TASK_ROLE_NAME 环境变量中获得taskRole名称，然后我们可以使用该名称获得其他的启动参数（如上述代码所示）。
值得注意的是如何获得分布式通信的端口号 。我们在提交分布式任务的页面上需要让OpenPAI为我们预留一个端口号（在高级设置中Ports设置中），然后我们可以可以通过设置的port label（上述的例子是"distributed_sync_port"）从环境变量中得到OpenPAI为我们预留的端口号（代码的第9行）。

Horovod分布式框架

Horovod 是Uber开源的一个分布式训练框架，主要目标是通过尽量少的代码修改将单个GPU的训练代码拓展到单机多卡或者多机多卡任务上，并且能够达到最快的训练速度。具体开源地址为：https://github.com/horovod/horovod。Horovod为目前主流的深度学习框架如pytorch、tensorflow、mxnet等提供了统一、简单的api接口，方便调用。其底层通信主要基于ring-allreduce 方法实现，依赖于open-mpi等开源库。horovod 与pytorch的DistributedDataParallel速度上并没太大差异，在实际使用中，两者取其一即可。horovod 具体安装方式参考 horovod github 中readme文档。

Pytorch使用示例

  1. import torch
  2. # 1.导入horovod.torch依赖库
  3. import horovod.torch as hvd
  5. # 初始化 Horovod
  6. hvd.init()
  8. # 将GPU绑定到每个线程
  9. torch.cuda.set_device(hvd.local_rank())
  11. # Define dataset...
  12. train_dataset = ...
  14. # 使用pytorch自带的 DistributedSampler用于多GPU训练
  15. train_sampler = torch.utils.data.distributed.DistributedSampler(
  16.     train_dataset, num_replicas=hvd.size(), rank=hvd.rank())
  18. train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)
  20. # Build model...
  21. model = ...
  22. model.cuda()
  24. optimizer = optim.SGD(model.parameters())
  26. # 使用horovod 的 DistributedOptimizer
  27. optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())
  29. # 将参数从第一个线程同步到其他所有线程
  30. hvd.broadcast_parameters(model.state_dict(), root_rank=0)
  32. for epoch in range(100):
  33.    for batch_idx, (data, target) in enumerate(train_loader):
  34.        optimizer.zero_grad()
  35.        output = model(data)
  36.        loss = F.nll_loss(output, target)
  37.        loss.backward()
  38.        optimizer.step()
  39.        if batch_idx % args.log_interval == 0:
  40.            print('Train Epoch: {} [{}/{}]\tLoss: {}'.format(
  41.                epoch, batch_idx * len(data), len(train_sampler), loss.item()))

在插入上述的关建代码之后，可以运行如下命令实现单机多卡或者多机多卡

  1. # 训练单机4卡
  2. $ horovodrun -np 4 python train.py
  4. # 训练双机8卡
  5. $ horovodrun -np 8 -H hostname1:4,hostname2:4 python train.py

OpenPAI horovod示例

OpenPAI提供了pytorch horovod的相应例程，地址为：https://github.com/microsoft/pai/blob/master/examples/pytorch_cifar10/yaml/Resnet18_horovod.yaml。详细配置如下：

  1. protocolVersion: 2
  2. name: Resnet18_horovod
  3. type: job
  4. jobRetryCount: 0
  5. prerequisites:
  6.   - type: dockerimage
  7.     uri: 'openpai/standard:python_3.6-pytorch_1.2.0-gpu'
  8.     name: docker_image_0
  9. taskRoles:
  10.   taskrole:
  11.     instances: 1
  12.     completion:
  13.       minFailedInstances: 1
  14.       minSucceededInstances: -1
  15.     taskRetryCount: 0
  16.     dockerImage: docker_image_0
  17.     resourcePerInstance:
  18.       gpu: 4
  19.       cpu: 12
  20.       memoryMB: 24576
  21.     commands:
  22.       - branch_name=pai-for-edu
  23.       - >-
  24.         wget
  25.         https://raw.githubusercontent.com/microsoft/pai/${branch_name}/contrib/edu-examples/pytorch_cifar10/src/cifar.py
  26.       - >-
  27.         wget
  28.         https://raw.githubusercontent.com/microsoft/pai/${branch_name}/contrib/edu-examples/pytorch_cifar10/src/horovod_cifar.py
  29.       - >-
  30.         wget
  31.         https://raw.githubusercontent.com/microsoft/pai/${branch_name}/contrib/edu-examples/pytorch_cifar10/src/init.sh
  32.       - bash init.sh
  33.       - >-
  34.         HOROVOD_GPU_ALLREDUCE=NCCL HOROVOD_GPU_BROADCAST=NCCL pip install
  35.         horovod
  36.       - >-
  37.         horovodrun -np 4 python horovod_cifar.py --epoch 200 --batchsize 512
  38.         --arch ResNet18
  39. defaults:
  40.   virtualCluster: default
  41. extras:
  42.   com.microsoft.pai.runtimeplugin:
  43.     - plugin: ssh
  44.       parameters:
  45.         jobssh: true

用户可以直接使用上述YAML配置文件运行OpenPAI提供的horovod例程，其核心运行命令为

  1. horovodrun -np 4 python horovod_cifar.py --arch ResNet18