学习使用DDP: DistributedDataParallel

发布时间:2023年12月18日

Introduction

  • “DistributedDataParalled” 是Pytorch中用于分布式训练的模块,相较与比较老的DataParallel更高效,易用(我在使用DataParallel时经常遇到参数和数据没有在一块卡的报错情况,非常烦人)。
  • 它允许在多个GPU甚至多个节点上并行,在每个GPU上分发参数并建立模型副本。每个进程都会加载不同的数据,DDP会自动处理跨GPU的梯度同步,每个GPU上的模型都会在自己的数据上进行前向反向传播,然后通过梯度同步机制更新模型参数。

名次定义

  • WORLD_SIZE:总进程数量,通常每个进程都会被分配一个唯一的 rank,而 world_size 就是所有进程的总数量。
  • MASTER_ADDRMASTER_PORT:这两个环境变量用于指定主进程的地址和端口号。在分布式训练中,主进程通常用于协调其他进程的工作。
  • LOCAL_RANKLOCAL_SIZE: 这两个环境变量是 torch.distributed.launch 工具特有的,用于指定当前进程在本地节点上的 local_rank 和节点上 GPU 的总数量 local_size。
  • RANKWORLD_SIZE(对于 torch.multiprocessing): 在使用 torch.multiprocessing 进行分布式训练时,RANKWORLD_SIZE 是用于指定进程全局排名和总进程数量的环境变量。

具体步骤

  1. 初始化分布式环境
	import torch.distributed as dist
	
	# 在 torch.distributed.launch 中设置 local_rank 和 rank
	local_rank = int(os.environ['LOCAL_RANK']) #这里的local_rank相当于告诉程序这是第几个进程
	rank = int(os.environ['RANK'])
	
	# 初始化分布式环境
	dist.init_process_group(backend='nccl', init_method='tcp://localhost:23456', world_size=world_size, rank=rank)
	
	# 使用 local_rank 指定当前进程在本地节点上的 GPU 设备
	torch.cuda.set_device(local_rank)
	device = torch.device("cuda", local_rank)
  • 这里我的理解是:在命令行中通过torch.distributed.launch启动了多个进程,每个进程都运行main.py, 同时通过分配local_rank来识别每个进程。

2.通过torch.nn.parallel.DistributedDataParallel定义模型。

	model = MyModel()
	model.to(device)
	model = torch.nn.parallel.DistributedDataParallel(model)
  • 创建你的神经网络模型,并将其放在 DistributedDataParallel 中。DistributedDataParallel 将模型的参数分发到每个 GPU 上,并在每个 GPU 上创建一个模型副本。
  1. 加载数据:
  • 确保每个进程都获得不同的数据切片。
	train_dataset = MyDataset()
	train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
	train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler)
  1. 训练
	for epoch in range(num_epochs):
	    # 新增2:设置sampler的epoch,DistributedSampler需要这个来维持各个进程之间的相同随机数种子
	    trainloader.sampler.set_epoch(epoch)
	    for data, label in trainloader:
			training...
  1. 销毁
    训练完成后需清理分布式环境,释放资源。
	torch.distributed.destroy_process_group()

启动

如果只是在一台机子上运行,

python -m torch.distributed.launch --nproc_per_node=n --batch-size=n main.py
  • –nproc_per_node specifies how many GPUs you would like to use. In the example above, it is 8.
  • –batch-size is now the Total batch-size. It will be divided evenly to each GPU. (这里的batch-size也可以不设置,在代码中设置每个进程的batch_size)

报错及解决方案

  • 如果是创建了多个文件夹,log文件等,则需给相关语句加上if 条件,只让local_rank == 0 时输出或者记录。
  • 保存模型时要保存 model.module
  • 如果出现报错有些参数在反向传播时未使用,则需要加上find_unused_parameters = True.
	 model = torch.nn.parallel.DistributedDataParallel(model, find_unused_parameters = True)
文章来源:https://blog.csdn.net/blueag1e/article/details/135047342
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。