PyTorch分布式训练中各节点如何通信

深度学习

---

---

前言

同学们在处理分布式训练时经常会遇到以下几个疑问：

1、分布式节点如何知道彼此并通信。

每个节点启动训练任务需要设置主节点的ip:port，通过跟主节点通信来启动服务。

2、是不是每个节点都要执行python程序。

是的，写好训练代码后，每个节点都需要执行python命令，但是一般平台会帮我们管理分布式服务，自动把脚本分发到各个节点，并执行python成勋，才会产生不需要每个节点执行python代码的错觉。

3、如果忘记在某个节点执行python命令，结果会怎样。
ddp会阻塞在初始化的位置，一直等到有world_size进程为止，这个就是靠主节点来跟踪进程数。

4、每个节点是不是都需要有完备的数据。

是的，如果使用ddp的DistributedSampler来获取数据，那么一定要保证每个节点都有完整的数据，要不然根据indics无法获取正确的数据。

5、每个节点获取的mini-batch数据如何保证不重合。

不同节点设置相同的随机数种子，这样采样的batch就是一摸一样的，然后再根据各自的rank来从batch中取出自己的mini-batch。1. pytorch如何初始化分布式训练

pytorch如何初始化分布式训练

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

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

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

怎么知道要使用哪几台机器进行训练的

假如在同一个局域网内有6台机器，其中三台机器训练bert，另外三台训练gpt，每台机器是如何知道其他节点跟自己是否训练的是同一个任务呢，更准确的是说不同进程之间是如何保证自己是同一个进程组的？同一个机器可以使用同一进程组进行标识，但是训练任务通常会分散在不同的机器上。一种非常简单的方法就是给同一个训练任务的机器分配唯一的id，id相同的时候大家进行通信，id不同时则不通信。pytorch使用的唯一标识是ip+port，也就是说同一个训练任务大家使用同一个通信地址ip:port 不同任务即使使用了同一个机器，即ip相同，也可以通过port来进行区分，可以看到192.168.1.2这台机器上跑了两个任务（这台机器有两个gpu，或者资源不够，用户就是想要指定两个进程来训练不同的任务），我们依然可以通过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.2	192.168.1.2:5004	python train.py --master_addr 192.168.1.2 --master_port 5004
gpt : slave	192.168.1.6	192.168.1.2:5004	python train.py --master_addr 192.168.1.2 --master_port 5004
gpt : slave	192.168.1.7	192.168.1.2:5004	python train.py --master_addr 192.168.1.2 --master_port 5004

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

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

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

tcp：直接指定tcp的ip和端口，init_method=‘tcp://192.168.1.2:5003’
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在使用多进程的时候，会先从本地拷贝一份环境变量，然后这份环境变量再分发给各个进程

可以试试，先启动一个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 #空

如何获取进程的唯一标识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	gpu1	gpu2	gpu3	gpu4
192.168.1.2	0	1	2	3
192.168.1.3	1*4+0	1*4+1	1*4+2	1*4+3
192.168.1.4	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

如何实现

import os
import argparse
import torch
import random
import numpy as np
import torch.distributed as distdef setup_new_process(local_rank, callee, args):rank = args.node_rank * args.nproc_per_node + local_rankworld_size = args.nnodes * args.nproc_per_noderandom.seed(rank)np.random.seed(rank)torch.manual_seed(rank)torch.cuda.manual_seed_all(rank)dist_backend = 'gloo'init_method = "env://"if args.use_env:os.environ['MASTER_ADDR'] = args.master_addros.environ['MASTER_PORT'] = args.master_portos.environ['RANK'] = str(rank)os.environ['WORLD_SIZE'] = str(world_size)else:init_method = f"tcp://{args.master_addr}:{args.master_port}"if torch.cuda.is_available():torch.cuda.set_device(local_rank)torch.cuda.empty_cache()# 通信后端，nvidia GPU推荐使用NCCLif torch.distributed.is_nccl_available():dist_backend = 'nccl'print(f'start init process: rank = {rank}')dist.init_process_group(backend=dist_backend,init_method=init_method,world_size=world_size,rank=rank)callee(args)import torch.nn as nn
import torch.nn.functional as F
import torchdef train(args):if dist.is_initialized():rank = dist.get_rank()print(f"rank = {rank} | strat train.......")def main(callee):parse = argparse.ArgumentParser()parse.add_argument('--use_env', action="store_true")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=0)parse.add_argument('--node_rank', type=int, default=0)parse.add_argument('--nnodes', type=int, default=1)args = parse.parse_args()if args.nproc_per_node == 0: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个数，启动多进程的时候，我们发现并没有传递local_rank这个参数，这是因为使用torch.multiprocessing.spawn会自动传入这个参数，并且是递增，，从0到nproc_per_node。

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 --nnodes=3
192.168.1.3	python train.py --master_addr 192.168.1.2 --master_port 5003 --node_rank=1 --nnodes=3
192.168.1.4	python train.py --master_addr 192.168.1.2 --master_port 5003 --node_rank=2 --nnodes=3

前面只是启动了分布式训练而已，我们创建一个小模型来试试ddp，直接使用DistributedDataParallel对模型进行进行一层包装即可使用dpp

torch.nn.parallel.DistributedDataParallel(module, device_ids=None, output_device=None, dim=0, broadcast_buffers=True, process_group=None, bucket_cap_mb=25, find_unused_parameters=False, check_reduction=False)

我们写一个小的demo

def train(args):model = nn.Linear(5, 1, bias=False).to(args.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上(同样要求在gpu)，当使用cpu训练的时候需要将device_ids和output_device都设置为None

device_ids
int 列表或 torch.device 对象，用于指定要并行的设备。对于数据并行，即完整模型放置于一个 GPU 上（single-device module）时，需要提供该参数，表示将模型副本拷贝到哪些 GPU 上。
对于模型并行的情况，即一个模型，分散于多个 GPU 上的情况（multi-device module），以及 CPU 模型，该参数比必须为 None，或者为空列表。
output_device
int 或者 torch.device，对于 single-device 的模型，表示结果输出的位置。
对于 multi-device module 和 GPU 模型，该参数必须为 None 或空列表。

def train(args):device = torch.device(f"cuda:{args.local_rank}" if torch.cuda.is_available() else "cpu")model = nn.Linear(5, 1, bias=False).to(device)for name, params in model.named_parameters():print(f'before dpp : rank = {args.rank}, name = {name}, params = {params.tolist()}')if torch.distributed.is_initialized():if torch.cuda.is_available():model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])else: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).to(device))loss = F.mse_loss(py,torch.rand(32,1).to(device))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的模型参数传给各个子节点，作为初始化的参数。这样可以保证每个节点拿到的模型参数都是一样的。训练的过程中由于梯度共享的原因，所以每一次迭代梯度也是相同的。

数据如何分发

接着我们探讨另一个问题，如何把数据分发到各个节点上，一个非常简单的想法就是，读取一个batch的数据，然后将这个batch分成n份，n=world_size，也就是进程的数量，这样每个进程就都有完全不同的数据了。

pytorch实现的也非常简单，定义了一个DistributedSampler来将batch拆分成world_size份。首先每个节点都有需要从远端下载一份数据，然后加载到dataloader中。一开始以为是主节点计算好每个进程要的数据，然后再分发到各个节点计算，当时的一个疑惑是，如果分发数据的话，图像该怎么办呢？图像数据那么大分发起来可就太慢了。后来发现并没有采样数据，而是获取采样下标，然后每个节点根据分发过来的indics自己读取真实的数据

看了代码发现，主节点并没有分发数据，而是每个节点直接采样torch.randperm下标的。这就更加困惑，各个节点调用随机函数采样，那数据不就会存在重合的现象了吗。直到有一天福至心灵，突然意识到如果随机种子相同，那么采样的数据是不会发生变化的。

这样，给不同的节点设置相同的随机数种子，那么就会得到相同的结果，然后在根据rank获取自己的数据就可以了。下面代码写的非常详细了。

class DistributedSampler(Sampler[T_co]):def __iter__(self) -> Iterator[T_co]:# deterministically shuffle based on epoch and seedg = torch.Generator()# 设置随机数种子，因为是根据epoch来的，所以每个节点的随机数种子是相同的g.manual_seed(self.seed + self.epoch)# 采样数据下标，seed一样，每个节点采样出来的下标也都一样indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore[arg-type]# 依据自己的rank，获取部分数据返回indices = indices[self.rank:self.total_size:self.num_replicas]return iter(indices)def __len__(self) -> int:return self.num_samples# 每一轮我们调用一下set_epoch函数，这样随机数种子就会更新了# 每一轮我们采样的结果才会发生变化，要不然随机数种子不变，每个epoch采样的顺序都不会变化def set_epoch(self, epoch: int) -> None:self.epoch = epoch

保存和加载模型

保存模型就很简单了，唯一需要注意的点是，我们需要只需要在某个节点保存模型就可以了，不需要每个节点都保存，因为ddp通过all_reduce保证每个节点的模型参数是一致的。pytorch模型保存很有意思，我们知道神经网络有自己的结构，每一层都有自己的名字和参数，更细粒度的，每一个节点都有自己的名字和参数，所以save的是一个dict，key就是每一层的名字，value就是每一层参数tensor。然后把dict序列化成pt或者pth。一个很明显的问题就是只保存了模型的参数，并没有保存模型的结构，所以我们在加载的时候需要使用之前定义的模型结构，然后通过模型加载dict。torch.save的本质就是调用pickle，把dict序列化成pt或pth，仅此而已，torch.load则是使用pickle把之前保存的pt再转换成dict，二model.load_state_dict则是根据dict的key找到模型对应的名字，对权重赋值，仅此而已。所以我们会发现即使大家都使用pytorch，相同的模型结构，加载的时候还是有可能找不到key的，就是因为有的人给网络参数起了不同的名字，这个时候就需要在load之后，把dict的key给重命名一下，然后再保存就可以了。

# save model dict to pt/pth
torch.save(model.state_dict(), PATH)
# load model
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval（)

还有一个小坑也需要注意一下，我们训练模型的时候通常会指定device，如果device=cuda:1，那么保存的时候也会把这个device的信息保存到模型中，直接使用torch.load默认会加载到cuda:1，如果我们只有cpu或者只有一个cuda:0，那就会报错，所以加载的时候需要把device重映射一下（本质就跟修改dict的key一样），

device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))
mode.to(device)
# cuda:1映射到cuda:0
device = torch.device('cuda:0')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))
mode.to(device)

import os
import torchclass SaveCheckpoint:def __init__(self, output_dir):self.output_dir = output_dirself.model_dir = os.path.join(output_dir, 'checkpoints')if not os.path.exists(self.model_dir):os.makedirs(self.model_dir)def save(self, global_step, model, optimizer=None, scheduler=None, arguments=None):if hasattr(model, 'module'):model = model.module  # extract model from a distributed/data-parallel wrappercheckpoint_dir = os.path.join(self.model_dir, f'{global_step}')if not os.path.exists(checkpoint_dir):os.makedirs(checkpoint_dir)checkpoint_path = os.path.join(checkpoint_dir, f"pytorch_model.bin")checkpoint = {'global_step': global_step, 'model_state_dict': model.state_dict()}if optimizer:checkpoint['optimizer_state_dict'] = optimizer.state_dict()if scheduler:checkpoint['scheduler_state_dict'] = scheduler.state_dict()if arguments:checkpoint['arguments'] = argumentstorch.save(checkpoint, checkpoint_path)return checkpoint_pathdef train(args):sb = SaveCheckpoint(args.output_dir)device = torch.device(f"cuda:{args.local_rank}" if torch.cuda.is_available() else "cpu")model = nn.Linear(5, 1, bias=False).to(device)if torch.distributed.is_initialized():if torch.cuda.is_available():model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])else:model = torch.nn.parallel.DistributedDataParallel(model)optimizer = torch.optim.Adam(model.parameters(),lr=0.01)step = 0for epoch in range(args.epochs):for data in data_loader():step += 1optimizer.zero_grad()py = model(torch.rand(32,5).to(device))loss = F.mse_loss(py,torch.rand(32,1).to(device))loss.backward() # 计算梯度optimizer.step() # 更新权重if args.rank == 0 and step % args.save_step == 0:sb.save(step, model)