ddp训练流程-pytorch教程
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地址 | 数字标识 | 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获取这个唯一标识的方式也有两种（其实是三种，文件系统共享我没用过）
-
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代码设置环境变量
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的模型参数传给各个子节点，作为初始化的参数。这样可以保证每个节点拿到的模型参数都是一样的。训练的过程中由于梯度共享的原因，所以每一次迭代梯度也是相同的。