ddp训练流程-pytorch教程

1. pytorch如何初始化分布式训练

核心函数如下，下面具体分析一下

    torch.distributed.init_process_group(backend=dist_backend,
                                         init_method=init_method,
                                         world_size=world_size,
                                         rank=rank)

1. backend就是通信协议，使用分布式时，在梯度汇总求平均的过程中，各主机之间需要进行通信。因此，需要指定通信的协议架构等。gpu是nccl，cpu是gloo。
2. init_method 指定当前进程组初始化方式，也就是获取其他节点的信息，进行同步
3. world_size是进程的个数，比如我们有3台机器，每台机器有2个gpu，那么就有3x2=6个进程
4. rank则表示进程的标识

2. 我们怎么知道要使用哪几台机器进行训练的？

假如在同一个局域网内有6台机器，其中三台机器训练bert，另外三台训练gpt，每台机器是如何知道其他节点跟自己是否训练的是同一个任务呢（更准确的是不同进程之间是如何保证）？一种非常简单的方法就是给同一个训练任务的机器分配唯一的id，id相同的时候大家进行通信，id不同时则不通信。pytorch使用的唯一标识是ip+port，细想觉得非常的巧妙，我们知道不同机器之间的通信主要靠的就是网络，网络服务就是通过ip+port来指定和区分。

训练任务	ip地址	数字标识	ip+port 标识
bert	192.168.1.2	1	192.168.1.2:5003
bert	192.168.1.3	1	192.168.1.2:5003
bert	192.168.1.4	1	192.168.1.2:5003
gpt	192.168.1.5	2	192.168.1.5:5004
gpt	192.168.1.6	2	192.168.1.5:5004
gpt	192.168.1.7	2	192.168.1.5:5004

如果不同任务即使使用了同一个机器，即ip相同，也可以通过port来进行区分，可以看到192.168.1.2这台机器上跑了两个任务（这台机器有两个gpu，或者资源不够，用户就是想要指定两个进程来训练不同的任务），我们可以通过port来区分这两个任务，我们也可以使用一个不同的ip来区分，bert是192.168.1.2:5003，gpt是192.168.1.6:5003，端口一样，但是ip不一样，当然也可以直接ip+port都不相同

训练任务	ip地址	port不同标识	ip不同标识	ip:port不同标识
bert	192.168.1.2	192.168.1.2:5003	192.168.1.2:5003	192.168.1.2:5003
bert	192.168.1.3	192.168.1.2:5003	192.168.1.2:5003	192.168.1.2:5003
bert	192.168.1.4	192.168.1.2:5003	192.168.1.2:5003	192.168.1.2:5003
gpt	192.168.1.2	192.168.1.2:5004	192.168.1.6:5003	192.168.1.6:5004
gpt	192.168.1.6	192.168.1.2:5004	192.168.1.6:5003	192.168.1.6:5004
gpt	192.168.1.7	192.168.1.2:5004	192.168.1.6:5003	192.168.1.6:5004

总而言之，不同任务之间是通过ip和port来作为唯一标识区分的。我们启动任务的时候指定这个ip+port，这个ip:port将会作为服务的主节点

训练任务	ip地址	ip+port 标识	启动
bert : master	192.168.1.2	192.168.1.2:5003	python train.py --master_addr 192.168.1.2 --master_port 5003
bert : slave	192.168.1.3	192.168.1.2:5003	python train.py --master_addr 192.168.1.2 --master_port 5003
bert : slave	192.168.1.4	192.168.1.2:5003	python train.py --master_addr 192.168.1.2 --master_port 5003
gpt : master	192.168.1.5	192.168.1.5:5004	python train.py --master_addr 192.168.1.5 --master_port 5004
gpt : slave	192.168.1.6	192.168.1.5:5004	python train.py --master_addr 192.168.1.5 --master_port 5004
gpt : slave	192.168.1.7	192.168.1.5:5004	python train.py --master_addr 192.168.1.5 --master_port 5004

3. 如何根据标识进行初始化（init_method）

torch获取这个唯一标识的方式也有两种（其实是三种，文件系统共享我没用过）

1. tcp：直接指定tcp的ip和端口，init_method='tcp://192.168.1.2:5003'

2. env : 我们获取到输入参数master_addr和master_port之后，设置环境变量

   os.environ['MASTER_ADDR'] = '192.168.1.2'
   os.environ['MASTER_PORT'] = '5003'
   

   然后就可以通过指定init_method="env://"来初始化服务了。很多博客都说要在环境变量中写入MASTER_ADDR和MASTER_PORT，我当时的理解是pytorch会直接把这两个变量写入到系统环境中

   export MASTER_ADDR="192.168.1.2"
   export MASTER_PORT="5003"
   

   然后就变得非常非常的困惑，有那么多的任务，如果大家都把自己的配置写到系统中不就存在了冲突了吗，最重要的是我在系统中根本就没看到这两个变量。后来才意识到，python会拷贝一份环境变量， os.environ['MASTER_ADDR'] = '192.168.1.2'其实添加的是拷贝环境变量的值，而不是真的在环境变量中指定了这个值。

   所谓的会自动从环境变量中获取MASTER_ADDR和MASTER_PORT，其实都是这个拷贝的环境变量，修改的也是这个拷贝值。

   可以试试，先启动一个python代码设置环境变量

   	import os
   	os.environ['MASTER_ADDR'] = '192.168.1.2'
   	print(os.environ['MASTER_ADDR']) # '192.168.1.2'
   

   再启动一个python代码读取环境变量

   import os
   print(os.environ['MASTER_ADDR']) # 空
   

   到系统环境中查看

   echo $master #空
   

4. 如何获取进程的唯一标识rank

world_size很好计算，这个是自己指定的，例如我们使用3台机器，每个节点有4个gpu，全部使用的话world_size=3*4=12，很直接world_size=nnodes * nproc_per_node。其中nnodes就是我们指定的节点个数，nproc_per_node就是单个节点执行的进程数，通常是每个机器gpu的数量。如果是cpu训练的话，就是cpu的个数，通常每台机器只有一个cpu。
上面讲了通过ip+port我们可以确定每个任务的唯一标识，通常一个任务我们会进行多几多卡训练，即启动多个进程。每个进程都有自己的唯一标识，这个就是rank。有趣的是，pytorch的进程id并不是根据全部机器或者world_size来分配的每个进程的rank的，假如我们有3个节点，每个节点4张卡，理想情况是我们执行pytorch的dpp初始化后，每个gpu都有一个rank值，依次递增到world_size-1

ip	gpu1	gpu2	gpu3	gpu4
192.168.1.2	0	1	2	3
192.168.1.3	4	5	6	7
192.168.1.4	8	9	10	11

但实际上，pytorch只会根据每个节点自身确定一个local_rank值，每次都是从0开始增加的

ip	gpu1	gpu2	gpu3	gpu4
192.168.1.2	0	1	2	3
192.168.1.3	0	1	2	3
192.168.1.4	0	1	2	3

所以为了获取全局的rank需要我们手动做一次转换rank=node_rank*n_gpu+local_rank

ip	node_rank	gpu1	gpu2	gpu3	gpu4
192.168.1.2	0	0	1	2	3
192.168.1.3	1	1*4+0	1*4+1	1*4+2	1*4+3
192.168.1.4	2	2*4+0	2*4+1	2*4+2	2*4+3

node_rank是我们给每个节点的编号。其实在这里有一个问题，可不可以一个节点使用2个gpu，一个节点使用3个gpu呢？这个时候该怎么获取每个进程的id呢？
还有一个问题需要关注，如果使用的init_method="env://"，那么也需要将rank和world_size也写入到环境变量中

	os.environ['RANK'] = rank
	os.environ['WORLD_SIZE'] = 12

5. 如何实现

import os
import argparse
import torch
import random
import numpy as np

def setup_new_process(local_rank, callee, args):
    args.local_rank = local_rank
    args.rank = args.node_rank * args.nproc_per_node + local_rank
    random.seed(args.rank)
    np.random.seed(args.rank)
    torch.manual_seed(args.rank)
    torch.cuda.manual_seed_all(args.rank)
    if args.init_method == "env://":
        os.environ['RANK'] = str(args.rank)

    if torch.cuda.is_available():
        torch.cuda.set_device(local_rank)
        torch.cuda.empty_cache()
    # 通信后端，nvidia GPU推荐使用NCCL
    dist_backend = 'nccl' if torch.distributed.is_nccl_available() else 'gloo'

    print(f'start init process: rank = {args.rank}')
    torch.distributed.init_process_group(backend=dist_backend,
                                         init_method=args.init_method,
                                         world_size=args.world_size,
                                         rank=args.rank)

    callee(args)

def train(args):
    if torch.distributed.is_initialized():
        model = torch.nn.parallel.DistributedDataParallel(model)
        print(f"rank = {args.rank} | strat train.......")    

def main(callee):
    parse = argparse.ArgumentParser()
    parse.add_argument('--init_method', type=str, default="env://")
    parse.add_argument('--master_addr', type=str, default="127.0.0.1")
    parse.add_argument('--master_port', type=str, default="5003")
    parse.add_argument('--nproc_per_node', type=int, default=-1)
    parse.add_argument('--node_rank', type=int, default=0)
    parse.add_argument('--world_size', type=int, default=4)
    parse.add_argument('--rank', type=int, default=4)
    parse.add_argument('--local_rank', type=int, default=4)

    args = parse.parse_args()
    if args.init_method == "env://":
        os.environ['WORLD_SIZE'] = str(args.world_size)
        os.environ['MASTER_ADDR'] = args.master_addr
        os.environ['MASTER_PORT'] = args.master_port
    else:
        args.init_method = f"tcp://{args.master_addr}:{args.master_port}"

    if args.nproc_per_node == -1:
        if torch.cuda.device_count() > 0:
            args.nproc_per_node = torch.cuda.device_count()
        else:
            args.nproc_per_node = os.cpu_count()
    torch.multiprocessing.spawn(setup_new_process, nprocs=args.nproc_per_node,
              args=(callee, args), join=True)


if __name__ == '__main__':
    main(train)

唯一没有解释的是torch.multiprocessing.spawn这个函数，这个函数用来启动分布式训练，本质就是创建多个线程。我的本地有4个cpu，所以我这里直接创建了4个进程来执行，注意nproc_per_node是单个节点进程数，也就是单机的gpu个数。

start init process: rank = 0
start init process: rank = 1
start init process: rank = 2
start init process: rank = 3
rank = 0 | strat train…
rank = 3 | strat train…
rank = 2 | strat train…
rank = 1 | strat train…

如果我们有多个节点，则需要在每个节点执行脚本

ip	命令
192.168.1.2	python train.py --master_addr 192.168.1.2 --master_port 5003 --node_rank=0 --world_size=4
192.168.1.3	python train.py --master_addr 192.168.1.2 --master_port 5003 --node_rank=1 --world_size=4
192.168.1.4	python train.py --master_addr 192.168.1.2 --master_port 5003 --node_rank=2 --world_size=4

前面只是启动了分布式训练而已，我们创建一个小模型来试试ddp

def train(args):
    rank = 0
    model = nn.Linear(5, 1, bias=False).to(rank)
    if torch.distributed.is_initialized():
        rank = torch.distributed.get_rank()
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank])
        print(f"rank = {rank} | start train.......")
    
    optimizer = torch.optim.Adam(model.parameters(),lr=0.01)

    for _  in range(10):
        optimizer.zero_grad()
        py = model(torch.rand(32,5).to(rank)
        loss = F.mse_loss(py,torch.rand(32,1).to(rank))
        print(loss)
        loss.backward()
        optimizer.step()

在单机CPU的模式下发现

raise ValueError(ValueError: DistributedDataParallel device_ids and output_device arguments only work with single-device/multiple-device GPU modules or CPU modules, but got device_ids [2], output_device None, and module parameters {device(type='cpu')}.

抛出了一个异常，上面这个代码主要是执行在gpu上的，to(rank)的意思就是把数据或模型加载到编号为rank的gpu上，我本地没有gpu，所以不能使用to(rank)，其次torch.nn.parallel.DistributedDataParallel也会把模型输出到某个device_id上，我们稍作修改让本地cpu可以跑起来

import torch.nn as nn
import torch.nn.functional as F
import torch

def train(args):
    model = nn.Linear(5, 1, bias=False)
    
    for name, params in model.named_parameters():
        print(f'before dpp : rank = {args.rank}, name = {name}, params = {params.tolist()}')
    
    if torch.distributed.is_initialized():
        model = torch.nn.parallel.DistributedDataParallel(model)
        print(f"rank = {args.rank} | strat train.......")
        
    for name, params in model.named_parameters():
        print(f'after dpp : rank = {args.rank}, name = {name}, params = {params.tolist()}')

    optimizer = torch.optim.Adam(model.parameters(),lr=0.01)

    for _  in range(10):
        optimizer.zero_grad()
        py = model(torch.rand(32,5))
        loss = F.mse_loss(py,torch.rand(32,1))
        loss.backward()
        optimizer.step()
    
    for name, params in model.named_parameters():
        print(f'finish dpp rank = {args.rank}, name = {name}, params = {params.tolist()}')

这里把模型的参数也打印出来了。代码中我们是直接使用随机数初始化每个网络的，因此可以看到每个进程的模型参数是不同的，但是训练结束之后可以看到，模型参数都变成了同一个

before dpp : rank = 0, name = weight, params = [[-0.0033482015132904053, 0.23990488052368164, -0.36807698011398315, -0.3291219472885132, -0.1722462773323059]]
before dpp : rank = 3, name = weight, params = [[-0.44340017437934875, -0.3527894914150238, -0.19154831767082214, -0.423104465007782, -0.025388896465301514]]
before dpp : rank = 1, name = weight, params = [[0.2304326891899109, -0.1973903477191925, -0.08669748902320862, 0.20990818738937378, -0.4210233688354492]]
before dpp : rank = 2, name = weight, params = [[0.10258638858795166, -0.10642534494400024, 0.12263882160186768, -0.022842705249786377, 0.1910441517829895]]
after dpp : rank = 2, name = module.weight, params = [[-0.0033482015132904053, 0.23990488052368164, -0.36807698011398315, -0.3291219472885132, -0.1722462773323059]]
after dpp : rank = 0, name = module.weight, params = [[-0.0033482015132904053, 0.23990488052368164, -0.36807698011398315, -0.3291219472885132, -0.1722462773323059]]
after dpp : rank = 3, name = module.weight, params = [[-0.0033482015132904053, 0.23990488052368164, -0.36807698011398315, -0.3291219472885132, -0.1722462773323059]]
after dpp : rank = 1, name = module.weight, params = [[-0.0033482015132904053, 0.23990488052368164, -0.36807698011398315, -0.3291219472885132, -0.1722462773323059]]
finish dpp rank = 0, name = module.weight, params = [[0.09621907025575638, 0.33785226941108704, -0.26929566264152527, -0.23034155368804932, -0.07334098219871521]]
finish dpp rank = 3, name = module.weight, params = [[0.09621907025575638, 0.33785226941108704, -0.26929566264152527, -0.23034155368804932, -0.07334098219871521]]
finish dpp rank = 1, name = module.weight, params = [[0.09621907025575638, 0.33785226941108704, -0.26929566264152527, -0.23034155368804932, -0.07334098219871521]]
finish dpp rank = 2, name = module.weight, params = [[0.09621907025575638, 0.33785226941108704, -0.26929566264152527, -0.23034155368804932, -0.07334098219871521]]

这是因为torch.nn.parallel.DistributedDataParallel(model)在加载模型的时候，会把rank=0的模型参数传给各个子节点，作为初始化的参数。这样可以保证每个节点拿到的模型参数都是一样的。训练的过程中由于梯度共享的原因，所以每一次迭代梯度也是相同的。