PyTorch中DistributedDataParallel 使用笔记

1. 为何使用DistributedDataParallel在PyTorch中的多GPU训练一般有2种DataParallel和DistributedDataParallelDataParallel是最简单的的单机多卡实现但是它使用多线程模型并不能够在多机多卡的环境下使用所以本文将介绍DistributedDataParallelDDP 基于使用多进程而不是使用多线程的DP并且存在GIL争用问题并且可以扩充到多机多卡的环境所以它是分布式多GPU训练的首选。2. 基本概念在使用DistributedDataParallel时有一些概率必须掌握多机多卡含义world_size代表有几台机器可以理解为几台服务器rank第几台机器即第几个服务器local_rank某台机器中的第几块GPU单机多卡含义world_size代表机器一共有几块GPUrank第几块GPUlocal_rank第几块GPU与rank相同3. 使用方法3.1. 修改主函数在运行的时候DistributedDataParallel会往你的程序中加入一个参数local_rank所以要现在你的代码中解析这个参数如parser.add_argument(--local_rank, typeint, default1, helpnumber of cpu threads to use during batch generation)3.2. 初始化torch.distributed.init_process_group(backendnccl) os.environ[CUDA_VISIBLE_DEVICES] 0, 1, 2 # 有几块GPU写多少3.3. 设定devicelocal_rank torch.distributed.get_rank() torch.cuda.set_device(local_rank) global device device torch.device(cuda, local_rank)我没用arg.local_rank新定义了一个local_rank变量是因为我更信任distributed.get_rank()这个函数这里用torch.device来写并且加了global是因为后面模型和数据都要用到这个device不会出错3.4. 模型加载到多gpumodel.to(device) # 这句不能少最好不要用model.cuda() model torch.nn.parallel.DistributedDataParallel(model, device_ids[local_rank], output_devicelocal_rank, find_unused_parametersTrue) # 这句加载到多GPU上3.5. 数据加载到gpu数据.to(device)3.6. 启动torchrun --nproc_per_node4 --rdzv_endpointlocalhost:12345 train_cylinder_asym.py参考文献Pytorch并行计算(二): DistributedDataParallel介绍_dist.barrier_harry_tea的博客-CSDN博客DistributedDataParallel多GPU分布式训练全过程总结 跟着做90%成功_BRiAq的博客-CSDN博客百度安全验证