[Connor-XY]

Depends on #7991, #8722, #8723, #8736, #8731, #8727, #8739, #8763, #8773 and #8775. Together they contain every other PR of this series except #8777, which sits beside this one on the same base. Please review the diff from their merge: scoped diff.

• The test uses Add test helpers to forbid global grid reads and build collections #8017's helpers that forbid the global grid, ParallelStateDeprecationWarning and its per-call-site record from Deprecate the parallel_state accessors of the global grid #8765, build_process_group_collection() from Extract build_process_group_collection from initialize_model_parallel #8039, get_registered_process_groups() from Add a public process-group registry to parallel_state #8747, and the optimizer's grad-stats fallback record from Warn on optimizer, TP and helper fallbacks to WORLD or the global grid #8734.
• Every phase runs with all process-group fallback and deprecation warnings raised as errors. So the test also needs the PRs of [TRACKING] Deprecate parallel_state: explicit process groups in megatron/core #6307 that pass the collection through model build, DDP, the optimizer, the schedules, checkpointing, RNG seeding and the inference entry points, and the fallback warnings of Warn on optimizer, TP and helper fallbacks to WORLD or the global grid #8734. That is why the base holds the whole series.
• The merges resolve the existing conflicts between these branches in the Transformer Engine wrappers, t5_model.py, multi_token_prediction.py, fully_sharded_optimizer.py, the Gloo accessor docstrings of parallel_state.py and the duplicated-mode linears of the attention and MoE modules, and regenerate the process-group allowlist. Each merge commit lists its resolutions.

Land this right before #8777, after everything it contains, #8772, #8773 and #8775 included.

Why. The tests that forbid the global grid each cover one component, mostly with stand-ins. No test ran model build, DDP and the optimizer, training steps, a torch_dist save and load, and the inference build together on a collection whose TP and DP groups differ from the global ones. A component that reads the global grid, or that passes no group (WORLD) to a collective, could hide between those tests.

What. New tests/unit_tests/test_process_group_collection_end_to_end.py. No library change.

• The global grid is TP=2 in the default rank order, built and registered by Utils.initialize_model_parallel. A second collection comes from build_process_group_collection(order="dp-tp-cp-ep-pp") and is not registered. Its TP and DP groups differ from the global ones on every rank.
• Each phase runs inside forbid_global_process_groups() with ProcessGroupFallbackWarning and ParallelStateDeprecationWarning as errors. The once-per-process records of fallbacks, accessor call sites and optimizer grad-stats fallbacks are emptied first, so a warning from an earlier test cannot hide a fallback. A failure names the phase, and the chained error names the accessor or fallback.
• The phases: build a GPT model with the training builder (Float16Module, DDP with the distributed optimizer), build the optimizer with get_megatron_optimizer, run two steps with get_forward_backward_func, save and reload a torch_dist checkpoint, and build the model for inference with get_model_for_inference and run it through GPTInferenceWrapper.
• Since main's refactor(training): Refactor Tokenizer Initialization #8391, get_model_for_inference reads the run config and passes the tokenizer config to the model builder. The inference test takes the run_config fixture, and its stub builder takes tokenizer_config.
• The checks:
  • Modules, model wrappers (DDP, Float16Module), DDP buffers and bucket groups, the optimizer and its grad-stats group, and the inference wrapper keep groups with the collection's ranks. None of them keeps WORLD or a group object of the global grid, directly or in a list, tuple, dict or collection. Several groups have the same ranks in both layouts, so only the object tells them apart. The inference context resolves the model's collection.
  • DP replicas are equal after the build and after training. The TP shards of the built model differ across the collection's TP group, which a broadcast over WORLD would break.
  • Per-DP-rank losses match the same run on the global grid. This catches a global group whose ranks differ from the collection's, but not WORLD, which spans the same ranks in both layouts.
  • A replica's loss and unreduced gradients do not change when the other replicas take other batches. A forward or backward collective over WORLD breaks this.
  • The checkpoint's replica ids, offsets and optimizer keys follow the collection. The checkpoint loads into a new model and optimizer, and into a model on the global grid.
  • The inference model reproduces the trained model's logits.
• Both runs are seeded with training's _set_random_seed(..., pg_collection=...), as initialize_megatron() does, so the CPU generator that initializes the learned position embeddings is seeded too.
• The unit-test conftest sets NVTE_FLASH_ATTN and NVTE_FUSED_ATTN before every test, and since Make MIMO vision encoder attention backend configurable #7362 set_attention_backend() refuses values that disagree with the config's backend. The test clears them before each model it builds.

Tests.

• Before the sync, on 4 CPU ranks with Gloo, the 2 CPU tests passed on every rank and the 8 CUDA tests skipped. With the record reset for the grad-stats fallback removed, test_the_block_rejects_global_reads_and_fallbacks failed on every rank.
• Before the sync, on a CPU stand-in for the CUDA phases (not committed), the identity check reported a module that keeps the global grid's cp group or WORLD, and the shard check reported a broadcast over WORLD.
• On 4 GPUs against its base, test_process_group_collection_end_to_end.py passes 10 of 10. The file is new, so all 10 are new tests. The same job runs 12 other test files of the series (data paths, dataset providers and ownership, get_batch, Megatron-FSDP MoE, optimizer parameter names, the high-level inference API, example process groups, owner linears and the parallel_state deprecation tests): 637 tests pass here and 637 on the base, and none fail. There are no regressions against the base.
• In the integration run (main 77391e7 with all 93 PRs of [TRACKING] Deprecate parallel_state: explicit process groups in megatron/core #6307 and 6 main-based fixes merged, against main), this PR's test files show no regressions against main. There test_process_group_collection_end_to_end.py passes 10 of 10 on 4 GPUs.
• Before the sync, a stacked version of this PR ran on GPUs. Its base also had Use the shared expert's TP group in the fused SwiGLU path #8024, Scatter LLaVA's language-model input over the model's TP group #8025 and a test-only follow-up of Take the FP8/FP4 amax reduction group from a process-group collection #8034, which this base does not have. The 4-GPU run above shows that the test passes without them. On the stacked version the test passed 10 of 10 on 4 GPUs and on 8 GPUs (two nodes), with and without -W error::megatron.core.process_groups_config.ProcessGroupFallbackWarning. Run in one process after test_emerging_optimizers.py, all 113 tests passed. Before the NVTE fix, 3 of the 10 failed there because of the conftest variables.

Compatibility.

• Test only. It needs CUDA and an even world size of at least 4, and skips otherwise. CI's catch-all unit bucket runs it on 8 GPUs, as TP2 x DP4.
• The tolerances assume repeatable kernels for identical shapes: rel 1e-5 for the step-0 loss across layouts, abs 1e-3 for the later losses, rel 1e-5 for a replica's loss between two passes, a relative-norm change of 1e-3 for its gradients, and 2e-2 for the inference logits.
• The result checks gather per-rank problem lists and fail on every rank. A phase that raises on only some ranks still leaves the others in a collective until the NCCL timeout.
• determinism-exempt is needed only while the PRs below this one are unlanded. CI diffs against main, and the base contains registered kernel files. This PR's own diff passes the coverage check.

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.

Later on 2026-11-01 the branch also merged main bfab98f. Its #8471 moves FullyShardedDataParallelV1 and FullyShardedDataParallelV2 from mcore_fsdp_adapter.py into mcore_adapter_v1.py and mcore_adapter_v2.py, and keeps mcore_fsdp_adapter.py as the factory that re-exports them. The merge carries its bases' edits to those classes where the code moved; this PR's own diff does not change. On 4 GPUs against its new base, the Megatron-FSDP adapter tests (7 files) show no regressions: 81 passed here, 81 on the base. The two CUDA-graph end-to-end cases of test_mcore_fully_sharded_data_parallel.py fail on both sides, as they do on main.

Part of #6307.

🤖 Generated with Claude Code

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

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