Properly checkpoint models trained with Apex

[JohnGiorgi]

Is your feature request related to a problem? Please describe.

#3866 enables mixed precision training with Apex, but pushes the changes required to properly save and load a model trained with amp down the line. As discussed, I am simply opening this issue so I don't forget and will (hopefully) make a PR in the next couple of days.

Describe the solution you'd like

Follow the instructions in the Apex docs to checkpoint a model trained with amp.

Additional context

Apex docs.

[schmmd]

Thanks @JohnGiorgi

[schmmd]

@JohnGiorgi are you still able to create a PR for this? We would love a contribution before our 1.0 release, otherwise this is unlikely to make it in.

[matt-gardner]

I'm also trying to give time estimates for the remaining issues in the 1.0 milestone. Looking at the docs that @JohnGiorgi linked to, this is probably a very fast change, perhaps as little as an hour. But I'm giving it a "day" label, as there might be some complexity here that I'm missing.

[JohnGiorgi]

Thanks for pinging me -- I do still plan on working on this, just got swamped lately. Hopefully I can finish this before the end of the week.

[schmmd]

Cool, that'd be great! We'd love to have this in 1.0.