[Connor-XY]

Depends on #6303 (which contains #6293), #8017 (which contains #6259), #8021 and #8049. Please review the diff from their merge: scoped diff.

• The checks use Add a dedicated fallback warning and process-group collection helpers #8049's require_process_groups, and the new tests use its ProcessGroupFallbackWarning and Add test helpers to forbid global grid reads and build collections #8017's helpers that forbid the global grid.
• The refit check needs the dp group that Require dp in the refit source collection instead of filling it #8021 puts in refit source collections. The new test_model_metadata.py tests also use the ProcessGroupCollection import that Require dp in the refit source collection instead of filling it #8021 adds there.
• The HybridModel pairing test uses the module-level ProcessGroupCollection import that Require an explicit process-group collection in LanguageModule #6303 adds to test_hybrid_model.py.
• The base merges all four in one commit. Its only conflict is the process-group allowlist total, which it sets to the number of merged entries. No entry changes.

Land this after #8748, which passes the collection in the same TEFusedMLP.as_mlp_submodule call where this PR passes the checked tp_group. Whichever lands second keeps both arguments. Land it after #8038 too: TEDotProductAttention drops its cp assert and relies on #8038's check that CP attention has a CP group.

Why. ProcessGroupCollection returns None for a declared field that was never set. So hasattr(pgc, name) is always true, and getattr(pgc, name, None) cannot tell an unset field from an axis that is off. Core code used both as presence checks, so a collection missing a field silently led to:

• an MLP sharded over the global TP group, with only a fallback warning;
• TEDotProductAttention getting no TP group;
• DSA attention running with TP and CP off;
• refit describing TP shards as whole tensors, and a None group reaching torch.distributed as WORLD;
• HybridModel choosing its logging ranks as if an unset group were an axis that is off;
• MIMO checkpoint checks that vanish under python -O.

What.

• TEDotProductAttention, MLP.as_mlp_submodule, TEFusedMLP.as_mlp_submodule, DSAttention, HybridModel's logging groups and refit's extract_param_metadata check the fields they read with require_process_groups. An unset field raises a ValueError that names it. None still means the axis is off.
• DSAttention checks tp and cp once at construction instead of reading them with getattr in every forward. The cuDNN full-fusion hook still treats an omitted collection as TP off.
• HybridModel's tp and dp_cp must still both be groups or both be None.
• Refit also rejects a None dp, and a missing ep for expert-parallel parameters.
• MIMO replica ids and sharded state dicts raise ValueError instead of asserting.
• Optional reads use vars(), with no behavior change: dp_cp_gtp_remat in the GPT and Hybrid MTP metrics, tp_cp in MTP, and gtp_remat, expt_gtp_remat and DDP's expt_dp in the DDP and optimizer setup.
• The allowlist does not change. vars() and require_process_groups are not counted reads.

Tests.

• New tests/unit_tests/test_process_group_field_presence.py:
  • DSAttention rejects an unset tp or cp and accepts None;
  • the MLP builders reject a collection without tp;
  • on an even number of ranks, TEDotProductAttention and the MLP builder use a separate communicator over the same TP ranks with global reads forbidden, and an omitted collection still warns.
• New tests/unit_tests/models/mimo/test_mimo_checkpoint_groups.py covers the replica ids and sharded state dicts.
• test_model_metadata.py adds six refit metadata tests.
• The HybridModel pairing test in test_hybrid_model.py builds real collections and expects unset fields to raise. The stand-in model in main's config_list_model_boundary fixture now sets tp and dp_cp to None next to pp, since HybridModel requires both.
• Three DSAttention tests in test_attention_variant_dsa.py built the module from an empty namespace and now declare both axes off.

On 4 GPUs against its base, with no regressions:

• test_process_group_field_presence.py passes all 9 new tests, including the 5 that need Transformer Engine;
• test_mimo_checkpoint_groups.py passes all 17 new tests;
• test_model_metadata.py has 50 passed here vs 41 on the base, and the six new tests run as 9 cases that all pass;
• test_hybrid_model.py has 84 passed and 4 skipped on both sides;
• test_attention_variant_dsa.py has 112 passed on both sides. The three TestIndexerTensorParallel tests fail on both sides because they need 8 ranks.

In the integration run, this PR's test files show no regressions against main. There the DSA file passes all 115 tests on 8 GPUs (two nodes), as on main, including the three TestIndexerTensorParallel tests. With all the #6307 changes applied, end-to-end training smokes on 4 GPUs give bit-identical losses to main and record no process-group fallback warning.

Before the sync with main, the branch ran on 4 CPU gloo ranks, one test file at a time, with no regressions:

• test_model_metadata.py passes all 50 tests;
• the new MIMO file passes all 17 tests;
• the new field-presence file passes 4 tests and skips 5 that need Transformer Engine;
• test_hybrid_model.py needs Transformer Engine to collect. Run on its own without it, the pairing test passes.

Compatibility.

• Callers that pass a collection must set the fields above, to None where the axis may be off. In-repo callers already do, and so do Megatron-Bridge's MIMO collections. Other callers get a new ValueError from:
  • TEDotProductAttention and the MLP builders, for an unset tp. On main the MLP took the global TP group with a warning, and TE got tp_group=None;
  • DSAttention, for an unset tp or cp;
  • refit, for an unset tp, pp or dp, a None dp, or a missing ep on an expert-parallel parameter;
  • MIMO checkpointing, whose messages also change.
• A duck-typed collection that keeps its groups as class attributes now reads None at the vars() sites. No in-repo caller passes one there.
• The new comments in process_groups_config.py say that an unset gtp_remat or expt_gtp_remat means the axis is off for DDP and optimizer setup. resolve_gtp_remat_group instead falls back to the global GTP group for an unset field. That mismatch predates this PR; Warn before resolve_gtp_remat_group reads the global grid #8769 makes the fallback warn.
• extract_param_metadata runs require_process_groups for every tensor, though once per component would do. The new Raises sections leave out the TypeError that require_process_groups raises for a MultiModuleProcessGroupCollection.
• Needs determinism-exempt: the kernel coverage check flags the Transformer Engine wrappers and the DSA files. Only process-group checks change.

Synced with main 77391e7 on 2026-11-01 by a merge commit that lists each conflict and its resolution, so existing commits and review links stay valid.

Part of #6307.

🤖 Generated with Claude Code

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

[Connor-XY]

Superseded by #8900, which carries this change as commit e8a55e6 under the same title, with "Originally #8756." in its message. Closing so review happens in one place. The branch stays on the fork.