From 0f8e077b7dfa94333a05d2e1d2e698023a5fc37b Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Wed, 12 Aug 2026 20:07:14 -0700 Subject: [PATCH 1/2] fix(gtp): gather GDP in_proj shards before the DCP split GatedDeltaProductMixer.sharded_state_dict asserted under GTP_remat: in_proj.weight is GTP-sliced on axis 0 and pad-aligned, so the local shard matches neither in_proj_dim nor the [z|V|K|Q|b|a] split boundaries. Port MambaMixer's treatment -- all-gather the shards, strip the pad, then split (checkpoint stays identical to a non-GTP_remat run); wrap merge_fn to re-pad and re-slice on load. Signed-off-by: Shiqing Fan --- .../core/generalized_tensor_parallel.md | 8 +- megatron/core/ssm/gated_delta_product.py | 85 +++++- .../test_gtp_dcp.py | 267 ++++++++++++++++++ 3 files changed, 356 insertions(+), 4 deletions(-) diff --git a/docs/api-guide/core/generalized_tensor_parallel.md b/docs/api-guide/core/generalized_tensor_parallel.md index 11c01b8a504..6f8d4c26341 100644 --- a/docs/api-guide/core/generalized_tensor_parallel.md +++ b/docs/api-guide/core/generalized_tensor_parallel.md @@ -182,7 +182,7 @@ GTP_remat runs under both the standard **Adam** `DistributedOptimizer` and **Muo - **Adam** shards optimizer state over the gtp_remat/egtp_remat-excluded replicate group, like any GTP_remat run (§3.2). - **Muon** keeps matrix params *whole* (Newton–Schulz needs the full 2D weight). A GTP_remat-replicated whole param (e.g. MoE router, latent-proj MLPs by default) then lands on one checkpoint key shared by all GTP_remat peers, so the LayerWise optimizer folds `gtp_rank` into its `replica_id` — exactly one peer writes (the optimizer-state analog of the model-side fold in §3.3). -- **Native-FP8 optimizer-state matching (Muon path).** The save-side dequantize (§3.3) hands DCP a *fresh* BF16 tensor, which breaks the id-based optimizer-param → model-`ShardedTensor` match for every native-FP8 GTP_remat weight. The dequantized copy carries a `_gtp_dequant_src` backlink to the live FP8 param, and `_backfill_gtp_sharded_param_map` reuses the model's **own** entry (backlink first, tagged-name second) — preserving its full offsets (expert axes included) and `replica_id`. Only truly-unmatched params (Mamba `in_proj`, a gathered+split factory) take the per-shard rebuild, which refuses expert-parallel params rather than emit EP-colliding shards. +- **Native-FP8 optimizer-state matching (Muon path).** The save-side dequantize (§3.3) hands DCP a *fresh* BF16 tensor, which breaks the id-based optimizer-param → model-`ShardedTensor` match for every native-FP8 GTP_remat weight. The dequantized copy carries a `_gtp_dequant_src` backlink to the live FP8 param, and `_backfill_gtp_sharded_param_map` reuses the model's **own** entry (backlink first, tagged-name second) — preserving its full offsets (expert axes included) and `replica_id`. Only truly-unmatched params (the SSM `in_proj` weights, gathered+split factories) take the per-shard rebuild, which refuses expert-parallel params rather than emit EP-colliding shards. Neither path adds a GTP_remat-specific checkpoint format or call site. @@ -543,7 +543,9 @@ Because the offsets reconstruct the global shape, the checkpoint is independent **Alignment padding & cross-topology reshard.** When `_gtp_slice_one_param` pads `out_features` to a multiple of `gtp_remat_size · pad_for_alignment`, the saved global describes the *padded* shape, so the helper sets `allow_shape_mismatch=True`. DCP then tolerates a load-side topology whose alignment yields a different padded size — the unpadded data overlaps and the tail pad rows are zeros GTP_remat recomputes. -> Note: Mamba's `in_proj` is a special case: it **all-gathers its GTP_remat shards** back to the logical TP-local size and strips the pad *before* saving, so its global is topology-independent and needs no `allow_shape_mismatch`. +> Note: the SSM `in_proj` weights — Mamba's (`mamba_mixer.py`, split `[z|x|B|C|dt]`) and gated-delta-product's (`gated_delta_product.py`, split `[z|V|K|Q|b|a]`) — are a special case: each **all-gathers its GTP_remat shards** back to the logical TP-local size and strips the pad *before* saving, so its global is topology-independent and needs no `allow_shape_mismatch`. This is required, not just tidier: the split-chunk boundaries do not line up with the GTP_remat slice boundaries, so a raw shard cannot be split at all. The checkpoint therefore matches a non-GTP_remat run byte-for-byte. +> +> On **load**, the split factory's `merge_fn` is wrapped to invert this: it cats the chunks back to the unpadded TP-local width, re-pads with zeros up to `gtp_remat_local_size · gtp_remat_size`, and slices by the GTP_remat rank — mirroring `_gtp_slice_one_param` so the tensor lands in the live shard's layout. `gtp_remat_size == 1` skips both the gather and the pad/slice. **Optimizer state.** The distributed optimizer's master/moment `ShardedObject`s are keyed by `dp_group_idx`. Under GTP_remat/EGTP_remat each peer owns a *different* master shard (the optimizer shards over the gtp_remat/egtp_remat-**excluded** replicate group), so the index is taken from the gtp_remat/egtp_remat-**merged** model-parallel group (`mp_group` for dense, `expt_tp_pp_with_egtp_remat_group` for expert) — giving every peer a distinct key while replicate-group ranks remain true replicas under that key. @@ -739,7 +741,7 @@ torchrun --nproc-per-node 4 -m pytest tests/unit_tests/generalized_tensor_parall | `test_gtp_grad_correctness.py` | Gradient + dist-opt + grad-norm numeric parity vs a DP baseline at replicate (DP) > 1. Also the fp32-accumulation reduce-scatter (§2.6): gtp_remat-axis and DDP-axis parity, plus the size-2 bypass. | | `test_gtp_cudagraph_grad.py` | Capture-step grad-norm guard (§1.2): `_backup_grads_before_capture`/`_restore_grads_after_capture` keep a graph capture from clobbering finalized `main_grad` (own params + cross-graph `next_w`, incl. routed-expert `weight_list`). | | `test_gtp_partial_cg.py` | Four-layer partial-CG loss and eager-vs-replay grad-norm parity with two-slot ring reuse across independently replayed graphs (§3.5). | -| `test_gtp_dcp.py` | DCP sharding metadata (§3.3): TP×GTP_remat offsets, pad reshard, `replica_id`, native-FP8 save/load. | +| `test_gtp_dcp.py` | DCP sharding metadata (§3.3): TP×GTP_remat offsets, pad reshard, `replica_id`, native-FP8 save/load. Also the SSM `in_proj` gather+split: the gated-delta-product mixer's factory build/merge at MXFP8 alignment, and a full DCP save→load roundtrip of that mixer. | | `test_gtp_muon_dcp.py` | Muon optimizer-state DCP roundtrip (§1.6): `replica_id` fold + native-FP8 backfill matching. | | `test_gtp_recompute_chain.py` | Recompute-chain buffers (§3.1): adjacent nodes never share a gather buffer, dense and grouped, plus dgrad/wgrad parity vs no-recompute. | | `test_gtp_mtp.py` | GTP_remat + MTP shared weights (§3.5), 14 cases over `mtp_use_repeated_layer` × dense/MoE. Both MTP hazards are silent, so each needs its own guard: the async reduce-scatter path is compared numerically against the sync path on an identical model/sharding/batch, and all-gathers issued are tallied against consumes to catch a consume reading a buffer nothing gathered into. | diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 62f86e842d3..7a0f011a775 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -29,14 +29,25 @@ get_cu_seqlens, ) from megatron.core.tensor_parallel import get_cuda_rng_tracker +from megatron.core.tensor_parallel.gtp_api import HAVE_GTP from megatron.core.transformer import TransformerConfig from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.utils import ( + ensure_metadata_has_dp_cp_group, make_sharded_tensors_for_checkpoint, sharded_state_dict_default, ) -from megatron.core.utils import deprecate_inference_params, is_using_quantization_scales +from megatron.core.utils import ( + deprecate_inference_params, + is_using_quantization_scales, + make_tp_sharded_tensor_for_checkpoint, +) + +if HAVE_GTP: + from megatron.core.tensor_parallel.gtp_api import is_gtp_param +else: + is_gtp_param = None try: from causal_conv1d import causal_conv1d_fn, causal_conv1d_update @@ -932,6 +943,9 @@ def _get_states_from_cache(self, inference_context, batch_size, *, inference_par def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): """Provide a sharded state dictionary for distributed checkpointing.""" + # Guard for cases metadata is not provided + metadata = ensure_metadata_has_dp_cp_group(metadata) + sharded_state_dict = {} # Parameters self._save_to_state_dict(sharded_state_dict, "", keep_vars=True) @@ -968,6 +982,41 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + (1 + self.num_householder) * self.ngroups_local_tp * self.d_state + self.nheads_local_tp * (1 + self.num_householder) ) + # Under GTP, in_proj.weight is GTP-sliced along axis 0. The [z|V|K|Q|b|a] split boundaries + # don't line up with GTP slice boundaries, so gather the shards back to TP-local size + # (strip the trailing pad rows from the gathered tail) and fall through to the same + # split path the non-GTP run uses — saved ckpt format matches a non-GTP run. + in_proj_gtp_remat_size = getattr(self.in_proj.weight, "gtp_remat_size", 1) + in_proj_is_gtp = ( + in_proj_gtp_remat_size > 1 and HAVE_GTP and is_gtp_param(self.in_proj.weight) + ) + if in_proj_is_gtp: + gtp_remat_group = self.in_proj.weight.group + # in_proj.weight was already built at the sharded size by the submodule + # sharded_state_dict above — and, for native-FP8 GTP, dequantized to BF16 there + # (make_tp_sharded_tensor_for_checkpoint). Gather those (BF16) shards back to the + # full TP-local size so the [z|V|K|Q|b|a] split below matches a non-GTP run. + local = sharded_state_dict[f"{prefix}in_proj.weight"].data.contiguous() + gathered = torch.empty( + (local.shape[0] * in_proj_gtp_remat_size,) + local.shape[1:], + dtype=local.dtype, + device=local.device, + ) + torch.distributed.all_gather_into_tensor(gathered, local, group=gtp_remat_group) + if gathered.shape[0] != in_proj_dim: + gathered = gathered[:in_proj_dim].contiguous() + # Gathered weight is replicated across full dp_cp; replica_id needs only the DP slot. + dp_cp_rank = torch.distributed.get_rank(metadata["dp_cp_group"]) + sharded_state_dict[f"{prefix}in_proj.weight"] = make_tp_sharded_tensor_for_checkpoint( + gathered, + f"{prefix}in_proj.weight", + tp_axis=0, + replica_id=(0, 0, dp_cp_rank), + prepend_offsets=sharded_offsets, + tp_group=self.pg_collection.tp, + dp_cp_group=metadata["dp_cp_group"], + ) + assert sharded_state_dict[f"{prefix}in_proj.weight"].data.size(0) == in_proj_dim, ( in_proj_dim, sharded_state_dict[f"{prefix}in_proj.weight"], @@ -999,6 +1048,40 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): sharded_state_dict[key], in_proj_split_sections, in_proj_split_names, 0 ) + # GTP load-side inverse of the save-time all-gather (see + # docs/api-guide/core/generalized_tensor_parallel.md §3.3, in_proj note): the checkpoint + # stores the FULL TP-local in_proj.weight (pad stripped) under the 6 split keys + # [z|V|K|Q|b|a], so the default merge_fn cats them back to ``in_proj_dim`` rows with no + # padding. To reload into the live GTP param we must mirror init + # (``_gtp_slice_one_param``): F.pad the merged tensor with zeros up to + # ``gtp_remat_local_size * gtp_remat_size``, then slice by ``gtp_remat_local_rank``. + # gtp_remat_size=1 has no pad/slice. + if in_proj_is_gtp: + factory = sharded_state_dict[f"{prefix}in_proj.weight"] + gtp_remat_local_rank = torch.distributed.get_rank(self.in_proj.weight.group) + gtp_remat_local_size = self.in_proj.weight.data.size(0) + original_merge_fn = factory.merge_fn + + @torch.no_grad() + def _gtp_slice_after_cat( + sub_state_dict, + _orig=original_merge_fn, + _rank=gtp_remat_local_rank, + _size=gtp_remat_local_size, + _gtp_remat_size=in_proj_gtp_remat_size, + ): + full = _orig(sub_state_dict) + aligned_total = _size * _gtp_remat_size + pad_rows = aligned_total - full.shape[0] + if pad_rows > 0: + full = torch.nn.functional.pad(full, (0, 0, 0, pad_rows)) + start = _rank * _size + return full[start : start + _size].contiguous() + + sharded_state_dict[f"{prefix}in_proj.weight"] = replace( + factory, merge_fn=_gtp_slice_after_cat + ) + conv_dim = ( self.d_inner_local_tp * self.num_householder + (1 + self.num_householder) * self.ngroups_local_tp * self.d_state diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py index da42a086a91..64fb2ae9ada 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py @@ -942,6 +942,261 @@ def _worker_mamba_inproj_optim_param_map(rank, world_size, port): GTPShardedParam._chain_state = {} +# --------------------------------------------------------------------------- +# Gated-delta-product (GDP) in_proj: gather+split under GTP_remat +# +# GDP's ``in_proj.weight`` is GTP-sliced along axis 0 and zero-padded to an alignment multiple, +# while the checkpoint splits it into 6 chunks [z|V|K|Q|b|a] whose boundaries do NOT line up with +# the GTP slice boundaries. ``GatedDeltaProductMixer.sharded_state_dict`` therefore all-gathers the +# shards back to the TP-local width and strips the pad before splitting (§3.3), and wraps the +# factory's merge_fn to re-pad + re-slice on load. The three workers below cover that contract. +# --------------------------------------------------------------------------- + +# in_proj width = d_inner(256)*4 + 4*ngroups(2)*d_state(128) + nheads(4)*4 = 2064. With +# pad_for_alignment=32 (what setup_gtp_remat_from_recipe picks for MXFP8) and gtp_remat_size=2 the +# alignment block is 64, so 48 pad rows fire -- the padded-shard case the split path must handle. +_GDP_HIDDEN_SIZE = 256 + + +def _build_gdp_mixer(required_pgs): + """Build a 1-layer GatedDeltaProductMixer. Returns ``(mixer, pg, in_proj_dim)``. + + Callers must have set ``update_gtp_config(pad_for_alignment=32)`` and initialized model + parallel with ``gtp_remat_size=2`` first. + """ + from megatron.core.models.hybrid.hybrid_layer_specs import gdp_stack_spec + from megatron.core.ssm.gated_delta_product import GatedDeltaProductMixer + + pg = ProcessGroupCollection.use_mpu_process_groups(required_pgs=required_pgs) + config = TransformerConfig( + num_layers=1, + hidden_size=_GDP_HIDDEN_SIZE, + num_attention_heads=4, + mamba_num_heads=4, + mamba_head_dim=64, + mamba_num_groups=2, + mamba_state_dim=128, + params_dtype=torch.bfloat16, + bf16=True, + ) + mixer = GatedDeltaProductMixer( + config, + gdp_stack_spec.submodules.mamba_layer.submodules.mixer.submodules, + config.hidden_size, + layer_number=1, + pg_collection=pg, + ).cuda() + in_proj_dim = ( + mixer.d_inner_local_tp * (1 + mixer.num_householder) + + (1 + mixer.num_householder) * mixer.ngroups_local_tp * mixer.d_state + + mixer.nheads_local_tp * (1 + mixer.num_householder) + ) + in_proj_w = mixer.in_proj.weight + assert isinstance(in_proj_w, GTPShardedParam), "in_proj.weight should be GTP_remat-sharded" + assert in_proj_w.data.size(0) * 2 > in_proj_dim, ( + f"expected GTP alignment padding to fire (got {in_proj_w.data.size(0)} * 2 == " + f"{in_proj_dim}); these tests must cover the strip-pad / re-pad path" + ) + return mixer, pg, in_proj_dim + + +def _gdp_valid_rows(in_proj_w, in_proj_dim): + """Rows of this rank's GTP shard that hold real weights (the rest are alignment pad).""" + local_rows = in_proj_w.data.size(0) + gtp_remat_rank = torch.distributed.get_rank(in_proj_w.group) + return max(0, min(local_rows, in_proj_dim - gtp_remat_rank * local_rows)) + + +def _worker_gdp_inproj_gather_split(rank, world_size, port): + """GatedDeltaProductMixer.sharded_state_dict under GTP_remat. + + Regression for the GDP save crash: the raw GTP shard neither matches ``in_proj_dim`` nor lines + up with the [z|V|K|Q|b|a] split boundaries -- the pre-fix code asserted here. Verify the mixer + gathers back to TP-local size, splits into the 6 chunks a non-GTP_remat run would write, and + that the load-side merge_fn re-pads + re-slices back to the live GTP shard. + """ + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=2 + ) + try: + model_parallel_cuda_manual_seed(42) + update_gtp_config(pad_for_alignment=32) # MXFP8 alignment + mixer, _, in_proj_dim = _build_gdp_mixer(['tp', 'cp', 'gtp_remat']) + in_proj_w = mixer.in_proj.weight + + metadata = {'dp_cp_group': ps.get_data_parallel_group(with_context_parallel=True)} + # Pre-fix this raised AssertionError((in_proj_dim, ShardedTensor(...))). + sd = mixer.sharded_state_dict(prefix='mixer.', metadata=metadata) + + factory = sd['mixer.in_proj.weight'] + assert isinstance(factory, ShardedTensorFactory), type(factory) + # Save side: the gathered tensor is the full TP-local width, pad stripped. + assert factory.data.size(0) == in_proj_dim, (factory.data.size(0), in_proj_dim) + + chunks = factory.build_fn(factory.key, factory.data, factory.replica_id, None) + assert [t.key.rsplit('.', 1)[-1] for t in chunks] == ["z", "V", "K", "Q", "b", "a"] + assert sum(t.data.size(0) for t in chunks) == in_proj_dim + + # Load side: cat the 6 chunks, re-pad, re-slice -> exactly this rank's live GTP shard. + merged = factory.merge_fn([t.data for t in chunks]) + assert tuple(merged.shape) == tuple(in_proj_w.data.shape), ( + tuple(merged.shape), + tuple(in_proj_w.data.shape), + ) + # The pad rows the last GTP rank carries are never written to the ckpt -> back as zeros. + n_valid = _gdp_valid_rows(in_proj_w, in_proj_dim) + torch.testing.assert_close(merged[:n_valid], in_proj_w.data[:n_valid], rtol=0, atol=0) + assert torch.equal(merged[n_valid:], torch.zeros_like(merged[n_valid:])) + finally: + ps.destroy_model_parallel() + ps.initialize_model_parallel() + GTPShardedParam._chain_state = {} + + +def _worker_gdp_save_load_roundtrip(rank, world_size, ckpt_base): + """End-to-end DCP save->load of a GatedDeltaProductMixer under GTP_remat. + + Companion to ``_worker_gdp_inproj_gather_split``, which only checks the factory build/merge + functions in isolation. This drives the real ``save``/``load`` so the load-side merge_fn + (re-pad + re-slice back to the live GTP shard) is exercised through DCP, and so a + duplicate-writer replica_id would surface as an 'Invalid access pattern'. + """ + from megatron.core.dist_checkpointing import load, save + from tests.unit_tests.dist_checkpointing import TempNamedDir + + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=2 + ) + try: + model_parallel_cuda_manual_seed(42) + update_gtp_config(pad_for_alignment=32) # MXFP8 alignment + mixer, pg, in_proj_dim = _build_gdp_mixer( + ['tp', 'cp', 'gtp_remat', 'dp_cp', 'dp_cp_gtp_remat'] + ) + in_proj_w = mixer.in_proj.weight + + # ``save_checkpoint_and_time`` threads the gtp_remat-INCLUSIVE group; using the + # gtp_remat-excluding pg.dp_cp here collides replica_ids across gtp_remat peers. + metadata = {'dp_cp_group': pg.dp_cp_gtp_remat} + golden = {k: v.detach().clone() for k, v in mixer.state_dict().items()} + + with TempNamedDir(ckpt_base / 'gdp_gtp_dcp_roundtrip', sync=True) as ckpt_dir: + save(mixer.sharded_state_dict(prefix='mixer.', metadata=metadata), ckpt_dir) + + # Scribble over every param so a no-op load cannot pass. + with torch.no_grad(): + for p in mixer.parameters(): + p.data.fill_(float(rank + 1)) + loaded = load(mixer.sharded_state_dict(prefix='mixer.', metadata=metadata), ckpt_dir) + + # in_proj comes back through the 6-way split + the GTP re-pad/re-slice merge_fn. + merged = loaded['mixer.in_proj.weight'] + assert tuple(merged.shape) == tuple(in_proj_w.data.shape), ( + tuple(merged.shape), + tuple(in_proj_w.data.shape), + ) + n_valid = _gdp_valid_rows(in_proj_w, in_proj_dim) + torch.testing.assert_close( + merged[:n_valid].cpu(), golden['in_proj.weight'][:n_valid].cpu(), rtol=0, atol=0 + ) + + # The rest of the mixer must round-trip too -- a colliding replica_id across gtp_remat + # peers would either fail the load or return another rank's data. + for name in ( + 'A_log', + 'dt_bias', + 'conv1d.weight', + 'norm.weight', + 'out_proj.weight', + 'in_proj.layer_norm_weight', + ): + key = f'mixer.{name}' + assert key in loaded, f"{key} missing from the loaded state dict: {sorted(loaded)}" + torch.testing.assert_close( + loaded[key].cpu(), golden[name].cpu(), rtol=0, atol=0, msg=f"{name} drifted" + ) + finally: + ps.destroy_model_parallel() + ps.initialize_model_parallel() + GTPShardedParam._chain_state = {} + + +def _worker_gdp_inproj_optim_param_map(rank, world_size, port): + """GDP ``in_proj`` must survive the optimizer id->ShardedTensor match (Muon path, §1.6). + + Same gap as the Mamba case (``_worker_mamba_inproj_optim_param_map``): the model entry for a + gathered+split ``in_proj`` exposes the *gathered* tensor, so it never id-matches the per-shard + GTP optimizer param and ``get_param_id_to_sharded_param_map`` drops it -> KeyError in + ``Float16OptimizerWithFloat16Params.sharded_state_dict``. Unlike that test, this one drives the + real production backfill (``_backfill_gtp_sharded_param_map``) rather than reproducing its + rebuild, so it also pins that GDP takes the per-shard rebuild branch, not the EP refusal. + """ + from megatron.core.dist_checkpointing.optimizer import ( + get_param_id_to_sharded_param_map, + make_sharded_optimizer_tensor, + ) + from megatron.core.optimizer.optimizer import _backfill_gtp_sharded_param_map + from megatron.core.tensor_parallel.generalized_tensor_parallelism import ( + tag_gtp_params_with_names, + ) + + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=2 + ) + try: + model_parallel_cuda_manual_seed(42) + update_gtp_config(pad_for_alignment=32) # MXFP8 alignment + mixer, pg, _ = _build_gdp_mixer(['tp', 'cp', 'gtp_remat', 'dp_cp', 'dp_cp_gtp_remat']) + tag_gtp_params_with_names(mixer) # sets _debug_name, mirrors production setup + in_proj_w = mixer.in_proj.weight + + metadata = {'dp_cp_group': pg.dp_cp_gtp_remat} + model_sd = mixer.sharded_state_dict(prefix='mixer.', metadata=metadata) + + # The gap: the gathered+split factory does not id-match the per-shard optimizer param. + id_map = get_param_id_to_sharded_param_map(model_sd, [in_proj_w]) + assert 0 not in id_map, "expected in_proj to be MISSING from the id map (the KeyError gap)" + + # The production backfill must fill it via the per-shard rebuild. An expert-parallel param + # would raise instead; in_proj is dense, so it must rebuild cleanly. + _backfill_gtp_sharded_param_map(id_map, [[in_proj_w]], model_sd) + assert 0 in id_map, "backfill did not restore in_proj" + entry = id_map[0] + # A plain per-shard ShardedTensor keyed by the tagged name -- NOT the model's gathered+split + # factory (reusing that would hand the optimizer the wrong shape). + assert isinstance(entry, ShardedTensor), type(entry) + assert entry is not model_sd['mixer.in_proj.weight'] + assert entry.key == in_proj_w._debug_name, (entry.key, in_proj_w._debug_name) + assert tuple(entry.local_shape) == tuple(in_proj_w.shape), ( + f"rebuilt local_shape {tuple(entry.local_shape)} != param shape " + f"{tuple(in_proj_w.shape)}" + ) + # The rebuilt entry must describe this rank's GTP slice of the PADDED global (what the + # optimizer shard actually is), not the gathered/pad-stripped width the model entry uses. + gtp_remat_rank = torch.distributed.get_rank(in_proj_w.group) + assert entry.global_offset[0] == gtp_remat_rank * in_proj_w.shape[0], ( + entry.global_offset, + gtp_remat_rank, + ) + assert entry.global_shape[0] == in_proj_w.shape[0] * in_proj_w.gtp_remat_size, ( + entry.global_shape, + in_proj_w.shape, + ) + + # make_sharded_optimizer_tensor must accept it for a same-shape optimizer state tensor. + opt_state = torch.zeros_like(in_proj_w) + osh = make_sharded_optimizer_tensor(entry, opt_state, prefix='optimizer.state.exp_avg') + assert osh is not None + assert tuple(osh.local_shape) == tuple(in_proj_w.shape) + finally: + ps.destroy_model_parallel() + ps.initialize_model_parallel() + GTPShardedParam._chain_state = {} + + def _worker_save_load_roundtrip_needs_gtp_inclusive_group(rank, world_size, ckpt_base): """Save->load roundtrip: save and load must use the gtp_remat-INCLUSIVE replica group. @@ -1022,6 +1277,18 @@ def test_mamba_inproj_optim_param_map(self): _require_world_size(4) _worker_mamba_inproj_optim_param_map(dist.get_rank(), 4, None) + def test_gdp_inproj_gather_split(self): + _require_world_size(4) + _worker_gdp_inproj_gather_split(dist.get_rank(), 4, None) + + def test_gdp_save_load_roundtrip(self, tmp_path_dist_ckpt): + _require_world_size(4) + _worker_gdp_save_load_roundtrip(dist.get_rank(), 4, tmp_path_dist_ckpt) + + def test_gdp_inproj_optim_param_map(self): + _require_world_size(4) + _worker_gdp_inproj_optim_param_map(dist.get_rank(), 4, None) + def test_replicated_param_needs_gtp_inclusive_dp_cp(self): _require_world_size(4) _worker_replicated_param_needs_gtp_inclusive_dp_cp(dist.get_rank(), 4, None) From 74b1ebf6f945c1cb532f9216c47555cd42a73e4c Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Wed, 12 Aug 2026 21:21:54 -0700 Subject: [PATCH 2/2] fix new added UTs and refine comments Signed-off-by: Shiqing Fan --- .../core/generalized_tensor_parallel.md | 2 +- megatron/core/ssm/gated_delta_product.py | 14 +++++++------- .../test_gtp_dcp.py | 19 +++++++++++++++---- 3 files changed, 23 insertions(+), 12 deletions(-) diff --git a/docs/api-guide/core/generalized_tensor_parallel.md b/docs/api-guide/core/generalized_tensor_parallel.md index 6f8d4c26341..eca9e2b7eb8 100644 --- a/docs/api-guide/core/generalized_tensor_parallel.md +++ b/docs/api-guide/core/generalized_tensor_parallel.md @@ -543,7 +543,7 @@ Because the offsets reconstruct the global shape, the checkpoint is independent **Alignment padding & cross-topology reshard.** When `_gtp_slice_one_param` pads `out_features` to a multiple of `gtp_remat_size · pad_for_alignment`, the saved global describes the *padded* shape, so the helper sets `allow_shape_mismatch=True`. DCP then tolerates a load-side topology whose alignment yields a different padded size — the unpadded data overlaps and the tail pad rows are zeros GTP_remat recomputes. -> Note: the SSM `in_proj` weights — Mamba's (`mamba_mixer.py`, split `[z|x|B|C|dt]`) and gated-delta-product's (`gated_delta_product.py`, split `[z|V|K|Q|b|a]`) — are a special case: each **all-gathers its GTP_remat shards** back to the logical TP-local size and strips the pad *before* saving, so its global is topology-independent and needs no `allow_shape_mismatch`. This is required, not just tidier: the split-chunk boundaries do not line up with the GTP_remat slice boundaries, so a raw shard cannot be split at all. The checkpoint therefore matches a non-GTP_remat run byte-for-byte. +> Note: the SSM `in_proj` weights — Mamba's (`mamba_mixer.py`, split `[z|x|B|C|dt]`) and gated-delta-product's (`gated_delta_product.py`, split householder-major into `z|V*|K*|Q|b*|a`) — are a special case: each **all-gathers its GTP_remat shards** back to the logical TP-local size and strips the pad *before* saving, so its global is topology-independent and needs no `allow_shape_mismatch`. This is required, not just tidier: the split-chunk boundaries do not line up with the GTP_remat slice boundaries, so a raw shard cannot be split at all. The checkpoint therefore matches a non-GTP_remat run byte-for-byte. > > On **load**, the split factory's `merge_fn` is wrapped to invert this: it cats the chunks back to the unpadded TP-local width, re-pads with zeros up to `gtp_remat_local_size · gtp_remat_size`, and slices by the GTP_remat rank — mirroring `_gtp_slice_one_param` so the tensor lands in the live shard's layout. `gtp_remat_size == 1` skips both the gather and the pad/slice. diff --git a/megatron/core/ssm/gated_delta_product.py b/megatron/core/ssm/gated_delta_product.py index 7a0f011a775..77c49198739 100644 --- a/megatron/core/ssm/gated_delta_product.py +++ b/megatron/core/ssm/gated_delta_product.py @@ -982,10 +982,10 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + (1 + self.num_householder) * self.ngroups_local_tp * self.d_state + self.nheads_local_tp * (1 + self.num_householder) ) - # Under GTP, in_proj.weight is GTP-sliced along axis 0. The [z|V|K|Q|b|a] split boundaries - # don't line up with GTP slice boundaries, so gather the shards back to TP-local size - # (strip the trailing pad rows from the gathered tail) and fall through to the same - # split path the non-GTP run uses — saved ckpt format matches a non-GTP run. + # Under GTP, in_proj.weight is GTP-sliced along axis 0. The [z|V*|K*|Q|b*|a] split + # boundaries don't line up with GTP slice boundaries, so gather the shards back to + # TP-local size (strip the trailing pad rows from the gathered tail) and fall through + # to the same split path the non-GTP run uses — saved ckpt matches a non-GTP run. in_proj_gtp_remat_size = getattr(self.in_proj.weight, "gtp_remat_size", 1) in_proj_is_gtp = ( in_proj_gtp_remat_size > 1 and HAVE_GTP and is_gtp_param(self.in_proj.weight) @@ -995,7 +995,7 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): # in_proj.weight was already built at the sharded size by the submodule # sharded_state_dict above — and, for native-FP8 GTP, dequantized to BF16 there # (make_tp_sharded_tensor_for_checkpoint). Gather those (BF16) shards back to the - # full TP-local size so the [z|V|K|Q|b|a] split below matches a non-GTP run. + # full TP-local size so the [z|V*|K*|Q|b*|a] split below matches a non-GTP run. local = sharded_state_dict[f"{prefix}in_proj.weight"].data.contiguous() gathered = torch.empty( (local.shape[0] * in_proj_gtp_remat_size,) + local.shape[1:], @@ -1050,8 +1050,8 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): # GTP load-side inverse of the save-time all-gather (see # docs/api-guide/core/generalized_tensor_parallel.md §3.3, in_proj note): the checkpoint - # stores the FULL TP-local in_proj.weight (pad stripped) under the 6 split keys - # [z|V|K|Q|b|a], so the default merge_fn cats them back to ``in_proj_dim`` rows with no + # stores the FULL TP-local in_proj.weight (pad stripped) under the per-householder split + # keys, so the default merge_fn cats them back to ``in_proj_dim`` rows with no # padding. To reload into the live GTP param we must mirror init # (``_gtp_slice_one_param``): F.pad the merged tensor with zeros up to # ``gtp_remat_local_size * gtp_remat_size``, then slice by ``gtp_remat_local_rank``. diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py index 64fb2ae9ada..af62304c5f8 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py @@ -946,7 +946,7 @@ def _worker_mamba_inproj_optim_param_map(rank, world_size, port): # Gated-delta-product (GDP) in_proj: gather+split under GTP_remat # # GDP's ``in_proj.weight`` is GTP-sliced along axis 0 and zero-padded to an alignment multiple, -# while the checkpoint splits it into 6 chunks [z|V|K|Q|b|a] whose boundaries do NOT line up with +# while the checkpoint splits it into householder-major chunks whose boundaries do NOT line up with # the GTP slice boundaries. ``GatedDeltaProductMixer.sharded_state_dict`` therefore all-gathers the # shards back to the TP-local width and strips the pad before splitting (§3.3), and wraps the # factory's merge_fn to re-pad + re-slice on load. The three workers below cover that contract. @@ -1011,7 +1011,7 @@ def _worker_gdp_inproj_gather_split(rank, world_size, port): """GatedDeltaProductMixer.sharded_state_dict under GTP_remat. Regression for the GDP save crash: the raw GTP shard neither matches ``in_proj_dim`` nor lines - up with the [z|V|K|Q|b|a] split boundaries -- the pre-fix code asserted here. Verify the mixer + up with the in_proj split boundaries -- the pre-fix code asserted here. Verify the mixer gathers back to TP-local size, splits into the 6 chunks a non-GTP_remat run would write, and that the load-side merge_fn re-pads + re-slices back to the live GTP shard. """ @@ -1034,11 +1034,22 @@ def _worker_gdp_inproj_gather_split(rank, world_size, port): # Save side: the gathered tensor is the full TP-local width, pad stripped. assert factory.data.size(0) == in_proj_dim, (factory.data.size(0), in_proj_dim) + from megatron.core.ssm.gated_delta_product import _get_in_proj_checkpoint_split_layout + + # The chunk names/sizes come from _get_in_proj_checkpoint_split_layout (householder-major: + # z, V0..V(M-1), K0..K(M-1), Q, b0..b(M-1), a). Derive the expectation from that helper so + # this pins "GTP splits exactly like a non-GTP_remat run" rather than a frozen key list. + _, expected_names = _get_in_proj_checkpoint_split_layout( + mixer.d_inner_local_tp, + mixer.ngroups_local_tp * mixer.d_state, + mixer.nheads_local_tp, + mixer.num_householder, + ) chunks = factory.build_fn(factory.key, factory.data, factory.replica_id, None) - assert [t.key.rsplit('.', 1)[-1] for t in chunks] == ["z", "V", "K", "Q", "b", "a"] + assert [t.key.rsplit('.', 1)[-1] for t in chunks] == expected_names assert sum(t.data.size(0) for t in chunks) == in_proj_dim - # Load side: cat the 6 chunks, re-pad, re-slice -> exactly this rank's live GTP shard. + # Load side: cat the chunks, re-pad, re-slice -> exactly this rank's live GTP shard. merged = factory.merge_fn([t.data for t in chunks]) assert tuple(merged.shape) == tuple(in_proj_w.data.shape), ( tuple(merged.shape),