open-mmlab/mmengine

[Feature] remove AMP wrap in train_step

Opened this issue · 5 comments

whlook commented

What is the feature?

  • 目前mmengine的行为在train_step上包装了AMP的context，（相比原来的mmcv是提供auto_fp16由用户自定义包装，且是module级别），这样的方式让AMP使用更简单，方便了用户，但是也引入了一个问题，参见:https://discuss.pytorch.org/t/autocast-and-torch-no-grad-unexpected-behaviour/93475 的讨论。
  • 由于在最上层包装了amp，那么在用户的模型内就无法正确使用no_grad进行forward操作，如以上链接的示例case。这个对用户的灵活度和问题排查也带来了一些问题（最近我也在该问题上排查了很久）。
  • 因此建议是否可以这种amp包装恢复成以前mmcv的方式？或者有另一种方案，就是提供“no amp wrap”版本的AmpOptimizerWrapper，相比之前的wrap，只是把amp的context去掉，而可以保留grad scale等操作

Any other context?

No response

zhouzaida commented

感谢你的反馈。我想具体了解一下，你希望提供“no amp wrap”版本的 AmpOptimizerWrapper，那 autocast 的操作你打算在哪里完成？

HAOCHENYE commented

你可以尝试使用 with autocast(enabled=False): 在里面的代码临时关闭 autocast

  • 👍1
whlook commented

感谢你的反馈。我想具体了解一下，你希望提供“no amp wrap”版本的 AmpOptimizerWrapper，那 autocast 的操作你打算在哪里完成？

我最初的想法是如果使用去掉autocast的AmpOptimizerWrapper，那么可以在自己的代码里手动使用torch.autocast等操作，不过这样对用户来说可能也不是很方便，目前看可能还是with autocast(enabled=False):更可行一些

whlook commented

你可以尝试使用 with autocast(enabled=False): 在里面的代码临时关闭 autocast

抱歉重新提起这个问题。目前发现使用with autocast(enabled=False):不是一个合适的解决方法。因为这个会导致内存突然上升，这个背离了使用amp的理由。所以为了兼顾amp使用，并且保证在训练代码中使用no_grad的操作不会破坏训练效果，我建议在mmengine中的amp包装上修改一下默认参数(cache_enabled=False)，默认为True:
image

whlook commented

你可以尝试使用 with autocast(enabled=False): 在里面的代码临时关闭 autocast

抱歉重新提起这个问题。目前发现使用with autocast(enabled=False):不是一个合适的解决方法。因为这个会导致内存突然上升，这个背离了使用amp的理由。所以为了兼顾amp使用，并且保证在训练代码中使用no_grad的操作不会破坏训练效果，我建议在mmengine中的amp包装上修改一下默认参数(cache_enabled=False)，默认为True: image

更新，mmengine也可以不改，在自己模型的代码里按需关闭cache也可以。顺便提醒路人了，这个还是挺隐蔽的坑