Add ability to save optimizer and resume while training

[coolkp]

This pr adds

1. Add optimizer state to checkpoint (save checkpoint path, whole state saved)
2. Load optimizer and step, overwrite optimizer in training loop (be careful of shardings) its inherited.
3. resume training from last checkpointed step

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid