[TRACKING] Deprecate parallel_state: explicit process groups in megatron/core
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
Tracking issue for deprecating `megatron.core.parallel_state` and moving `megatron/core` to explicit process-group passing.
## Why
`parallel_state` holds the process groups for a **single, global** parallel grid. For a job with one grid, reading them is merely deprecated. For a job that builds **independent** grids — a vision encoder and an LLM with different parallelism, GTP, MIMO — the accessors return a group belonging to the *wrong* grid. Nothing raises. The collective runs on the wrong ranks, or a rank is seeded with the wrong RNG offset, and the job produces wrong numbers.
This is not hypothetical. Three gradient-counting bugs trace to this mechanism:
- #5916 (merged) — gradient-norm undercounting with EP + TP
- #6080 — gradient reduction when EP=1, EP=TP, EGTP != GTP
- #6099 — *"gradient-norm duplicate filtering fell back to the global TP group, incorrectly excluding some expert parameters when ETP differed from TP"*
Baseline on `main` today: **207 global process-group reads across 54 files** in `megatron/core` (178 accessor reads + 29 `use_mpu_process_groups` calls), measured by the checker in #6258.
## Open PRs
Suggested landing order — #6257 first, since the later descriptions lean on its vocabulary.
| PR | What | Size |
|---|---|---|
| #6257 | Document the deprecation and its replacement (docs only) | 4 files, +414/-45 |
| #6258 | CI ratchet: fail on *new* global reads, 198 grandfathered | 4 files, +686/-0 |
| #6259 | Give gloo process groups a home in `ProcessGroupCollection` | 4 files, +91/-29 |
| #6293 | Remove global reads from `megatron/core` (207 → 155) | 44 files, +492/-215 |
| #6303 | Require an explicit collection in `LanguageModule` | 62 files, +402/-119 |
#6257 and #6258 carry no behavioural risk. #6259 changes a core datastructure's contract. #6293 and #6303 are **breaking for external callers** that construct core classes without a collection — see below.
Each change was validated with a **paired run against unmodified `main` at the same commit, in the same container**, comparing failure *sets* rather than counts.
## Decisions needed
1. **Breaking change** — #6293 and #6303 make ~10 public classes assert instead of silently defaulting to the global grid. Downstream impact measured: Megatron-Bridge needs **2** call-site edits (its real model paths already thread `pg_collection`); nemo-rl-internal needs none. Needs a release-note decision.
2. **Make #6258 a required check?** Without it, reads regrow — `main` added 9 in the 302 commits while this work was in flight.
3. **Coordination with #6234** (GTP `pg_collection` cleanup) — 7 files overlap with #6293. The concerns are complementary (that PR threads `gtp_remat`; this one covers pp/cp/tp/dp) but they will conflict textually. Happy to rebase behind it.
## Still to do
- **~155 remaining reads.** Largest remaining blocks: `TransformerBlock` (17 callers), `TransformerLayer` (9), `cuda_graphs.py` (16, needs `pg_collection` plumbed into `MambaLayer` first).
- **Runtime fallback reporting** — instrumentation that warns once per site and counts every occurrence, so real reliance on global state can be measured before more fallbacks are removed. Built, not yet submitted.
- **Tier-4 globals** — the leftovers in `parallel_state` that are not process groups. Mostly resolved: the virtual-pipeline *rank* migration to `vp_stage` was already complete; virtual-pipeline *size* has a home in `TransformerConfig`; `GlobalMemoryBuffer` is a shared GPU scratch allocator that should stay a singleton but move out of `parallel_state`. NCCL options were never global. Branch pending.
- **Deleting `parallel_state.py` outright is explicitly out of scope.** Megatron-Bridge holds 785 references — more than `megatron/core`'s 207. The realistic end state is demotion to bootstrap-only (`initialize_model_parallel` / `destroy_model_parallel` / `is_initialized`).
## Guidance for contributors, now
Until this lands, the rule from #megatron-core-developments applies:
- **New features must not** call `parallel_state.get_*_group()` / `get_*_rank()` / `get_*_world_size()` in `megatron/core`. Accept a `ProcessGroupCollection` or an explicit `torch.distributed.ProcessGroup` and pass it through.
- **`ProcessGroupCollection.use_mpu_process_groups()` is not a migration target** — it reads the same global state, so swapping one for the other is a lateral move.
- **Bug fixes may leave existing calls alone.** Process-group plumbing belongs in its own PR.
Full write-up in `docs/developer/parallel-state-deprecation.md` (added by #6257).
Contributor guide
Research direction
Start with docs/developer/parallel-state-deprecation.md and review the open PRs #6257, #6258, #6259, #6293, and #6303 to understand the migration plan. The remaining work is spread across TransformerBlock, TransformerLayer, and cuda_graphs.py, with MambaLayer plumbing still needed; completion requires coordinated removal of the listed global reads rather than one self-contained change.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 15/100