fix(ckpt): support MFSDP v2 MoE checkpoint resume

[wujingyue]

What does this PR do?

Enable MFSDP v2 MoE checkpoint save/resume in Bridge and verify that resumed training matches uninterrupted training. Bridge currently rejects these checkpoints during config validation. After allowing them, loading into the unwrapped model bypasses the MFSDP adapter's BF16 compute-weight refresh, causing the first resumed loss to differ.

Depends on NVIDIA/Megatron-LM#8230. That fix must be available in Bridge's Core dependency before this PR merges. This PR does not change the Core submodule pin.

Changelog

• Allow MFSDP v2 checkpoint save/load with fsdp_dtensor.
• Restore model state through the original wrapper, matching Core's loader, so the MFSDP v2 adapter refreshes BF16 compute weights from the restored main weights.
• Add an EP=2 regression to the existing two-GPU MFSDP functional suite. Reuse its small MoE model and training helpers; save after steps 1–2, rebuild and load, then compare resumed steps 3–4 exactly against uninterrupted training. Check the executed step numbers and require the checkpoint to exist.

GitHub Actions CI

Initial local validation used Core commit dafb72047b05dac29dacc2be4ecebdafcf3ee21b from NVIDIA/Megatron-LM#8230, TE 2.18.0+27486e03, and two A6000 GPUs:

• New BF16 EP=2 save/resume test: passed on both ranks, with the precision-aware optimizer enabled and exact loss comparisons.
• Latest EP=2 rerun with Core bb229ffa31bd4566cd3edfff942965de7b789bf2: passed on both ranks using the default test-data fixture, with GH_TOKEN supplied for asset downloads.
• Existing dense MFSDP save/resume test: passed on both ranks.
• Existing checkpoint loading and resume unit tests: 102 passed.
• Against Core before Fix MFSDP v2 checkpointing for expert-parallel MoE models NVIDIA/Megatron-LM#8230, the new test fails at the first save with the original EP NotImplementedError.
• With Fix MFSDP v2 checkpointing for expert-parallel MoE models NVIDIA/Megatron-LM#8230 but without the Bridge wrapper-loading change, the new test fails at the step-3 loss comparison.
• uv run pre-commit run --all-files: passed.
• Strict mypy with --follow-imports=skip reports 235 typing errors across the three changed files; type checking is not clean.

Run the new regression with a Core build containing the fix:

uv run python -m torch.distributed.run --nproc_per_node=2 -m pytest \
  tests/functional_tests/test_groups/megatron_fsdp/test_megatron_fsdp.py::TestMegatronFSDP::test_fsdp_v2_moe_ep2_save_resume

Before your PR is "Ready for review"

• Read and followed the contributor guidelines.
• Added the necessary functional regression test.
• Reviewed documentation needs; no documentation changes included.
• Affects optional components: no new optional dependencies or imports.

Additional Information

The test runs automatically through the existing H100 L0 MFSDP launch script. No workflow changes are needed.

[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.