Repository navigation
Conversation
Contributor
Author
|
/ok to test 9a16919 |
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
sharded_state_dict() called without metadata filled in the global get_data_parallel_group(with_context_parallel=True), and replica ids were computed over it. A model on its own grid, saved by a tool, a converter or a test that passes no metadata, then claimed replicas by another grid's DP ranks: several ranks or none wrote each shard. ensure_metadata_has_dp_cp_group takes an optional dp_cp_group that it uses instead of the global group when the metadata has none; a group already in the metadata still wins. MegatronModule, MLP and TransformerBlock pass their collection's dp_cp_gtp_remat (dp_cp when the collection has no GTP-inclusive group). That is the group the global default stands for, so replica ids on the standard grid do not change. MLP now keeps its pg_collection. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
LanguageModule.sharded_state_dict fills missing metadata for GPT, hybrid and BERT models before their submodules can, so a model built on its own collection still got the global DP x CP group for every replica id below it. Pass the model collection's dp_cp_gtp_remat (dp_cp without GTP), the group the global default stands for. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
make_sharded_tensors_for_checkpoint, make_tp_sharded_tensor_for_checkpoint and make_sharded_tensor_for_checkpoint fell back to the global groups only when both groups were None. A tp_group passed without its dp_cp_group made the DP x CP group a singleton, so every data-parallel rank claimed to be the main replica of each shard. That is now a ValueError. No in-repo caller passes a TP group alone. A dp_cp_group without a tp_group still means the tensors are not tensor-parallel: the GTP checkpoint helper, MIMO's sharded_state_dict for plain submodules and the GTP checkpoint tests rely on it, so the docstrings now say so instead of raising. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
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-b2-ckpt-replica-groups
branch
from
October 26, 2026 22:26
9a16919 to
482ceca
Compare
Contributor
Author
|
/ok to test 482ceca |
Connor-XY
added a commit
to Connor-XY/Megatron-LM
that referenced
this pull request
Nov 1, 2026
The SSM, MoE and TE layers that follow pass their own DP x CP group to ensure_metadata_has_dp_cp_group, which takes that argument from NVIDIA#8044, and their tests reuse NVIDIA#8044's sharded-state DP x CP test helpers, so merge NVIDIA#8044 here. NVIDIA#8044's head already contains main, so the merge is clean and its tree is NVIDIA#8044's. 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
NVIDIA#8044 lets modules fill missing sharded-state metadata with their own DP x CP group and rejects a TP group passed without one to the checkpoint helpers. The owners pass that group, and the checkpoint helpers are to warn when they fall back to the global groups, in the code NVIDIA#8044 changes, so merge NVIDIA#8044 here. The merge is clean. 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 one-warning tests that follow check that a MoELayer built without a collection gives its experts the GTP groups of the global grid, and read them from experts.pg_collection. NVIDIA#8732's 645cd0d makes SequentialMLP and TEGroupedMLP keep their collection; without it that test case fails with "'SequentialMLP' object has no attribute 'pg_collection'". NVIDIA#8732 brings NVIDIA#8044. 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 1, 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.
Why.
sharded_state_dict()called without metadata fillsmetadata['dp_cp_group']from the global DP x CP group. A model on its own grid, as built by tools, converters, tests or exports, then computes replica ids from another grid's ranks, so several ranks or none write each shard. Separately, atp_grouppassed to the sharded-tensor helpers without adp_cp_groupmakes every data-parallel rank a main replica.What.
ensure_metadata_has_dp_cp_group(metadata, dp_cp_group=None): a group already in the metadata still wins; otherwise the given group, then the global default.MegatronModule,MLP(which keeps itspg_collection),TransformerBlockandLanguageModulepass their collection'sdp_cp_gtp_remat, ordp_cpwhen unset: the group the global default stands for. One global read is removed.make_sharded_tensors_for_checkpoint,make_tp_sharded_tensor_for_checkpointandmake_sharded_tensor_for_checkpointraiseValueErrorfor atp_groupwithout adp_cp_group. Adp_cp_groupalone still means "not tensor-parallel". No in-repo caller is affected; the commit can be dropped if the restriction is unwanted.Tests. With the global DP x CP accessor patched to raise, new tests compare
(key, global_offset, replica_id)with the global default for an MLP, a transformer block and a GPT model at TP=2 and on a GTP=2 grid, and check the new error. Before the sync they failed on main and passed here on 4 GPUs, and distributed-checkpointing GPT and MLP, MIMO checkpoint and GTP DCP tests matched main. In the integration run on 4 GPUs, the two new test files give 8 passed and 1 skipped:test_sharded_state_dict_dp_cp_group.py6 passed,test_gpt_sharded_state_dict_dp_cp_group.py2 passed and 1 skipped. The skipped case is the GPT model with GTP-sharded weights, because the container's Transformer Engine is older than GTP needs. Neither file exists on main. In the integration run, this PR's test files show no regressions against main.Compatibility. Replica ids on the standard grid do not change. MoE experts, hybrid/SSM mixers and TE layers still fill the global group; they follow per owner.
Synced with main 77391e7 on 2026-11-01 by a merge commit, so existing commits and review links stay valid. The merge had no conflicts; its message records the one change made for main.
Part of #6307.
🤖 Generated with Claude Code