Give gloo process groups a home in ProcessGroupCollection

[Connor-XY]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Today, explicit process groups and gloo process groups are mutually exclusive — and the distributed optimizer is on every training run.

ProcessGroupCollection.setup_process_groups_for_optimizer branches on pg_collection is None:

• the None branch builds gloo groups from parallel_state.get_data_parallel_group_gloo() / get_expert_data_parallel_group_gloo()
• the else branch raises (process_groups_config.py):

if use_gloo_process_groups:
    raise ValueError(
        "Gloo process groups are not supported when pg_collection is "
        "provided. Please set use_gloo_process_groups to False."
    )

So any job passing an explicit collection cannot use gloo groups at all. This blocks the parallel_state migration for the optimizer path.

Root cause

get_data_parallel_group_gloo and get_expert_data_parallel_group_gloo were the only accessors in use_mpu_process_groups's mapping without a check_initialized parameter, so they assert rather than returning None when gloo groups were never created (initialize_model_parallel(create_gloo_process_groups=False)). A collection therefore could not carry them.

Changes

• Add check_initialized to both gloo accessors, matching every other accessor in that mapping.
• Add intra_dp_cp_gloo / intra_expt_dp_gloo to ProcessGroupCollection, named for their existing intra_dp_cp / intra_expt_dp siblings, and populate them in use_mpu_process_groups.
• setup_process_groups_for_optimizer reads them instead of refusing. Absent → None (the legitimate no-gloo case); present → passed through.
• With that, the optimizer's fallback hoists to the edge: three resolution sites inside private helpers become one at get_megatron_optimizer.

That last part is the shape proposed in #6099: "move all global-process-group fallbacks to the edge, around get_megatron_optimizer, and make all internal code require explicit process groups."

Test change worth reviewing

test_optimizer.py "Test 6" asserted the old refusal (pytest.raises(ValueError, match="Gloo process groups are not supported")). It now asserts the new contract: a collection without gloo groups yields None; one carrying them passes them through.

Verification

Paired runs on 4×GPU against unmodified main at the same commit, same container:

Tree	Result
this PR	108 passed, 41 skipped
pristine main	108 passed, 41 skipped

Suites: test_optimizer.py, test_emerging_optimizers.py, test_process_groups_config.py.

Scope

This covers only the gloo groups the optimizer needs. Virtual-pipeline rank/size, GlobalMemoryBuffer and NCCL options are not process groups and still have no home in ProcessGroupCollection; they need a separate design.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• I have added relevant documentation
• I have run the autoformatter on my PR

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

[Connor-XY]

Part of #6307 — tracking issue for the parallel_state deprecation, with the landing order, the decisions needed, and what remains.