Remove global process-group reads from megatron/core (207 -> 155) - #6293
Draft
Connor-XY wants to merge 25 commits into
Draft
Remove global process-group reads from megatron/core (207 -> 155)#6293Connor-XY wants to merge 25 commits into
Connor-XY wants to merge 25 commits into
Conversation
Six sites in megatron/core called parallel_state accessors even though the correct group was already stored on the object (self.tp_group / self.cp_group). Pure substitution: no signature changes, no new arguments, no caller impact. SelfAttention.run_realtime_tests 3 sites -> self.tp_group SelfAttention.get_query_key_value_tensors 1 site -> self.tp_group MultiTokenPredictionLayer 1 site -> self.tp_group LLaVAModel 1 site -> self.cp_group Global process-group reads in megatron/core: 163 -> 157. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
…scope Two more sites where the correct group was reachable without any API change. transformer_engine.py: in the cp_comm_type == 'a2a+p2p' branch the code read the global hierarchical CP groups, even though the caller-supplied branch directly above asserts hasattr(pg_collection, 'hcp') for exactly that case. The assert was already guaranteeing the group the code then declined to use. language_module.py: the word-embedding replica id derived its DP-CP rank from the global grid while the sharded tensor nine lines later is built against metadata['dp_cp_group']. Now both use the same group, so replica id and sharding agree under a caller-supplied grid. Global process-group reads in megatron/core: 157 -> 155. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
megatron/core/models/multimodal/context_parallel.py hard-coded the global CP grid in 11 places, across module-level helpers with no group parameter. This is the vision-encoder path, so it is exactly where an independent encoder grid silently used the LLM's groups instead. cp_group is now a required argument on the public helpers and is threaded through the private ones and the autograd Function. get_pg_size/get_pg_rank preserve the previous behaviour when CP is disabled (size 1, rank 0). The call graph is closed inside this repo: LLaVAModel already holds self.cp_group, and the only other callers are unit tests. The module is not re-exported and has no users in Megatron-Bridge or nemo-rl. Global process-group reads in megatron/core: 155 -> 144. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
_gather_from_tensor_parallel_region read the global TP grid for its size, rank and both collectives. Both callers are subclasses that already hold the group (self.tp_group on ColumnParallelLinear, self._tp_group on the TE variant), and the mapping helpers it calls already accept an explicit group. Global process-group reads in megatron/core: 144 -> 141. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
… in tests Caught by a distributed run on draco: changing the autograd Function signature missed one direct .apply() call site in the unit tests. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
hybrid_context_parallel_forward_backward read the global TP and DP-CP grids for its broadcast source rank, broadcast group, rank checks and barrier. Its only caller is forward_backward_no_pipelining, which already holds pg_collection. dp_cp is asserted in the hybrid-CP branch rather than unconditionally, matching how TEDotProductAttention asserts hcp only for cp_comm_type == 'a2a+p2p'; the legacy schedule collection already requests dp_cp. get_tensor_model_parallel_src_rank() becomes torch.distributed.get_global_rank(tp_group, 0), which is its definition. hybrid_cp_schedule.py no longer imports parallel_state at all. Global process-group reads in megatron/core: 141 -> 136. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
run_realtime_tests compared Q/K layernorm parameters across the global DP grid. The DP group is required only by this debug check, so it is asserted there rather than in __init__, which asks use_mpu_process_groups for tp/cp only -- the same idiom TEDotProductAttention uses for hcp. megatron/core/transformer/attention.py no longer references parallel_state. Global process-group reads in megatron/core: 136 -> 133. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
reduce_loss_in_tracker read the global PP and DP grids for its two all-reduces. The chain is a single hop: training.py -> track_indexer_metrics -> reduce_loss_in_tracker, all within this repo, so the groups are supplied at the edge where reading parallel_state is allowed. Global process-group reads in megatron/core: 127 -> 125. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
mtp_on_this_rank read the global PP rank to decide whether MTP layers live on this stage. Both model callers (GPTModel, HybridModel) already hold self.pp_group from LanguageModule.__init__, which runs before the call. Global process-group reads in megatron/core: 125 -> 124. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
RotaryEmbedding and MultimodalRotaryEmbedding fell back to the global CP grid via a ternary when cp_group was not passed. 10 of the 22 constructor calls in this repo already passed one; the remaining 12 (3 in core, 9 in tests) now do too, so the fallback is gone rather than merely unused. megatron/core/models/common/embeddings/rotary_pos_embedding.py no longer references parallel_state. Global process-group reads in megatron/core: 124 -> 122. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
…nd T5Model Both classes fell back to ProcessGroupCollection.use_mpu_process_groups() when no collection was passed. The fallback is now an assertion naming the missing argument, which is the 'fails loudly rather than silently' half of the goal. DotProductAttention: already dead in-repo -- Attention always passes pg_collection when building core_attention, and the one direct test construction passes it positionally. T5Model: 4 of 5 constructors did not pass one (3 tests, 1 example); they now do. This is a deliberate breaking change for external callers that construct these directly without a collection, which is the intent of the migration. Global process-group reads in megatron/core: 122 -> 120 (shims 20 -> 18). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
…lock The use_mpu_process_groups fallback becomes an assertion naming the missing argument. All 7 constructors without one live in a single test module and now pass a collection. Global process-group reads in megatron/core: 120 -> 119 (shims 18 -> 17). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
LLaVAModel is the case the migration exists for: a vision encoder and an LLM may run on independent parallel grids, so falling back to the global grid silently picks the wrong one. The fallback becomes an assertion that says exactly that. All 8 constructors without a collection (7 tests, 1 example) now pass one. Global process-group reads in megatron/core: 119 -> 118 (shims 17 -> 16). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Attention, MultiLatentAttention (x2) and AbsorbedMLASelfAttention each fell back to use_mpu_process_groups when no collection was passed. All are dead in-repo: TransformerLayer sets attention_optional_kwargs['pg_collection'] unconditionally before build_module, and every direct constructor already passes one (verified by AST across megatron/, tests/ and examples/). Global process-group reads in megatron/core: 118 -> 114 (shims 16 -> 12). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Global process-group reads in megatron/core: 114 -> 113 (shims 12 -> 11). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
These were based on a faulty inventory. My constructor scan filtered on
len(node.args) < 2, which silently skipped every call site passing arguments
positionally, so 'no caller is missing pg_collection' was wrong.
A paired draco run showed 85 real failures:
63x AssertionError: MultiLatentAttention requires an explicit pg_collection
22x AssertionError: Attention requires an explicit pg_collection
10x TypeError: DSAIndexer.__init__() got multiple values for 'pg_collection'
(my inserted keyword duplicated an existing positional argument)
Reverting to the last validated state. The earlier fallback removals
(DotProductAttention, T5Model, MultiTokenPredictionBlock, LLaVAModel) were each
validated against a pristine baseline and are retained.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Redo of the reverted attempt, this time with a signature-aware inventory: both DSAIndexer constructors already pass pg_collection positionally, so the fallback is dead and the tests need no change. The previous attempt added a redundant keyword and produced 'got multiple values for argument pg_collection'. Global process-group reads in megatron/core: 118 -> 117 (shims 16 -> 15). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Second attempt, this time driven by the signature-aware scanner: 15 of 16 SelfAttention constructors were missing pg_collection (all in test modules) and now pass one. The earlier attempt used a scan that skipped positional call sites and had to be reverted. Global process-group reads in megatron/core: 117 -> 116 (shims 15 -> 14). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
…sses Second attempt, driven by the signature-aware scanner. MultiLatentAttention (x2) and AbsorbedMLASelfAttention no longer fall back to use_mpu_process_groups. 18 of 19 MLASelfAttention constructors were missing pg_collection, all in one test module, and now pass one; AbsorbedMLASelfAttention's single constructor already did. Global process-group reads in megatron/core: 116 -> 113 (shims 14 -> 11). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Paired draco run: 24 failed / 43 passed against a pristine 20 / 57, with 110 'requires an explicit pg_collection' assertions. The failures are concentrated in TestFusedMLA*, i.e. a MultiLatentAttention subclass my scan never covered -- I scanned MLASelfAttention and AbsorbedMLASelfAttention by name and assumed those were the only constructors reaching the base __init__. The Attention removal in the preceding commit validated clean (5F/121P on both trees) and is retained. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
…sses Third attempt. The previous two failed because my inventory was incomplete: the first scan skipped positional call sites, the second scanned MLASelfAttention and AbsorbedMLASelfAttention by name and missed FusedMLASelfAttention, which also reaches MultiLatentAttention.__init__ (110 assertion failures, all in TestFusedMLA*). This time the subclass list was derived programmatically from the class hierarchy: MLASelfAttention (18 sites), FusedMLASelfAttention (6 sites), AbsorbedMLASelfAttention (0). All now pass a collection. Global process-group reads in megatron/core: 116 -> 113 (shims 14 -> 11). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
The only construction path is GroupedMLP.__init__ (experts.py), which dereferences pg_collection.ep and .expt_tp unconditionally before passing it on -- so it cannot be None there. No direct constructors exist anywhere in the repo. Global process-group reads in megatron/core: 113 -> 112 (shims 11 -> 10). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
dynamic_engine and abstract_model_inference_wrapper each derived a collection from global state when the inference config carried none. Both now call InferenceConfig.resolve_pg_collection(), so the fallback exists once at the config boundary and consumers can assume a collection. Resolution is lazy rather than in __post_init__, since configs are commonly built before initialize_model_parallel. Also drops a stale comment in attention.py referring to a fallback that no longer exists. Global process-group reads in megatron/core: 110 -> 109 (shims 10 -> 9). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
…ries Two independent reductions of the same shape as the MoE one: many global reads become one, with no signature or caller changes. distributed_trtllm_model_weights_converter: five parallel_state reads (pp size, tp size, tp rank, pp rank, tp group) become one use_mpu_process_groups call with sizes and ranks derived from it. The file no longer imports parallel_state. inference/communication_utils: broadcast_tensor fetched the model-parallel source rank and the group separately; the source rank is rank 0 of that group, so one lookup gives both. broadcast_list called get_model_parallel_src_rank twice for the same value. Global process-group reads in megatron/core: 102 -> 96 (accessors 92 -> 85). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Rebasing onto current main surfaced three FusedMLASelfAttention constructors added since this work began. They reach MultiLatentAttention.__init__, which this branch makes require an explicit collection. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Contributor
Author
|
Part of #6307 — tracking issue for the |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Removes reads of global process-group state from
megatron/core, so a model built on independent parallel grids gets the right groups instead of silently inheriting the global ones.207 → 155 global reads in
megatron/core(accessor178 → 136,use_mpu_process_groups29 → 19), measured with the checker from #6258.Motivation: three gradient-counting bugs — #5916 (merged), #6080, #6099 — all trace to a global-grid fallback. #6099 states it directly: "gradient-norm duplicate filtering fell back to the global TP group."
Two kinds of change
1. Thread the group where the call graph already has it. No signature change; the caller already holds a collection.
models/multimodal/context_parallel.py(11 sites) — the vision-encoder CP path, i.e. exactly where an independent encoder grid silently borrowed the LLM's groups.cp_groupis now required on the public helpers and threaded through the private ones and the autogradFunction.pipeline_parallel/hybrid_cp_schedule.py(5) — file no longer importsparallel_statetransformer/attention.py(3) — file no longer referencesparallel_statetransformer/heterogeneous/linear_replacements.py(3),multi_token_prediction.py,extensions/transformer_engine.py, and others2. Require an explicit collection where the fallback was already dead in-repo. Each becomes an assertion naming the missing argument — the "fail loudly" half of the goal:
DotProductAttention,T5Model,MultiTokenPredictionBlock,LLaVAModel,DSAIndexer,TEGroupedLinear,Attention,MultiLatentAttention,AbsorbedMLASelfAttention, rotary embeddings.These are breaking changes for external code constructing those classes directly without a collection. In-repo and example call sites are updated here.
Verification
Every change was validated with a paired run against unmodified
mainat the same commit, in the same container — comparing failure sets, not just counts.Latest, 8×GPU:
mainZero regressions; none of the added assertions fire. Suites:
test_attention.py,test_multi_latent_attention.py,test_multimodal_context_parallel.py,test_llava_model.py.Note for reviewers
Rebasing onto current
maindropped five commits as superseded — upstream independently added_build_default_pg_collection(),_get_te_ops_tensor_parallel_context(), andwith_gtp_rematplumbing. It also surfaced threeFusedMLASelfAttentionconstructors added since this work began, fixed here.Overlaps #6234 (GTP
pg_collectioncleanup) in 7 files. The concerns are complementary — that PR threadsgtp_rematgroups; this one covers pp/cp/tp/dp — but they will conflict textually. Happy to rebase behind it.Contribution process
Pre-checks