PyTorch 分布式训练

简介

随着神经网络的参数数量和计算需求的增长，在许多节点和许多GPU上有效并行化神经网络训练变得越来越重要，因为等待大型网络训练数月会减慢实验速度并限制进一步的发展。本文介绍 PyTorch 上进行多机多卡分布式训练的原理和操作方法，以便于训练神经网络时获得显著的性能提升。

在将神经网络的训练并行化到许多GPU上时，必须选择如何将不同的操作分配到可用的不同GPU上。 在这里，我们重点介绍一种称为数据并行随机梯度下降（SGD）的技术。 与标准SGD中一样，梯度下降是对数据子集（mini-batch）进行的，需要进行多次迭代才能遍历整个数据集。 但是，在数据并行训练中，每个GPU都具有整个神经网络模型的完整副本，并且对于每次迭代，mini-batch中仅分配了样本的子集。 对于每次迭代，每个GPU都会在其数据上进行网络的正向传播和反向传播，以计算相对于网络参数的损耗梯度。

最后，GPU彼此通信以平均不同GPU计算的梯度，将平均梯度应用于权重以获得新的权重。 所有GPU都以锁定步长(lock-step)进行迭代，一旦GPU完成迭代，它就必须等待所有其他GPU完成迭代，才能正确地更新权重。 这等效于在单个GPU上执行SGD，但是我们通过在多个GPU之间分配数据并并行进行计算来获得加速。

两种架构

TreeAllReduce

这种架构如图所示。采用这种计算模型的分布式，通常会遇到网络的问题，随着 worker 数量的增加，其加速比会迅速地恶化，需要借助其他辅助技术。

由于某一个 GPU 需要接收其他所有 GPU 的梯度，并求平均以及 broadcast 回去，通信成本随 GPU 的数量线性增加。

RingAllReduce

这种架构的特点是：通信成本是恒定的，并且与系统中GPU的数量无关，并且仅有系统中GPU之间最慢的连接确定。

如图所示，在 RingAllRedce 中，GPU 集群被组织成一个逻辑环，每个 GPU 只从左邻居接受数据、并发送数据给右邻居，即每次同步每个 GPU 只获得部分梯度更新，等一个完整的 Ring 完成，每个 GPU 都获得了完整的参数。

该算法分两个步骤进行：首先是 scatter reduce，然后是 all gather。在 scatter reduce 步骤中，GPU 将交换数据，以使每个 GPU 获得最终结果的一部分。在 all gather 步骤中，GPU 将交换这些块，以使所有 GPU 获得完整的最终结果。

详细流程

The Scatter-Reduce

为了简单起见，我们假设目标是按元素逐个汇总一个大型浮点数数组中的所有元素； 系统中有N个 GPU，每个 GPU 都有一个相同大小的数组，并且在all reduce结束，每个 GPU 都应具有相同大小的数组，其中包含原始数组中数字的总和。

首先，GPU 将阵列划分为N个较小的块（其中N是环中 GPU 的数量）。

接下来，GPU 将执行 scatter reduce 的 N-1 次迭代； 在每次迭代中，GPU 都会将其块之一发送到其右邻居，并从其左邻居接收一个块，并累积到该块中。 每个GPU发送和接收的块在每次迭代中都不同。 第 N 个GPU从发送块N 和接收块 N – 1 开始，然后从那里向后进行，每次迭代都发送在上一次迭代中接收到的块。

例如，在第一次迭代中，上图中的五个 GPU 将发送和接收以下块：

在第一个发送和接收完成之后，每个 GPU 将具有一个块，该块由两个不同 GPU 上相同块的总和组成。 例如，第二个 GPU 上的第一个块将是来自第二个 GPU 和第一个 GPU 的那个块中的值之和。

在接下来的迭代中，该过程继续进行，到最后，每个 GPU 将具有一个块，其中包含所有 GPU 中该块中所有值的总和。 下图显示了 scatter reduce 完成后的结果。

The Allgather

ring allgather 与 scatter reduce 相同（发送和接收进行 N-1 次迭代）相同，除了它们不累积 GPU 接收的值，而是简单地覆盖块。 第 N 个 GPU 首先发送第 N + 1 个块并接收第 N 个块，然后在以后的迭代中始终发送它刚接收的块。

例如，在我们的五个 GPU 设置的第一个迭代中，GPU 将发送和接收以下块：

在接下来的迭代中，该过程继续进行，最后，每个GPU将具有整个阵列的完全累加值。下图显示了 Allgather 完成后的结果。

在深度学习上的应用

为了最小化通信开销，我们可以利用神经网络的结构。 在每次迭代中，每个GPU都运行正向传播以计算误差，然后运行反向传播以计算神经网络的每个参数的梯度。 反向传播计算从输出层开始并向输入层移动，这意味着输出层参数的梯度在较早层的梯度之前明显可用。 由于Allreduce一次可以对网络参数的一个子集进行操作，因此我们可以在输出层参数上开始Allreduce，而其他梯度仍在计算中。 这样做会使通信与反向传播步骤中的其余计算重叠，从而减少了每个GPU最终等待通信完成的总时间。

PyTorch 分布式训练

基本概念

概念	含义
group	即进程组。默认情况下，只有一个组，一个 job 即为一个组，也即一个 world。
world size	表示全局进程个数。
rank	表示进程序号，用于进程间通讯，表征进程优先级。rank = 0 的主机为 master 节点。
local_rank	进程内，GPU 编号，非显式参数，由 torch.distributed.launch 内部指定。比方说， rank = 3，local_rank = 0 表示第 3 个进程内的第 1 块 GPU。

基本使用流程

PyTorch 中分布式的基本使用流程如下：

1. 在使用 distributed 包的任何其他函数之前，需要使用 init_process_group 初始化进程组，同时初始化 distributed 包。
2. 如果需要进行小组内集体通信，用 new_group 创建子分组
3. 创建分布式并行模型 DDP(model, device_ids=device_ids)
4. 为数据集创建 Sampler

重要 API

过时 API：DataParallel

这个wrapper能够很方便的使用多张卡，而且将进程控制在一个。唯一的问题就在于，DataParallel只能满足一台机器上gpu的通信，而一台机器一般只能装8张卡，对于一些大任务，8张卡就很吃力了。

PyTorch 最新版本多机多卡训练的 API 如下：

• torch.nn.parallel.DistributedDataParallel

​ 这个包是实现多机多卡分布训练最核心东西，它可以帮助我们在不同机器的多个模型拷贝平均梯度。

• torch.utils.data.distributed.DistributedSampler

​ 在多机多卡情况下分布式训练数据的读取也是一个问题，不同的卡读取到的数据应该是不同的。dataparallel的做法是直接将batch切分到不同的卡，这种方法对于多机来说不可取，因为多机之间直接进行数据传输会严重影响效率。于是有了利用sampler确保dataloader只会load到整个数据集的一个特定子集的做法。DistributedSampler就是做这件事的。它为每一个子进程划分出一部分数据集，以避免不同进程之间数据重复。

# 分布式训练示例
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel

dataset = your_dataset()
datasampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
dataloader = DataLoader(dataset, batch_size=batch_size_per_gpu, sampler=datasampler)
model = your_model()
model = DistributedDataPrallel(model, device_ids=[local_rank], output_device=local_rank)

其他部分就和正常训练代码无异了。

注意要点：

1. 和dataparallel不同，dataparallel需要将batchsize设置成n倍的单卡batchsize，而distributedsampler使用的情况下，batchsize设置与单卡设置相同。
2. 这里有几个新的参数：world size, rank, local rank, rank。world size指进程总数，在这里就是我们使用的卡数；rank指进程序号，local_rank指本地序号，两者的区别在于前者用于进程间通讯，后者用于本地设备分配。

Slurm 集群训练示例

多线程分布式后端 有 gloo, npi, nccl 三种。gloo 基本只支持 cpu，mpi 需要在本地重新编译 PyTorch，nccl 对 GPU 支持良好还不需要重新编译，官方推荐 nccl。

关于获取节点信息的详细代码：

import os
os.environ['SLURM_NTASKS']          #可用作world size
os.environ['SLURM_NODEID']          #node id
os.environ['SLURM_PROCID']          #可用作全局rank
os.environ['SLURM_LOCALID']         #local_rank
os.environ['SLURM_STEP_NODELIST']   #从中取得一个ip作为通讯ip

完整示例：

import torch
torch.multiprocessing.set_start_method('spawn')

import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel
import os

def dist_init(host_addr, rank, local_rank, world_size, port=23456):
    host_addr_full = 'tcp://' + host_addr + ':' + str(port)
    torch.distributed.init_process_group("nccl", init_method=host_addr_full,
                                         rank=rank, world_size=world_size)
    num_gpus = torch.cuda.device_count()
    torch.cuda.set_device(local_rank)
    assert torch.distributed.is_initialized()

rank = int(os.environ['SLURM_PROCID'])
local_rank = int(os.environ['SLURM_LOCALID'])
world_size = int(os.environ['SLURM_NTASKS'])
# get_ip函数自己写一下 不同服务器这个字符串形式不一样
# 保证所有task拿到的是同一个ip就成
ip = get_ip(os.environ['SLURM_STEP_NODELIST'])

dist_init(ip, rank, local_rank, world_size)


# 接下来是写dataset和dataloader，这个网上有很多教程
# 我这给的也只是个形式，按自己需求写好就ok
dataset = your_dataset()  #主要是把这写好
datasampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
dataloader = DataLoader(dataset, batch_size=batch_size_per_gpu, sampler=source_sampler)

model = your_model()     #也是按自己的模型写
model = DistributedDataPrallel(model, device_ids=[local_rank], output_device=local_rank)

# 此后训练流程与普通模型无异

Reference

• [1] Bringing HPC Techniques to Deep Learning

• [2] Bringing HPC Techniques to Deep Learning (翻译)

• [3] Getting Started With Distributed Data Parallel (PyTorch Tutorial)

• [4] Pytorch多机多卡分布式训练