torch分布式训练

1. 基本概念

• rank
  进程号，在多进程上下文中，我们通常假定rank 0是第一个进程或者主进程，其它进程分别具有1，2，3不同rank号，这样总共具有4个进程

• node
  物理节点，可以是一个容器也可以是一台机器，节点内部可以有多个GPU；nnodes指物理节点数量， nproc_per_node指每个物理节点上面进程的数量

• local_rank
  指在一个node上进程的相对序号，local_rank在node之间相互独立

• WORLD_SIZE
  全局进程总个数，即在一个分布式任务中rank的数量

• Group
  进程组，一个分布式任务对应了一个进程组。只有用户需要创立多个进程组时才会用到group来管理，默认情况下只有一个group

如下图所示，共有3个节点(机器)，每个节点上有4个GPU，每台机器上起4个进程，每个进程占一块GPU，那么图中一共有12个rank，nproc_per_node=4，nnodes=3，每个节点都一个对应的node_rank。

注意
rank与GPU之间没有必然的对应关系，一个rank可以包含多个GPU；一个GPU也可以为多个rank服务（多进程共享GPU），在torch的分布式训练中习惯默认一个rank对应着一个GPU，因此local_rank可以当作GPU号

• backend
  通信后端，可选的包括：nccl（NVIDIA推出）、gloo（Facebook推出）、mpi（OpenMPI）。一般建议GPU训练选择nccl，CPU训练选择gloo

• master_addr与master_port
  主节点的地址以及端口，供init_method 的tcp方式使用。 因为pytorch中网络通信建立是从机去连接主机，运行ddp只需要指定主节点的IP与端口，其它节点的IP不需要填写。 这个两个参数可以通过环境变量或者init_method传入

  方式1：

  os.environ[‘MASTER_ADDR’] = ‘localhost’
  os.environ[‘MASTER_PORT’] = ‘12355’
  dist.init_process_group(“nccl”,
  rank=rank,
  world_size=world_size)

  方式2：

  dist.init_process_group(“nccl”,
  init_method=“tcp://localhost:12355”,
  rank=rank,
  world_size=world_size)

2. 使用分布式训练模型

使用DDP分布式训练，一共就如下个步骤：

1. 初始化进程组 dist.init_process_group
2. 设置分布式采样器 DistributedSampler
3. 使用DistributedDataParallel封装模型
4. 使用torchrun 或者 mp.spawn 启动分布式训练

2.1 初始化进程组

进程组初始化如下:

torch.distributed.init_process_group(backend, 
									init_method=None, 
									world_size=-1, 
									rank=-1, 
									store=None,
									...)

• backend: 指定分布式的后端，torch提供了NCCL, GLOO,MPI三种可用的后端，通常CPU的分布式训练选择GLOO, GPU的分布式训练就用NCCL即可
• init_method：初始化方法，可以是TCP连接、File共享文件系统、ENV环境变量三种方式
  • init_method='tcp://ip:port'： 通过指定rank 0（即：MASTER进程）的IP和端口，各个进程进行信息交换。 需指定 rank 和 world_size 这两个参数
  • init_method='file://path'：通过所有进程都可以访问共享文件系统来进行信息共享。需要指定rank和world_size参数
  • init_method=env://：从环境变量中读取分布式的信息(os.environ)ÿ