open-mmlab/mmengine

[Feature] remove AMP wrap in train_step

Opened this issue · 5 comments

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

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

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

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

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

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

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

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

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

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