pytorch DistributedDataParallel 分布式训练踩坑记录

目录

    • 一、几个比较常见的概念:
    • 二、踩坑记录
      • [2.1 dist.init_process_group初始化](#2.1 dist.init_process_group初始化)
      • [2.2 spawn启动(rank怎么来的)](#2.2 spawn启动(rank怎么来的))
      • [2.3 loss backward](#2.3 loss backward)
      • [2.4 model cuda设置](#2.4 model cuda设置)
      • [2.5 数据加载](#2.5 数据加载)

一、几个比较常见的概念:

  • rank: 多机多卡时代表某一台机器,单机多卡时代表某一块GPU

  • world_size : 多机多卡时代表有几台机器,单机多卡时代表有几块GPU

    python 复制代码
    world_size = torch.cuda.device_count()
  • local_rank : 多机多卡时代表某一块GPU, 单机多卡时代表某一块GPU
    单机多卡的情况要比多机多卡的情况常见的多。

  • DP:适用于单机多卡(=多进程)训练。算是旧版本的DDP

  • DDP:适用于单机多卡训练、多机多卡。

二、踩坑记录

2.1 dist.init_process_group初始化

这一步就是设定一个组,这个组里面设定你有几个进程(world_size),现在是卡几(rank)。让pycharm知道你要跑几个进程,包装在组内,进行通讯这样模型参数会自己同步,不需要额外操作了。

python 复制代码
import os
import torch.distributed as dist
 
def ddp_setup(rank,world_size):
    os.environ['MASTER_ADDR'] = 'localhost' #rank0 对应的地址
    os.environ['MASTER_PORT'] = '6666' #任何空闲的端口
    dist.init_process_group(
            backend='nccl',  #nccl Gloo #nvidia显卡的选择nccl
            world_size=world_size, 
            init_method='env://',
            rank=rank
        ) #初始化默认的分布进程组
 
    dist.barrier() #等到每块GPU运行到这再继续往下走

2.2 spawn启动(rank怎么来的)

rank是自动分配的。怎么分配呢?这里用的是spawn也就一行代码。

python 复制代码
import torch.multiprocessing as mp
def main (rank:int,world_size:int,args):
    pass
    #训练代码 主函数
 
mp.spawn(main,args=(args.world_size,args), nprocs=args.world_size)

注意,调用spawn的时候,没有输入main的其中一个参数rank,rank由代码自动分配。将代码复制两份在两张卡上同时跑,你可以print(rank),会发现输出 0 1。两份代码并行跑。

另外,nprocs=args.world_size。如果不这么写,代码会卡死在这,既不报错,也不停止。

2.3 loss backward

shell 复制代码
one of the variables needed for gradient computation has been modified by an inplace operation

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [2048]] is at version 4; expected version 3 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

经过调试发现,当使用nn.DataParallel并行训练或者单卡训练均可正常运行;另外如果将两次模型调用集成到model中,即通过out1, out2 = model(input0, input1) 的方式在分布式训练下也不会报错。

在分布式训练中,如果对同一模型进行多次调用则会触发以上报错,即nn.parallel.DistributedDataParallel方法封装的模型,forword()函数和backward()函数必须交替执行,如果执行多个(次)forward()然后执行一次backward()则会报错。

解决此问题可以聚焦到nn.parallel.DistributedDataParallel接口上,通过查询PyTorch官方文档发现此接口下的两个参数:

  • find_unused_parameters: 如果模型的输出有不需要进行反向传播的,此参数需要设置为True;若你的代码运行后卡住不动,基本上就是该参数的问题。
  • broadcast_buffers: 该参数默认为True,设置为True时,在模型执行forward之前,gpu0会把buffer中的参数值全部覆盖到别的gpu上。
shell 复制代码
model = nn.parallel.DistributedDataParallel(model, 
                                            device_ids=[args.local_rank], 
                                            broadcast_buffers=False, 
                                            find_unused_parameters=True)

2.4 model cuda设置

shell 复制代码
RuntimeError: NCCL error in: ../torch/lib/c10d/ProcessGroupNCCL.cpp:859, invalid usage, NCCL version 21.1.1
ncclInvalidUsage: This usually reflects invalid usage of NCCL library (such as too many async ops, too many collectives at once, mixing streams in a group, etc).

*这是因为model和local_rank所指定device不一致引起的错误。

shell 复制代码
model.cuda(args.local_rank)
model = torch.nn.parallel.DistributedDataParallel(model, 
                                                  device_ids=[args.local_rank],
                                                  broadcast_buffers=False,
                                                  find_unused_parameters=True)

2.5 数据加载

使用distributed加载数据集,需要使用DistributedSampler自动为每个gpu分配数据,但需要注意的是sampler和shuffle=True不能并存。

shell 复制代码
train_sampler = DistributedSampler(trainset)
train_loader = torch.utils.data.DataLoader(trainset,
                                           batch_size=args.train_batch_size,
                                           num_workers=args.train_workers,
                                           sampler=train_sampler)
相关推荐
阿里云大数据AI技术几秒前
【AAAI2026】阿里云人工智能平台PAI视频编辑算法论文入选
人工智能
玄同7653 分钟前
我的 Trae Skill 实践|使用 UV 工具一键搭建 Python 项目开发环境
开发语言·人工智能·python·langchain·uv·trae·vibe coding
Yorlen_Zhang13 分钟前
Python Tkinter Text 控件完全指南:从基础编辑器到富文本应用
开发语言·python·c#
苍何24 分钟前
腾讯重磅开源!混元图像 3.0 图生图真香!
人工智能
千里马也想飞27 分钟前
人工智能在医疗领域的应用与研究论文写作实操:AI辅助快速完成框架+正文创作
人工智能
Rorsion32 分钟前
PyTorch实现二分类(单特征输出+单层神经网络)
人工智能·pytorch·分类
勾股导航40 分钟前
K-means
人工智能·机器学习·kmeans
liliangcsdn40 分钟前
Diff2Flow中扩散和流匹配的对齐探索
人工智能
HAPPY酷41 分钟前
C++ 和 Python 的“容器”对决:从万金油到核武器
开发语言·c++·python
SmartBrain1 小时前
战略洞察:以AI为代表的第四次工业革命
人工智能·语言模型·aigc