总结
这边博客并不是教会如何使用dpp，只是记录一下自己使用dpp时遇到的一些困惑
- 分布式节点如何知道彼此并通信。每个节点启动训练任务需要设置主节点的ip:port，通过跟主节点通信来启动服务。
- 是不是每个节点都要执行python程序。是的，写好训练代码后，每个节点都需要执行python命令，但是一般平台会帮我们管理分布式服务，自动把脚本分发到各个节点，并执行python成勋，才会产生不需要每个节点执行python代码的错觉
- 如果忘记在某个节点执行python命令，结果会怎样。ddp会阻塞在初始化的位置，一直等到有world_size进程为止，这个就是靠主节点来跟踪进程数（我猜的）
- 每个节点是不是都需要有完备的数据。是的，如果使用ddp的DistributedSampler来获取数据，那么一定要保证每个节点都有完整的数据，要不然根据indics无法获取正确的数据。
- 每个节点获取的mini-batch数据如何保证不重合。不同节点设置相同的随机数种子，这样采样的batch就是一摸一样的，然后再根据各自的rank来从batch中取出自己的mini-batch。
1. 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则表示进程的标识
2. 我们怎么知道要使用哪几台机器进行训练的？
假如在同一个局域网内有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将会作为服务的主节点
3. 如何根据标识进行初始化（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 #空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'] = 125. 如何实现
import os
import argparse
import torch
import random
import numpy as np
import torch.distributed as dist
def setup_new_process(local_rank, callee, args):
rank = args.node_rank * args.nproc_per_node + local_rank
world_size = args.nnodes * args.nproc_per_node
random.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_addr
os.environ['MASTER_PORT'] = args.master_port
os.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推荐使用NCCL
if 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 torch
def 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 seed
g = 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 torch
class SaveCheckpoint:
def __init__(self, output_dir):
self.output_dir = output_dir
self.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 wrapper
checkpoint_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'] = arguments
torch.save(checkpoint, checkpoint_path)
return checkpoint_path
def 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 = 0
for epoch in range(args.epochs):
for data in data_loader():
step += 1
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() # 更新权重
if args.rank == 0 and step % args.save_step == 0:
sb.save(step, model)
