学习使用DDP: DistributedDataParallel

简介

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

名词定义

• WORLD_SIZE:总进程数量，通常每个进程都会被分配一个唯一的 rank，而 world_size 就是所有进程的总数量。
• MASTER_ADDR 和 MASTER_PORT：这两个环境变量用于指定主进程的地址和端口号。在分布式训练中，主进程通常用于协调其他进程的工作。
• LOCAL_RANK 和 LOCAL_SIZE: 这两个环境变量是 torch.distributed.launch 工具特有的，用于指定当前进程在本地节点上的 local_rank 和节点上 GPU 的总数量 local_size。
• RANK 和 WORLD_SIZE（对于 torch.multiprocessing）: 在使用 torch.multiprocessing 进行分布式训练时，RANK 和 WORLD_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 上创建一个模型副本。

3. 加载数据：

• 确保每个进程都获得不同的数据切片，这里的作用是保证每个进程都平分不同的数据，而不是加载相同的数据。假设有一个1万数据量的dataset，开N个线程，则每个线程都会在一个epoch中加载1万数据，这些数据一样而造成冗余，进而使用DistributedSampler就可以让16个线程平分这1万数据,共同完成一个epoch的训练。

	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)

4. 训练

	for epoch in range(num_epochs):
	    # 新增2：设置sampler的epoch，DistributedSampler需要这个来维持各个进程之间的相同随机数种子
	    train_loader.sampler.set_epoch(epoch)
	    for data, label in trainloader:
			training...

5. 销毁
   训练完成后需清理分布式环境，释放资源。

	torch.distributed.destroy_process_group()

启动

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

python -m torch.distributed.launch --nproc_per_node=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.state_dict()
• 如果出现报错有些参数在反向传播时未使用，则需要加上find_unused_parameters = True.

    model = torch.nn.parallel.DistributedDataParallel(model, find_unused_parameters = True)

• 从不并行改到并行代码时也要注意，多个进程会平分数据集，单个进程在计算剩余时间时也会出错。$剩余时间=epochnum\ast datanum/batchsize\ast steptime$,这里的batch_size应该是多个线程总的batch_size, 而非单线程的batch_size.
• 还有warning说模型某个参数在优化前被改变，不影响训练，但有可能影响效果。这个报错只在训练一开始出现，我暂时直接忽略了。
• 还遇到了一些关于数据集Dataloader的问题，可以参考：DataLoader的教程