pytorch 分布式训练 distributed parallel 笔记

此文章仅供学习使用，请勿用作商业用途。

在GPU上进行训练

查看自己的GPU并赋值device

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
inputs, labels = inputs.to(device), labels.to(device)

多GPU并行编程

普通情况下的单个GPU：

device = torch.device("cuda:0")
model.to(device)
mytensor = my_tensor.to(device) #注意这个操作会返回在GPU上的一个tensor，之前的那个tensor并没有被销毁，所以需要重新赋值

python在默认情况下只使用一个GPU，在多个GPU的情况下就需要使用pytorch提供的DataParallel

参考内容，由简单到复杂：

data_parallel_tutorial

parallelism_tutorial

pytroch分布式

在pytorch中分布式分为两种，一种是一机多卡，另一种是多机多卡。

MXNet之ps-lite及parameter server原理讲解了多级分布式训练的种种基础问题。

对于一机多卡:

对于多机多卡：

分布式训练

pytorch分布式训练通信的后端使用的是gloo或者NCCL，一般来说使用NCCL对于GPU分布式训练，使用gloo对CPU进行分布式训练，MPI需要从源码重新编译。

默认情况下通信方式采用的是Collective Communication也就是通知群体，当然也可以采用Point-to-Point Communication即一个一个通知。

在分布式多个机器中，需要某个主机是主要节点。首先主机地址应该是一个大家都能访问的公共地址,主机端口应该是一个没有被占用的空闲端口。

其次，需要定义world-size代表全局进程个数（一般一个GPU上一个进程），rank代表进程的优先级也是这个进程的编号，rank=0的主机就是主要节点。

同时，rank在gpu训练中，又表示了gpu的编号/进程的编号。

torch.nn.parallel.DistributedDataParallel()不同于torch.nn.DataParallel()``````torch.distributed也不同于 torch.multiprocessing后者是对单机多卡的情况进行处理的.即使是在单机的情况下，DistributedDataParallel的效果也要好于DataParallel，因为：

每个流程维护自己的优化器，并在每次迭代中执行完整的优化步骤。
虽然这可能看起来是多余的，因为梯度已经收集在一起并跨进程平均，
因此每个进程都是相同的，这意味着不需要参数广播步骤，从而减少节点之间传输张量的时间。
每个进程都包含一个独立的Python解释器，消除了额外的解释器开销和“GIL-thrashing”，
即从单个Python进程中驱动多个执行线程、模型副本或gpu。
这对于大量使用Python运行时的模型(包括具有循环层或许多小组件的模型)尤其重要。

为了进行分布式训练，多个机器之间必须可以进行网络通信，且每个机器都需要各自运行训练的代码.通信可以使用各种后端，其中对于多机多卡GPU一般使用NCCL。在实际分布式运行起来的时候会涉及到物理网络端口使用的问题，使用的时候一般会出现很多问题，一般来说需要配置环境变量：

export NCCL_SOCKET_IFNAME=eth0 # 设置NCCL_SOCKET_IFNAME
export GLOO_SOCKET_IFNAME=eth0 # 设置GLOO_SOCKET_IFNAME

其他NCCL的环境变量

对执行训练的进程进行初始化

torch.distributed.init_process_group(backend, init_method='env://', timeout=datetime.timedelta(0, 1800), **kwargs)能够对分布式通信进行初始化，其中：

backend str/Backend 是通信所用的后端，可以是"ncll" "gloo"或者是一个torch.distributed.Backend类（Backend.GLOO）
init_method str 这个URL指定了如何初始化互相通信的进程
world_size int 执行训练的所有的进程数
rank int this进程的编号，也是其优先级
timeout timedelta 每个进程执行的超时时间，默认是30分钟，这个参数只适用于gloo后端
group_name str 进程所在group的name

其他的一系列函数：

torch.distributed.get_backend(group=group) # group是可选参数，返回字符串表示的后端 group表示的是ProcessGroup类
torch.distributed.get_rank(group=group) # group是可选参数，返回int，执行该脚本的进程的rank
torch.distributed.get_world_size(group=group) # group是可选参数,返回全局的整个的进程数
torch.distributed.is_initialized() # 判断该进程是否已经初始化
torch.distributed.is_mpi_avaiable() # 判断MPI是否可用
torch.distributed.is_nccl_avaiable() # 判断nccl是否可用

使用TCP

初始化：

import torch.distributed as dist
dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456',
                        rank=args.rank, world_size=4)

需要指定第0个机器的地址和端口号，以及指定world size

共享文件系统

初始化：

import torch.distributed as dist
dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile',
                        world_size=4, rank=args.rank)

共享文件系统需要指定file://开头的URL，各个进程在共享文件系统中通过这个文件来实现同步或者异步，所以每次在开始训练的时候保证这个文件被清空。这个文件也可以是不存在的，会自动创建，但是目录需要是存在的。

环境变量

默认情况下使用的都是环境变量来进行分布式通信，也就是指定init_method="env://"，这个进程会自动从本机的环境变量中读取如下数据：

MASTER_PORT: rank0上机器的一个空闲端口
MASTER_ADDR: rank0机器的地址
WORLD_SIZE: 这里可以指定，在init函数中也可以指定
RANK: 本机的rank，也可以在init函数中指定

GROUP

一个group就是一个world，默认情况下是只有一个group，如果需要更加精细更加特殊的通信，不同的机器进行不同的操作，这个时候就可以划分多个group。在torch.distributed.init_process_group()之后就可以调用torch.distributed.new_group(ranks=None, timeout=datetime.timedelta(0, 1800))创建新的group，所有的group都必须使用相同的backend。参数ranks是一个list，包含了在这个group中的进程的rank。

点对点通信

torch.distributed.send(tensor, dst, group=group, tag=0) # 向dst（rank）进程发送tensor，也可以指定group，tag用于对方recv进行对应
torch.distributed.recv(tensor, src=None, group=group, tag=0) # 从src处接受tensor，如果src为None表示接受任意源发来的tensor，tag用于对应
# 下面两个函数返回的是distributed request objects
torch.distributed.isend(tensor, dst, group=<object object>, tag=0)
torch.distributed.irecv(tensor, src, group=<object object>, tag=0)

同步和异步的collective operations

每个collective operations也就是群体操作都支持同步和异步的方式。

同步collective operations

默认的情况下都是同步的模式，async_op被设置成False。同步表示这个collective operations返回的时候就一定执行了，并且collective function不会返回任何东西。

异步collective operations

async_op被设置为True，调用collective function返回一个distributed request object，这个object有如下两个方法用于控制：

is_completed() 判断是否执行完毕
wait() 使用这个方法来阻塞这个进程，直到调用的collective function执行完毕

Collective functions

这些op需要好好看看，怎么使用。

torch.distributed.broadcast(tensor, src, group=group, async_op=False) # 将tensor从src(rank)广播到group中

torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=group, async_op=False) # 对tensor进行原地in-pllace的reduce，op是torch.distributed.ReduceOp中的一个，指定了某种确定的element-wise的操作

torch.distributed.reduce(tensor, dst, op=ReduceOp.SUM, group=<object object>, async_op=False)

torch.distributed.all_gather(tensor_list, tensor, group=<object object>, async_op=False) # 将group中的tensor集中到tensor_list中

torch.distributed.gather(tensor, gather_list, dst, group=<object object>, async_op=False) # 将group中的tensor集中到dst(rank)处

torch.distributed.scatter(tensor, scatter_list, src, group=<object object>, async_op=False)

torch.distributed.barrier(group=<object object>, async_op=False)

多GPU的collective functions

如果每个主机上有超过一个GPU，并且使用NCCL的通信后端的时候，可以使用broadcast_multigpu() all_reduce_multigpu() reduce_multigpu() all_gather_multigpu()，这些函数可以轻松使用，要注意到每个进程上的tensor list长度都必须相同。

假设有两个节点，每个节点上有8个GPU：

# 第0个节点上执行的代码
import torch
import torch.distributed as dist

dist.init_process_group(backend="nccl",
                        init_method="file:///distributed_test",
                        world_size=2,
                        rank=0)
tensor_list = []
for dev_idx in range(torch.cuda.device_count()):
    tensor_list.append(torch.FloatTensor([1]).cuda(dev_idx))

dist.all_reduce_multigpu(tensor_list)

# 第1个节点上执行的代码
import torch
import torch.distributed as dist

dist.init_process_group(backend="nccl",
                        init_method="file:///distributed_test",
                        world_size=2,
                        rank=1)
tensor_list = []
for dev_idx in range(torch.cuda.device_count()):
    tensor_list.append(torch.FloatTensor([1]).cuda(dev_idx))

dist.all_reduce_multigpu(tensor_list)

在执行完毕之后，两个节点上每个都会产生8个tensor，并且会都减去16.

多GPU的分布式collective functions：

torch.distributed.broadcasr_multigpu(tensor_list, group=<object object>, async_op=False, src_tensor=0) # 将tensor广播到group中每个节点的每个GPU(进程)

torch.distributed.all_reduce_multigpu(tensor_list, op=ReduceOp.SUM, group=<object object>, async_op=False) # 减少所有机器上的张量数据，从而得到最终的结果。这个函数减少了每个节点上的张量数量，而每个张量位于不同的gpu上。因此，张量列表中的输入张量需要是GPU张量。另外，张量列表中的每个张量都需要驻留在不同的GPU上。

torch.distributed.reduce_multigpu(tensor_list, dst, op=ReduceOp.SUM, group=<object object>, async_op=False, dst_tensor=0) # 减少所有机器上多个gpu上的张量数据。tensor_list中的每个张量应该位于一个单独的GPU上只有有秩dst的进程上tensor_list[dst_张量]的GPU会收到最终结果。目前只支持nccl后端张量，应该只支持GPU张量

torch.distributed.all_gather_multigpu(output_tensor_lists, input_tensor_list, group=<object object>, async_op=False) # 从列表中的整个组收集张量。tensor_list中的每个张量应该位于一个单独的GPU上目前只支持nccl后端张量，应该只支持GPU张量

启动辅助工具 Launch utility

torch.distributed也提供了一个辅助启动工具torch.distributed.launch，这个工具可以辅助在每个节点上启动多个进程process，支持Python2 和 Python3.

这个工具可以用作CPU或者GPU，如果被用于GPU，每个GPU产生一个进程Process。这个工具对于具有多个直接支持gpu的Infiniband接口的系统尤其有利，因为所有这些接口都可以用于聚合通信带宽。

单节点多进程(GPU)分布式训练

python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE train.py [--arg1 --arg2 ...]

多节点多进程(GPU)分布式训练

假设两个节点，节点0和节点1。

节点0：192.168.1.1: 1234

python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=2 --node_rank=0 --master_addr="192.168.1.1" --master_port=1234 train.py [--arg1 --arg2 ...]

节点1：

python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=2 --node_rank=1 --master_addr="192.168.1.1" --master_port=1234 train.py [--arg1 --arg2 ...]

帮助

python -m torch.distributed.launch --help

Notices

这个工具在NCCL上才能发挥最好的性能，NCCL也是被推荐用于分布式GPU训练的。
在训练的train.py中必须要解析--local_rank=LOCAL_PROCESS_RANK这个命令行参数，这个命令行参数是由torch.distributed.launch提供的，指定了每个GPU在本地的rank。

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()

解析之后在代码中设置device(GPU)

torch.cuda.set_device(arg.local_rank) # 在开始训练之前要设置

或者是下面的形式：

with torch.cuda.device(arg.local_rank):
    train()

在代码中训练开始之前还需要调用init_process_group:

torch.distributed.init_process_group(backend='nccl', init_method='end://')

在训练的代码中使用常规的分布式函数调用或者是torch.nn.parallel.DistributedDataParallel()。如果使用了GPU训练并且使用了torch.nn.parallel.DistributedDataParallel()，使用如下的方法来进行配置:

model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[arg.local_rank], output_device=arg.local_rank)

device_ids需要是唯一的一个GPU编号，这个训练脚本在这个GPU上执行。device_ids需要是[args.local_rank]并且output_device需要是args.local_rank才能使用这个工具。

local_rank是一个局部的id，在每个机器上GPU的id，如果使用不当会出现如下问题.

孵化工具 Spawn utility

这个工具可以产生多个进程torch.multiprocessing.spawn()

总地来说，进一步的使用torch.distributed请看ImageNet的大规模训练。