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>
Conflicts: - hybrid_cp_schedule.py was deleted on main (dynamic CP replaced hybrid CP), so the hybrid-CP scheduler change and its tests are dropped. - schedules.py, multi_token_prediction.py checkpoint_handler, pretrain_gpt.py and pretrain_hybrid.py packing path: take main's version, which already passes the collection's pipeline group to mtp_on_this_rank. - test_get_batch.py: keep main's test and also assert that get_batch passes pg_collection.pp to mtp_on_this_rank. - test_multi_token_prediction.py: take main's test and build the MTP block with an explicit collection. - examples/t5/pretrain_t5.py: drop the imports main removed. - Regenerate the process-group allowlist from the merged tree. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> 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>
DistributedTRTLLMModelWeightsConverter takes the tensor- and pipeline-parallel groups of the grid the model was loaded on from the global parallel state, and TRTLLMHelper cannot pass other groups. A model exported from its own grid gets another grid's TP and PP sizes and ranks in its TRT-LLM mapping, and the vocabulary gather runs over the wrong group. The converter and TRTLLMHelper's get_trtllm_pretrained_config_and_model_weights now take an optional pg_collection, which the helper forwards through its distributed-setting method. Without one, the converter warns once with a ProcessGroupFallbackWarning (deprecated in 0.21, removed in 0.23) and reads tp and pp from parallel_state through resolve_process_groups, as before. A collection that leaves tp or pp unset raises ValueError, and one that sets either to None raises the existing RuntimeError, so the converter never exports a TP1/PP1 mapping by accident. The distributed export example passes its model's collection. Depends on NVIDIA#8049. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Record the converter shim read that became a guarded resolve_process_groups call. The total stays at 173, now with 1 guarded read. 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-c-export-converter-groups
branch
from
October 26, 2026 22:26
e39d7fc to
a03e035
Compare
Contributor
Author
|
/ok to test a03e035 |
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
The checker classification that follows requires every allowlisted read in megatron/core to be warned or guarded, and its test lists the only two reads allowed to stay unwarned. NVIDIA#8059 lets the TRT-LLM distributed converter take the model's process groups and warns before its use_mpu_process_groups() fallback. Without it, DistributedTRTLLMModelWeightsConverter.__init__'s shim read is unwarned and the classification test fails. The merge conflicts only in the process-group usage allowlist, which is regenerated on the merged tree. 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, and its test lists the only two reads allowed to stay unwarned. NVIDIA#8752 warns before get_megatron_optimizer's use_mpu_process_groups() fallback, before clip_qk's global data-parallel group (with NVIDIA#8043, which NVIDIA#8752 brings) and describes the GTP shards in _backfill_gtp_sharded_param_map with the optimizer's own groups. Without it those four reads are unwarned and the classification test fails. megatron/core/optimizer/__init__.py conflicts on its imports: NVIDIA#8045's e09f5de (which NVIDIA#8749 brings) imports get_ep_rank_and_size next to get_global_unique_param_name, and NVIDIA#8752's 2d4d3cc adds resolve_process_groups to the process_groups_config import. The import takes both. NVIDIA#8752 also carries NVIDIA#6303, which this branch does not have yet, so tests/unit_tests/export/trtllm/test_trtllm_distributed_gpu_converter.py conflicts on its imports: NVIDIA#6303 imports ProcessGroupCollection and NVIDIA#8059 adds ProcessGroupFallbackWarning to the same import. The import takes both. 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>
This was referenced Nov 1, 2026
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>
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.
DistributedTRTLLMModelWeightsConvertertakes its TP and PP groups from the global parallel state, andTRTLLMHelpercannot pass others. A model loaded on its own grid exports another grid's TP/PP sizes and ranks in its TRT-LLM mapping, and the vocabulary gather runs over the wrong group.What.
TRTLLMHelper.get_trtllm_pretrained_config_and_model_weightstake an optionalpg_collection; the helper forwards it to the converter.ProcessGroupFallbackWarning(deprecated in 0.21, removed in 0.23) and readstpandppfrom the global grid as before, throughresolve_process_groups.tporppraisesValueError;Nonegroups keep the existingRuntimeError, so nothing exports a TP1/PP1 mapping by accident.Tests. On 4 GPUs: a TP=2
HyperCommGridcollection of distinct communicators converts, with every global read patched to raise, to the same weights as the global path; the fallback warns once; incomplete collections raise; the helper forwards the collection without needing TensorRT-LLM. These fail on #8049's head. The existing export tests match. In the integration run on 4 GPUs,test_trtllm_distributed_gpu_converter.pyhas 6 passed vs 1 on main: this PR's 4 new test cases pass, and so doestest_requires_model_parallel_groupsfrom #6293.test_trtllm_helper.pyhas 1 passed and 1 skipped: the new forwarding test passes, andtest_exceptionsskips because TensorRT-LLM is not installed. In the integration run, this PR's test files show no regressions against main.Compatibility. New optional arguments only. Callers that pass none get one warning.
The branch was synced with main 77391e7 on 2026-11-01 by a merge commit, which had no conflicts, so existing commits and review links stay valid.
Part of #6307.
🤖 Generated with Claude Code