总结
这边博客并不是教会如何使用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'] = 12
5. 如何实现
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)