1. 基本概念

  • rank
    进程号,在多进程上下文中,我们通常假定rank 0是第一个进程或者主进程,其它进程分别具有1,2,3不同rank号,这样总共具有4个进程

  • node
    物理节点,可以是一个容器也可以是一台机器,节点内部可以有多个GPU;nnodes指物理节点数量, nproc_per_node指每个物理节点上面进程的数量

  • local_rank
    指在一个node上进程的相对序号,local_rank在node之间相互独立

  • WORLD_SIZE
    全局进程总个数,即在一个分布式任务中rank的数量

  • Group
    进程组,一个分布式任务对应了一个进程组。只有用户需要创立多个进程组时才会用到group来管理,默认情况下只有一个group

  • 如下图所示,共有3个节点(机器),每个节点上有4个GPU,每台机器上起4个进程,每个进程占一块GPU,那么图中一共有12个rank,nproc_per_node=4,nnodes=3,每个节点都一个对应的node_rank。

    注意
    rank与GPU之间没有必然的对应关系,一个rank可以包含多个GPU;一个GPU也可以为多个rank服务(多进程共享GPU),在torch的分布式训练中习惯默认一个rank对应着一个GPU,因此local_rank可以当作GPU号

  • backend
    通信后端,可选的包括:nccl(NVIDIA推出)、gloo(Facebook推出)、mpi(OpenMPI)。一般建议GPU训练选择nccl,CPU训练选择gloo

  • master_addr与master_port
    主节点的地址以及端口,供init_method 的tcp方式使用。 因为pytorch中网络通信建立是从机去连接主机,运行ddp只需要指定主节点的IP与端口,其它节点的IP不需要填写。 这个两个参数可以通过环境变量或者init_method传入

  • # 方式1:
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("nccl", 
                            rank=rank, 
                            world_size=world_size)
    # 方式2:
    dist.init_process_group("nccl", 
                            init_method="tcp://localhost:12355",
                            rank=rank, 
                            world_size=world_size)
    

    2. 使用分布式训练模型

    使用DDP分布式训练,一共就如下个步骤:

    1. 初始化进程组 dist.init_process_group
    2. 设置分布式采样器 DistributedSampler
    3. 使用DistributedDataParallel封装模型
    4. 使用torchrun 或者 mp.spawn 启动分布式训练

    2.1 初始化进程组

    进程组初始化如下:

    torch.distributed.init_process_group(backend, 
    									init_method=None, 
    									world_size=-1, 
    									rank=-1, 
    									store=None,
    									...)
    
  • backend: 指定分布式的后端,torch提供了NCCL, GLOO,MPI三种可用的后端,通常CPU的分布式训练选择GLOO, GPU的分布式训练就用NCCL即可
  • init_method:初始化方法,可以是TCP连接、File共享文件系统、ENV环境变量三种方式
  • init_method='tcp://ip:port': 通过指定rank 0(即:MASTER进程)的IP和端口,各个进程进行信息交换。 需指定 rank 和 world_size 这两个参数
  • init_method='file://path':通过所有进程都可以访问共享文件系统来进行信息共享。需要指定rank和world_size参数
  • init_method=env://:从环境变量中读取分布式的信息(os.environ),主要包括 MASTER_ADDR, MASTER_PORT, RANK, WORLD_SIZE。 其中,rank和world_size可以选择手动指定,否则从环境变量读取
  • tcp和env两种方式比较类似(其实env就是对tcp的一层封装),都是通过网络地址的方式进行通信,也是最常用的初始化方法

    case 1

    import os, argparse
    import torch
    import torch.distributed as dist
    
    parse = argparse.ArgumentParser()
    parse.add_argument('--init_method', type=str)
    parse.add_argument('--rank', type=int)
    parse.add_argument('--ws', type=int)
    args = parse.parse_args()
    
    if args.init_method == 'TCP':
    	dist.init_process_group('nccl', init_method='tcp://127.0.0.1:28765', rank=args.rank, world_size=args.ws)
    elif args.init_method == 'ENV':
        dist.init_process_group('nccl', init_method='env://')
    
    rank = dist.get_rank()
    print(f"rank = {rank} is initialized")
    # 单机多卡情况下,localrank = rank. 严谨应该是local_rank来设置device
    torch.cuda.set_device(rank)
    tensor = torch.tensor([1, 2, 3, 4]).cuda()
    print(tensor)
    

    假设单机双卡的机器上运行,则开两个终端,同时运行下面的命令

    # TCP方法
    python3 test_ddp.py --init_method=TCP --rank=0 --ws=2
    python3 test_ddp.py --init_method=TCP --rank=1 --ws=2
    # ENV方法
    MASTER_ADDR='localhost' MASTER_PORT=28765 RANK=0 WORLD_SIZE=2 python3 test_gpu.py --init_method=ENV
    MASTER_ADDR='localhost' MASTER_PORT=28765 RANK=1 WORLD_SIZE=2 python3 test_gpu.py --init_method=ENV
    

    如果开启的进程未达到 word_size 的数量,则所有进程会一直等待,直到都开始运行,可以得到输出如下:

    # rank0 的终端:
    rank 0 is initialized
    tensor([1, 2, 3, 4], device='cuda:0')
    # rank1的终端
    rank 1 is initialized
    tensor([1, 2, 3, 4], device='cuda:1')
    

    在初始化DDP的时候,能够给后端提供主进程的地址端口、本身的RANK,以及进程数量即可。初始化完成后,就可以执行很多分布式的函数了,比如dist.get_rank, dist.all_gather等等

    2.2 分布式训练数据加载

    DistributedSampler把所有数据分成N份(N为worldsize), 并能正确的分发到不同的进程中,每个进程可以拿到一个数据的子集,不重叠,不交叉

    torch.utils.data.distributed.DistributedSampler(
    				dataset,
    				num_replicas=None, 
    				rank=None, 
    				shuffle=True, 
    				seed=0, 
    				drop_last=False)
    
  • dataset: 需要加载的完整数据集
  • num_replicas: 把数据集分成多少份,默认是当前dist的world_size
  • rank: 当前进程的id,默认dist的rank
  • shuffle:是否打乱
  • drop_last: 如果数据长度不能被world_size整除,可以考虑是否将剩下的扔掉
  • seed:随机数种子。这里需要注意,从源码中可以看出,真正的种子其实是 self.seed+self.epoch 这样的好处是,不同的epoch每个进程拿到的数据是不一样,因此需要在每个epoch开始前设置下:sampler.set_epoch(epoch)
  • case 2

    sampler = DistributedSampler(dataset)
    loader = DataLoader(dataset, sampler=sampler)
    for epoch in range(start_epoch, n_epochs):
      sampler.set_epoch(epoch) # 设置epoch 更新种子
      train(loader)
    

    2.3 模型分布式封装

    将单机模型使用torch.nn.parallel.DistributedDataParallel 进行封装

    torch.cuda.set_device(local_rank)
    model = Model().cuda()
    model = DistributedDataParallel(model, device_ids=[local_rank])
    

    注意: 要调用model内的函数或者属性,使用model.module.xxxx

    这样在多卡训练时,每个进程有一个model副本和optimizer,使用自己的数据进行训练,之后反向传播计算完梯度的时候,所有进程的梯度会进行all-reduce操作进行同步,进而保证每个卡上的模型更新梯度是一样的,模型参数也是一致的。

    在save和load模型时候,为了减小所有进程同时读写磁盘,一般处理方法是以主进程为主

    case 3

    model = DistributedDataParallel(model, device_ids=[local_rank])
    CHECKPOINT_PATH ="./model.checkpoint"
    if rank == 0:
      torch.save(ddp_model.state_dict(), CHECKPOINT_PATH)
    # barrier()其他保证rank 0保存完成
    dist.barrier()
    map_location = {"cuda:0": f"cuda:{local_rank}"}
    model.load_state_dict(torch.load(CHECKPOINT_PATH, map_location=map_location))
    # 后面正常训练代码
    optimizer = xxx
    for epoch:
      for data in Dataloader:
          model(data)
          xxx
        # 训练完成 只需要保存rank 0上的即可
        # 不需要dist.barrior(), all_reduce 操作保证了同步性
      if rank == 0:
         torch.save(ddp_model.state_dict(), CHECKPOINT_PATH)
    

    2.4 启动分布式训练

    如case1所示我们手动运行多个程序,相对繁琐。实际上本身DDP就是一个python 的多进程,因此完全可以直接通过多进程的方式来启动分布式程序。 torch提供了以下两种启动工具来更加方便的运行torch的DDP程序。

    2.4.1 mp.spawn

    使用torch.multiprocessing(python的multiprocessing的封装类) 来自动生成多个进程

    mp.spawn(fn, args=(), nprocs=1, join=True, daemon=False)
    
  • fn: 进程的入口函数,该函数的第一个参数会被默认自动加入当前进*程的rank, 即实际调用: fn(rank, *args)
  • nprocs: 进程数量,即:world_size
  • args: 函数fn的其他常规参数以tuple的形式传递
  • case 4

    import torch
    import torch.distributed as dist
    import torch.multiprocessing as mp
    
    def fn(rank, ws, nums):
        dist.init_process_group('nccl', init_method='tcp://127.0.0.1:28765',
                                rank=rank, world_size=ws)
        rank = dist.get_rank()
        print(f"rank = {rank} is initialized")
        torch.cuda.set_device(rank)
        tensor = torch.tensor(nums).cuda()
        print(tensor)
    
    if __name__ == "__main__":
        ws = 2
        mp.spawn(fn, nprocs=ws, args=(ws, [1, 2, 3, 4]))
    

    直接执行一次命令 python3 test_ddp.py 即可,输出如下:

    rank = 0 is initialized
    rank = 1 is initialized
    tensor([1, 2, 3, 4], device='cuda:1')
    tensor([1, 2, 3, 4], device='cuda:0')
    

    这种方式同时适用于TCP和ENV初始化

    2.4.2 launch/run

    使用torch提供的 torch.distributed.launch工具,可以以模块的形式直接执行

    python3 -m torch.distributed.launch --配置 train.py --args参数
    

    常用配置有:

  • --nnodes: 使用的机器数量,单机的话,就默认是1了
  • --nproc_per_node: 单机的进程数,即单机的worldsize
  • --master_addr/port: 使用的主进程rank0的地址和端口
  • --node_rank: 当前的进程rank
  • 在单机情况下, 只有--nproc_per_node 是必须指定的,--master_addr/portnode_rank都是可以由launch通过环境自动配置

    case5 test_dist.py

    import torch
    import torch.distributed as dist
    import torch.multiprocessing as mp
    import os
    
    dist.init_process_group('nccl', init_method='env://')
    
    rank = dist.get_rank()
    local_rank = os.environ['LOCAL_RANK']
    master_addr = os.environ['MASTER_ADDR']
    master_port = os.environ['MASTER_PORT']
    print(f"rank = {rank} is initialized in {master_addr}:{master_port}; local_rank = {local_rank}")
    torch.cuda.set_device(rank)
    tensor = torch.tensor([1, 2, 3, 4]).cuda()
    print(tensor)
    

    输入如下命令

    python3 -m torch.distribued.launch --nproc_per_node=2 test_dist.py
    

    得到如下输出

    rank = 0 is initialized in 127.0.0.1:29500; local_rank = 0
    rank = 1 is initialized in 127.0.0.1:29500; local_rank = 1
    tensor([1, 2, 3, 4], device='cuda:1')
    tensor([1, 2, 3, 4], device='cuda:0')
    

    注意:torch1.10开始用终端命令torchrun来代替torch.distributed.launch,具体来说,torchrun实现了launch的一个超集,不同的地方在于:

  • 完全使用环境变量配置各类参数,如RANK,LOCAL_RANK, WORLD_SIZE等,尤其是local_rank不再支持用命令行隐式传递的方式
  • 能够更加优雅的处理某个worker失败的情况,重启worker。需要代码中有load_checkpoint(path)和save_checkpoint(path) 这样有worker失败的话,可以通过load最新的模型,重启所有的worker接着训练。具体参考 imagenet-torchrun
  • 训练的节点数目可以弹性变化
  • 上面的命令可以写成如下

    torchrun --nproc_per_node=2 test_dist.py
    

    torchrun或者launch对上面ENV的初始化方法支持最完善,TCP初始化方法的可能会出现问题,因此尽量使用env来初始化dist

    3. 分布式做evaluation

    分布式做evaluation的时候,一般需要先所有进程的输出结果进行gather,再进行指标的计算,两个常用的函数:

  • dist.all_gather(tensor_list, tensor) : 将所有进程的tensor进行收集并拼接成新的tensorlist返回
  • dist.all_reduce(tensor, op) 这是对tensor的in-place的操作, 对所有进程的某个tensor进行合并操作,op可以是求和等
  • case 6 test_ddp.py

    import torch
    import torch.distributed as dist
    
    dist.init_process_group('nccl', init_method='env://')
    rank = dist.get_rank()
    torch.cuda.set_device(rank)
    
    tensor = torch.arange(2) + 1 + 2 * rank
    tensor = tensor.cuda()
    print(f"rank {rank}: {tensor}")
    
    tensor_list = [torch.zeros_like(tensor).cuda() for _ in range(2)]
    dist.all_gather(tensor_list, tensor)
    print(f"after gather, rank {rank}: tensor_list: {tensor_list}")
    
    dist.barrier()
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    print(f"after reduce, rank {rank}: tensor: {tensor}")
    

    通过torchrun --nproc_per_node=2 test_ddp.py 输出结果如下:

    rank 1: tensor([3, 4], device='cuda:1')
    rank 0: tensor([1, 2], device='cuda:0')
    after gather, rank 1: tensor_list: [tensor([1, 2], device='cuda:1'), tensor([3, 4], device='cuda:1')]
    after gather, rank 0: tensor_list: [tensor([1, 2], device='cuda:0'), tensor([3, 4], device='cuda:0')]
    after reduce, rank 0: tensor: tensor([4, 6], device='cuda:0')
    after reduce, rank 1: tensor: tensor([4, 6], device='cuda:1')
    

    在evaluation的时候,可以拿到所有进程中模型的输出,最后统一计算指标,基本流程如下

    pred_list = []
    for data in Dataloader:
        pred = model(data)
        batch_pred = [torch.zeros_like(label) for _ in range(world_size)]
        dist.all_gather(batch_pred, pred)
        pred_list.extend(batch_pred)
    pred_list = torch.cat(pred_list, 1)
    # 所有进程pred_list是一致的,保存所有数据模型预测的值
    

    4. 常用函数

  • torch.distributed.get_rank(group=None) 获取当前进程的rank
  • torch.distributed.get_backend(group=None) 获取当前任务(或者指定group)的后端
  • data_loader_train = torch.utils.data.DataLoader(dataset=data_set, batch_size=32,num_workers=16,pin_memory=True)
  • num_workers: 加载数据的进程数量,默认只有1个,增加该数量能够提升数据的读入速度。(注意:该参数>1,在低版本的pytorch可能会触发python的内存溢出)
  • pin_memory: 锁页内存,加快数据在内存上的传递速度。 若数据加载成为训练速度的瓶颈,可以考虑将这两个参数加上
  • 进程内指定显卡,很多场景下使用分布式都是默认一张卡对应一个进程,所以通常,我们会设置进程能够看到卡数
    # 方式1:在进程内部设置可见的device
    torch.cuda.set_device(args.local_rank)
    # 方式2:通过ddp里面的device_ids指定
    ddp_model = DDP(model, device_ids=[rank]) 
    # 方式3:通过在进程内修改环境变量
    os.environ['CUDA_VISIBLE_DEVICES'] = loac_rank
    
  • 参考

    分布式情感分析-github
    Pytorch DDP分布式训练介绍
    PyTorch分布式DPP涉及的基本概念与问题

    来源:orangerfun

    物联沃分享整理
    物联沃-IOTWORD物联网 » torch分布式训练

    发表评论