NVIDIA / NVIDIA/Megatron-LM

[TRACKING] Deprecate parallel_state: explicit process groups in megatron/core

Open
#6,307 0 comments 1 reaction 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.