[deepakn94]

Summary

Under GTP, GatedDeltaProductMixer.sharded_state_dict all-gathered every layer's in_proj.weight back to the full TP-local size so it could split it into the [z|V*|K*|Q|b*|a] checkpoint chunks, whose boundaries don't line up with the GTP shard boundaries. Every GTP rank kept those gathered copies in the state dict until the save was staged. For a large hybrid model with a 61,760 × 10,240 bf16 in_proj (≈ 1.2 GiB per layer) across 72 GDP layers, that is ~85 GiB live per rank, and it grew reserved GPU memory by ~48 GB at each save (allocated memory unchanged, so the copies were freed but the allocator kept the new segments).

Each split chunk is a contiguous block of TP-local rows, so a GTP rank's shard covers parts of one or a few chunks. Each rank now saves those parts directly, as views of its shard, at their row offsets within each chunk (uneven shards, stored by DCP as-is via CheckpointableShardedTensor). Pad rows are not saved. On load, the same pieces are read and the merge rebuilds the padded shard, zero-filling the pad rows. A rank whose shard holds only padding keeps a zero-element LocalNonpersistentObject so the factory key and device survive load preprocessing. No rank gathers or materializes the full tensor.

The keys and global shapes are exactly what a non-GTP run writes, so:

• GTP and non-GTP runs can load each other's checkpoints, as before;
• checkpoints written by the old gather-based save load into this code.

The model entry is a fresh view, so the Muon optimizer still takes its per-shard rebuild (_backfill_gtp_sharded_param_map) and the optimizer-state checkpoint layout for in_proj is unchanged.

Mamba's in_proj uses the same gather pattern and is not changed here; the GTP doc now says so.

Testing

• test_gtp_dcp.py:
  • test_gdp_inproj_shard_split (replaces the gather+split test): the factory's pieces are views of the shard, tile every chunk exactly once across the GTP group with non-GTP keys and global shapes, and merge back to the padded shard.
  • New test_gdp_inproj_cross_gtp_load: saves at GTP 2 and loads at GTP 1, then the reverse.
  • New test_gdp_gtp_padding_save_load_preprocessing (CPU and CUDA, bf16 and fp32): drives save_preprocess/load_preprocess/apply_factory_merges for a fully real, a partially padded and an all-padding shard, checking that real rows round-trip, pad rows come back as zeros, and nothing leaks into common state.
  • test_gdp_save_load_roundtrip and test_gdp_inproj_optim_param_map are kept, with comments updated.
• The same change on an earlier base passed test_gtp_dcp.py -k "gdp or mamba or roundtrip" on 4× GB300 (the GDP tests, the two Mamba tests and the other save/load round trips). Not yet re-run on this branch.

🤖 Generated with Claude Code

[copy-pr-bot]

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.

[Phlip79]

/review

[svcnvidia-nemo-ci]

Found a local-checkpoint loading regression, a focused test-coverage gap, and one minor style violation. Static review only; tests were not executed.

Restore padding-only factories through the local-checkpoint load path

The new padding-only branch works through serialization.load, but not through MCoreTensorAwareStateDict.to_state_dict, which local non-persistent checkpoint loading calls.

Failure: For the documented 61,760-row, GTP=64, alignment=32 configuration, rank 63 builds only a LocalNonpersistentObject. In to_state_dict, that placeholder is merged into recreated_state_dict, while apply_factory_merges receives only sharded_part. The GDP factory key is absent there, so loading raises ValueError: Different dict keys encountered in apply_factory_merges instead of restoring a zero-filled shard. The previous gather-based factory supplied checkpoint chunks even on this rank, so this failure is introduced by the new placeholder path.

Correction: In MCoreTensorAwareStateDict.to_state_dict, combine the loaded shards with common/nonpersistent state before applying factory merges, matching the ordering used by serialization.load. Replace the block at lines 386–389 with:

        with debug_time("apply_factory_merges", logger):
            merge(recreated_state_dict, sharded_part)
            recreated_state_dict = apply_factory_merges(recreated_state_dict, sh_ten_factories)

The invariant is that every factory merges against its complete loaded subtree, including locally reconstructed nonpersistent values. This is a suggested patch based on source inspection; verify it before adopting it.

Preservation and acceptance: Extend the proposed padding regression to round-trip through MCoreTensorAwareStateDict.from_state_dict/to_state_dict, with NVRX available, for both atomic and fully_parallel. The padding-only rank must return a (992, 2) tensor of zeros on the requested device and dtype, without saving the placeholder or padding. Full-real and partially padded ranks must retain their real values and restore zero padding. Also retain an ordinary split factory and unrelated common/nonpersistent values in the setup: those must still round-trip unchanged. Keep the existing global-DCP preprocessing and cross-GTP load assertions.

Exercise the replacement factory at TP > 1

The changed GDP workers initialize TP=1, and the synthetic preprocessing test passes tp_rank=0, tp_size=1. Consequently, neither the new factory’s TP-rank offset nor its TP-size global-shape calculation is exercised. Removing tp_rank * section_size would make no difference to these tests. The existing TP2→TP1 semantic-order regression calls only the old _split_tensor_factory.

Add a focused regression comparing assembled new-factory chunks with the non-GTP factory. For example, the following proposed test can be added alongside that existing semantic-order regression; its existing imports provide defaultdict, torch, ShardedTensor, and the layout/non-GTP helpers:

def test_gtp_inproj_tp2_matches_non_gtp_chunks():
    from megatron.core.ssm.gated_delta_product import _gtp_split_tensor_factory

    sections, names = _get_in_proj_checkpoint_split_layout(10, 3, 2, 3)
    expected, actual = defaultdict(list), defaultdict(list)
    for tp_rank in range(2):
        full = torch.arange(120).reshape(60, 2) + 1000 * tp_rank
        ref = ShardedTensor.from_rank_offsets("w", full, (0, tp_rank, 2))
        for chunk in _split_tensor_factory(ref, sections, names, 0).build():
            expected[chunk.key].append(chunk)
        padded = torch.nn.functional.pad(full, (0, 0, 0, 4))
        for gtp_rank in range(2):
            shard = padded[gtp_rank * 32 : (gtp_rank + 1) * 32]
            original = ShardedTensor.from_rank_offsets(
                "w", shard, (0, tp_rank * 2 + gtp_rank, 4)
            )
            factory = _gtp_split_tensor_factory(
                original, sections, names, gtp_rank, tp_rank, 2
            )
            chunks = factory.build()
            torch.testing.assert_close(factory.merge_fn([c.data for c in chunks]), shard)
            for chunk in chunks:
                actual[chunk.key].append(chunk)

    assert actual.keys() == expected.keys()
    for key, refs in expected.items():
        chunks = sorted(actual[key], key=lambda c: c.global_offset)
        shape, end = refs[0].global_shape, 0
        for chunk in chunks:
            assert chunk.global_shape == shape
            assert chunk.global_offset == (end, 0)
            end += chunk.local_shape[0]
        assert end == shape[0]
        torch.testing.assert_close(
            torch.cat([c.data for c in chunks]), torch.cat([c.data for c in refs])
        )

Here each TP-local tensor has 60 real rows in two 32-row GTP shards, and the GTP boundary cuts V2. Acceptance is identical checkpoint keys, global shapes and assembled values, complete nonoverlapping coverage, and correct local merges including four zero tail rows. Retain the existing TP=1 cross-GTP loads, storage-view assertions, and full/partial/all-padding cases. This suggested test has not been executed; verify it before adopting it.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex. This was a full review. Request a full review with /review reset=true.

[svcnvidia-nemo-ci]

Use double-quoted literals in the newly added _worker_gdp_inproj_cross_gtp_load. The repository style guide explicitly requires double quotes; Black preserves quote style, so formatting will leave this violation unchanged.

Suggested change

	key = 'mixer.in_proj.weight'
	key = "mixer.in_proj.weight"

Apply the same delimiter-only correction to the new worker’s group names, metadata keys, prefixes, and directory names. Preserve every string value and assertion: checkpoint keys and test behavior must not change. Acceptance is a static check that the new worker uses double quotes with identical literal contents. No test execution is claimed.

[deepakn94]

/ok to test 8fb2feb

[balasaajay]

/ok to test 7f474a1

[asolergi-nv]

Why this compute capability gate?

[deepakn94]

Not sure? Let's punt this to a future PR?

[jaredcasper]

/ok to test 303857d

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/41208466380