Adopt checkpoint bundles in training and losses

[relh]

Summary

• Switch training save/load to checkpoint bundles.
• Unwrap CheckpointPolicy in kickstarter losses.
• Update training tests and fast train tool to use bundles.

Changes

• CheckpointManager + Checkpointer now use bundle dirs.
• Losses unwrap CheckpointPolicy to get underlying policy.
• Tests updated for bundle-based resume/load.

Tests

• not run (stack recut)

[relh]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• Update docs and recipes for checkpoint bundles #4468
• Move MptPolicy compatibility into metta #4467
• Build submission zips from checkpoint bundles #4466
• Adopt checkpoint bundles in training and losses #4465 👈 (View in Graphite)
• Resolve checkpoint bundle URIs and load specs #4464
• Add checkpoint bundle primitives #4463
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.