[FEATURE] save optimizer and amp state into checkpoint
daden-ms opened this issue · comments
Description
Currently, in the common.py for transformer models, a checkpoint only saves model state and the optimizer and amp state info is not saved. We can consider saving this info like in
https://github.com/NVIDIA/apex#checkpointing