Support resuming training after interruption

[norabelrose]

Optimizer and LR scheduler state are now saved to disk, along with the global step