Add and Remove ZeRO 3 Hooks (#5658)
Gives the ability to add and remove the forward hooks in ZeRO 3 by using
a context manager. These code changes were taken from a Huggingface
[PR](https://github.com/huggingface/trl/pull/1617) and integrated for
direct support in DeepSpeed.
This is useful in the inference case and the speedup can be observed
[here](https://github.com/huggingface/trl/pull/1483).
---------
Co-authored-by: root <root@deepspeed-c000004.2d1icxc5dsxehnpuwt3ifc34ph.gvxx.internal.cloudapp.net>
Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Heyang Qin <heyangqin@microsoft.com>
Co-authored-by: Logan Adams <loadams@microsoft.com>
