torchrun分布式训练的配置踩坑

土豆丝yu

已于 2023-09-12 21:04:55 修改

阅读量507

点赞数

文章标签： pytorch

于 2023-08-03 21:46:10 首次发布

本文链接：https://blog.csdn.net/zaixiadouing/article/details/132093064

版权

分布式训练初始化：

需要区分是否正确使用rank和local_rank;
local_rank为某一个机器上进程的相对序号，在单机多卡训练中一般为0，多机器训练时，才会有其他索引；
rank为全局的进程序号

参考定义：

在这里插入图片描述
初始化单机多卡时，每个进程对应唯一的rank，一般设置：

import torch.distributed as dist

rank = int(os.environ['RANK'])
num_gpus = torch.cuda.device_count()
torch.cuda.set_device(rank % num_gpus) 
dist.init_process_group(backend='nccl', **kwards)

若没有正确初始化，可能会卡在dist.init_process_group处，形成所谓的死锁

数据加载的分布式并行

模型的DDP封装

dist.barrier()同步所有进程

终端启动方式推荐

单机多卡：

# 配置可能需要的环境
export OMP_NUM_THREADS=1
export MASTER_PORT=$((12000 + $random % 20000))

CUDA_VISIBLE_DEVICES=0,1,2,3 \ 
torchrun --proc_per_node=4 --master_port=MASTER_PORT\
train.py

分布式训练的最终目的是优化：

数据很多时少用append，预先分配好整个容器大小，每次用索引去修改内容；
DP， DDP ，MP分布式训练，单机多卡，多机多卡训练；

来自GPT的汇总：

使用torchrun进行多卡训练模型时，你需要在代码中进行以下配置：

初始化进程组：使用torch.distributed.init_process_group函数来初始化GPU通信方式（例如，‘nccl’）和参数的获取方式（例如，‘env://’）。同时，你需要设置当前进程使用的GPU¹。

parse.add_argument('--local_rank',type=int)
args=parser.parse_args()
local_rank=args.local_rank
torch.cuda.set_device(local_rank)
torch.distributed.init_process_group('nccl', init_method='env://')
device = torch.device(f'cuda:{args.local_rank}')

处理Dataloader：使用torch.utils.data.distributed.DistributedSampler来打乱数据，因为一个batch被分配到了好几个进程中，要确保不同的GPU拿到的不是同一份数据¹。

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, shuffle=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

初始化模型：使用torch.nn.parallel.DistributedDataParallel来包装模型。这个函数能帮助我们为不同GPU上求得的梯度进行allreduce（即汇总不同GPU计算所得的梯度，并同步计算结果）¹。

model=torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

设备转移：将inputs、labels转移到相应的设备上¹。

参考文献

(1) PyTorch 分布式训练实现(DP/DDP/torchrun/多机多卡) - 知乎. https://zhuanlan.zhihu.com/p/489011749.
(2) torchrun (Elastic Launch) — PyTorch 2.0 documentation. https://pytorch.org/docs/stable/elastic/run.html.
(3) Pytorch - 多机多卡极简实现（附源码） - 知乎 - 知乎专栏. https://zhuanlan.zhihu.com/p/486130584.