refactor mtp

src/mcore_bridge/__init__.py

@@ -10,15 +10,15 @@
 if TYPE_CHECKING:
     from .bridge import GPTBridge
     from .config import MLAModelConfig, ModelConfig, hf_to_mcore_config
-    from .model import get_mcore_model
+    from .model import get_mcore_model, resolve_mtp_config, resolve_mtp_num_layers
     from .tuners import LoraParallelLinear
     from .utils import get_logger, set_random_seed, split_cp_inputs, unwrap_model
     from .version import __release_datetime__, __version__
 else:
     _import_structure = {
         'bridge': ['GPTBridge'],
         'config': ['MLAModelConfig', 'ModelConfig', 'hf_to_mcore_config'],
-        'model': ['get_mcore_model'],
+        'model': ['get_mcore_model', 'resolve_mtp_config', 'resolve_mtp_num_layers'],
         'tuners': ['LoraParallelLinear'],
         'utils': ['get_logger', 'set_random_seed', 'split_cp_inputs', 'unwrap_model'],
         'version': ['__release_datetime__', '__version__'],

src/mcore_bridge/bridge/gpt_bridge.py

@@ -1931,8 +1931,14 @@ def _convert(self, mg_models, hf_state_dict, hf_prefix: str, to_mcore: bool, tqd
             if to_mcore and self.pp_rank > 0:
                 self._set_state_dict(lm_model, 'embedding.word_embeddings.weight', hf_state_dict, self.hf_embed_key,
                                      to_mcore)
+            from mcore_bridge.model.mtp import get_physical_mtp_count
+            if not to_mcore:
+                phys_count = self._reduce_mtp_physical_count(lm_model)
+            else:
+                # Import: only the owning (last) stage reaches here.
+                phys_count = get_physical_mtp_count(lm_model)
             layer_idx = 0
-            while layer_idx < self.config.mtp_num_layers:
+            while layer_idx < phys_count:
                 res = self._convert_mtp_layer(lm_model, hf_state_dict, f'{self.hf_mtp_prefix}.', layer_idx, to_mcore)
                 layer_idx += 1
                 if to_mcore:
@@ -1968,7 +1974,19 @@ def _convert_mtp_extra(self, mtp_layer, hf_state_dict, to_mcore, origin_hf_state
         self._set_state_dict(mtp_layer, 'final_layernorm.weight', hf_state_dict, self.hf_mtp_final_layernorm_key,
                              to_mcore)
 
+    def _reduce_mtp_physical_count(self, lm_model) -> int:
+        """PP export: obtain the physical MTP layer count from the owning stage
+        so every rank iterates the same number of times (collective safety)."""
+        from mcore_bridge.model.mtp import get_physical_mtp_count
+        local = get_physical_mtp_count(lm_model)
+        if self.pp_size <= 1:
+            return local
+        t = torch.tensor([local], dtype=torch.long, device='cuda')
+        dist.all_reduce(t, op=dist.ReduceOp.MAX, group=self.pp_group)
+        return int(t.item())
+
     def _convert_mtp_layer(self, lm_model, hf_state_dict, hf_prefix: str, layer_idx: int, to_mcore: bool):
+        from mcore_bridge.model.mtp import get_mtp_inner_layer
         mtp_layer = lm_model.mtp.layers[layer_idx] if hasattr(lm_model, 'mtp') else None
         if self.hf_mtp_prefix == self.hf_layers_prefix:
             hf_layer_idx = layer_idx + self.config.num_layers
@@ -1981,8 +1999,14 @@ def _convert_mtp_layer(self, lm_model, hf_state_dict, hf_prefix: str, layer_idx:
             if len(hf_state_dict) == 0:
                 if self._peft_format:
                     return {}
-                logger.info(f'MTP Layer {mtp_layer.layer_number} safetensors weights not found, '
-                            'this part will be randomly initialized.')
+                expected = f'{hf_prefix}*'
+                if not self.config.mtp_allow_random_init:
+                    raise ValueError(f'MTP weights not found under {expected} (expected '
+                                     f'{self.config.mtp_num_layers} physical set(s) after the {self.config.num_layers} '
+                                     'decoder layers). Set mtp_shared_weights=True to reuse one set across depths, or '
+                                     'mtp_allow_random_init=True to train MTP from random init.')
+                logger.warning(f'MTP Layer {mtp_layer.layer_number} weights not found under {expected}; '
+                               'randomly initializing (mtp_allow_random_init=True).')
                 for param in mtp_layer.parameters():
                     if param.ndim == 2:
                         mtp_layer.config.init_method(param.data)
@@ -1991,7 +2015,7 @@ def _convert_mtp_layer(self, lm_model, hf_state_dict, hf_prefix: str, layer_idx:
             origin_hf_state_dict = {}
             hf_state_dict = {}
         self._convert_mtp_extra(mtp_layer, hf_state_dict, to_mcore, origin_hf_state_dict)
-        transformer_layer = None if mtp_layer is None else mtp_layer.transformer_layer
+        transformer_layer = None if mtp_layer is None else get_mtp_inner_layer(mtp_layer)
         self._convert_mtp_embeds(lm_model, hf_state_dict, to_mcore)
         hf_state_dict.update(self._set_layer_attn(transformer_layer, hf_state_dict, -1, to_mcore))
         hf_state_dict.update(self._set_layer_mlp(transformer_layer, hf_state_dict, -1, to_mcore, is_mtp=True))

src/mcore_bridge/config/model_config.py

@@ -250,6 +250,7 @@ class ModelConfig(TransformerConfig):
     csa_compress_rotary_base: float = 40000.0
     o_groups: int = 8
     o_lora_rank: int = 1024
+
     # Canonical names consumed by TransformerConfig.__post_init__ since the main->dev sync (0.20+);
     # `o_groups` / `o_lora_rank` remain the dev-era aliases and are reconciled against these.
     output_projection_groups: int = 8
@@ -287,7 +288,10 @@ class ModelConfig(TransformerConfig):
 
     # mtp
     mtp_decoder_input_detach: bool = False
-    mtp_shared_weights: bool = False
+    # Allow a full-finetune checkpoint that carries no MTP weights to random-init them.
+    # Off by default: a missing MTP key set is usually a wrong path/name, not intent.
+    mtp_allow_random_init: bool = False
+    mtp_shared_weights: Optional[bool] = None
 
     # visual
     language_model_only: bool = False
@@ -399,12 +403,42 @@ def __post_init__(self):
             self.apply_query_key_layer_scaling = self.fp16
         if self.apply_query_key_layer_scaling:
             os.environ['NVTE_APPLY_QK_LAYER_SCALING'] = '1'
+        # mtp_num_layers is always the *logical* prediction depth.
+        # mtp_shared_weights is a deprecated alias for mtp_use_repeated_layer.
+        if isinstance(self.mtp_num_layers, str):
+            raise TypeError('mtp_num_layers must be an int or None; '
+                            'use resolve_mtp_config() to resolve "auto".')
         if self.mtp_shared_weights:
-            assert self.mtp_num_layers is not None
-            self.mtp_unroll_steps = self.mtp_num_layers
-            self.mtp_num_layers = 1
-        else:
-            self.mtp_unroll_steps = self.mtp_num_layers
+            assert self.mtp_num_layers is not None and self.mtp_num_layers > 0, \
+                'mtp_shared_weights requires a positive mtp_num_layers'
+            if not hasattr(self, 'mtp_use_repeated_layer'):
+                # COMPAT(0.16): the field does not exist yet; the block provides the shim.
+                object.__setattr__(self, 'mtp_use_repeated_layer', True)
+            else:
+                self.mtp_use_repeated_layer = True
+        elif not hasattr(self, 'mtp_use_repeated_layer'):
+            object.__setattr__(self, 'mtp_use_repeated_layer', False)
+        if (self.mtp_use_repeated_layer and self.pipeline_model_parallel_layout is not None
+                and (self.mtp_num_layers or 0) > 1
+                and 'mtp_use_repeated_layer' not in TransformerConfig.__dataclass_fields__):
+            # MCore 0.16 only: its layout assert counts MTP entries against
+            # config.mtp_num_layers (the LOGICAL depth) while the bridge
+            # repeated-layer shim physically builds ONE layer -- with depth>1
+            # no consistent layout mapping exists.  Depth 1 is conflict-free
+            # (logical == physical == 1, the pre-migration behaviour) and stays
+            # allowed.  Note for 0.17+: upstream supports repeated layers
+            # natively but its layout assert still requires the custom layout
+            # to declare the LOGICAL depth in MTP entries; layouts written for
+            # the old shared semantics (one physical entry) must be updated
+            # accordingly -- upgrading alone does not remap them.
+            raise ValueError('mtp_shared_weights with mtp_num_layers > 1 and a custom '
+                             'pipeline_model_parallel_layout is not supported on megatron-core 0.16: '
+                             'the upstream layout assert counts MTP entries against the logical depth '
+                             'while the repeated-layer shim builds one physical layer. Use depth 1, '
+                             'drop the custom layout, or use megatron-core >= 0.17 with a layout that '
+                             'declares mtp_num_layers MTP entries.')
+        # Legacy alias kept for downstream code that still reads it
+        self.mtp_unroll_steps = self.mtp_num_layers
         # ``num_nextn_predict_layers`` counts the draft layers for both DeepSeek's standard MTP and
         # V4.1's DSpark, so it alone cannot tell them apart: on a plain V3/V4 checkpoint it means MTP
         # and there is no DSpark at all. ``dspark_block_size`` is what actually marks a DSpark

src/mcore_bridge/model/__init__.py

@@ -3,4 +3,5 @@
 from .constant import ModelType
 from .gpt_model import GPTModel
 from .mm_gpt_model import MultimodalGPTModel
+from .mtp_detect import resolve_mtp_config, resolve_mtp_num_layers
 from .register import MODEL_MAPPING, get_mcore_model, get_mcore_model_type, get_model_meta