fix(gtp): reshardable optimizer checkpoints under GTP by xuwchen · Pull Request #7162 · NVIDIA/Megatron-LM

xuwchen · 2026-09-09T09:15:10Z

What does this PR do?

dist_ckpt_optim_fully_reshardable is the only distributed-optimizer checkpoint format whose state survives a change of parallel layout; every other format pins its entries to the DP grouping in force at save time. Under GTP this format could neither be written nor read back, so a GTP run could not hand its optimizer state to a differently-sharded run in either direction.

Two independent defects, each sufficient to block it:

Save aborted. sharded_param_state_fully_reshardable matches a grad-buffer parameter to its model ShardedTensor by object identity. GTP puts a different tensor in the model entry -- a dequantized BF16 copy for native-FP8 weights, and a factory over the gathered tensor for the fused SSM in_proj -- so the lookup failed and the save stopped.
What it wrote could not be read. The model side gathers in_proj to TP-local rows before splitting it into per-segment keys; the optimizer did not follow, and keyed the same state under the fused name.

A third, unrelated dist_checkpointing bug sits on the same path: _ChunkStorageMetadata builds shard offsets spanning prepend_axis_num + ndim axes but sizes spanning only ndim, so any shard with a prepend axis (every expert-parallel or GTP weight) describes a chunk whose offset and size ranks disagree. DCP's global-plan validation rejects it, or mis-slices when validation is waived. GTP only made this easy to reach.

Fix

checkpointable.py: give offsets and sizes the same rank -- _dcp_chunk_size() prepends one per prepend axis and _dcp_shard_view() adds the matching singleton dimensions.
distrib_optimizer.py: _resolve_gtp_sharded_metadata() follows the _gtp_dequant_src backlink or rebuilds the per-shard ShardedTensor, the way the distributed-Muon path already does, and refuses the expert-parallel rebuild that would write duplicate shards across EP groups. For factory-backed weights the optimizer state is gathered across the GTP group to the width the factory expects, with the alignment pad stripped, and the model's own factory is reused so the keys match a non-GTP run.
When the match fails, the error now names the parameter and what the identity-keyed map actually holds.

Also in this PR

With fp32 params the optimizer's main param is a view of the model param, so loading optimizer state rewrites model weights over the shard this rank owns. That is a no-op while the checkpoint's two copies of the weights agree; when they do not, what the run starts from depends on the parallel layout. This PR does not change what is loaded -- it counts the disagreement device-side during the copy that already runs and warns once per load.

Verified / not verified

Verified on a Qwen3.5-VL proxy: the checkpoint saves without dp_group_idx keys, a same-topology resume continues, and a 3D-saved checkpoint loads into a GTP run. The reverse direction -- a GTP-saved checkpoint loaded by a non-GTP run -- follows from the same key mapping but was not exercised. Alignment padding > 0 on fused weights is handled in a follow-up PR together with the model-side padding fix.

copy-pr-bot · 2026-09-09T09:15:14Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

copy-pr-bot · 2026-09-10T16:37:02Z

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.

fanshiqing · 2026-09-11T07:28:46Z

/claude strict-review

Place the gather and load-slice helpers directly in tensor_parallel/gtp_ckpt.py and re-export them from ssm.utils. Mamba and GDP save logical projection sections and restore contiguous physical GTP shards through these helpers. Run model gathers under no_grad so live Parameters do not enter autograd. Include the GTP rank in checkpoint replica IDs, reject flattened model factory merges, and remove Mamba's unused dataclasses.replace import. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Gather the TP-local GDN input projection before splitting q, k, v, z, beta, and alpha, then restore padding and the local GTP slice on load. Reuse the shared helpers so section boundaries and writer election are independent of the physical GTP shard boundaries. Exercise TP2/GTP2 with four and eight ranks, including DP2 replicas at the CI world size. Check semantic sections, alignment padding, merge round trips, complete GTP/DP replica IDs, and exactly one checkpoint writer. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Gather the full TP-local gate/up weight before applying its checkpoint factory, then restore the physical GTP row slice on load using the common checkpoint helpers. Document the fallback when optional GTP imports are unavailable and guard unsupported grouped EGTP gated weights. Add elementwise forward and gradient checks at GTP4 and GTP8. Check TP2/GTP2 checkpoint sections, nonzero padding, merge round trips, and unique writer election with DP1 or DP2 so the new cases run at the default CI world size. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Centralize the GTP gather, semantic section split, and load-side pad/slice used by GDN, GDP, and Mamba. Derive the logical row count from the live shard and padding metadata to retain section-size validation. Reuse the shared splitter for GDP convolution parameters and preserve householder checkpoint keys and replica election. Cover live BF16/FP32 Parameters, padding, cross-topology reconstruction, complete source and destination DCP metadata, invalid sections, and OOM fallback. Fold removed GTP ranks into destination DP and keep one owned CPU default process group across checkpoint test modules. Prior validation of the unchanged production and test logic (the current history restores copyright headers and clarifies fixture documentation): - CPU/Gloo: 57 passed per rank on eight ranks. - CW H100, full test_gtp_dcp.py: 23 passed / 2 skipped per rank on four ranks; 2 passed / 23 skipped per rank on eight ranks. - CW H100, full test_gtp_basics.py: 60 passed / 2 skipped per rank on four ranks; 34 passed / 28 skipped per rank on eight ranks. New GDN/SwiGLU DCP and fused-layout runtime cases passed on every rank at both sizes. MXFP8 requires compute capability >= 10 and was skipped on H100. Eight-rank runs also skip legacy four-rank-only cases. GPU validation used cached CI 26.04 with Torch 2.12 / CUDA 13.2 and TE 2.19.0.dev0+8260f496 built inside that image. Current CI pins 26.06 / TE 2.18, so this does not claim exact current-CI parity. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Carry the NVIDIA#7161 review updates through the optimizer checkpoint stack without rewriting any previously reviewed commits. Follow the gtp_utils module rename in optimizer backlink resolution, grouped experts, and padding and checkpoint regression imports. Preserve logical padding geometry, untrimmed checkpoint buffers, native-FP8 backlinks, grouped EGTP factories, and optimizer companion behavior. The shared projection factory and its optimizer tests follow the merged NVIDIA#7162 logical-shape and parameter-detection interface unchanged. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Use the GTP parameter flag and logical-shape property when constructing fused projection checkpoints. Rename the shared helper module to gtp_utils and remove the unused MLP detection fallback function. Access optional GTP symbols through the API module behind HAVE_GTP so a missing detector needs no conditional import or local None placeholder. Document guarded optional-symbol access in the facade. Align CPU test parameters with normal size-one initialization and the logical-shape interface, retaining independent invalid-section checks. Isolate test mocks to the SSM consumer's API reference. Keep the four reviewed commits intact and add this follow-up separately. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Carry the folded NVIDIA#7161 helper review through the optimizer and padding checkpoint stack. Follow the gtp_utils module rename in optimizer backlink resolution, grouped experts, and padding and checkpoint regression imports. Preserve logical padding geometry, untrimmed checkpoint buffers, native-FP8 backlinks, grouped EGTP factories, and optimizer companion behavior. The shared projection factory and its optimizer tests retain the validated logical-shape and guarded GTP API module interface. Consolidate the two prior synchronization merges without changing the final tree or any original functional commit. This stack contains no main merge. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Mock the input-projection factory at the GDP call site after extracting its implementation into the common SSM helper. Keep the convolution splitter mock and verify the live weight, explicit TP/DP groups, and nonempty sharding offsets reach the new factory. The existing GDP checkpoint test file now passes all four CPU tests; before this change the wrapper test reached the real ShardedTensor type check with its SimpleNamespace test double. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Share the real GDN and SwiGLU checkpoint setup. Cover native-FP8 GDN through the shared factory and model restore, including alignment padding, and pin the EGTP grouped gated-fc1 rejection with one small test. Focus helper tests on GTP boundary crossings, padding, GTP2/4/1 restore and writer election. Remove duplicate layout/dtype combinations, redundant GPU writer checks, the GDP alias OOM case, and unrelated Linear/out_proj coverage. Retain the GDP fixture fix and shorten shared fixture docs. Place new checkpoint cases before existing upstream blocks to avoid the append conflict with main. Production code is unchanged. The PR's net UT addition falls from 734 to 326 lines. Validation: seven focused tests passed on each of four OCI GPU ranks with TE 2.20.0.dev0, without skips. The three helper tests also previously passed on eight CPU/Gloo ranks. Syntax, formatting and diff checks passed. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Share the SwiGLU section-split call and discover grouped weights from registered parameters instead of a fixed index range. Clarify the grouped EGTP guard and the separate optimizer checkpoint compatibility limits. Extend the existing SwiGLU checkpoint regression with one real TE fc1 forward against independently seeded logical weights. Check grouped weight indices beyond the old scan limit without restoring the larger microbatch test suite. Validation: eight focused tests passed on each of four OCI GB200 ranks with TE 2.20.0.dev0. A temporary row-permutation injection failed at the new forward assertion on all four ranks. Pylint, Black, isort, syntax, and diff checks passed. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Gathered projection factories and dequantized native-FP8 checkpoint entries cannot be matched to live GTP parameters by object identity. Distributed Adam's fully_reshardable and legacy fully_sharded_model_space formats fail even when the GTP topology is unchanged. Reject these combinations before reading optimizer state or exchanging buffers, and recommend dp_reshardable through either the CLI or checkpoint metadata API. Extend the existing fused-projection regressions to check both rejections and the supported format, and document the limitation. Clarify the Muon helper docstring accordingly. GPU validation before the diagnostic-text update: 31 GTP DCP, 2 Muon DCP, and 3 ordinary Adam fully-reshardable tests passed per rank on four OCI GB200 GPUs with the 26.07 v10 container (TE 2.20.0.dev0), without skips. The follow-up changes only the diagnostic and docs; AST equivalence, Black, isort, Pylint, syntax, and PR diff checks passed. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Resolve the test_gtp_dcp.py insertion conflict by retaining the upstream cross-GTP checkpoint helpers and the PR's GDN/SwiGLU test classes. Validation: the merged tree differs from the prior GPU-tested revision only in the optimizer diagnostic and documentation. Executable ASTs match after excluding docstrings and normalizing the diagnostic. Prior tests passed all 31 GTP DCP, 2 Muon DCP, and 3 ordinary Adam fully-reshardable cases per rank on four OCI GB200 GPUs with the 26.07 v10 container, without skips. Black, isort, Pylint, syntax, and PR diff checks passed. Signed-off-by: xuwenc <xuwenc@nvidia.com>

) Replay of NVIDIA/Megatron-LM PR NVIDIA#7162 (branch gtp-reshardable-optimizer-ckpt @ dcabf80) onto dev, squashed into one commit: the tree difference between NVIDIA#7161's 146bae1 and the NVIDIA#7162 merge dcabf80. The PR's commits beyond NVIDIA#7161, in order: 9b9ea3d fix(checkpoint): map fused optimizer fragments to logical sections 51e9d1b fix(dcp): represent flat optimizer fragments as rectangular shards 1534ede fix(optimizer): resolve dequantized checkpoint entries by source identity e564bcc fix(gdn): checkpoint projection bias in semantic sections e4dec75 fix(dcp): align prepended-axis chunk sizes and tensor views 974b560 fix(optimizer): restore FP32 parameter views without autograd 4cf822b feat(checkpoint): warn about model and FP32 master disagreement tests/unit_tests/ssm/test_in_proj_factory.py is taken from the NVIDIA#7162 merge state: NVIDIA#7161's later commit f648f67 rewrote this file around simpler fixtures, and NVIDIA#7162's optimizer factory tests are written against the earlier structure. The NVIDIA#7162 version is the superset. Signed-off-by: xuwenc <xuwenc@nvidia.com>

…#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>

…#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>

) Replay of NVIDIA/Megatron-LM PR NVIDIA#7162 (branch gtp-reshardable-optimizer-ckpt @ dcabf80) onto dev, squashed into one commit: the tree difference between NVIDIA#7161's 146bae1 and the NVIDIA#7162 merge dcabf80. The PR's commits beyond NVIDIA#7161, in order: 9b9ea3d fix(checkpoint): map fused optimizer fragments to logical sections 51e9d1b fix(dcp): represent flat optimizer fragments as rectangular shards 1534ede fix(optimizer): resolve dequantized checkpoint entries by source identity e564bcc fix(gdn): checkpoint projection bias in semantic sections e4dec75 fix(dcp): align prepended-axis chunk sizes and tensor views 974b560 fix(optimizer): restore FP32 parameter views without autograd 4cf822b feat(checkpoint): warn about model and FP32 master disagreement tests/unit_tests/ssm/test_in_proj_factory.py keeps NVIDIA#7161's later simplified fixtures (f648f67) and ports NVIDIA#7162's optimizer-factory coverage onto them: the companion keeps the parameter identity for an ordinary weight, and physical GTP/DP optimizer slices (full and flat DP fragments) tile each semantic section exactly once and restore the padded shard. The fixtures keep TP fixed at one, so the TP>1 optimizer-offset case of the original test is not carried here. Signed-off-by: xuwenc <xuwenc@nvidia.com>

…#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>

) Replay of NVIDIA/Megatron-LM PR NVIDIA#7162 (branch gtp-reshardable-optimizer-ckpt @ 23fcbe5, rebased on NVIDIA#7161 @ 83eac91) onto dev, squashed into one commit: the tree difference between the NVIDIA#7161 head and the NVIDIA#7162 head. The PR's commits beyond NVIDIA#7161, in order: 8772fa3 fix(dcp): map optimizer fragments to rectangular logical shards 336b1e2 fix(gtp): map optimizer state through source-bound companions 49c3948 fix(gdn): checkpoint projection bias in semantic sections cd43082 fix(optimizer): restore FP32 parameter views without autograd 23fcbe5 feat(checkpoint): warn about model and FP32 master disagreement Signed-off-by: xuwenc <xuwenc@nvidia.com>

…#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>

) Replay of NVIDIA/Megatron-LM PR NVIDIA#7162 (branch gtp-reshardable-optimizer-ckpt @ 23fcbe5, rebased on NVIDIA#7161 @ 83eac91) onto dev, squashed into one commit: the tree difference between the NVIDIA#7161 head and the NVIDIA#7162 head. The PR's commits beyond NVIDIA#7161, in order: 8772fa3 fix(dcp): map optimizer fragments to rectangular logical shards 336b1e2 fix(gtp): map optimizer state through source-bound companions 49c3948 fix(gdn): checkpoint projection bias in semantic sections cd43082 fix(optimizer): restore FP32 parameter views without autograd 23fcbe5 feat(checkpoint): warn about model and FP32 master disagreement dev-only adaptation: the guard test's fake optimizer sets _optimizer_state_offloader=None, because dev's DistributedOptimizer.sharded_state_dict consults the chunked optimizer-state offloader before reaching the guard (main has no such code). Signed-off-by: xuwenc <xuwenc@nvidia.com>

…#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>

) Replay of NVIDIA/Megatron-LM PR NVIDIA#7162 (branch gtp-reshardable-optimizer-ckpt @ 23fcbe5, rebased on NVIDIA#7161 @ 83eac91) onto dev, squashed into one commit: the tree difference between the NVIDIA#7161 head and the NVIDIA#7162 head. The PR's commits beyond NVIDIA#7161, in order: 8772fa3 fix(dcp): map optimizer fragments to rectangular logical shards 336b1e2 fix(gtp): map optimizer state through source-bound companions 49c3948 fix(gdn): checkpoint projection bias in semantic sections cd43082 fix(optimizer): restore FP32 parameter views without autograd 23fcbe5 feat(checkpoint): warn about model and FP32 master disagreement dev-only adaptations: test_gtp_checkpoint_optimizer calls apply_swiglu_sharded_factory without tp_group/dp_group (dev's signature lacks them; main added them for FSDP2 DP-sharded params only); and the guard test's fake optimizer sets _optimizer_state_offloader=None, because dev's DistributedOptimizer.sharded_state_dict consults the chunked optimizer-state offloader before reaching the guard (main has no such code). Signed-off-by: xuwenc <xuwenc@nvidia.com>

…#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>

…#7168) Expose weight rematerialization alignment override (--gtp-remat-pad-for-alignment), keep alignment padding outside checkpoint coordinates (logical-layout fix), and wire the shared projection and optimizer factories for EGTP grouped gated fc1. Rebased from review/gtp-7168-on-7162-20260922 @ ff480df (tree diff above NVIDIA#7162 squash). Dev adaptations: --gtp-remat-pad-for-alignment registered in _add_distributed_args; training.py reads gtp_remat_pad_for_alignment attr via getattr. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Represent flat DP fragments as rectangular logical shards, including prepended expert axes. Keep partial boundaries and trailing padding explicit for DCP. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Attach source-bound optimizer companions to gathered model factories. Preserve keys, expert offsets and outer replica updates without gathering optimizer state. Elect DP writers from the live parameter replica group even when training metadata includes GTP. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Apply the projection section layout to GDN biases. Preserve same-TP legacy loading and reject ambiguous legacy resharding until a semantic checkpoint is saved. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Copy checkpoint masters and moments without recording autograd operations, including FP32 shards that alias live model parameters. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Count FP32 model/master disagreements before restoring optimizer masters and report one rank-zero warning after reducing the count. Preserve legacy loads that do not initialize the counter. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Cover fragment geometry, source identity, replica ownership, GTP-inclusive metadata, native FP8, GDN bias migration and FP32 restoration. Consolidate shared test setup and keep DCP disk round trips and independent value checks. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Exclude alignment tails from checkpoint shapes and offsets. Preserve distinct native-FP8 identity and data backlinks, validate optimizer padding and restore physical buffers on load. Reuse the rectangular fragment mapper from NVIDIA#7162. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Update the existing dev port from upstream PR NVIDIA#7162 at 23fcbe5 to 1ceb6b0. Normalize writer replicas from the live parameter before resolving optimizer companions, and report model/master disagreements within each optimizer's DP group so pipeline stages without that optimizer do not join a world collective. Include the consolidated restoration tests, subgroup warning regression, GTP-inclusive writer metadata coverage and test cleanup. Preserve dev's three-argument SwiGLU factory API and optimizer-offload test stub. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Update the NVIDIA#7168 dev port to 9d94f9a, stacked on NVIDIA#7162 at 1ceb6b0. Consolidate BF16/native-FP8 padded checkpoint round trips while preserving both dtype cases, exact loaded values, optimizer restoration, physical padding and live-parameter identity checks. Simplify identity-preserving checkpoint data selection and group the expert projection helper imports. Adapt grouped fc1 to dev's three-argument SwiGLU factory API; retain explicit process groups on the GTP gather. Signed-off-by: xuwenc <xuwenc@nvidia.com>

) Port NVIDIA/Megatron-LM PR NVIDIA#7162 at 1ceb6b0, based on NVIDIA#7161 at 83eac91. Map flat optimizer fragments to rectangular logical shards and bind companions to live parameters. Normalize writer replicas before propagating model metadata, checkpoint GDN projection biases by semantic section, and restore FP32 parameter views without autograd. Report model/master mismatches within each optimizer's DP group so unrelated pipeline stages do not join the collective. Include the consolidated checkpoint regressions. Dev adaptations: retain the three-argument SwiGLU factory API and the optimizer-offload field in the test stub. The multi-bucket optimizer resharding prerequisite (NVIDIA#5077) remains a separate preceding commit. Signed-off-by: xuwenc <xuwenc@nvidia.com>

…#7168) Port NVIDIA/Megatron-LM PR NVIDIA#7168 at 9d94f9a, stacked on NVIDIA#7162 at 1ceb6b0. Expose weight-rematerialization alignment, keep physical padding outside logical checkpoint coordinates, and support grouped gated fc1 under EGTP. Use the shared BF16/native-FP8 checkpoint round trips while retaining exact loaded values, optimizer states, physical padding and parameter identity. Dev adaptations: register --gtp-remat-pad-for-alignment in distributed arguments, read the optional training attribute with getattr, and call the three-argument SwiGLU factory. Keep explicit groups on the GTP gather. Signed-off-by: xuwenc <xuwenc@nvidia.com>

) Port NVIDIA/Megatron-LM PR NVIDIA#7162 at 1ceb6b0, based on NVIDIA#7161 at 83eac91. Map flat optimizer fragments to rectangular logical shards and bind companions to live parameters. Normalize writer replicas before propagating model metadata, checkpoint GDN projection biases by semantic section, and restore FP32 parameter views without autograd. Report model/master mismatches within each optimizer's DP group so unrelated pipeline stages do not join the collective. Include the consolidated checkpoint regressions. Dev adaptations: retain the three-argument SwiGLU factory API and the optimizer-offload field in the test stub. The multi-bucket optimizer resharding prerequisite (NVIDIA#5077) remains a separate preceding commit. Normalize pre-existing dev formatting required by the PR-wide Black check; these whitespace and layout edits do not change execution or checkpoint data. Signed-off-by: xuwenc <xuwenc@nvidia.com>

…#7168) Port NVIDIA/Megatron-LM PR NVIDIA#7168 at 9d94f9a, stacked on NVIDIA#7162 at 1ceb6b0. Expose weight-rematerialization alignment, keep physical padding outside logical checkpoint coordinates, and support grouped gated fc1 under EGTP. Use the shared BF16/native-FP8 checkpoint round trips while retaining exact loaded values, optimizer states, physical padding and parameter identity. Dev adaptations: register --gtp-remat-pad-for-alignment in distributed arguments, read the optional training attribute with getattr, and call the three-argument SwiGLU factory. Keep explicit groups on the GTP gather. Signed-off-by: xuwenc <xuwenc@nvidia.com>

Save rank-local intersections with logical projection sections instead of all-gathering GDP weights. Use the shared fragment implementation for KDA input projections and dense, shared, and routed GLU projections. Keep logical keys, expert offsets, and replica ownership. Omit alignment padding from checkpoint coordinates, reconstruct zero tails on load, and retain all-padding ranks through local nonpersistent metadata. Preserve native-FP8 checkpoint buffers separately from live-parameter identity. Bind Kimi optimizer factories to the source parameters so FP32 masters and moments use the same logical fragments, including partial-row flat DP slices. Preserve master indices and accept legacy list containers. GDP retains its unsplit optimizer fallback; Mamba retains its gather path. Older physical GTP optimizer layouts are not automatically converted. Consolidate the GDP implementation and its Kimi extension so review shows the final shared helper directly. Adapt only the required fixes from NVIDIA#7161/NVIDIA#7162/NVIDIA#7168; do not import the complete PR stack. Keep coverage, padding, cross-layout model/master/momentum roundtrips, continued updates, real Kimi wiring, and flat Adam regressions together. Source-Commit: d8fae99 Source-Commit: 65360b3 Signed-off-by: xuwenc <xuwenc@nvidia.com>

Save GDP and Kimi projection weights as rank-local logical fragments without GTP all-gathers. Keep section keys, expert offsets, replica ownership, zero padding, native-FP8 buffers, and partial-row optimizer slices intact. Use source-bound fragments for Kimi FP32 masters and optimizer moments; retain the GDP optimizer fallback and the existing Mamba gather path. Elect one dense-TP writer for shared MLP activation state, including ETP < TP. Keep fragment interfaces, layout and padding regressions, cross-layout model/master/momentum loading, real-module wiring, and writer ownership together. Simplify the real-module test's single-entry container without removing test cases, parameter combinations, or assertions. Adapt the required parts of NVIDIA#7161/NVIDIA#7162/NVIDIA#7168 without importing their entire stacks. Older physical GTP optimizer layouts are not automatically converted. Source-Commit: 5232c0d Source-Commit: 8c1331e Source-Commit: da64550 Signed-off-by: xuwenc <xuwenc@nvidia.com>

xuwchen requested a review from fanshiqing September 9, 2026 09:15

xuwchen mentioned this pull request Sep 9, 2026

fix(gtp): logical checkpoint layout for padded shards and EGTP gated fc1 #7168

Draft

xuwchen linked an issue Sep 9, 2026 that may be closed by this pull request

Support Qwen3.5-VL with GTP #7164

Open

xuwchen force-pushed the gtp-reshardable-optimizer-ckpt branch from 84418d5 to 7baaf9b Compare September 10, 2026 16:36

xuwchen force-pushed the gtp-reshardable-optimizer-ckpt branch 2 times, most recently from ac20a24 to 16ce491 Compare September 11, 2026 06:55

xuwchen added 4 commits September 11, 2026 00:34

xuwchen force-pushed the gtp-reshardable-optimizer-ckpt branch from 16ce491 to 4cf822b Compare September 11, 2026 07:42

xuwchen force-pushed the gtp-reshardable-optimizer-ckpt branch from c82944d to dcabf80 Compare September 11, 2026 10:08

Merge branch 'main' into gtp-fused-projection-ckpt-layout

df4e6d0

lhb8125 mentioned this pull request Sep 14, 2026

Megatron-LM Monthly Activity Report — September 2026 (Final, through September 30, 2026) #7100

Open

xuwchen added 5 commits September 13, 2026 23:07

xuwchen force-pushed the gtp-reshardable-optimizer-ckpt branch from dcabf80 to 23fcbe5 Compare September 22, 2026 17:24

xuwchen added 6 commits September 24, 2026 20:20

feat(dcp): support rectangular optimizer checkpoint fragments

b13429f

Represent flat DP fragments as rectangular logical shards, including prepended expert axes. Keep partial boundaries and trailing padding explicit for DCP. Signed-off-by: xuwenc <xuwenc@nvidia.com>

fix(gdn): checkpoint projection biases by semantic section

515cec1

Apply the projection section layout to GDN biases. Preserve same-TP legacy loading and reject ambiguous legacy resharding until a semantic checkpoint is saved. Signed-off-by: xuwenc <xuwenc@nvidia.com>

fix(optimizer): restore FP32 parameter views without autograd

e4ebd1a

Copy checkpoint masters and moments without recording autograd operations, including FP32 shards that alias live model parameters. Signed-off-by: xuwenc <xuwenc@nvidia.com>

xuwchen force-pushed the gtp-reshardable-optimizer-ckpt branch from 23fcbe5 to 1ceb6b0 Compare September 25, 2026 13:15

lhb8125 mentioned this pull request Oct 1, 2026

Megatron-LM Monthly Activity Report — September 2026 #7773

Open

xuwchen mentioned this pull request Oct 7, 2026

fix(gtp): checkpoint fused projections in logical layout #7161

Open

Navigation Menu

Repository navigation

fix(gtp): reshardable optimizer checkpoints under GTP - #7162

fix(gtp): reshardable optimizer checkpoints under GTP#7162
xuwchen wants to merge 17 commits into
NVIDIA:mainfrom
xuwchen:gtp-reshardable-optimizer-ckpt

xuwchen commented Sep 9, 2026

Uh oh!

copy-pr-bot Bot commented Sep 9, 2026

Uh oh!

copy-pr-bot Bot commented Sep 10, 2026

Uh oh!

fanshiqing commented Sep 11, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

xuwchen commented Sep 9, 2026

What does this PR do?

Fix

Also in this PR

Verified / not verified

Uh oh!

copy-pr-bot Bot commented Sep 9, 2026

Uh oh!

copy-pr-bot Bot commented Sep 10, 2026

Uh oh!

fanshiqing commented Sep 11, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants