Summary

• Add per-rank optimizer state saving/loading to checkpoint utilities
• Fixes loss spikes and training instability when resuming from checkpoints

Problem

The current checkpoint system saves:

• backbone.pt - model weights
• scheduler.pt - LR scheduler state
• PFC module weights per rank

But it does NOT save optimizer.state_dict(), which contains:

• exp_avg (1st moment / momentum)
• exp_avg_sq (2nd moment / RMS accumulator)
• step count for bias correction

When resuming training without optimizer state, AdamW starts fresh, causing:

• Transient loss spikes (especially severe if resuming mid-training with decayed LR)
• Training instability lasting 1k-10k steps (due to beta2=0.999)
• Wasted compute to recover the optimization trajectory

Solution

• Add optimizer parameter to save_checkpoint() and load_checkpoint()
• Save optimizer state per-rank since it includes moments for local PFC shards
• Backward compatible: old checkpoints without optimizer_*.pt will warn and continue

Files Changed

• training/checkpoint_utils.py - Core save/load logic
• training/train.py - Pass optimizer to checkpoint functions

Testing

• Existing training should work unchanged (optimizer param defaults to None)
• New checkpoints will include optimizer_{rank:03d}.pt files
• Resume from new checkpoints will restore full optimizer state