一、解决在分布式训练中,如果对同一模型进行多次调用的报错
报错如下:
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [256)] is at version 4; expected version 3 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).
参考知乎文章《【PyTorch踩坑】一个排查了一下午的坑》
经过一些调试发现,只有当某些特定情况下才会触发此报错。下面结合一个对比学习的例子(并不是完整的脚本)来简单描述:
python
import torch
import torch.nn as nn
from torchvision.models import resnet50
def main():
model = resnet50(num_classes=256).cuda()
model = nn.parallel.DistributedDataParallel(model,
device_ids=[args.local_rank],
find_unused_parameters=True)
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(),
lr=0.001,
momentum=0.99,
weight_decay=1e-4)
for i in range(10):
input0 = torch.randn((4, 3, 224, 224), dtype=torch.float32).cuda()
input2 = torch.randn((4, 3, 224, 224), dtype=torch.float32).cuda()
out1 = model(input0)
out2 = model(input1)
loss = criterion(out1, out2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if __name__ == '__main__':
main()
经过调试发现,当使用nn.DataParallel
并行训练或者单卡训练均可正常运行;另外如果将两次模型调用集成到model中,即通过out1, out2 = model(input0, input1)
的方式在分布式训练下也不会报错。
由此可以猜测:在分布式训练中,如果对同一模型进行多次调用则会触发以上报错,即nn.parallel.DistributedDataParallel
方法封装的模型,forword()
函数和backward()
函数必须交替执行,如果执行多个(次)forward()
然后执行一次backward()
则会报错。
那么解决此问题的入手点则可以聚焦到nn.parallel.DistributedDataParallel
接口上。 通过查询PyTorch官方文档发现此接口下的两个参数:
- find_unused_parameters: 如果模型的输出有不需要进行反向传播的,此参数需要设置为True;若你的代码运行后卡住不动,基本上就是该参数的问题。
- broadcast_buffers: 该参数默认为True,设置为True时,在模型执行forward之前,gpu0会把buffer中的参数值全部覆盖到别的gpu上。
问题基本可以定位出来了,即broadcast_buffers=True
导致参数被覆盖修改。解决办法:
model = nn.parallel.DistributedDataParallel(model,
device_ids=[args.local_rank],
broadcast_buffers=False,
find_unused_parameters=True)
参考
distributed: https://pytorch.org/docs/stable/distributed.html
Inplace error if DistributedDataParallel module that contains a buffer is called twice