[Connor-XY]

This PR finishes the core side of the parallel_state deprecation (#6307). The last two unwarned reads of the global grid in Megatron Core move behind their fallback warnings, and resolve_gtp_remat_group warns before it reads the grid. The process-group checker now tells warned reads from unwarned ones, and a new test fails once the package version reaches an announced removal version. Pytest turns a ParallelStateDeprecationWarning raised from this repository's own code outside the tests into an error. The docs describe the CI check, how to build a ProcessGroupCollection without the globals, and common ways a component ends up on the wrong groups. It is for callers who build layers without a collection, and for maintainers of Megatron Core, MoE, the transformer layers and CI. It consolidates four draft PRs unchanged. Each member is one commit, so review commit by commit. A final commit regenerates tools/process_group_usage_allowlist.json. It builds on #6257 and on the units for FP8 amax, CUDA graphs and 1F1B (#8038), attention and layers (#8729), MoE plumbing (#8724), DDP and Megatron-FSDP (#8725), transformer and MoE warnings (#8748, #8732, #8057, #8060, #8073), DDP/FSDP warnings (#8735, #8749), the registry and tools (#8059), checkpoint-helper warnings (#8756), optimizer warnings (#8752, #8734), core helper warnings (#8755) and the accessor deprecation (#8767). It lands after all of them.

Review

The unit's own commits: Connor-XY/Megatron-LM@4df70b1...61b02bb

The full PR diff also shows the earlier units until they merge.

Members

• docs: document the process-group CI check and building a ProcessGroupCollection #7991: docs: document the process-group CI check and building a ProcessGroupCollection. Rewrites the deprecation guide for the landed code, adds pages on building a ProcessGroupCollection from a HyperCommGrid and on 11 failure patterns, and updates the Core QuickStart to pass groups explicitly and train with get_megatron_optimizer, with no runtime change.
• Warn before resolve_gtp_remat_group reads the global grid #8769: Warn before resolve_gtp_remat_group reads the global grid. The fallback warns once per owner and argument before it reads the grid, also with GTP off, so an error filter stops the read; SequentialMLP passes its collection to each expert MLP.
• Make the process-group checks enforce what the deprecation notes claim #8770: Make the process-group checks enforce what the deprecation notes claim. The checker classifies each core read as warned or unwarned (:warned allowlist suffix) and counts the GTP and setup helpers as guarded, the guard outside core follows aliases, getattr, importlib and star imports, and the PR adds a blocking removal-version test and a pytest error filter; the filter is a separate pyproject.toml change that can be dropped on its own.
• Warn on the last unwarned core reads and once per omitted collection #8772: Warn on the last unwarned core reads and once per omitted collection. MoEMetricsTracker._sync_metrics and mtp_on_this_rank read inside their warned branches, leaving 0 unwarned core reads, and Attention, the MTP block and MoE's default collection add the GTP fields to their fallback collections, so an owner built without a collection gives 1 warning instead of 3 to 5.

Behavior changes for callers

• A linear built without pg_collection, or with a collection that lacks gtp_remat or expt_gtp_remat, now emits ProcessGroupFallbackWarning even with GTP off, and raises under --strict-process-group-fallbacks. Set both fields, to None when an axis is off, to avoid it.
• A MoELayer built with pg_collection=None, Transformer Engine grouped GEMM and expert GTP above 1 now shards its experts over the global expert GTP group. Before, it did not. With expert GTP off nothing changes.
• In tests, a ParallelStateDeprecationWarning whose caller is repository code outside the tests becomes an error. Other packages stay ignored.
• The allowlist format gains the :warned suffix. --update converts an old allowlist.
• The removal-version test fails the CI job once megatron/core/package_info.py reaches a version that a warning announces as the removal version.

Teams

ci, core-adlr, core-nemo, mixture-of-experts-adlr, mixture-of-experts-devtech, transformer

Testing

Each member was GPU-tested against its base before consolidation. GTP-on paths were checked separately with a Transformer Engine 2.19 development build: 33 test files, no regressions against main. The consolidated chain (every unit in landing order, on current main) was run against main on 4 GPUs: 258 test files, and no test fails on the chain that passes on main. Tests in 28 files fail on both, mostly tests written for 8 GPUs, and one file times out on both.

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