[ckpt] fix: honor disabled RNG saves

[cuichenx]

Summary

What

• collect checkpoint RNG state only when CheckpointConfig.save_rng is enabled
• make low-memory conversion saves omit RNG state consistently with their
  existing save_rng=False configuration
• cover enabled and disabled normal saves plus the low-memory conversion path

Root cause

save_megatron_model() explicitly constructs conversion checkpoint state with
save_rng=False, because a weight-conversion checkpoint does not need training
RNG state. Both Bridge save paths ignored that contract:

• save_checkpoint() called get_rng_state() unconditionally
• the low-memory conversion path separately called get_rng_state() before
  generating its prebuilt state dictionary

On a genuine CPU-only node this attempted torch.cuda.get_rng_state() after all
20.9B GPT-OSS parameters had loaded. Exact-public DFW job 15100682 failed at
checkpoint save with RuntimeError: Found no NVIDIA driver on your system.
No local patch was used in that run and no checkpoint output was produced.

Normal training behavior is unchanged: when save_rng=True, RNG state is still
collected and included as before.

Validation

• normal checkpoint save with save_rng=True
• normal checkpoint save with save_rng=False
• low-memory conversion save with RNG disabled
• three focused tests passed in a driverless 26.06.01 container
• uv run --no-sync pre-commit run --all-files
• unchanged full GPT-OSS CPU import from a fresh exact-public integration
  head containing this fix

Full CPU conversion remains unverified until the original workload passes and
the produced checkpoint passes its later GPU/TE load gate.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[cuichenx]

/ok to test 3c5b4bb

[claude]

LGTM - clean, well-scoped change that avoids collecting RNG state (and the associated all_gather_object collective plus CUDA init) when save_rng=False.

Verified:

• checkpointing.py: generate_state_dict (line ~2075) already gated the rng_state entry on ckpt_cfg.save_rng, so when save_rng=False the RNG state was previously gathered and then discarded. Skipping get_rng_state upfront produces identical saved checkpoints while eliminating wasted work. Since save_rng is a rank-uniform config value, all ranks skip the collective together - no risk of a collective-mismatch hang.
• model_load_save.py: the low-memory conversion path called get_rng_state (which touches CUDA via torch.cuda.get_rng_state()) despite save_rng=False discarding the result; passing rng_state=None directly is correct and avoids needless CUDA init during CPU-side conversion. pg_collection is still needed downstream (line ~765), so leaving its retrieval in place is right.
• Test coverage is solid: test_save_checkpoint_global is now parametrized over save_rng with explicit assert_called_once/assert_not_called checks, and the new test_low_memory_save_omits_rng_collection asserts get_rng_state is not called and rng_state is None is passed through.

Suggested test cases:

• tests/unit_tests/training/test_checkpointing.py::TestSaveCheckpoint::test_save_checkpoint_global[True]
• tests/unit_tests/training/test_checkpointing.py::TestSaveCheckpoint::test_save_checkpoint_global[False]
• tests/unit_tests/training/test_model_load_save.py::TestSaveMegatronModel::test_low_memory_save_omits_rng_collection

No perf tests impacted.