Pytorch DDP实践

关于使用多GPU进行分布式训练的详细教程可见https://zhuanlan.zhihu.com/p/113694038

实现一个distributed data parallel的重点是管理各个进程之间的通信问题。首先我们需要创建进程组。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

from torch.nn.parallel import DistributedDataParallel as DDP

# On Windows platform, the torch.distributed package only
# supports Gloo backend, FileStore and TcpStore.
# For FileStore, set init_method parameter in init_process_group
# to a local file. Example as follow:
# init_method="file:///f:/libtmp/some_file"
# dist.init_process_group(
#    "gloo",
#    rank=rank,
#    init_method=init_method,
#    world_size=world_size)
# For TcpStore, same way as on Linux.

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

这里的communication backend有四个选项：gloo（facebook），MPI（OpenMPI），NCCL（nvidia）与TCP。他们属于不同的分布式框架。一般GPU训练选择nccl，CPU训练选gloo。

单机多卡

首先，检测GPU数目。
可以使用torch.cuda.device_count()来打印可用的GPU数目。

多卡训练要考虑通信开销的，是个trade off的过程，不见得四块卡一定比两块卡快多少，可能是训练到四块卡的时候通信开销已经占了大头。

怎样将原来的单GPU训练代码改为多卡distributed训练代码？

模型侧，只需要用DistributedDataParallel包装一下原来的model即可，在backend端它会支持梯度的All-Reduce操作。数据侧，我们nn.utils.data.DistributedSampler来给各个进程切分数据，只需要在dataloader中使用这个sampler就好，值得注意的一点是你要训练循环过程的每个epoch开始时调用train_sampler.set_epoch(epoch)，（主要是为了保证每个epoch的划分是不同的）其它的训练代码都保持不变。

怎样查看本地ip与通信端口？

配置分布式任务的第一步，是确保机器之间能够互相正常通信：

多机多卡