fix(checkpointing): include optimizer scaffold while loading

[cuichenx]

Summary

• include optimizer sharded-state scaffolding during checkpoint load even when the next checkpoint should not save optimizer state
• add a regression unit test for load-time optimizer scaffold generation

Blast Radius / Test Assessment

• Touches shared training checkpointing and resume behavior.
• L0/unit coverage is required and included for generate_state_dict.
• Broader manual validation should include resuming from a distributed optimizer checkpoint with save_optim=False; no L1/L2 CI labels are added.

Validation

• uv run --no-sync pre-commit run --files src/megatron/bridge/training/checkpointing.py tests/unit_tests/training/test_checkpointing.py
• uv run --no-sync python -m py_compile src/megatron/bridge/training/checkpointing.py tests/unit_tests/training/test_checkpointing.py

[claude]

Light Code Review - The fix is correct and well-scoped. The is_loading flag is already set on both load paths (lines 2629 and 2694) and is never present on the save path (line 1177), so there is no risk of accidentally including optimizer state during saves. The unit test covers the positive case (scaffold included when save_optim=False and is_loading=True). One gap: Missing negative-case test -- no test asserting that save_optim=False without is_loading still excludes optimizer state. Consider adding a companion test with optim_sd_kwargs={} and asserting optimizer not in result. No bugs, typos, or doc issues found. LGTM. Suggested test cases: No perf tests impacted.

[cuichenx]

/claude review

[claude]

LGTM — clean, well-scoped fix with a solid regression test.

The logic change at checkpointing.py:1848 correctly decouples the load-time optimizer scaffold from the save_optim flag by checking is_loading in optim_sd_kwargs. Both call sites (_load_dist_checkpoint at line 2629 and the FSDP path at line 2696) already pass is_loading=True, so this is a safe, targeted fix.

The concurrent-access test fix in test_checkpoint_utils.py (replacing the global builtins.open mock with a real temp file) is the right approach for thread safety.

One minor observation: there is no negative test asserting that save_optim=False without is_loading still excludes optimizer state from the result. The existing tests happen to pass optimizer=None in that scenario, which skips the branch anyway. A test with save_optim=False, no is_loading key, and a non-None optimizer would strengthen the coverage, but this is not blocking.

Suggested test cases: No perf tests impacted.

[cuichenx]

/ok to test 4472834