学习使用DDP: DistributedDataParallel

简介

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

基本原理

1. 启动多个进程，每个进程在一张卡上加载一个模型（参数在数值上相同）。
2. 训练过程中，通过Ring-Reduce的方法通讯，交换梯度，每个进程获取所有进程梯度。
3. 用平均后的体度更新参数，保证更新后的参数也都相同。

• 这篇大神博文写的特别好，有兴趣深入研究的同学可以参考。知乎

名词定义

• 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'])
	world_size = int(os.environ['WOERLD_SIZE'])
	#似乎也可以：
	##################################################################
	parser.add_argument('--world_size', default=-1, type=int,
                    help='number of nodes for distributed training')
	parser.add_argument('--local_rank', default=-1, type=int,
                    help='node rank for distributed training')
    ##################################################################
	# 初始化分布式环境
	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.cuda.empty_cache()

• 这里我的理解是：在命令行中通过torch.distributed.launch启动了多个进程，每个进程都运行main.py, 同时通过分配local_rank来识别每个进程。

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

	model = MyModel()
	model.to(device)
	model = torch.nn.parallel.DistributedDataParallel(model,device_ids=[local_rank])

• 创建你的神经网络模型，并将其放在 DistributedDataParallel 中。DistributedDataParallel 将模型的参数分发到每个 GPU 上，并在每个 GPU 上创建一个模型副本，把parameter，buffer从master节点传到其他节点，使所有进程上的状态一致。
  注释：DDP通过这一步保证所有进程的初始状态一致。所以，请确保在这一步之后，你的代码不会再修改模型的任何东西了，包括添加、修改、删除parameter和buffer！

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)

报错及解决方案

• 注意，一定要在构建DDP模型前加载预训练参数，否则每个进程模型上的参数是不一样的！！！
  
  这张图中的不同进程的loss差距就极大。
• 另一方面，parameters要在DDP后传入optimizer.
• 如果是创建了多个文件夹，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说模型某个参数在优化前被改变，不影响训练，但有可能影响效果。这个报错只在训练一开始出现，我暂时直接忽略了。这个问题的改进方法找到了：在模型中的permute()加上.contiguous()。
• 还遇到了一些关于数据集Dataloader的问题，可以参考：DataLoader的教程
• 如果0节点总是比其他节点占用的内存大，可以指定加载权重时使用CPU:

  checkpoint = torch.load("checkpoint.pth", map_location='cpu')
  model.load_state_dict(checkpoint["state_dict"])
  

• 我还发现即便在

  train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
  

  设置shuffle = True（默认就是True）每次训练的数据依旧没有打乱，查阅网上帖子发现是因为Sampler中的随机seed与self.epoch挂钩，每个epoch运行：train_loader.sampler.set_epoch(epoch) 但是seed与epoch标号一起从0到num_epochs,每次训练不同epoch是打乱了，但是整体训练一次都是相同的。