Refactor state dict canonicalization

Do not modify `sharded_state_dict` method, instead use
`generate_state_dict` to apply the transformation only when
loading/saving.

Author: janEbert

megatron/core/models/hybrid/hybrid_block.py

@@ -13,7 +13,7 @@
 from torch import Tensor, nn
 
 from megatron.core.dist_checkpointing.mapping import ShardedStateDict
-from megatron.core.dist_checkpointing.utils import apply_prefix_mapping, replace_prefix_for_sharding
+from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
 from megatron.core.enums import Fp8Recipe
 from megatron.core.extensions.transformer_engine import TENorm
 from megatron.core.fp4_utils import get_fp4_context
@@ -32,25 +32,6 @@
 from megatron.core.transformer.utils import sharded_state_dict_default
 from megatron.core.utils import WrappedTensor, deprecate_inference_params, make_viewless_tensor
 
-# Canonical slot name each layer symbol uses in its stand-alone
-# sharded-state-dict output – i.e., the attribute path under which the
-# primitive's weights live when the block contains just that symbol. For a
-# fused `[XY]` block, the same primitives sit under `self_attention` (for the
-# sequence mixer X) and `mlp` (for the channel mixer Y); `sharded_state_dict`
-# uses this table to rewrite fused keys back into the stand-alone layout so
-# fused and unfused patterns produce the same checkpoint keys. Only Mamba
-# needs an intra-block rename (`self_attention.` -> `mixer.`); every other
-# sequence mixer is stand-alone-hosted in a TransformerLayer whose
-# `self_attention` slot already matches the fused layout.
-_CANONICAL_SLOT_FOR_SYMBOL: dict[str, str] = {
-    LayerSymbols.MAMBA: "mixer",
-    LayerSymbols.GDN: "self_attention",
-    LayerSymbols.ATTENTION: "self_attention",
-    LayerSymbols.DS_ATTENTION: "self_attention",
-    LayerSymbols.MLP: "mlp",
-    LayerSymbols.MOE: "mlp",
-}
-
 
 @dataclass
 class HybridStackSubmodules:
@@ -103,13 +84,6 @@ class HybridStack(GraphableMegatronModule, MegatronModule):
         pp_layer_offset (int, optional): the global layer offset for this pipeline
             segment, measured in physical blocks (fused groups count as one).
             Defaults to 0.
-        sub_layer_offset (int, optional): the global sub-layer offset for this
-            pipeline segment, measured in unfused sub-layers (a fused `[XY]`
-            contributes 2). Used only to emit canonical-form sharded-state-dict
-            keys so fused and unfused patterns produce the same checkpoint
-            layout. When None (default), falls back to `pp_layer_offset`, which
-            is correct whenever no earlier pipeline segment contained a fusion
-            group.
         post_layer_norm (bool, optional): whether to include a final layer norm.
             Defaults to True.
         post_process (bool, optional): whether to include an output layer.
@@ -128,7 +102,6 @@ def __init__(
         pre_process: bool = True,
         layer_type_list: Optional[list[str]] = None,
         pp_layer_offset: int = 0,
-        sub_layer_offset: Optional[int] = None,
         post_layer_norm: bool = True,
         post_process: bool = True,
         device=None,
@@ -156,12 +129,6 @@ def __init__(
             "--hybrid-layer-pattern by HybridModel."
         )
         self.layer_type_list = layer_type_list
-        # Sub-layer offset defaults to pp_layer_offset, which is correct for
-        # any pattern whose earlier pipeline segments contain no fusion groups
-        # (i.e., physical-block count == sub-layer count up to that point).
-        self.sub_layer_offset = (
-            sub_layer_offset if sub_layer_offset is not None else pp_layer_offset
-        )
 
         # Build layers from the pre-selected segment
         self.layers = nn.ModuleList()
@@ -479,70 +446,21 @@ def sharded_state_dict(
         sharded_state_dict = {}
         layer_prefix = f'{prefix}layers.'
 
-        # Sub-layer cursor tracks the running sub-layer index across this
-        # segment; combined with `sub_layer_offset` it yields the same global
-        # sub-layer indices whether the current pattern contains fused groups
-        # or not. That is what gives checkpoints a fusion-independent layout –
-        # a pattern saved unfused can be loaded into a fused model (and vice
-        # versa) without any external translation step.
-        sub_layer_cursor = self.sub_layer_offset
+        for local_layer_idx, layer in enumerate(self.layers):
 
-        for local_layer_idx, (layer_type, layer) in enumerate(
-            zip(self.layer_type_list, self.layers)
-        ):
+            global_layer_offset = layer.layer_number - 1  # self.layer_number starts at 1
             state_dict_prefix = (
                 f'{layer_prefix}{local_layer_idx}.'  # module list index in HybridStack
             )
+
+            sharded_prefix = f'{layer_prefix}{global_layer_offset}.'
             sharded_pp_offset = []
+
             layer_sharded_state_dict = layer.sharded_state_dict(
                 state_dict_prefix, sharded_pp_offset, metadata
             )
 
-            if len(layer_type) == 1:
-                # Stand-alone block: one physical block == one sub-layer, and
-                # the block's attribute layout already matches the canonical
-                # stand-alone layout. Only the outer block index needs to move
-                # from the local module-list index to the global sub-layer
-                # index.
-                canonical_prefix = f'{layer_prefix}{sub_layer_cursor}.'
-                replace_prefix_for_sharding(
-                    layer_sharded_state_dict, state_dict_prefix, canonical_prefix
-                )
-                sub_layer_cursor += 1
-            else:
-                # Fused block `[XY]`: split the single physical block's keys
-                # into two sub-layer prefixes so the checkpoint looks exactly
-                # as it would for stand-alone `X` followed by stand-alone `Y`.
-                # Norms attached to X (e.g. `input_layernorm` for DSA) attach
-                # to X's sub-layer index; norms attached to Y (e.g.
-                # `pre_mlp_layernorm` for MoE) attach to Y's sub-layer index.
-                x_sym, y_sym = layer_type[0], layer_type[1]
-                canonical_x_prefix = f'{layer_prefix}{sub_layer_cursor}.'
-                canonical_y_prefix = f'{layer_prefix}{sub_layer_cursor + 1}.'
-                slot_for_x = _CANONICAL_SLOT_FOR_SYMBOL[x_sym]
-                slot_for_y = _CANONICAL_SLOT_FOR_SYMBOL[y_sym]
-
-                # Order matters: `apply_prefix_mapping` picks the first
-                # matching prefix, so list each specific sub-prefix before
-                # the bare block prefix fallback.
-                prefix_map = {
-                    f'{state_dict_prefix}input_layernorm.': (
-                        f'{canonical_x_prefix}input_layernorm.'
-                    ),
-                    f'{state_dict_prefix}self_attention.': (f'{canonical_x_prefix}{slot_for_x}.'),
-                    f'{state_dict_prefix}self_attn_bda.': (f'{canonical_x_prefix}self_attn_bda.'),
-                    f'{state_dict_prefix}pre_mlp_layernorm.': (
-                        f'{canonical_y_prefix}pre_mlp_layernorm.'
-                    ),
-                    f'{state_dict_prefix}mlp.': f'{canonical_y_prefix}{slot_for_y}.',
-                    f'{state_dict_prefix}mlp_bda.': f'{canonical_y_prefix}mlp_bda.',
-                    # Fallback for any stray top-level fused-block keys (e.g.,
-                    # `_extra_state` attached to the TransformerLayer itself);
-                    # attach them to X's sub-layer index by convention.
-                    state_dict_prefix: canonical_x_prefix,
-                }
-                apply_prefix_mapping(layer_sharded_state_dict, prefix_map)
-                sub_layer_cursor += 2
+            replace_prefix_for_sharding(layer_sharded_state_dict, state_dict_prefix, sharded_prefix)
 
             sharded_state_dict.update(layer_sharded_state_dict)
 

megatron/core/models/hybrid/hybrid_layer_fusion.py

@@ -16,6 +16,8 @@
 
 from typing import TYPE_CHECKING
 
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.dist_checkpointing.utils import apply_prefix_mapping
 from megatron.core.extensions.transformer_engine import TENorm
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols as LayerSymbols
@@ -222,3 +224,129 @@ def build_fused_layer(
         build_kwargs["pp_layer_offset"] = pp_layer_offset
 
     return build_module(fused_spec, **build_kwargs)
+
+
+# Canonical slot name each layer symbol uses in its stand-alone
+# sharded-state-dict output – i.e., the attribute path under which the
+# primitive's weights live when the block contains just that symbol. For a
+# fused `[XY]` block, the same primitives sit under `self_attention` (for the
+# sequence mixer X) and `mlp` (for the channel mixer Y); canonicalization
+# uses this table to rewrite fused keys back into the stand-alone layout so
+# fused and unfused patterns produce the same checkpoint keys. Only Mamba
+# needs an intra-block rename (`self_attention.` -> `mixer.`); every other
+# sequence mixer is stand-alone-hosted in a TransformerLayer whose
+# `self_attention` slot already matches the fused layout.
+_CANONICAL_SLOT_FOR_SYMBOL: dict[str, str] = {
+    LayerSymbols.MAMBA: "mixer",
+    LayerSymbols.GDN: "self_attention",
+    LayerSymbols.ATTENTION: "self_attention",
+    LayerSymbols.DS_ATTENTION: "self_attention",
+    LayerSymbols.MLP: "mlp",
+    LayerSymbols.MOE: "mlp",
+}
+
+
+def canonicalize_hybrid_sharded_state_dict(
+    sharded_state_dict: ShardedStateDict,
+    layer_prefix: str,
+    layer_type_list: list[str],
+    physical_offset: int = 0,
+    sub_layer_offset: int = 0,
+) -> None:
+    """Rewrite HybridStack layer keys into the canonical (unfused) layout, in place.
+
+    `HybridStack.sharded_state_dict` emits keys indexed by global physical
+    block position within the model (a fused `[XY]` group still occupies a
+    single physical block). Fused blocks are realized as `TransformerLayer`s
+    whose `self_attention` slot holds the sequence mixer and `mlp` slot
+    holds the channel mixer, so their keys do not match what a stand-alone
+    `X` followed by stand-alone `Y` would produce. This function rewrites
+    each fused block's keys into two sub-layer-indexed prefixes that
+    do match: `layers.{sub_layer_offset + i}.mixer.*` for mamba sub-layers,
+    `layers.{sub_layer_offset + i}.mlp.*` for MLP, etc. Stand-alone blocks
+    are simply re-indexed from physical to sub-layer index.
+
+    The resulting keys are fusion-independent – a checkpoint written with
+    `[*-]M` and one written with `*-M` end up with the same set of keys, so
+    the dist_checkpointing layer can load either into either.
+
+    Args:
+        sharded_state_dict: The sharded state dict to rewrite in place. Only
+            entries whose keys start with `layer_prefix` are touched.
+        layer_prefix: The full prefix up to and including `"layers."`, e.g.
+            `"decoder.layers."`. Keys outside this prefix are left alone.
+        layer_type_list: The per-physical-block layer-type symbols for this
+            pipeline segment (a single `"M"`, `"*"`, etc. for a stand-alone
+            block; a two-char string like `"*-"` for a fused group).
+        physical_offset: The global physical-block index at which this
+            pipeline segment starts (i.e., the value `HybridStack` uses to
+            derive each layer's `layer_number`). Defaults to 0, which is
+            correct for non-pipeline-parallel runs.
+        sub_layer_offset: The global sub-layer index at which this pipeline
+            segment starts. Accounts for sub-layers contributed by earlier
+            pipeline segments so that fused groups in those segments are
+            correctly counted. When the model is not pipeline-parallel (or
+            no earlier segment contains a fusion group), this is equal to
+            `physical_offset` and may be left at its default.
+
+    Notes:
+        - Build up a combined prefix map across all layers and apply it in
+          one `apply_prefix_mapping` pass. This keeps the rewrite narrow
+          (sibling keys outside `layer_prefix` are untouched) and lets the
+          function safely run on the full model state dict without tripping
+          on embedding or output-layer entries.
+        - Order matters inside the combined prefix map: `apply_prefix_mapping`
+          picks the first matching prefix, so the specific sub-prefixes
+          (e.g. `"input_layernorm."`) are inserted before the bare block
+          prefix fallback.
+        - Norms attached to the X sub-layer (e.g. `input_layernorm` for DSA)
+          stay with X's sub-layer index; norms attached to Y (e.g.
+          `pre_mlp_layernorm` for MoE) attach to Y's sub-layer index.
+    """
+    prefix_map: dict[str, str] = {}
+    sub_layer_cursor = sub_layer_offset
+
+    for local_layer_idx, layer_type in enumerate(layer_type_list):
+        physical_prefix = f'{layer_prefix}{physical_offset + local_layer_idx}.'
+
+        if len(layer_type) == 1:
+            # Stand-alone block: one physical block == one sub-layer, and the
+            # block's attribute layout already matches the canonical
+            # stand-alone layout. Only the outer block index needs to move
+            # from the local module-list index to the global sub-layer index.
+            canonical_prefix = f'{layer_prefix}{sub_layer_cursor}.'
+            prefix_map[physical_prefix] = canonical_prefix
+            sub_layer_cursor += 1
+        else:
+            # Fused block `[XY]`: split the single physical block's keys into
+            # two sub-layer prefixes so the checkpoint looks exactly as it
+            # would for stand-alone `X` followed by stand-alone `Y`.
+            x_sym, y_sym = layer_type[0], layer_type[1]
+            canonical_x_prefix = f'{layer_prefix}{sub_layer_cursor}.'
+            canonical_y_prefix = f'{layer_prefix}{sub_layer_cursor + 1}.'
+            slot_for_x = _CANONICAL_SLOT_FOR_SYMBOL[x_sym]
+            slot_for_y = _CANONICAL_SLOT_FOR_SYMBOL[y_sym]
+
+            # Specific sub-prefixes before the bare block prefix fallback.
+            prefix_map[f'{physical_prefix}input_layernorm.'] = (
+                f'{canonical_x_prefix}input_layernorm.'
+            )
+            prefix_map[f'{physical_prefix}self_attention.'] = (
+                f'{canonical_x_prefix}{slot_for_x}.'
+            )
+            prefix_map[f'{physical_prefix}self_attn_bda.'] = (
+                f'{canonical_x_prefix}self_attn_bda.'
+            )
+            prefix_map[f'{physical_prefix}pre_mlp_layernorm.'] = (
+                f'{canonical_y_prefix}pre_mlp_layernorm.'
+            )
+            prefix_map[f'{physical_prefix}mlp.'] = f'{canonical_y_prefix}{slot_for_y}.'
+            prefix_map[f'{physical_prefix}mlp_bda.'] = f'{canonical_y_prefix}mlp_bda.'
+            # Fallback for any stray top-level fused-block keys (e.g.,
+            # `_extra_state` attached to the TransformerLayer itself);
+            # attach them to X's sub-layer index by convention.
+            prefix_map[physical_prefix] = canonical_x_prefix
+            sub_layer_cursor += 2
+
+    if prefix_map:
+        apply_prefix_mapping(sharded_state_dict, prefix_map)

megatron/core/models/hybrid/hybrid_model.py

@@ -194,12 +194,16 @@ def __init__(
             first_stage_layers=self.config.num_layers_in_first_pipeline_stage,
             last_stage_layers=self.config.num_layers_in_last_pipeline_stage,
         )
-        # Sub-layer offset mirrors `layer_offset` but counts each character in a
-        # fused `[XY]` group separately. `HybridStack` uses it to emit
-        # sharded-state-dict keys in the canonical (unfused) layout so
-        # checkpoints saved under one fusion configuration load cleanly under
-        # another.
-        sub_layer_offset = get_sub_layer_offset(main_pattern, layer_offset)
+        # Read at checkpoint save/load time by
+        # `megatron.training.checkpointing._apply_hybrid_canonicalization_if_applicable`,
+        # which rewrites the decoder's sharded keys into a fusion-independent
+        # layout so checkpoints saved under one fusion configuration load
+        # cleanly under another. `_decoder_physical_offset` mirrors `layer_offset`
+        # (physical-block index where this pipeline segment starts);
+        # `_decoder_sub_layer_offset` is its sub-layer counterpart, counting
+        # each character of a fused `[XY]` group separately.
+        self._decoder_physical_offset = layer_offset
+        self._decoder_sub_layer_offset = get_sub_layer_offset(main_pattern, layer_offset)
 
         # Determine if MTP is needed (based on pattern parsing)
         self.mtp_process = (
@@ -263,7 +267,6 @@ def __init__(
             pre_process=self.pre_process,
             layer_type_list=layer_type_list,
             pp_layer_offset=layer_offset,
-            sub_layer_offset=sub_layer_offset,
             post_process=self.post_process,
             dtype=config.params_dtype,
             pg_collection=self.pg_collection,