[Naeemkh]

Summary

• Muon checkpoints saved before Add an exact Newton-Schulz option for sharded Muon weights #281 fail to resume on main: orthogonalize is in every Muon param group, and
  DCP requires every key of the load template (Missing key in checkpoint state_dict: optimizer.param_groups.0.orthogonalize).
• Muon.LEGACY_GROUP_DEFAULTS = {"orthogonalize": "local"}: Muon.__setstate__ fills listed keys into loaded
  groups that lack them, and step reads group["orthogonalize"].
• CheckpointManager.load leaves a listed key out of the DCP template only where the checkpoint lacks it, and
  logs it. Any other missing key still fails the load.
• A resume keeps the mode its checkpoint was saved with; a checkpoint from before the option resumes with
  "local", the mode it was trained with.

Testing

• uv run ruff check kempnerforge/ tests/ passes
• uv run ruff format --check kempnerforge/ tests/ scripts/ passes
• uv run pyright kempnerforge/ passes (0 errors)
• uv run pytest tests/unit/ -v --timeout=60 passes (CPU: 2324 passed, 3 skipped, 1 xfailed)
• If distributed code changed: uv run torchrun --nproc_per_node=4 -m pytest tests/distributed/ -v
• If training loop / parallelism / optimizers changed: uv run pytest tests/e2e/ --e2e -v

New tests in TestMuonResume: a checkpoint without orthogonalize resumes as "local" and matches an
uninterrupted run bitwise; it stays "local" with "exact" configured; a checkpoint with the key keeps its saved
mode; a checkpoint missing ns_steps still fails; the oldest layout (no inner AdamW, no orthogonalize)
resumes Muon.

Checked with real checkpoints from older code: one saved at fb23328 resumes with step 3 bitwise equal to the
uninterrupted run; one saved at 8776509 (before #258) resumes, with only the inner AdamW's parameters
differing, as that layout's warning says. Both fail on 000e803.

Closes #388

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
kempnerforge/checkpoint/manager.py	96.61% <100.00%> (+0.11%)	⬆️
kempnerforge/training/optimizer.py	92.15% <100.00%> (+0.22%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[timothyngo]

LGTM, see the comment I added: a short warning would be helpful when something is missing.