Repository navigation
Conversation
Signed-off-by: Yan Xu <yxu1@nvidia.com>
A collection without an embd field, or tied pipeline embedding/output stages without a usable embedding group, is a caller error. Report it with ValueError instead of an assert, which is skipped under python -O. Signed-off-by: Yan Xu <yxu1@nvidia.com>
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: 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>
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 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 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>
Rebasing onto main (c9b53d0) brought in nine constructor calls written against the old fallbacks, which this series has removed: - YarnRotaryEmbedding x2 in test_mla_yarn_rope_apply.py: cp_group - SelfAttention and RotaryEmbedding in test_attention_packed_seq.py: pg_collection / cp_group (the build-time CP group the test already reads) - MLASelfAttention in test_multi_latent_attention.py: pg_collection - MultiTokenPredictionBlock x4 in test_multi_token_prediction.py: pg_collection with cp/tp/pp, since the block now dereferences pg_collection.pp as well Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Two call sites the series left behind: - pretrain_vlm.py built LLaVAModel without a collection although model_provider already accepts pg_collection. Resolve the global-grid fallback at the script boundary and forward it, as examples/multimodal/model.py does. - The hybrid CP scheduler tests monkeypatched hybrid_cp_schedule.parallel_state, which the scheduler no longer imports. Hand them a fake pg_collection with tp and dp_cp groups instead and pass it through; the TP broadcast source now comes from torch.distributed.get_global_rank, so map that to the group rank too. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
TEGroupedLinear now requires an explicit process-group collection. The mixed-split determinism test builds one directly; pass the global collection there as the other constructions in this file do. Signed-off-by: Yan Xu <yxu1@nvidia.com>
Fallbacks to the global process groups warn through warn_global_process_group_fallback, but they emit a plain FutureWarning with one hard-coded deprecation window, and each component validates its collection in its own way. ProcessGroupCollection.__getattr__ returns None for a field that was never set, so a component cannot tell an absent field from an axis that is off, and hasattr() is always true. - ProcessGroupFallbackWarning subclasses FutureWarning, so existing filters and tests keep matching it, and CI can turn only the fallbacks into errors. - warn_global_process_group_fallback takes deprecated_in and removed_in per call; the defaults keep every existing call site unchanged. It warns once per (owner, argument) pair, so a caller that omits two arguments of one owner hears about both. A warning that an error filter raises is not recorded, so strict mode raises on every fallback, not only the first. - require_process_groups() returns the named fields in order. An absent field or NON_GROUP_MEMBER raises a ValueError that names the owner and every offending field; an explicit None (axis off, or not a member) is allowed. A MultiModuleProcessGroupCollection raises TypeError, because only the caller knows which module's groups apply. - resolve_process_groups() is the one sanctioned fallback: when the collection is omitted it warns, builds only the required fields from parallel_state, and validates them. An explicit collection is validated and returned without a warning. No call site changes here. Depends on NVIDIA#6293. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
resolve_process_groups() reads the global grid when the caller omits its groups, but that read happens inside process_groups_config.py, which the ratchet exempts. Wrapping an existing fallback in the helper would therefore look like a removed read. Count each call in megatron/core as a "guarded" entry, keyed by its enclosing scope like the shim. A guarded call that replaces a removed accessor or shim read in the same file and scope is a conversion: --update swaps the entry and the total does not grow. Any other new guarded call fails like a new read, and a silent read cannot replace a guarded call. --stats lists the guarded count next to the accessor and shim counts. The tree has no guarded calls yet, so the allowlist is unchanged (173 reads: 152 accessor, 21 shim). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
MoELayer builds its process groups from the global grid when the caller passes no collection, through get_default_pg_collection(). The MoE metrics tracker and the expert-bias update fall back to global groups as well. None of these fallbacks warns, so callers get no signal before the fallbacks are removed. A collection that lacks a field the MoE layer reads is not noticed at construction either: the layer, its router and its token dispatchers read None, which becomes a size-1 group, a crash in the forward pass, or a collective over the default group. get_default_pg_collection() goes through resolve_process_groups with the fields it set before: the seven that MoELayer, its router and its token dispatchers read, plus the experts' expt_dp and expt_dp_gtp_remat. The shim reads them with the same accessors and flags, so the default path selects the same groups. The fallback warns once per owner; MoELayer passes its class name. It is deprecated in 0.21 and removed in 0.23. The shim returns None for a group that parallel_state has not created, where the accessors asserted, so the fallback raises a RuntimeError in that case to keep failing at construction. MoELayer checks the seven layer fields with require_process_groups, so a collection without one of them raises a ValueError at construction. The GTP-inclusive groups are not required, because the inference shard collections do not set them. Three existing tests built MoE layers from collections without some of the layer fields. They set them: groups of the standard grid, or None where the test's dummy router and dispatcher read no group. MoEMetricsTracker.report and get_updated_expert_bias warn when they use the global groups, and keep the groups they use today. Depends on NVIDIA#8049. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
get_default_pg_collection() reads the MoE groups through one guarded resolve_process_groups call instead of nine accessor calls. The checker records the conversion: 173 reads become 165 (accessor 152 to 143, guarded 0 to 1). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Connor-XY
force-pushed
the
yx/pstate-b7-moe-defaults
branch
from
October 26, 2026 22:26
ef2b75a to
56cdf4f
Compare
Contributor
Author
|
/ok to test 56cdf4f |
The fallback warnings said they were deprecated since Megatron Core 0.20 and would be removed in 0.22. No release emits them yet: neither core_r0.20.0 nor main has the warning code, and package_info is at 0.21.0, so they first ship in 0.21. Removing them in 0.22 would leave users a single release of warnings, short of the N+2 policy in docs/api-backwards-compatibility-check.md that the comment cites. Set the defaults to 0.21 and 0.23, say in the comment why, and check the dated text in the warning test instead of comparing the message with the constants it is built from. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
The device fixture read torch.cuda.current_device(), which is cuda:0 on every rank until Utils.initialize_distributed binds the rank's GPU. pytest resolves device before vision_cp_group, which is what first initializes distributed when the file runs on its own. So the first vision context-parallel test put every rank's tensors on cuda:0, and NCCL failed with "Multiple ranks detected using the same GPU", then aborted the communicators for the tests after it. The fixture now initializes distributed before it reads the device. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…anch NVIDIA#6293 now dates the global process-group fallbacks from Megatron Core 0.21, the first release that emits them, with removal in 0.23. Take its defaults and keep this branch's note that a fallback which starts to warn in a later release passes its own versions. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
NVIDIA#6293 now dates the fallback defaults from 0.21, the first release that emits them, with removal in 0.23. The docstring of warn_global_process_group_fallback still said that a fallback which starts to warn after 0.20 passes its own versions; that cohort now starts after 0.21. Two tests checked that explicit versions reach the message by passing 0.21 and 0.23. Those now equal the defaults, so the tests would pass even if the arguments were ignored. Pass 0.22 and 0.24 instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
NVIDIA#8049 now dates the global process-group fallbacks from Megatron Core 0.21, with removal in 0.23, and aligns its docstring and explicit-version tests with those dates. Merge its current head so this branch builds on it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
This was referenced Oct 30, 2026
Connor-XY
added a commit
to Connor-XY/Megatron-LM
that referenced
this pull request
Nov 1, 2026
NVIDIA#8073 makes the MoE layer warn when it falls back to the global process groups, validates MoELayer collections, and adds test_moe_process_groups.py. The GTP remat-group warning that follows edits that test, so merge NVIDIA#8073 here. Only the process-group usage allowlist conflicts, on its total. It is regenerated on the merged tree: 162 reads (140 accessor, 19 shim, 3 guarded). It keeps NVIDIA#8073's guarded get_default_pg_collection read and the reads that NVIDIA#8057 and NVIDIA#8060 remove or guard. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Connor-XY
added a commit
to Connor-XY/Megatron-LM
that referenced
this pull request
Nov 1, 2026
…#8734 NVIDIA#8734 makes the optimizer, seeding, duplicate-filter and silent group helpers warn when they fall back to the global grid, and adds the owner= fallback warning of resolve_gtp_remat_group with its TestGlobalFallbackWarnings tests. The change that follows rewrites that warning and those tests, so merge it here. layers.py merges without a textual conflict, but NVIDIA#8060 removes the global TP rank read of _initialize_affine_weight_cpu and the branch removes the one of param_is_not_tensor_parallel_duplicate, so neither side uses get_tensor_model_parallel_rank any more. Drop it from the parallel_state import; pylint's unused-import check would flag it otherwise. The process-group usage allowlist conflicts on its total. It is regenerated on the merged tree: 152 reads (128 accessor, 18 shim, 6 guarded). It keeps the branch's guarded seeding, duplicate-filter and gtp_replica_rank reads and NVIDIA#6259's get_megatron_optimizer shim, and the reads that NVIDIA#8057, NVIDIA#8060 and NVIDIA#8073 remove or guard. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Connor-XY
added a commit
to Connor-XY/Megatron-LM
that referenced
this pull request
Nov 1, 2026
The checker classification that follows requires every allowlisted read in megatron/core to be warned or guarded. NVIDIA#8769 warns before resolve_gtp_remat_group reads the global grid, and brings NVIDIA#8734 (optimizer, duplicate-filter and tensor-parallel seeding fallback warnings), NVIDIA#8057, NVIDIA#8060 and NVIDIA#8073. Five files conflict. Each is resolved as follows: - optimizer/optimizer.py, process_groups_config import: take NVIDIA#8752's require_process_groups and resolve_process_groups (92d2030) and NVIDIA#8734's _FALLBACK_WARNING_SKIP_PREFIXES, ProcessGroupFallbackWarning, resolve_gtp_remat_group and warn_global_process_group_fallback (ec801ee). - generalized_tensor_parallelism.py: keep NVIDIA#7029's TeardownStage and register_model_parallel_teardown import, which NVIDIA#8767 brings, and NVIDIA#8734's resolve_process_groups import. - inference_layers.py: InferenceLinear keeps NVIDIA#8035's comment and pg_collection=pg_collection, which NVIDIA#8755 brings, and passes NVIDIA#8734's owner=type(self).__name__ to resolve_gtp_remat_group. - tensor_parallel/layers.py: NVIDIA#8734 stops calling the four parallel_state rank and size getters, so their import goes; NVIDIA#8734's process_groups_config import stays, and get_global_memory_buffer comes from megatron.core.utils as NVIDIA#8730, which NVIDIA#8767 brings, has it. - tests/unit_tests/ssm/test_hybrid_block.py get_pg_collection: the required groups are the union of NVIDIA#8035's embd, tp_cp and tp_dp_cp and NVIDIA#8073's ep, expt_tp and tp_ep, with both comments. The process-group usage allowlist is regenerated on the merged tree. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Conflicts: - examples/t5/pretrain_t5.py: Main adds `cfg = get_run_config()` to model_provider and drops the `build_tokenizer` import in favour of `get_tokenizer()`. The PR moves that import for isort and passes a ProcessGroupCollection into T5Model. The result keeps main's `cfg` line and the PR's pg_collection fallback and argument. It drops the `build_tokenizer` import, which would otherwise be unused. - megatron/core/models/common/embeddings/rotary_pos_embedding.py: Main imports `register_capture_constant` for its shared rotary cache. The PR imports `warn_global_process_group_fallback`. Both imports are kept, along with main's cache refactor and the PR's warning cp_group fallback. - pretrain_gpt.py, pretrain_hybrid.py: Main (NVIDIA#8341, NVIDIA#8592) moves `is_dataset_built_on_rank` into megatron/training/datasets/providers.py and the dataset config builder into argument_utils. The PR adds `pp_group` to the `mtp_on_this_rank` call in the local predicate. The result takes main's layout and keeps the PR's `pp_group` in get_batch. Adapted to main outside the conflicts: - megatron/training/datasets/providers.py: The PR's predicate edit is re-applied in main's shared `is_dataset_built_on_rank_explicit`. It passes `pp_group=mpu.get_pipeline_model_parallel_group()` to `mtp_on_this_rank`, so middle pipeline stages do not hit the fallback warning. The function's signature is unchanged. The process-group allowlist records the reads main added since this PR was written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
The process-group allowlist records the reads main added since this PR was written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Adapted to main outside the conflicts: - tests/unit_tests/models/test_hybrid_mhc.py: NVIDIA#7521 adds test_full_recompute_real_moe_matches_forward_backward, which builds real MoE layers from a process-group collection without tp_cp and tp_dp_cp. This PR makes MoELayer require those groups, so the test now requests them, matching the sibling test test_real_moe_raw_branch_forward_backward. The process-group allowlist records the reads main added since this PR was written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
This was referenced Nov 3, 2026
Contributor
Author
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
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.
Depends on #8049 (and #6293). Please review the diff from #8049's head: scoped diff.
Why. Without a
pg_collection,MoELayersilently builds its groups fromparallel_state, andMoEMetricsTracker.reportandget_updated_expert_biassilently use global groups. A collection that omits a field the layer, router or token dispatchers read yieldsNone: a size-1 group, a crash in forward, or a collective over the default group.What.
get_default_pg_collection(owner=...)goes throughresolve_process_groups, warns once per owner (deprecated in 0.21, removed in 0.23), and builds the same nine fields with the same accessor flags, so the default groups are unchanged. Ifparallel_statewas never initialized it raises, as the accessors did.MoELayervalidatesep,tp,cp,expt_tp,tp_ep,tp_cpandtp_dp_cpwithrequire_process_groups. GTP-inclusive groups stay optional.dp_cp_gtp_remat) andget_updated_expert_bias(no group) warn and keep their groups.test_full_recompute_real_moe_matches_forward_backward, came from main (Add mHC precision controls, mean output contract, and full activation recompute with mHC #7521) and now requeststp_cpandtp_dp_cplike its siblingtest_real_moe_raw_branch_forward_backward.Tests. New
test_moe_process_groups.py: identity with the global groups on five grids (including GTP and EGTP), one warning per owner, an explicit collection with every global accessor patched to raise, missing fields, and the logging and expert-bias fallbacks. Before the sync, 21 of its 23 tests failed on #8049's head. On the pre-sync branch, the MoE layer, router, dispatcher, aux-loss, expert and hybrid suites matched #8049's head on 4 GPUs. In the integration run (main 77391e7 with all #6307 PRs merged), this PR's test files show no regressions against main. On 4 GPUs,test_hybrid_mhc.pypasses 43 tests,test_hybrid_block.py67 andtest_latent_moe_layer.py15, the same counts as main.test_moe_process_groups.pypasses all 23 tests there; the file is new, so main has no results for it.Compatibility. Omitting the collection warns. A collection without a layer field raises; Megatron Bridge's builders set all seven.
The branch was synced with main 77391e7 on 2026-11-01 by a merge commit that brings in #8049 after its own update to main and lists each conflict and its resolution, so existing commits and review links stay valid.
Part of #6307.
🤖 Generated with Claude Code