[TRTLLM-13394][feat] Support loading MTP weights from standalone checkpoint

tensorrt_llm/_torch/models/modeling_nemotron_h.py

@@ -912,7 +912,17 @@ def __init__(
             model_nextn = self.config.num_nextn_predict_layers
             ckpt_nextn = self.config.num_nextn_predict_layers
             self.num_hidden_layers = self.config.num_hidden_layers
-            assert ckpt_nextn > 0, "There are not MTP modules in the checkpoint."
+            has_external_mtp = (
+                model_config.spec_config.loads_mtp_from_separate_checkpoint)
+            assert ckpt_nextn > 0 or has_external_mtp, (
+                "There are not MTP modules in the checkpoint. "
+                "Set speculative_config.speculative_model to a separate MTP "
+                "heads checkpoint, or use a target checkpoint that embeds MTP.")
+            if ckpt_nextn == 0 and has_external_mtp:
+                # Neither checkpoint declares a head count: fall back to a
+                # single shared head, matching MTPForCausalLM's MTP-Eagle
+                # default.
+                ckpt_nextn = model_nextn = 1
             if ckpt_nextn == 1 and not model_config.spec_config.use_mtp_vanilla:
                 pass
             else:
@@ -976,6 +986,13 @@ def load_weights(self,
                      weights: dict,
                      weight_mapper: BaseWeightMapper,
                      allow_partial_loading: bool = False):
+        from tensorrt_llm._torch.speculative.utils import (
+            filter_mtp_checkpoint_weights, loads_mtp_from_speculative_model)
+
+        if loads_mtp_from_speculative_model(self.model_config.spec_config):
+            # Filter before preprocess: mapper remaps mtp.layers.* ->
+            # model.layers.{N}.* and would otherwise load embedded MTP heads.
+            weights = filter_mtp_checkpoint_weights(weights)
         new_weights = weight_mapper.preprocess_weights(weights)
         super().load_weights(weights=new_weights,
                              weight_mapper=weight_mapper,

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -2122,20 +2122,101 @@ def forward(
 
         return logits
 
+    def mtp_head_module_names(self) -> List[str]:
+        """Names of the MTP heads under every alias they are reachable by.
+
+        One-model MTP registers the same head objects twice: under
+        ``draft_model.mtp_layers.{h}`` and, after the target model extends its
+        layer list, under ``model.layers.{num_hidden_layers + h}``. A load that
+        wants to leave the heads untouched has to exclude both aliases.
+        """
+        mtp_layers = getattr(self.draft_model, "mtp_layers", None)
+        if not mtp_layers:
+            return []
+        head_ids = {id(layer) for layer in mtp_layers}
+        return [
+            name for name, module in self.named_modules(remove_duplicate=False)
+            if name and id(module) in head_ids
+        ]
+
     def load_weights(self,
                      weights: Dict,
                      weight_mapper: Optional[BaseWeightMapper] = None,
                      params_map: Optional[Dict[str, str]] = None,
                      allow_partial_loading: bool = False):
+        from tensorrt_llm._torch.speculative.utils import (
+            filter_mtp_checkpoint_weights, loads_mtp_from_speculative_model)
+
+        skip_modules = ["draft_model"]
+        if loads_mtp_from_speculative_model(self.spec_config):
+            # The heads come from speculative_model in a second pass
+            # (load_draft_weights), so exclude them here. They must be
+            # *skipped* rather than tolerated via allow_partial_loading:
+            # partial loading suppresses process_weights_after_loading() on
+            # every quantized Linear/MoE it touches, which would leave the
+            # target model's quant scales (NVFP4 alphas, MoE input scales)
+            # uninitialized.
+            weights = filter_mtp_checkpoint_weights(weights)
+            skip_modules.extend(self.mtp_head_module_names())
         super().load_weights(weights=weights,
                              weight_mapper=weight_mapper,
-                             skip_modules=["draft_model"],
+                             skip_modules=skip_modules,
                              params_map=params_map,
                              allow_partial_loading=allow_partial_loading)
 
     def load_draft_weights(self,
                            weights: Dict,
                            weight_mapper: Optional[BaseWeightMapper] = None):
+        from tensorrt_llm._torch.models.modeling_utils import \
+            _load_weights_impl_v2
+        from tensorrt_llm._torch.speculative.utils import (
+            loads_mtp_from_speculative_model,
+            remap_preprocessed_mtp_weights_for_draft_model,
+            select_mtp_checkpoint_weights,
+            skip_modules_for_separate_mtp_checkpoint)
+
+        if loads_mtp_from_speculative_model(self.spec_config):
+            # Load MTP heads into draft_model only, and verify every non-shared
+            # MTP parameter has a matching tensor. The previous parent-model
+            # load used allow_partial_loading=True, which silently left MTP
+            # modules at random init when keys did not bind.
+            n_total = len(weights)
+            weights = select_mtp_checkpoint_weights(weights)
+            if not weights:
+                raise ValueError(
+                    "speculative_model was set for MTP but no 'mtp.*' weights "
+                    f"were found in {self.spec_config.speculative_model!r}. "
+                    "Expected keys like 'mtp.layers.0.*'.")
+            n_dropped = n_total - len(weights)
+            if n_dropped:
+                logger.warning(
+                    "Ignoring %d non-mtp.* tensors from speculative_model while "
+                    "loading MTP heads (kept %d mtp.* tensors).", n_dropped,
+                    len(weights))
+            if weight_mapper is None:
+                raise ValueError(
+                    "weight_mapper is required to load separate MTP heads")
+            weights = weight_mapper.preprocess_weights(weights)
+            num_hidden_layers = self.config.num_hidden_layers
+            num_mtp_layers = len(self.draft_model.mtp_layers)
+            weights = remap_preprocessed_mtp_weights_for_draft_model(
+                weights,
+                num_hidden_layers=num_hidden_layers,
+                num_mtp_layers=num_mtp_layers,
+            )
+
+            # Skip optional modules (e.g. shared_head) only when absent from
+            # this checkpoint; architectures that ship those tensors still load
+            # them under allow_partial_loading=False.
+            _load_weights_impl_v2(
+                self.draft_model,
+                weights,
+                weight_mapper,
+                skip_modules=skip_modules_for_separate_mtp_checkpoint(weights),
+                allow_partial_loading=False,
+            )
+            return
+
         args = inspect.getfullargspec(self.draft_model.load_weights).args
         if "weight_mapper" in args:
             self.draft_model.load_weights(weights=weights,

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -412,6 +412,11 @@ def load_config_and_apply_defaults(
             config_kwargs['mapping'] = llm_args.parallel_config.to_mapping()
 
         if llm_args.speculative_config:
+            from tensorrt_llm._torch.speculative.utils import \
+                resolve_mtp_checkpoint_source
+
+            resolve_mtp_checkpoint_source(llm_args.speculative_config,
+                                          checkpoint_dir)
             config_kwargs['spec_config'] = llm_args.speculative_config
 
         config = checkpoint_loader.load_config(checkpoint_dir, **config_kwargs)
@@ -555,8 +560,7 @@ def load(
 
             loads_draft_weights = (
                 self.spec_config is not None
-                and (self.spec_config.spec_dec_mode.need_load_draft_weights()
-                     or self.spec_config._use_shared_kv_cache))
+                and self.spec_config.needs_separate_draft_weights)
             speculative_mode = self._speculative_mode_name(self.spec_config)
             post_transform_qualification = self._qualify_post_transform_profile(
                 model,
@@ -708,19 +712,7 @@ def init_meta_tensor(t: torch.Tensor):
                                             self.weight_mapper)
 
                 if loads_draft_weights:
-                    weights = checkpoint_loader.load_weights(
-                        self.spec_config.speculative_model,
-                        mapping=self.mapping)
-
-                    draft_model_arch = model.draft_config.pretrained_config.architectures[
-                        0]
-                    draft_weight_mapper = AutoCheckpointMapper.get(
-                        checkpoint_loader.checkpoint_format, draft_model_arch)
-                    draft_weight_mapper.init_model_and_config(
-                        model.draft_model, model.draft_config)
-
-                    self._call_load_weights(model.load_draft_weights, weights,
-                                            draft_weight_mapper)
+                    self._load_separate_draft_weights(model, checkpoint_loader)
 
             elif load_format == LoadFormat.GMS:
                 # GPU Memory Service path: weight tensors live in a
@@ -843,21 +835,8 @@ def init_meta_tensor_in_pool(t: torch.Tensor):
                                     "pool.")
 
                             if loads_draft_weights:
-                                draft_weights = checkpoint_loader.load_weights(
-                                    self.spec_config.speculative_model,
-                                    mapping=self.mapping)
-
-                                draft_model_arch = model.draft_config.pretrained_config.architectures[
-                                    0]
-                                draft_weight_mapper = AutoCheckpointMapper.get(
-                                    checkpoint_loader.checkpoint_format,
-                                    draft_model_arch)
-                                draft_weight_mapper.init_model_and_config(
-                                    model.draft_model, model.draft_config)
-
-                                self._call_load_weights(
-                                    model.load_draft_weights, draft_weights,
-                                    draft_weight_mapper)
+                                self._load_separate_draft_weights(
+                                    model, checkpoint_loader)
 
                             # Run post_load hooks INSIDE the pool so any
                             # tensors they create or rebind (fused QKV,
@@ -1117,6 +1096,32 @@ def _speculative_mode_name(
             return "unknown"
         return mode_name.lower()
 
+    def _load_separate_draft_weights(
+            self, model: DecoderModelForCausalLM,
+            checkpoint_loader: BaseCheckpointLoader) -> None:
+        """Load draft/MTP weights from ``speculative_model`` into the one-engine model.
+
+        Eagle3 / external drafters use a draft-specific mapper and ``draft_config``.
+        One-model MTP with separate heads reuses the target architecture mapper
+        because MTP modules are already attached under the target model.
+        """
+        draft_weights = checkpoint_loader.load_weights(
+            self.spec_config.speculative_model, mapping=self.mapping)
+
+        if model.draft_config is not None:
+            draft_model_arch = model.draft_config.pretrained_config.architectures[
+                0]
+            draft_weight_mapper = AutoCheckpointMapper.get(
+                checkpoint_loader.checkpoint_format, draft_model_arch)
+            draft_weight_mapper.init_model_and_config(model.draft_model,
+                                                      model.draft_config)
+        else:
+            # MTP one-model + separate MTP checkpoint: no draft HF architecture.
+            draft_weight_mapper = self.weight_mapper
+
+        self._call_load_weights(model.load_draft_weights, draft_weights,
+                                draft_weight_mapper)
+
     @classmethod
     def _qualify_post_transform_profile(
             cls,
@@ -1359,6 +1364,12 @@ def _load_and_validate_config(
             self, checkpoint_dir: str,
             checkpoint_loader: BaseCheckpointLoader) -> ModelConfig:
         """Loads and validates the model configuration."""
+        from tensorrt_llm._torch.speculative.utils import (
+            loads_mtp_from_speculative_model, resolve_mtp_checkpoint_source,
+            update_spec_config_from_model_config)
+
+        resolve_mtp_checkpoint_source(self.spec_config, checkpoint_dir)
+
         load_config_kwargs = dict(
             checkpoint_dir=checkpoint_dir,
             trust_remote_code=self.llm_args.trust_remote_code,
@@ -1403,6 +1414,14 @@ def _load_and_validate_config(
 
         config = checkpoint_loader.load_config(**load_config_kwargs)
 
+        if loads_mtp_from_speculative_model(self.spec_config):
+            # `load_config_and_apply_defaults` already ran this, but against a
+            # config object it then discards. The MTP heads' structure fields
+            # (head count, block pattern) come from `speculative_model` and
+            # have to reach the config the model is actually built from.
+            update_spec_config_from_model_config(self.spec_config,
+                                                 config.pretrained_config)
+
         # Store nvfp4 config in extra_attrs for Linear layer access
         config.extra_attrs[
             'nvfp4_gemm_allowed_backends'] = config.nvfp4_gemm_allowed_backends