【PyTorch distributed】多卡训练原理简析与RuntimeError: Expected to mark a variable ready only once.报错分析

引言

单机多卡使用torch.distributed训练模型时碰上了报错，于是趁机研究了一下这个torch.distributed的使用逻辑。
报错信息如下展示，本文仅记录报错信息，复现报错，并简单讨论其背后的成因，为此也要简要的讨论torch.distributed是如何完成多卡训练的操作。

如果对你有帮助，点赞让我知道！嘿嘿

报错信息

RuntimeError: Expected to mark a variable ready only once. This error is caused by one of the following reasons: 1) Use of a module parameter outside the forward function. Please make sure model parameters are not shared across multiple concurrent forward-backward passes. or try to use _set_static_graph() as a workaround if this module graph does not change during training loop.2) Reused parameters in multiple reentrant backward passes. For example, if you use multiple checkpoint functions to wrap the same part of your model, it would result in the same set of parameters been used by different reentrant backward passes multiple times, and hence marking a variable ready multiple times. DDP does not support such use cases in default. You can try to use _set_static_graph() as a workaround if your module graph does not change over iterations.

背景原理

先扯远一点，聊聊天torchdistributed的一些原理性的东西，最后解决这个报错。

梳理一下torch.distributed的工作流程

使用一个简单的案例，该案例可以重现上述的报错信息：

import os;import sys;import torch
import torch.distributed as dist
import torch.nn as nn;import torch.optim as optim
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
def cleanup():
    dist.destroy_process_group()

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(1, 1,bias=False)
        self.net2 = nn.Linear(1, 1,bias=False)
        with torch.no_grad():
            self.net1.weight.fill_(20.)
            self.net2.weight.fill_(20.)
    def forward(self, x,rank):
        if rank==0:
            return self.net1(x)
        else:
            return self.net2(x)
            
# 使用这个demo_fn会复现之前的报错信息
def demo_r(rank,world_size):
	print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)
    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank],find_unused_parameters=True,)
    optimizer = optim.SGD(ddp_model.parameters(), lr=1.)
    optimizer.zero_grad()
    k = ddp_model(torch.tensor(2.).reshape(1,1).to(rank),rank)
    if rank==0:
        k += ddp_model.module.net2(torch.tensor(32.).reshape(1,1).to(rank))
    else:
        k += ddp_model.module.net1(torch.tensor(41.).reshape(1,1).to(rank))
    k.backward()
    cleanup()

def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             join=True)

if __name__=='__main__':
    os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
    # 使用这个demo_fn会复现之前的报错信息
    run_demo(demo_r,2)

rank, world_size,torch.distributed本质上是工作逻辑，每个进程可以分配到一个rank标记，相当于一个进程id，world_size表示一共有多少个进程。以如上一个小demo为例。

首先主进程调用run_demo，住进程将创建rank=0和rank=1两个进程，俩个进程分别执行setup任务，确定通信的主机。

然后俩个进程分别初始化Toymodel 并将其转移到对应的GPU显存中，然后再用DDP包裹模型，生成DDPModel。这里有个值得注意的问题，DDPModel构建的时候完成了两件事情， 一是保证在不同的进程中，模型的各个副本参数完全一致。其二是每个进程都创造一个自己的Reducer，Reducer会处理各个进程中模型参数的梯度同步问题。

如何保证不同进程Model参数初始相同

一般来说，通过spawn生成的子进程其各个进程状态是一样，那么Toymodel会产生相同的随机情况，从而实现各个rank的model参数一致。但这个是得不到保证的，如果每个进程执行的内容不同（例如，强制修改不同的随机种子）从而导致Toymodel()得到不同的随机参数。那这样分布式训练的不同进程的出发点会不会不同呢？ 设计一个小实验，修改Toymodel的初始化，删掉对weight的指定，同时修改demo_fn为，其中为不同的rank采用了不同的随机种子，观察在toymodel初始化和包裹的过程中其参数的变化。

def demo_observe_init(rank,world_size):
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)
    torch.manual_seed(rank)
    torch.cuda.manual_seed_all(rank)
    print(f"rank:{rank},seed:{torch.seed()}")
    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    print(f'rank {rank}, layer 1 data before ddp,  {model.net1.weight.data}')
    print(f'rank {rank}, layer 2 data before ddp,  {model.net2.weight.data}')
    ddp_model = DDP(model, device_ids=[rank],find_unused_parameters=True,)
    print(f'rank {rank}, layer 1 data after ddp,  {model.net1.weight.data}')
    print(f'rank {rank}, layer 2 data after ddp,  {model.net2.weight.data}')
    print(f'rank {rank}, layer 1 ddpdata,  {ddp_model.module.net1.weight.data}')
    print(f'rank {rank}, layer 2 ddpdata,  {ddp_model.module.net2.weight.data}')

执行之后得到输出

Running basic DDP example on rank 0.
Running basic DDP example on rank 1.
rank:1,seed:2869173717154807797
rank:0,seed:9643613638370635189
rank 1, layer 1 data before ddp, tensor([[-0.2182]], device=‘cuda:1’)
rank 1, layer 2 data before ddp, tensor([[0.7092]], device=‘cuda:1’)
rank 0, layer 1 data before ddp, tensor([[-0.2300]], device=‘cuda:0’)
rank 0, layer 2 data before ddp, tensor([[0.4428]], device=‘cuda:0’)
rank 0, layer 1 data after ddp, tensor([[-0.2300]], device=‘cuda:0’)
rank 0, layer 2 data after ddp, tensor([[0.4428]], device=‘cuda:0’)
rank 0, layer 1 ddpdata, tensor([[-0.2300]], device=‘cuda:0’)
rank 0, layer 2 ddpdata, tensor([[0.4428]], device=‘cuda:0’)
rank 1, layer 1 data after ddp, tensor([[-0.2300]], device=‘cuda:1’)
rank 1, layer 2 data after ddp, tensor([[0.4428]], device=‘cuda:1’)
rank 1, layer 1 ddpdata, tensor([[-0.2300]], device=‘cuda:1’)
rank 1, layer 2 ddpdata, tensor([[0.4428]], device=‘cuda:1’)

观察输出可以直接得到结论：

1. 不同的进程确实产生了不同的ToyModel。
2. 当用DDP包裹原模型时，rank0会广播broadcast自己的模型参数到其他进程的模型上实现同步。 所有其他进程会copy，rank0对应的权重参数。

梯度如何同步

Reducer会把所有的参数梯度分成很多组，每一个组叫做一个bucket，一个bucket里面包含在所有进程中相同的一部分参数的梯度，如下图，有process0,process1两个进程，bucket1里面包含了模型俩个副本中param0和param1的梯度。

Reducer同时还会为每个梯度注册一个hook，起先每个梯度都会标记为unready状态，一旦在反向传播的过程中（ loss.backward() ），某个参数的梯度被计算出来，那么此时这个参数对应的hook被触发，该梯度也就变成了准备状态。当一个bucket中的所有hook都被触发，意味着这个bucket中的变量的梯度都已经计算得到，Reducer便立刻开启横跨所有进程的异步的梯度平均。例如bucket1中的hook都ready了，那么grad0会被重置为1/2*(process0.grad0 + process1.grad0)。 这个过程被称作是allreduce操作，下图展示了各种相关的操作。

当一个bucket中的所有grad都被平均了之后，这个bucket就被reduce了。Reducer处理一个bucket是异步的非blocking的，但是当所有的bucket都ready了，Reducer就会Block，开始阻塞等待，直到每个bucket都异步allreduce都完成。当所有bucket都被reduce了，那么梯度就在所有进程上统一了。

废话有点多，总结一下，在调用backward()的过程中，获得了梯度的参数会被标记为ready状态，当一个bucket的所有参数都ready了，这个bucket就ready了，同时也会被Reducer进行异步的allreduce，一旦所有bucket都ready了，Reducer将会阻塞等待所有allreduce的完成， 也就是说，模型梯度的同步会在 backward() 时完成，这也有一点join的味道了。

梯度同步的问题

看上去梯度同步也不复杂，就上面的做法就能通过Reducer实现同步，但是这里面存在问题，让我们回到我的简单例子上，模型的前向代码是：

def forward(self, x,rank):
        if rank==0:
            return self.net1(x)
        else:
            return self.net2(x)

此时rank0上的模型在前向时只使用net1的参数，rank1只使用net2，如果每个进程进行操作：

model = ToyModel().to(rank)
ddp_model = DDP(model, device_ids=[rank],find_unused_parameters=False,)
k = ddp_model(torch.tensor(2.).reshape(1,1).to(rank),rank)
k.backward()

那么在调用的backward的时候会出现

	net1.grad.status	net2.grad.status
rank0	ready	unready
rank1	unready	ready

此时由于不同rank的计算图不一样，没有任何一个bucket会ready，自然任何一个bucket的allreduce也不会被触发，Reducer虽然会在所有bucket变成ready后等待各自完成allreduce操作，但是当所有进程的backward()已经完成梯度的计算，Reducer将不会等待还未进入ready状态的bucket。这个时候就会导致严重的错误！！ ，梯度平均将不会发生，如果此时在不同进程上进行了参数更新，会直接导致不同进程的参数不一样，造成严重的问题。

find_unused_parameters解决梯度同步问题

实际情况是，确实有很多任务针对不同的数据输入会产生不同的计算图，在这样的条件下我们要如何确保梯度平均依旧能够触发，从而保证模型参数的一致性？答案是在DDP中设置find_unsued_prameters=True

ddp_model = DDP(model, device_ids=[rank],find_unused_parameters=True)

一旦find_unused_parameters=True，此时在调用forward的函数时，DDP模型会在本地rank的模型上花费额外的算力将未使用的params的梯度直接设置为ready状态。

如果进行操作

model = ToyModel().to(rank)
ddp_model = DDP(model, device_ids=[rank],find_unused_parameters=True,)
k = ddp_model(torch.tensor(2.).reshape(1,1).to(rank),rank)
k.backward()

在调用backward之前梯度的状态就已经为，forward过程已经将rank0.net2.grad.status和rank1.net1.grad.status设置为True

	net1.grad.status	net2.grad.status
rank0	unready	ready
rank1	ready	unready

因此在调用backward()之后，所有的param都会进入ready状态，从而正确触发所有的梯度平均操作。

报错信息分析

可以使用TORCH_DISTRIBUTED_DEBUG=DETAIL python yourscripts.py运行你的程序，这样在报错信息的结束处，会友善的提示你，到底是那个模块的哪个参数出了问题，例如：

RuntimeError: Expected to mark a variable ready only once … … and hence marking a variable ready multiple times. DDP does not support such use cases in default. You can try to use _set_static_graph() as a workaround if your module graph does not change over iterations.
Parameter at index 0 with name net1.weight has been marked as ready twice. This means that multiple autograd engine hooks have fired for this particular parameter during this iteration.

说了这么多，其实我们已经可以非常清晰的确定这个报错的原因了，我们的worker函数长这样：

# 使用这个demo_fn会复现之前的报错信息
def demo_r(rank,world_size):
	print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)
    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank],find_unused_parameters=True,)
    optimizer = optim.SGD(ddp_model.parameters(), lr=1.)
    optimizer.zero_grad()
    k = ddp_model(torch.tensor(2.).reshape(1,1).to(rank),rank)
    if rank==0:
        k += ddp_model.module.net2(torch.tensor(32.).reshape(1,1).to(rank))
    else:
        k += ddp_model.module.net1(torch.tensor(41.).reshape(1,1).to(rank))
    k.backward()
    cleanup()

仔细看，男主是小帅，由于我们设置了 find_unused_parameters=True，以rank0为例，在前向传播的过程中

k = ddp_model(torch.tensor(2.).reshape(1,1).to(rank),rank)

net1的网络参数会被使用，net2没被使用，因此，net2.weight.grad会被标记为ready状态。但在前向传播之外，我们又单独调用了net2的网络参数，进行运算，并把运算结果保留在了变量k中

 k += ddp_model.module.net2(torch.tensor(32.).reshape(1,1).to(rank))

此时调用k.backward()，由于net2的网络参数参与了k的计算，因此当net2.weight.grad被计算得到时，对应的hook会被激活标记net2参数为ready。但是如上所述，前向传播中net2已经被设置为了ready，这样的冲突导致了报错。

解决方案

经过上述的讨论我们知道find_unused_parameters=True是非常重要的参数，尤其是针对计算图不固定的场景，其决定了模型能不能在多卡上同步。
当你对你的模型有充足的把握时（例如此处），我们可以让find_unused_parameters=False，由于我们确定该计算图会涵盖所有参数。

参考资料

原理分析写的过于简陋，细看如下文档必有收获。
ddp_tutorial
ddp_communication&training
some_inner_design_of_ddp

DDP实践中额外要注意的几个点：

1. dist.set_epoch(epoch)肥肠重要
2. 有人纠结
   model = ToyModel().to(rank)
   ddp_model = DDP(model, device_ids=[rank],find_unused_parameters=True,)
   optimizer = optim.SGD(ddp_model.parameters(), lr=1.)
   代码中optimizer的parameters一定得是ddp_model.parameters()吗，可以是model.parameters()吗？
   答案是不必纠结，通过[i is j for i,j in zip(model.parameters(),ddp_model.parameters())]，测试一下就会发现，当DDP包裹model的时候，其实二者的参数是一样的了（只是parameters()传回的一样），但二者不能一概而论，例如model.net1和ddp_model.module.net1，二者还是得用不同的方法索取模型内容。
3. 如果使用我上面例子的方式spawn出多进程，那么脚本中一定要包含if __name__=='__main__':区别主进程和由run_demo函数调用mp.spawn生成的子进程。（可以尝试去掉，会产生报错信息。这个报错原理大概和进程的生成有关，我没仔细研究，猜想一下，可能是如果没有’main’限定，spawn出的子进程会再次run_demo…，如果我说的不对欢迎指正，我也想知道！！）
4. dist.barrier()这个函数是否是必要使用的？ 在之前的阐述中表明了其实每次的backward操作都触发了DDP model中的hook，这些hook通过bucket进行了同步，那么dist.barrier()是否有存在的必要？（确实有不用它的AI项目），官方的描述是：dist.barrier(group): Blocks all processes in group until each one has entered this function.[ddp_communication&training]。类似于join的操作，其中group是一堆rank的集合，可以通过group = dist.new_group([rank0,rank1,....])创建。