Repository navigation
Fix GTP checkpoint load failing across different GTP degrees - #7293
Merged
fanshiqing merged 5 commits intoSep 15, 2026
Merged
fanshiqing merged 5 commits into
fanshiqing merged 5 commits into
Conversation
svcnvidia-nemo-ci
marked this pull request as draft
September 14, 2026 05:38
Contributor
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
fanshiqing
marked this pull request as ready for review
September 14, 2026 05:38
fanshiqing
force-pushed
the
shiqingf/fix_gtp_ckpt_load_across_different_gtp_degree
branch
from
September 14, 2026 05:52
60a2ce4 to
f8592c4
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
fanshiqing
force-pushed
the
shiqingf/fix_gtp_ckpt_load_across_different_gtp_degree
branch
from
September 14, 2026 06:04
f8592c4 to
df1fc42
Compare
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Phlip79
reviewed
Sep 15, 2026
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Member
Author
|
/ok to test feeaaa2 |
Phlip79
approved these changes
Sep 15, 2026
deepakn94
approved these changes
Sep 15, 2026
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/34933415391 |
fanshiqing
deleted the
shiqingf/fix_gtp_ckpt_load_across_different_gtp_degree
branch
September 15, 2026 07:46
xuwchen
added a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Sep 22, 2026
…TP correctly (TP2 vs TP1)
xuwchen
added a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Sep 22, 2026
…TP correctly (TP2 vs TP1)
xuwchen
added a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Sep 22, 2026
…ffset With gtp_remat_pad_for_alignment > 0 a GTP weight is allocated padded to pad_for_alignment * gtp_remat_size rows, and every shard -- padding included -- was described to the checkpoint as if it were real data. ShardedTensor.from_rank_offsets computes global_offset[axis] = axis_rank_offset * local_axis_shape, so passing the PADDED shard size shifted every tp_rank > 0 by pad_length: those ranks read their neighbour's rows and their own tail came back as zeros. allow_shape_mismatch, set to absorb the inflated global shape, suppressed the validation that would have caught it. The symptom was a grad norm of 7.75e9 at pad_for_alignment=128; at the default of 16 it still corrupted every weight whose dim0 is not already aligned (vision fc1 pads 24 rows). Describe the checkpoint in LOGICAL layout instead. _make_gtp_logical_sharded_tensor sizes and offsets the global tensor from the unpadded row count and trims this rank's pad tail, so padding never reaches the checkpoint and allow_shape_mismatch is no longer needed. A shard that is entirely padding (vision linear_proj: dim0 1152, alignment 512) contributes an empty slice rather than raising. Trimming changes object identity, which the optimizer depends on -- it maps optimizer params to model entries by id(entry.data). The trimmed shard therefore carries a gtp_pad_src backlink to the live param, resolved by gtp_entry_backlink() in both the distributed-Adam and distributed-Muon lookups. The backlink lives on the ShardedTensor and not on its data, because the torch strategy does sh_ten.data = sh_ten.data.detach(), which drops tensor attributes. A full-length shard is handed back as the SAME object so identity still holds. A padded native-FP8 param points its backlink at the live param rather than at the dequantized copy, or the trimmed rank resolves to a copy the _gtp_dequant_src lookup can no longer see and its optimizer state silently vanishes. Two consumers need the untrimmed shard: an all-gather over the GTP group sizes itself off the short trailing rank and HANGS otherwise (untrimmed_gtp_shard()), and the fully-reshardable optimizer load expresses flat param ranges over the padded shard, so it zero-fills the excluded rows before applying the range. Covered by test_padded_shard_logical_offsets, test_padded_ep_egtp_offsets, test_padded_optimizer_state_mapping, test_padded_fused_projection_gather, test_padded_save_load_roundtrip and test_non_gtp_checkpoint_into_padded_gtp -- the last of which is the discriminating one: it writes with a non-GTP layout and reads into a padded GTP model, and is the only one that fails on the pre-fix mapping (want rows [10..15], got [12..17], +2 = pad_length). The symmetric round-trip cannot catch a self-consistently wrong mapping. Run them with scripts/gtp_dcp_unit_test.sh: the file gates on world_size == 4 and needs the TE >= 2.19 overlay, and getting either wrong skips the whole file while the command still exits 0. Rebase note (dev @ 0e704ed): the base now carries NVIDIA#7293, which handled the same symptom by keeping the padded global and waiving the shape check at load time (grant_shape_mismatch_for_gtp_padding). This commit supersedes that mechanism -- the checkpoint is logical, so a current save never mismatches -- and keeps the grant only as a reader for checkpoints written before the logical layout. Such a checkpoint is recoverable only at TP=1 (data rows, then the pad tail); the logical shard is tagged gtp_tp_size and the grant refuses TP>1, where the old grid shifted every TP rank past 0 by pad_length. NVIDIA#7293's tests were adapted: the cross-degree round trip now asserts that no waiver is needed, the grant-decision test writes its padded checkpoint by hand as a legacy artifact, and the "padded GTP save -> non-GTP load" case became an exact round trip instead of a hard failure. Signed-off-by: xuwenc <xuwenc@nvidia.com>
xuwchen
added a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Sep 22, 2026
…ffset With gtp_remat_pad_for_alignment > 0 a GTP weight is allocated padded to pad_for_alignment * gtp_remat_size rows, and every shard -- padding included -- was described to the checkpoint as if it were real data. ShardedTensor.from_rank_offsets computes global_offset[axis] = axis_rank_offset * local_axis_shape, so passing the PADDED shard size shifted every tp_rank > 0 by pad_length: those ranks read their neighbour's rows and their own tail came back as zeros. allow_shape_mismatch, set to absorb the inflated global shape, suppressed the validation that would have caught it. The symptom was a grad norm of 7.75e9 at pad_for_alignment=128; at the default of 16 it still corrupted every weight whose dim0 is not already aligned (vision fc1 pads 24 rows). Describe the checkpoint in LOGICAL layout instead. _make_gtp_logical_sharded_tensor sizes and offsets the global tensor from the unpadded row count and trims this rank's pad tail, so padding never reaches the checkpoint and allow_shape_mismatch is no longer needed. A shard that is entirely padding (vision linear_proj: dim0 1152, alignment 512) contributes an empty slice rather than raising. Trimming changes object identity, which the optimizer depends on -- it maps optimizer params to model entries by id(entry.data). The trimmed shard therefore carries a gtp_pad_src backlink to the live param, resolved by gtp_entry_backlink() in both the distributed-Adam and distributed-Muon lookups. The backlink lives on the ShardedTensor and not on its data, because the torch strategy does sh_ten.data = sh_ten.data.detach(), which drops tensor attributes. A full-length shard is handed back as the SAME object so identity still holds. A padded native-FP8 param points its backlink at the live param rather than at the dequantized copy, or the trimmed rank resolves to a copy the _gtp_dequant_src lookup can no longer see and its optimizer state silently vanishes. Two consumers need the untrimmed shard: an all-gather over the GTP group sizes itself off the short trailing rank and HANGS otherwise (untrimmed_gtp_shard()), and the fully-reshardable optimizer load expresses flat param ranges over the padded shard, so it zero-fills the excluded rows before applying the range. Covered by test_padded_shard_logical_offsets, test_padded_ep_egtp_offsets, test_padded_optimizer_state_mapping, test_padded_fused_projection_gather, test_padded_save_load_roundtrip and test_non_gtp_checkpoint_into_padded_gtp -- the last of which is the discriminating one: it writes with a non-GTP layout and reads into a padded GTP model, and is the only one that fails on the pre-fix mapping (want rows [10..15], got [12..17], +2 = pad_length). The symmetric round-trip cannot catch a self-consistently wrong mapping. Run them with scripts/gtp_dcp_unit_test.sh: the file gates on world_size == 4 and needs the TE >= 2.19 overlay, and getting either wrong skips the whole file while the command still exits 0. Rebase note (dev @ 0e704ed): the base now carries NVIDIA#7293, which handled the same symptom by keeping the padded global and waiving the shape check at load time (grant_shape_mismatch_for_gtp_padding). This commit supersedes that mechanism -- the checkpoint is logical, so a current save never mismatches -- and keeps the grant only as a reader for checkpoints written before the logical layout. Such a checkpoint is recoverable only at TP=1 (data rows, then the pad tail); the logical shard is tagged gtp_tp_size and the grant refuses TP>1, where the old grid shifted every TP rank past 0 by pad_length. NVIDIA#7293's tests were adapted: the cross-degree round trip now asserts that no waiver is needed, the grant-decision test writes its padded checkpoint by hand as a legacy artifact, and the "padded GTP save -> non-GTP load" case became an exact round trip instead of a hard failure. Signed-off-by: xuwenc <xuwenc@nvidia.com>
xuwchen
added a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Sep 22, 2026
…#7168) Replay of NVIDIA/Megatron-LM PR NVIDIA#7168 (branch gtp-alignment-pad-ckpt-layout @ 5a837a9) onto dev, squashed into one commit: the tree difference between the NVIDIA#7162 merge dcabf80 and the NVIDIA#7168 merge 5a837a9. The PR's commits beyond NVIDIA#7162, in order: 0768d59 feat(gtp): expose weight rematerialization alignment 019618e fix(gtp): checkpoint alignment padding in logical coordinates d279514 feat(egtp): checkpoint grouped gated projections in logical layout dev-only adaptations: - configure_gtp_remat_from_recipe: dev (via NVIDIA#7293) already resolves the recipe default in resolve_gtp_pad_for_alignment; an explicit --gtp-remat-pad-for-alignment wins over it. - The flag is registered in _add_distributed_args, since dev derives the other GTP flags from the config dataclasses and has no hand-written --gtp-remat-* group. - dist_checkpointing/strategies/torch.py: import get_torch_version at module level; the uneven-sharding error message uses it outside the function that imports it locally. - Reconciliation with NVIDIA#7293 (on dev, not on main when NVIDIA#7168 was written): the checkpoint is logical, so a current save never mismatches; grant_shape_mismatch_for_gtp_padding is kept only as a reader for checkpoints written before the logical layout. Such a checkpoint is recoverable only at TP=1; logical shards carry gtp_tp_size and the grant refuses TP>1, where the old grid shifted every TP rank past 0 by pad_length. NVIDIA#7293's tests were adapted: the cross-degree round trip asserts that no waiver is needed, the grant-decision test writes its padded checkpoint by hand as a legacy artifact, and the "padded GTP save -> non-GTP load" case became an exact round trip instead of a hard failure. Signed-off-by: xuwenc <xuwenc@nvidia.com>
xuwchen
added a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Sep 22, 2026
…#7168) Replay of NVIDIA/Megatron-LM PR NVIDIA#7168 (branch gtp-alignment-pad-ckpt-layout @ 5a837a9) onto dev, squashed into one commit: the tree difference between the NVIDIA#7162 merge dcabf80 and the NVIDIA#7168 merge 5a837a9. The PR's commits beyond NVIDIA#7162, in order: 0768d59 feat(gtp): expose weight rematerialization alignment 019618e fix(gtp): checkpoint alignment padding in logical coordinates d279514 feat(egtp): checkpoint grouped gated projections in logical layout dev-only adaptations: - configure_gtp_remat_from_recipe: dev (via NVIDIA#7293) already resolves the recipe default in resolve_gtp_pad_for_alignment; an explicit --gtp-remat-pad-for-alignment wins over it. - The flag is registered in _add_distributed_args, since dev derives the other GTP flags from the config dataclasses and has no hand-written --gtp-remat-* group. - dist_checkpointing/strategies/torch.py: the uneven-sharding error message uses get_torch_version outside the function that imports it locally; add a local import there (a module-level one is circular: megatron.core.utils imports dist_checkpointing). - Reconciliation with NVIDIA#7293 (on dev, not on main when NVIDIA#7168 was written): the checkpoint is logical, so a current save never mismatches; grant_shape_mismatch_for_gtp_padding is kept only as a reader for checkpoints written before the logical layout. Such a checkpoint is recoverable only at TP=1; logical shards carry gtp_tp_size and the grant refuses TP>1, where the old grid shifted every TP rank past 0 by pad_length. NVIDIA#7293's tests were adapted: the cross-degree round trip asserts that no waiver is needed, the grant-decision test writes its padded checkpoint by hand as a legacy artifact, and the "padded GTP save -> non-GTP load" case became an exact round trip instead of a hard failure. Signed-off-by: xuwenc <xuwenc@nvidia.com>
xuwchen
added a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Sep 22, 2026
…#7168) Replay of NVIDIA/Megatron-LM PR NVIDIA#7168 (branch gtp-alignment-pad-ckpt-layout @ 5a837a9) onto dev, squashed into one commit: the tree difference between the NVIDIA#7162 merge dcabf80 and the NVIDIA#7168 merge 5a837a9. The PR's commits beyond NVIDIA#7162, in order: 0768d59 feat(gtp): expose weight rematerialization alignment 019618e fix(gtp): checkpoint alignment padding in logical coordinates d279514 feat(egtp): checkpoint grouped gated projections in logical layout dev-only adaptations: - configure_gtp_remat_from_recipe: dev (via NVIDIA#7293) already resolves the recipe default in resolve_gtp_pad_for_alignment; an explicit --gtp-remat-pad-for-alignment wins over it. - The flag is registered in _add_distributed_args, since dev derives the other GTP flags from the config dataclasses and has no hand-written --gtp-remat-* group. - dist_checkpointing/strategies/torch.py: the uneven-sharding error message uses get_torch_version outside the function that imports it locally; add a local import there (a module-level one is circular: megatron.core.utils imports dist_checkpointing). - Reconciliation with NVIDIA#7293 (on dev, not on main when NVIDIA#7168 was written): the checkpoint is logical, so a current save never mismatches; grant_shape_mismatch_for_gtp_padding is kept only as a reader for checkpoints written before the logical layout. Such a checkpoint is recoverable only at TP=1; logical shards carry gtp_tp_size and the grant refuses TP>1, where the old grid shifted every TP rank past 0 by pad_length. NVIDIA#7293's tests were adapted: the cross-degree round trip asserts that no waiver is needed, the grant-decision test writes its padded checkpoint by hand as a legacy artifact, and the "padded GTP save -> non-GTP load" case became an exact round trip instead of a hard failure. Signed-off-by: xuwenc <xuwenc@nvidia.com>
xuwchen
added a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Sep 22, 2026
…#7168) Replay of NVIDIA/Megatron-LM PR NVIDIA#7168 (branch gtp-alignment-pad-ckpt-layout @ 5a837a9) onto dev, squashed into one commit: the tree difference between the NVIDIA#7162 merge dcabf80 and the NVIDIA#7168 merge 5a837a9. The PR's commits beyond NVIDIA#7162, in order: 0768d59 feat(gtp): expose weight rematerialization alignment 019618e fix(gtp): checkpoint alignment padding in logical coordinates d279514 feat(egtp): checkpoint grouped gated projections in logical layout dev-only adaptations: - configure_gtp_remat_from_recipe: dev (via NVIDIA#7293) already resolves the recipe default in resolve_gtp_pad_for_alignment; an explicit --gtp-remat-pad-for-alignment wins over it. - The flag is registered in _add_distributed_args, since dev derives the other GTP flags from the config dataclasses and has no hand-written --gtp-remat-* group. - dist_checkpointing/strategies/torch.py: the uneven-sharding error message uses get_torch_version outside the function that imports it locally; add a local import there (a module-level one is circular: megatron.core.utils imports dist_checkpointing). - Reconciliation with NVIDIA#7293 (on dev, not on main when NVIDIA#7168 was written): the checkpoint is logical, so a current save never mismatches; grant_shape_mismatch_for_gtp_padding is kept only as a reader for checkpoints written before the logical layout. Such a checkpoint is recoverable only at TP=1; logical shards carry gtp_tp_size and the grant refuses TP>1, where the old grid shifted every TP rank past 0 by pad_length. NVIDIA#7293's tests were adapted: the cross-degree round trip asserts that no waiver is needed, the grant-decision test writes its padded checkpoint by hand as a legacy artifact, and the "padded GTP save -> non-GTP load" case became an exact round trip instead of a hard failure. Signed-off-by: xuwenc <xuwenc@nvidia.com>
xuwchen
added a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Sep 22, 2026
…#7168) Replay of NVIDIA/Megatron-LM PR NVIDIA#7168 (branch gtp-alignment-pad-ckpt-layout @ 5a837a9) onto dev, squashed into one commit: the tree difference between the NVIDIA#7162 merge dcabf80 and the NVIDIA#7168 merge 5a837a9, applied on top of the rebased NVIDIA#7162 (23fcbe5). The PR's commits beyond NVIDIA#7162, in order: 0768d59 feat(gtp): expose weight rematerialization alignment 019618e fix(gtp): checkpoint alignment padding in logical coordinates d279514 feat(egtp): checkpoint grouped gated projections in logical layout dev-only adaptations: - NVIDIA#7162 (rebased) indexes optimizer metadata once in _get_param_id_to_sharded_metadata; the pad-trimmed shard's gtp_pad_src backlink joins the dequantized-source hop there via gtp_entry_backlink, instead of the per-parameter lookup NVIDIA#7168 was written against. - configure_gtp_remat_from_recipe: dev (via NVIDIA#7293) already resolves the recipe default in resolve_gtp_pad_for_alignment; an explicit --gtp-remat-pad-for-alignment wins over it. - The flag is registered in _add_distributed_args, since dev derives the other GTP flags from the config dataclasses and has no hand-written --gtp-remat-* group. - dist_checkpointing/strategies/torch.py: the uneven-sharding error message uses get_torch_version outside the function that imports it locally; add a local import there (a module-level one is circular: megatron.core.utils imports dist_checkpointing). - Reconciliation with NVIDIA#7293 (on dev, not on main when NVIDIA#7168 was written): the checkpoint is logical, so a current save never mismatches; grant_shape_mismatch_for_gtp_padding is kept only as a reader for checkpoints written before the logical layout. Such a checkpoint is recoverable only at TP=1; logical shards carry gtp_tp_size and the grant refuses TP>1, where the old grid shifted every TP rank past 0 by pad_length. NVIDIA#7293's tests were adapted: the cross-degree round trip asserts that no waiver is needed, the grant-decision test writes its padded checkpoint by hand as a legacy artifact, and the "padded GTP save -> non-GTP load" case became an exact round trip instead of a hard failure. Signed-off-by: xuwenc <xuwenc@nvidia.com>
xuwchen
added a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Sep 22, 2026
…#7168) Replay of NVIDIA/Megatron-LM PR NVIDIA#7168 (branch gtp-alignment-pad-ckpt-layout @ 5a837a9) onto dev, squashed into one commit: the tree difference between the NVIDIA#7162 merge dcabf80 and the NVIDIA#7168 merge 5a837a9, applied on top of the rebased NVIDIA#7162 (23fcbe5). The PR's commits beyond NVIDIA#7162, in order: 0768d59 feat(gtp): expose weight rematerialization alignment 019618e fix(gtp): checkpoint alignment padding in logical coordinates d279514 feat(egtp): checkpoint grouped gated projections in logical layout dev-only adaptations: - NVIDIA#7162 (rebased) indexes optimizer metadata once in _get_param_id_to_sharded_metadata; _resolve_gtp_sharded_metadata, which this PR's tests call, is kept as a thin per-parameter view over that index; the pad-trimmed shard's gtp_pad_src backlink joins the dequantized-source hop there via gtp_entry_backlink, instead of the per-parameter lookup NVIDIA#7168 was written against. - configure_gtp_remat_from_recipe: dev (via NVIDIA#7293) already resolves the recipe default in resolve_gtp_pad_for_alignment; an explicit --gtp-remat-pad-for-alignment wins over it. - The flag is registered in _add_distributed_args, since dev derives the other GTP flags from the config dataclasses and has no hand-written --gtp-remat-* group. - dist_checkpointing/strategies/torch.py: the uneven-sharding error message uses get_torch_version outside the function that imports it locally; add a local import there (a module-level one is circular: megatron.core.utils imports dist_checkpointing). - Reconciliation with NVIDIA#7293 (on dev, not on main when NVIDIA#7168 was written): the checkpoint is logical, so a current save never mismatches; grant_shape_mismatch_for_gtp_padding is kept only as a reader for checkpoints written before the logical layout. Such a checkpoint is recoverable only at TP=1; logical shards carry gtp_tp_size and the grant refuses TP>1, where the old grid shifted every TP rank past 0 by pad_length. NVIDIA#7293's tests were adapted: the cross-degree round trip asserts that no waiver is needed, the grant-decision test writes its padded checkpoint by hand as a legacy artifact, and the "padded GTP save -> non-GTP load" case became an exact round trip instead of a hard failure. Signed-off-by: xuwenc <xuwenc@nvidia.com>
xuwchen
added a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Sep 22, 2026
…#7168) Replay of NVIDIA/Megatron-LM PR NVIDIA#7168 (branch gtp-alignment-pad-ckpt-layout @ 5a837a9) onto dev, squashed into one commit: the tree difference between the NVIDIA#7162 merge dcabf80 and the NVIDIA#7168 merge 5a837a9, applied on top of the rebased NVIDIA#7162 (23fcbe5). The PR's commits beyond NVIDIA#7162, in order: 0768d59 feat(gtp): expose weight rematerialization alignment 019618e fix(gtp): checkpoint alignment padding in logical coordinates d279514 feat(egtp): checkpoint grouped gated projections in logical layout dev-only adaptations: - NVIDIA#7162 (rebased) indexes optimizer metadata once in _get_param_id_to_sharded_metadata; _resolve_gtp_sharded_metadata, which this PR's tests call, is kept as a thin per-parameter view over that index; the pad-trimmed shard's gtp_pad_src backlink joins the dequantized-source hop there via gtp_entry_backlink, instead of the per-parameter lookup NVIDIA#7168 was written against. - configure_gtp_remat_from_recipe: dev (via NVIDIA#7293) already resolves the recipe default in resolve_gtp_pad_for_alignment; an explicit --gtp-remat-pad-for-alignment wins over it. - The flag is registered in _add_distributed_args, since dev derives the other GTP flags from the config dataclasses and has no hand-written --gtp-remat-* group. - dist_checkpointing/strategies/torch.py: the uneven-sharding error message uses get_torch_version outside the function that imports it locally; add a local import there (a module-level one is circular: megatron.core.utils imports dist_checkpointing). - Reconciliation with NVIDIA#7293 (on dev, not on main when NVIDIA#7168 was written): the checkpoint is logical, so a current save never mismatches; grant_shape_mismatch_for_gtp_padding is kept only as a reader for checkpoints written before the logical layout. Such a checkpoint is recoverable only at TP=1; logical shards carry gtp_tp_size and the grant refuses TP>1, where the old grid shifted every TP rank past 0 by pad_length. NVIDIA#7293's tests were adapted: the cross-degree round trip asserts that no waiver is needed, the grant-decision test writes its padded checkpoint by hand as a legacy artifact, and the "padded GTP save -> non-GTP load" case became an exact round trip instead of a hard failure. Signed-off-by: xuwenc <xuwenc@nvidia.com>
This branch was previously deployed
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.
Summary
A checkpoint saved with GTP (Generalized Tensor Parallelism) alignment padding at one
gtp_remat_sizecould fail to load at a differentgtp_remat_size, or when loaded by a non-GTP process — the checkpoint's declared shape (padded) didn't match the loader's expected shape (possibly unpadded or padded differently), andallow_shape_mismatchwas never set to tolerate the difference.Here are two typical cases this MR should fix:
Example-1: GTP-save-with-padding -> GTP-load-w/o-padding
Example-2: GTP-save --> GTP load but with different padding requirement
Coverage matrix
gtp_remat_size)declared == dim0_unpadded, e.g. non-tile-aligned hidden dims)All "Fixed" cells are validated for plain TP×DP. Pipeline parallelism (PP > 1) is a known gap, not covered by this fix — see below.
Fix
Add
infer_gtp_allow_shape_mismatch()(megatron/core/utils.py), called right beforedist_checkpointing.load(). It reads the checkpoint's real on-disk shape and setsallow_shape_mismatchper tensor — but only when the declared-vs-expected difference is actually explainable as GTP padding (bounded by the recipe's alignment size), not for every mismatch.pad_for_alignment(the recipe's alignment tile size) is resolved via a newmegatron/core/tensor_parallel/gtp_recipe.py— a small, zero-dependency module so it can be computed during checkpoint load without importing TE-dependent GTP modules.Testing
test_cross_gtp_degree_save_load_roundtripdoes a genuinedist_checkpointing.save()/load()roundtrip for the originally-reported case — GTP4-padded save (dim0=192→ 256) loaded by GTP2 (no padding needed,192 % 64 == 0). The load-side buffer is pre-scribbled with a sentinel value before loading (so a no-op or partial load can't accidentally pass), and the result is checked bit-exact (rtol=0, atol=0) against an independently-constructed golden tensor. This confirms the fix doesn't just flip a boolean — the actual loaded weight data is correct.allow_shape_mismatchfrom an unrelated reason (must not be clobbered), and the bf16 (pad_for_alignment=1) case (must not vacuously accept any size).tests/unit_tests/generalized_tensor_parallel/suite (all 18 files, ~200+ cases) passes cleanly, including this one.Known limitations (not in the scope of this MR)
SwiGLU linear_fc1.weight + GTPis unreconciled by this fixalignment/exact-match.
Contribution process
Pre-checks