多卡训练Runtime Error: Function BroadcastBackward returned an invalid gradient at index XXX

代码在单卡训练时没有问题，但是在多卡训练（DP或者DDP模式）时，会在loss.backward()处报错，但是报错并不详细，只停留在loss.backward()，再往后就是pytorch后端C++代码了。报错如下：

Runtime Error: Function BroadcastBackward returned an invalid gradient at index 770- got [0] but expected shape compatible with [0,8,3,3]

可见是反向传播时shape不一致。但是经过检查，我的loss并没有出错，单卡训练测试也没问题，这个就比较棘手了。

以下是我的排查步骤。

1. 打印计算图

先安装计算图可视化的包：

pip install torchviz
pip install graphviz
sudo apt-get install graphviz

然后在loss处插入如下代码：

from torchviz import make_dot
# 先计算出loss
loss = model_loss(output, gt)
# 创建计算图
dot = make_dot(loss, params=dict(model.named_parameters()))
# 保存图像
dot.render("model_graph", format="svg")

loss.backward()

我并没有从计算图中发现什么问题。

2. 寻找没有参与反向传播的参数

使用单卡调试，在loss.backward()之后、optimizer.step()之前加入如下代码：

for name, param in model.named_parameters():
    if param.grad is None:
        print(name)

打印出来的参数就是没有参与loss运算的部分，它们梯度为None。

在这里，我发现有一些模块没有参与loss运算，原因是我定义了一些用于实验的模块，但是后续forward当中没有使用。例如，我定义了一个modulelist，但是只有前面3个模块被使用了，后面2个模块定义了但是没有使用，却也作为参数传入了（因为整个modulelist都传入了）。

经过修改，bug修复！我认为出现“单卡可以运行但是多卡报错”的原因是分布式计算时需要保证计算图在梯度回传时都要有梯度。