fix(ckpt): initialize model parallel for builder loads

[yaoyu-33]

What changed

Initialize and acquire MCore model-parallel process groups at the shared Bridge model-construction boundary in build_and_load_model().

Both supported native checkpoint paths now use the same ProcessGroupCollection:

• legacy ModelProviderMixin receives it through provide_distributed_model(pg_collection=...)
• new ModelConfig / builder receives it through build_distributed_models(...)

Supported trigger and impact

A collective caller can initialize torch.distributed and then call the documented AutoBridge.load_megatron_model() API on a builder-backed checkpoint. Muse Glimmer makes this path supported after #5479. The load path previously treated the default process group as proof that MCore groups also existed, then ProcessGroupCollection.use_mpu_process_groups() raised before model construction or checkpoint I/O.

This blocks native checkpoint load/refit/export for builder-backed models embedded in framework-owned distributed processes unless callers perform undocumented MCore initialization.

Root cause and fix

AutoBridge.load_megatron_model() skips its temporary distributed context when torch.distributed is initialized. The builder path consumed MCore process groups directly, while the legacy provider path had its own fallback initialization.

The shared helper now initializes missing MPU state from the already-normalized TP, PP, VPP, CP, EP, and expert-TP configuration, initializes the model-parallel CUDA RNG tracker when CUDA is available, and returns the MPU-backed collection. Model construction then branches only after this common setup.

Regression evidence

The tests cover:

• a real one-rank CPU/Gloo default process group with MPU initially absent
• initialization parameters and CUDA RNG seeding
• reuse of an already-initialized MPU state without reseeding
• propagation of the same returned collection to both legacy provider and builder paths

Validation:

uv run --no-sync python -m pytest tests/unit_tests/training/test_model_load_save.py -q

Result: 64 passed.

uv run --no-sync pre-commit run --all-files

Result: all hooks passed.

Scope

No public API, checkpoint format, dependency, workflow, or upstream MCore changes. The standalone temporary-context path and already-initialized MPU behavior remain intact. No GPU, convergence, performance, or full-suite claims are made.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[copy-pr-bot]

/ok to test dd4f9cd8c2bfb1576f0ed91081b66f41e353be2c

@yaoyu-33, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[yaoyu-33]

/ok to test dd4f9cd

[claude]

Light review - LGTM

Minimal, well-scoped fix. The builder branch now initializes MCore model-parallel state when the caller owns only the default PyTorch process group before use_mpu_process_groups() consumes those groups. It mirrors the sibling ModelProviderMixin.initialize_model_parallel (model_provider.py:323-330), including the or-1 None-handling for CP/EP. The isinstance(model_cfg, ModelConfig) guard makes the direct TP/PP/VPP/CP/EP/expert-TP attribute access safe.

Observations (non-blocking):

• On this skip-temp-context path the newly-initialized MPU state persists after the call returns (unlike temporary_distributed_context, which destroys it in a finally block). Consistent with the existing provider branch and arguably needed for the returned model to stay usable, so not a regression.
• The RNG tracker is not set up here, but load_megatron_model sets perform_initialization=False (weights come from the checkpoint), matching the sibling path.

Test coverage is good: the new test exercises the default-PG-only scenario with a real gloo group, asserts MPU gets initialized, and cleans up in a finally block.

Suggested test cases:

• tests/unit_tests/training/test_model_load_save.py::TestLoadMegatronModel::test_load_builder_config_initializes_model_parallel_with_existing_default_group
• tests/unit_tests/training/test_model_load_save.py::TestLoadMegatronModel

No perf tests impacted.

[yaoyu-33]

/ok to test a27a605

[yaoyu-33]

/ok to test 1b5cb0b

[yaoyu-33]

/ok to test 3dc48f6