pytorch distributed常见错误

本文主要介绍了Pytorch分布式的一些常见错误，避免大家踩坑

Distributed

加载参数

有一个坑是使用分布式计算的时候，每张卡的内存分配都应该是均匀的，但是有时候会出现0卡占用更多内存的情况，这个坑在知乎上有讨论：链接

分布式本身的内存分配应该是均匀的（左图），但是有时候会出现另一种情况（有图）

这是load模型的时候导致的，当用下面句子load模型时，torch.load会默认把load进来的数据放到0卡上，这样四个进程全部会在0卡占用一部分显存，解决方法就是将load进来的数据map到cpu上

# 在GPU上加载
checkpoint = torch.load("checkpoint.pth")
model.load_state_dict(checkpoint["state_dict"])

# 改为在CPU上加载模型
checkpoint = torch.load("checkpoint.pth", map_location=torch.device('cpu'))
model.load_state_dict(checkpoint["state_dict"])

计算图

RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by (1) passing the keyword argument `find_unused_parameters=True` to `torch.nn.parallel.DistributedDataParallel`; (2) making sure all `forward` function outputs participate in calculating loss. If you already have done the above two steps, then the distributed data parallel module wasn't able to locate the output tensors in the return value of your module's `forward` function. Please include the loss function and the structure of the return value of `forward` of your module when reporting this issue (e.g. list, dict, iterable).

出现这种错误有三个可能的原因：

1. 模型的输出的某个值没有参与loss计算。例如你在输出时输出了target和一个temp，这个temp可能是你想可视化的一个结果图，但是你计算loss的时候temp没有参与loss计算，就会报错
2. 模型某一个操作在前向传播时没有进行计算。比如你定义了三个卷积层conv1，conv2以及conv3，在正向传播时只用了conv1以及conv2，就会报错
3. 模型在反向传播时某一个操作没有获得梯度。这个与原因2类似，2是定义了conv3没有参与计算，不参与计算就不会获得梯度，conv3可能会参与计算，但是在conv3之后的参数反向传播时没有传播到conv3这里，就会报错

为了更清晰的明白上述三个问题，下面举了一个实例。我们定义了fc1，fc2，fc3，fc4程序中的运行过程为

• 错误一对应于我们在y1, y4 = model(input)输出了三个变量y1，y4，但是计算loss时只用了y1，没有用y4
• 错误二对应于我们在正向传播时没有用到fc3
• 错误三在这里没有举例，可以理解为在你定义的模型中有个conv参与了计算，但是是无效操作，即这个conv的输出没有作为任何函数的输入，也没有进行反向传播，除了浪费了计算时间和资源以外没有任何作用

class model(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.fc1 = nn.Linear(dim, dim*4)
		self.fc2 = nn.Linear(dim*4, dim)
		self.fc3 = nn.Linear(dim, dim)
		self.fc4 = nn.Linear(dim, dim)
    def forward(self, x):
		y1 = self.fc1(x)
		y2 = self.fc2(y1)
		y4 = self.fc4(y2)
		return y1, y4
...
optimizer.zero_grad()
y1, y4 = model(input)
loss = loss_l1(y1, x)
loss.backward()
optimizer.step()	

对于第一个错误，如果你的观察变量不会影响结果，你就可以将分布式中的torch.nn.parallel.DistributedDataParallel的参数find_unused_parameters=True，就不会报错了，他会忽略掉这个错误

多啰嗦几句

model可学习的参数类型有哪几种？

这里简单讲解一下pytorch中模型参数更新的方式，模型需要更新的参数一共分为两种，一种是parameter类型，他是有梯度才能更新的；另外一种是buffer类型，他更新的方式是根据每一个iteration的data数据，比如batchnorm中的running_mean以及running_var参数就是buffer类型的

model中的参数是什么时候更新的？

在参数更新时，对于parameter参数，也就是conv中的weight和bias等参数，我们执行完loss.backward()函数以后，他们会计算parameter的梯度grad，然后我们执行optimizer.step()以后会根据学习率来更新我们的parameter；而对于buffer类型的参数，他是没有梯度的，拿batchnorm举例，他是根据每一个iteration提取的batch data来更新的，也就是批归一化，在我们执行完output=model(input)之后，buffer类型的参数就自动更新了（不需要等到optimizer.step()之后）

分布式中对哪些参数进行判断？

在上面提到的问题中，如果网络参数没有更新或者没有获得梯度，分布式训练就会报错，对model的可学习参数了解之后，我们知道分布式检查的是parameter参数，不检查buffer参数（因为buffer类型的参数本身就没有梯度，不然只要有BN分布式一定报错），这里如果报错了以后给大家提供两个定位错误位置的方法

• 从参数的角度：在output = model(input)之前和optimizer.step()之后将模型的参数保存下来，然后运行下面的命令，检查哪些参数没有更新，从而定位错误位置。这个方法有一个问题，当你的梯度很小（更新paramter）或者batch data的位于0均值（更新buffer）的时候，模型的参数可能不变，因为变化真的太小了，所以从参数的角度可能会失效，而且还可能输出一堆running_mean的参数（因为下面代码会对buffer进行比较）

state_dict1 = torch.load('before.pth', map_location='cpu')
state_dict2 = torch.load('after.pth', map_location='cpu')
for i in state_dict1:
    if state_dict1[i].equal(state_dict2[i]):
        print(i)

• 从梯度的角度：上面讲过，没获得梯度分布式也会报错，我们在optimizer.step()之后（或者loss.backward()之后也可以）加入如下语句。代码表示查看哪些参数没有获得梯度，因为正常的反向传播parameter是要获得梯度的，如果没有获得梯度说明此参数被排除在了计算图之外。除此之外，这行代码能够有效避免对buffer的排查（因为下面代码不会对buffer进行比较）

if rank==0:
    for name, param in net.named_parameters():
        if param.requires_grad and param.grad is None:
            print(name)