diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 8516a3dbdb1e..b755f387ffee 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/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, diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index a991e266965d..3cd27ee0ae05 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/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, diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 2a6ea23076aa..c9e42c67559f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/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 diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index a8d92fed4f46..d3f1c041e864 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import json +import os from bisect import bisect_left from dataclasses import dataclass from typing import TYPE_CHECKING, Dict, Optional @@ -42,6 +44,247 @@ "Gemma4ForConditionalGeneration", ) +# MTP structure fields copied from a separate MTP-head checkpoint onto the +# target pretrained config when ``speculative_model`` is set. +# Prefer writable HF fields: NemotronHConfig exposes mtp_hybrid_override_pattern +# as a read-only property derived from mtp_layers_block_type. +_MTP_STRUCTURE_FIELDS_FROM_DRAFT = ( + "num_nextn_predict_layers", + "mtp_layers_block_type", + "mtp_block_configs", +) + +_MTP_PATTERN_TO_LAYER = { + "M": "mamba", + "E": "moe", + "*": "attention", + "-": "mlp", +} + + +def _set_pretrained_config_attr(model_config, + name: str, + value, + *, + required: bool = True) -> bool: + """Set a config field, tolerating read-only properties / strict dataclasses. + + Each write is verified by reading the value back: a class-level property + shadows ``__dict__``, so writing through ``vars()`` can appear to succeed + while the config keeps reporting its old value. When ``required`` is False, + failures are logged and ignored (used for optional fields like + ``mtp_block_configs``). + """ + writes = ( + lambda: setattr(model_config, name, value), + lambda: vars(model_config).__setitem__(name, value), + ) + for write in writes: + try: + write() + except (TypeError, AttributeError): + continue + if getattr(model_config, name, None) == value: + return True + + message = (f"Unable to set MTP config field '{name}' on " + f"{type(model_config).__name__}") + if required: + raise AttributeError(message) + logger.warning("%s; keeping the target checkpoint's value.", message) + return False + + +def _pattern_to_mtp_layers_block_type(pattern: str) -> list: + try: + return [_MTP_PATTERN_TO_LAYER[char] for char in pattern] + except KeyError as exc: + raise ValueError( + f"Invalid mtp_hybrid_override_pattern {pattern!r}: " + f"expected characters in {sorted(_MTP_PATTERN_TO_LAYER)}") from exc + + +def _is_mtp_checkpoint_weight_key(key: str) -> bool: + """Return True for checkpoint keys that belong to MTP heads.""" + return key.startswith("mtp.") or key.startswith("mtp/") + + +def filter_mtp_checkpoint_weights(weights: dict) -> dict: + """Drop ``mtp.*`` keys so embedded MTP heads do not override a separate MTP checkpoint.""" + return { + k: v + for k, v in weights.items() if not _is_mtp_checkpoint_weight_key(k) + } + + +def select_mtp_checkpoint_weights(weights: dict) -> dict: + """Keep only ``mtp.*`` keys from a (possibly full) checkpoint dict. + + Separate MTP-head checkpoints may still ship unrelated tensors (or a full + target copy). Loading those into the one-engine model would overwrite the + already-loaded target backbone and corrupt generation. + """ + return { + k: v + for k, v in weights.items() if _is_mtp_checkpoint_weight_key(k) + } + + +def remap_preprocessed_mtp_weights_for_draft_model( + weights: dict, + num_hidden_layers: int, + num_mtp_layers: int, +) -> dict: + """Map ``model.layers.{{N[+h]}}.*`` keys onto ``mtp_layers.{{h}}.*``. + + Nemotron preprocess rewrites ``mtp.layers.*`` onto the target module path + ``model.layers.{{num_hidden_layers}}.*``. For a strict draft-only load we + re-home those keys under ``draft_model.mtp_layers``. + """ + remapped: dict = {} + unused: list[str] = [] + for key, value in weights.items(): + matched = False + for head_idx in range(num_mtp_layers): + prefix = f"model.layers.{num_hidden_layers + head_idx}." + if key.startswith(prefix): + remapped[f"mtp_layers.{head_idx}.{key[len(prefix):]}"] = value + matched = True + break + if not matched: + unused.append(key) + if unused: + sample = ", ".join(unused[:8]) + more = "" if len(unused) <= 8 else f" (+{len(unused) - 8} more)" + raise ValueError( + "After MTP preprocess, expected keys under " + f"'model.layers.{{{num_hidden_layers}+h}}.*' for " + f"h in [0, {num_mtp_layers}), but found unmatched keys: " + f"{sample}{more}") + return remapped + + +def skip_modules_for_separate_mtp_checkpoint(weights: dict) -> list[str]: + """Modules to skip when loading a separate MTP checkpoint into draft_model. + + ``shared_head`` is optional across MTP architectures: + - Nemotron ``mtp.*`` checkpoints omit it (final_layernorm on the last + sublayer is the trained norm; ``shared_head`` only wraps ``lm_head``). + - DeepSeek / Qwen / Exaone / Step3 ship ``shared_head.norm`` (and Step3 + also ships a dedicated ``shared_head.output``). + + Skip only when the remapped weight dict has no ``shared_head`` keys so + strict loading stays architecture-agnostic. + """ + skip: list[str] = [] + if not any("shared_head" in key for key in weights): + skip.append("shared_head") + return skip + + +def loads_mtp_from_speculative_model(spec_config) -> bool: + """True when one-model MTP should load heads from ``speculative_model``.""" + if spec_config is None: + return False + return spec_config.loads_mtp_from_separate_checkpoint + + +def _refers_to_same_checkpoint(lhs, rhs) -> bool: + """True when two model references point at the same checkpoint.""" + if lhs is None or rhs is None: + return False + if str(lhs) == str(rhs): + return True + try: + return os.path.samefile(str(lhs), str(rhs)) + except OSError: + # At least one side is not an existing local directory (e.g. a Hub + # model id), so string equality above was the only comparison. + return False + + +def resolve_mtp_checkpoint_source(spec_config, checkpoint_dir) -> None: + """Keep one-model MTP on the target checkpoint when both paths match. + + Before separate MTP checkpoints were supported, ``speculative_model`` was + ignored for one-model MTP and the heads always came from the target + weights. Configs that point ``speculative_model`` at the target checkpoint + keep that behavior instead of switching to the separate-heads load path, + which the target checkpoint's key layout may not even satisfy. + """ + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + if not isinstance(spec_config, MTPDecodingConfig): + return + if spec_config.speculative_model is None: + return + if not _refers_to_same_checkpoint(spec_config.speculative_model, + checkpoint_dir): + return + if not spec_config._mtp_heads_in_target_checkpoint: + logger.info( + "speculative_model points at the target checkpoint " + f"({checkpoint_dir}); loading MTP heads from the target weights.") + spec_config._mtp_heads_in_target_checkpoint = True + + +def _load_speculative_model_config_dict(spec_config) -> Optional[dict]: + """Read ``config.json`` from ``spec_config.speculative_model``, if present.""" + draft_dir = getattr(spec_config, "speculative_model", None) + if not draft_dir: + return None + try: + cfg_path = os.path.join(str(draft_dir), "config.json") + if not os.path.isfile(cfg_path): + return None + with open(cfg_path) as f: + return json.load(f) + except (OSError, ValueError, TypeError, AttributeError) as exc: + logger.warning( + f"Unable to read speculative_model config from {draft_dir}: {exc}") + return None + + +def _merge_mtp_fields_from_speculative_model(spec_config, + model_config) -> Optional[int]: + """Overlay MTP structure fields from ``speculative_model`` onto ``model_config``. + + Returns the MTP layer count from the draft checkpoint when available. + """ + draft_cfg = _load_speculative_model_config_dict(spec_config) + if not draft_cfg: + return None + + draft_nextn = draft_cfg.get("num_nextn_predict_layers") + if draft_nextn is None: + draft_nextn = draft_cfg.get("mtp_num_hidden_layers") + + for field in _MTP_STRUCTURE_FIELDS_FROM_DRAFT: + if field in draft_cfg and draft_cfg[field] is not None: + _set_pretrained_config_attr( + model_config, + field, + draft_cfg[field], + required=(field != "mtp_block_configs"), + ) + + # HF NemotronHConfig: mtp_hybrid_override_pattern is a read-only property + # derived from mtp_layers_block_type. Convert the pattern when the draft + # checkpoint only provides the legacy string form. + if (draft_cfg.get("mtp_layers_block_type") is None + and draft_cfg.get("mtp_hybrid_override_pattern") is not None): + _set_pretrained_config_attr( + model_config, + "mtp_layers_block_type", + _pattern_to_mtp_layers_block_type( + draft_cfg["mtp_hybrid_override_pattern"]), + ) + + if draft_nextn is not None: + _set_pretrained_config_attr(model_config, "num_nextn_predict_layers", + draft_nextn) + return int(draft_nextn) + return None + def _is_effective_dynamic_tree(spec_config) -> bool: # At dynamic_tree_max_topK == 1 the tree collapses to a linear chain; route @@ -557,6 +800,15 @@ def update_spec_config_from_model_config(spec_config, model_config): and architectures[0] in _GEMMA4_SHARED_KV_TARGET_ARCHITECTURES): spec_config._use_shared_kv_cache = ( spec_config.spec_dec_mode.is_mtp_eagle_one_model()) + + # When MTP heads live in a separate checkpoint, prefer that checkpoint's + # layer count / pattern over the target model's (which may have no MTP or + # an older embedded MTP head that will be overridden at weight load). + draft_nextn = None + if loads_mtp_from_speculative_model(spec_config): + draft_nextn = _merge_mtp_fields_from_speculative_model( + spec_config, model_config) + # Read the MTP layer count from the model's pretrained config. This # determines the actual MTP layer count in the checkpoint and drives the # spec_dec_mode decision (EAGLE vs vanilla MTP). Different checkpoints expose @@ -564,13 +816,16 @@ def update_spec_config_from_model_config(spec_config, model_config): # `num_nextn_predict_layers`, while Qwen3Next-style configs (including # Qwen3.5) use `mtp_num_hidden_layers`. Fall back to a single shared MTP / # EAGLE layer when neither field is present. - num_nextn_predict_layers = getattr(model_config, "num_nextn_predict_layers", - None) - if num_nextn_predict_layers is None: + if draft_nextn is not None: + num_nextn_predict_layers = draft_nextn + else: num_nextn_predict_layers = getattr(model_config, - "mtp_num_hidden_layers", None) - if num_nextn_predict_layers is None: - num_nextn_predict_layers = 1 + "num_nextn_predict_layers", None) + if num_nextn_predict_layers is None: + num_nextn_predict_layers = getattr(model_config, + "mtp_num_hidden_layers", None) + if num_nextn_predict_layers is None: + num_nextn_predict_layers = 1 spec_config.num_nextn_predict_layers = num_nextn_predict_layers is_vanilla = spec_config.spec_dec_mode.is_mtp_vanilla() diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 55dbfa0f2ab8..f2cf73c93d66 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1731,7 +1731,9 @@ class DecodingBaseConfig(StrictBaseModel): "speculative_model_dir"), description= "The speculative (draft) model. Accepts either (1) a HuggingFace Hub model ID (e.g. 'yuhuili/EAGLE3-LLaMA3.1-Instruct-8B'), " - "which will be automatically downloaded, or (2) a local filesystem path to a downloaded model directory." + "which will be automatically downloaded, or (2) a local filesystem path to a downloaded model directory. " + "For MTP, when set to a checkpoint other than the target model, loads MTP heads from it instead of any " + "embedded mtp.* weights in the target; pointing it at the target model keeps the embedded heads." ) max_concurrency: Optional[PositiveInt] = Field( @@ -1811,6 +1813,9 @@ class DecodingBaseConfig(StrictBaseModel): _allow_separate_draft_kv_cache: bool = PrivateAttr(True) # If set, the draft model attends directly over the target model KV cache. _use_shared_kv_cache: bool = PrivateAttr(False) + # If set, speculative_model resolves to the target checkpoint, so one-model + # MTP loads its heads from the target weights instead of a separate file. + _mtp_heads_in_target_checkpoint: bool = PrivateAttr(False) # Internal: true when draft_len_schedule was auto-translated from max_concurrency. _translated_from_max_concurrency: bool = PrivateAttr(False) @@ -1921,6 +1926,30 @@ def supports_backend(self, backend: str) -> bool: """ return True + @property + def loads_mtp_from_separate_checkpoint(self) -> bool: + """Whether one-model MTP heads come from ``speculative_model``. + + False when ``speculative_model`` resolves to the target checkpoint: + the heads are then loaded from the target weights, as they were + before separate MTP checkpoints were supported. + """ + return (self.spec_dec_mode.is_mtp_one_model() + and self.speculative_model is not None + and not self._mtp_heads_in_target_checkpoint) + + @property + def needs_separate_draft_weights(self) -> bool: + """Whether draft weights must be loaded from ``speculative_model``. + + True for Eagle3 one-model / external drafters, Gemma4 shared-KV, and + one-model MTP when MTP heads live in a separate checkpoint. + """ + if (self.spec_dec_mode.need_load_draft_weights() + or self._use_shared_kv_cache): + return True + return self.loads_mtp_from_separate_checkpoint + @property def spec_dec_mode(self): # spec_dec_mode has more functionality than the raw decoding_mode string. @@ -2567,8 +2596,9 @@ class MTPDecodingConfig(DecodingBaseConfig): default=None, init=False, description="Number of MTP layers in the model checkpoint. " - "Auto-populated from the model's pretrained config. Do not set manually." - ) + "Auto-populated from the target pretrained config, or from " + "speculative_model's config.json when MTP heads are loaded separately. " + "Do not set manually.") begin_thinking_phase_token: NonNegativeInt = Field( default=128798, diff --git a/tests/unittest/_torch/executor/test_model_loader_gms.py b/tests/unittest/_torch/executor/test_model_loader_gms.py index 3d073d95bf90..8a315f95f3c2 100644 --- a/tests/unittest/_torch/executor/test_model_loader_gms.py +++ b/tests/unittest/_torch/executor/test_model_loader_gms.py @@ -204,7 +204,7 @@ def is_post_transform_weights_preloaded(self) -> bool: def _spec_config_needing_draft_weights(): return SimpleNamespace( - spec_dec_mode=SimpleNamespace(need_load_draft_weights=lambda: True), + needs_separate_draft_weights=True, speculative_model="/draft", ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py new file mode 100644 index 000000000000..a5be718ad5af --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py @@ -0,0 +1,422 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +from types import SimpleNamespace + +import torch + +from tensorrt_llm._torch.models.checkpoints.hf.nemotron_h_weight_mapper import ( + NemotronHHfWeightMapper, +) +from tensorrt_llm._torch.speculative.utils import ( + filter_mtp_checkpoint_weights, + loads_mtp_from_speculative_model, + resolve_mtp_checkpoint_source, + select_mtp_checkpoint_weights, + skip_modules_for_separate_mtp_checkpoint, + update_spec_config_from_model_config, +) +from tensorrt_llm.llmapi.llm_args import Eagle3DecodingConfig, MTPDecodingConfig + + +def test_needs_separate_draft_weights_for_mtp_with_speculative_model(): + cfg = MTPDecodingConfig(max_draft_len=3, speculative_model="/path/to/mtp") + assert cfg.needs_separate_draft_weights is True + + cfg_no_draft = MTPDecodingConfig(max_draft_len=3) + assert cfg_no_draft.needs_separate_draft_weights is False + + +def test_needs_separate_draft_weights_still_true_for_eagle3(): + cfg = Eagle3DecodingConfig(max_draft_len=3, speculative_model="/path/to/eagle3") + assert cfg.needs_separate_draft_weights is True + + +def test_loads_mtp_from_speculative_model_helper(): + assert ( + loads_mtp_from_speculative_model( + MTPDecodingConfig(max_draft_len=3, speculative_model="/path/to/mtp") + ) + is True + ) + assert loads_mtp_from_speculative_model(MTPDecodingConfig(max_draft_len=3)) is False + assert loads_mtp_from_speculative_model(None) is False + + +def test_speculative_model_equal_to_target_keeps_embedded_mtp(tmp_path): + """speculative_model == the target checkpoint is the pre-feature API usage.""" + target_dir = tmp_path / "target" + target_dir.mkdir() + + cfg = MTPDecodingConfig(max_draft_len=3, speculative_model=str(target_dir)) + assert loads_mtp_from_speculative_model(cfg) is True + + resolve_mtp_checkpoint_source(cfg, str(target_dir)) + assert loads_mtp_from_speculative_model(cfg) is False + assert cfg.needs_separate_draft_weights is False + # The user-provided value is left untouched. + assert cfg.speculative_model == str(target_dir) + + +def test_speculative_model_equal_to_target_matches_equivalent_paths(tmp_path): + target_dir = tmp_path / "target" + target_dir.mkdir() + link_dir = tmp_path / "target_link" + link_dir.symlink_to(target_dir, target_is_directory=True) + + cfg = MTPDecodingConfig(max_draft_len=3, speculative_model=str(link_dir)) + resolve_mtp_checkpoint_source(cfg, str(target_dir) + "/") + assert loads_mtp_from_speculative_model(cfg) is False + + +def test_separate_mtp_checkpoint_survives_resolution(tmp_path): + target_dir = tmp_path / "target" + target_dir.mkdir() + mtp_dir = tmp_path / "mtp_heads" + mtp_dir.mkdir() + + cfg = MTPDecodingConfig(max_draft_len=3, speculative_model=str(mtp_dir)) + resolve_mtp_checkpoint_source(cfg, str(target_dir)) + assert loads_mtp_from_speculative_model(cfg) is True + assert cfg.needs_separate_draft_weights is True + + +def test_resolution_does_not_affect_eagle3(tmp_path): + """Eagle3 always loads its draft from speculative_model, same dir or not.""" + target_dir = tmp_path / "target" + target_dir.mkdir() + + cfg = Eagle3DecodingConfig(max_draft_len=3, speculative_model=str(target_dir)) + resolve_mtp_checkpoint_source(cfg, str(target_dir)) + assert cfg.needs_separate_draft_weights is True + + +def test_filter_and_select_mtp_checkpoint_weights(): + weights = { + "backbone.layers.0.mixer.weight": torch.ones(2), + "mtp.layers.0.enorm.weight": torch.ones(4), + "mtp.layers.1.norm.weight": torch.ones(4), + "lm_head.weight": torch.ones(3), + } + filtered = filter_mtp_checkpoint_weights(weights) + assert "backbone.layers.0.mixer.weight" in filtered + assert "lm_head.weight" in filtered + assert "mtp.layers.0.enorm.weight" not in filtered + assert "mtp.layers.1.norm.weight" not in filtered + + selected = select_mtp_checkpoint_weights(weights) + assert set(selected) == { + "mtp.layers.0.enorm.weight", + "mtp.layers.1.norm.weight", + } + + +def test_update_spec_config_prefers_speculative_model_mtp_fields(tmp_path): + mtp_dir = tmp_path / "mtp_heads" + mtp_dir.mkdir() + (mtp_dir / "config.json").write_text( + json.dumps( + { + "num_nextn_predict_layers": 1, + "mtp_hybrid_override_pattern": "*E", + "mtp_block_configs": [{"block_type": "moe", "num_experts": 8}], + } + ) + ) + + # Target has no MTP (or stale MTP metadata). + model_config = SimpleNamespace( + architectures=["NemotronHForCausalLM"], + num_nextn_predict_layers=0, + mtp_layers_block_type=None, + mtp_block_configs=None, + ) + spec_config = MTPDecodingConfig(max_draft_len=5, speculative_model=str(mtp_dir)) + + update_spec_config_from_model_config(spec_config, model_config) + + assert spec_config.num_nextn_predict_layers == 1 + assert model_config.num_nextn_predict_layers == 1 + # Legacy pattern string is converted to the writable HF field. + assert model_config.mtp_layers_block_type == ["attention", "moe"] + assert model_config.mtp_block_configs == [{"block_type": "moe", "num_experts": 8}] + # User-set max_draft_len is preserved for Eagle-style replay. + assert spec_config.max_draft_len == 5 + + +def test_update_spec_config_uses_mtp_layers_block_type_when_present(tmp_path): + mtp_dir = tmp_path / "mtp_heads" + mtp_dir.mkdir() + (mtp_dir / "config.json").write_text( + json.dumps( + { + "num_nextn_predict_layers": 1, + "mtp_layers_block_type": ["attention", "moe"], + } + ) + ) + + model_config = SimpleNamespace( + architectures=["NemotronHForCausalLM"], + num_nextn_predict_layers=0, + mtp_layers_block_type=None, + ) + spec_config = MTPDecodingConfig(max_draft_len=3, speculative_model=str(mtp_dir)) + update_spec_config_from_model_config(spec_config, model_config) + assert model_config.mtp_layers_block_type == ["attention", "moe"] + + +def test_remap_preprocessed_mtp_weights_for_draft_model(): + from tensorrt_llm._torch.speculative.utils import remap_preprocessed_mtp_weights_for_draft_model + + weights = { + "model.layers.52.layers.0.enorm.weight": torch.ones(4), + "model.layers.52.layers.1.norm.weight": torch.ones(4), + } + remapped = remap_preprocessed_mtp_weights_for_draft_model( + weights, num_hidden_layers=52, num_mtp_layers=1 + ) + assert remapped == { + "mtp_layers.0.layers.0.enorm.weight": weights["model.layers.52.layers.0.enorm.weight"], + "mtp_layers.0.layers.1.norm.weight": weights["model.layers.52.layers.1.norm.weight"], + } + + +def test_skip_modules_for_separate_mtp_checkpoint_shared_head(): + # Nemotron-style: no shared_head tensors -> skip so strict load does not + # demand an absent module. + nemotron_keys = { + "mtp_layers.0.layers.0.enorm.weight": torch.ones(4), + "mtp_layers.0.layers.1.final_layernorm.weight": torch.ones(4), + } + assert skip_modules_for_separate_mtp_checkpoint(nemotron_keys) == ["shared_head"] + + # DeepSeek / Qwen / Exaone: shared_head.norm is present -> load it. + deepseek_keys = { + "mtp_layers.0.enorm.weight": torch.ones(4), + "mtp_layers.0.shared_head.norm.weight": torch.ones(4), + } + assert skip_modules_for_separate_mtp_checkpoint(deepseek_keys) == [] + + # Step3: shared_head also owns an output projection -> still load. + step3_keys = { + "mtp_layers.0.shared_head.norm.weight": torch.ones(4), + "mtp_layers.0.shared_head.output.weight": torch.ones(4, 4), + } + assert skip_modules_for_separate_mtp_checkpoint(step3_keys) == [] + + +def _make_one_engine_stub(spec_config, num_hidden_layers: int = 52): + """A bare SpecDecOneEngineForCausalLM with just the module tree we need. + + ``__init__`` builds a whole target model, so construct the instance + directly and register only the two aliases of a single MTP head: the + target's ``model.layers[N]`` and ``draft_model.mtp_layers[0]``. + """ + from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM + + class _OneEngineStub(SpecDecOneEngineForCausalLM): + # The real ``config`` is a read-only property over + # ``model_config.pretrained_config``, which this stub never builds. + config = SimpleNamespace(num_hidden_layers=num_hidden_layers) + + model = object.__new__(_OneEngineStub) + torch.nn.Module.__init__(model) + + head = torch.nn.Module() + inner = torch.nn.Module() + inner.layers = torch.nn.ModuleList([torch.nn.Module(), head]) + model.model = inner + draft = torch.nn.Module() + draft.mtp_layers = torch.nn.ModuleList([head]) + model.draft_model = draft + model.spec_config = spec_config + return model + + +def _capture_parent_load_weights(monkeypatch) -> dict: + """Intercept the base-class load to inspect the dispatch arguments.""" + from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM + + captured = {} + + def fake_load_weights( + self, + weights, + weight_mapper=None, + skip_modules=(), + params_map=None, + allow_partial_loading=False, + ): + captured["skip_modules"] = list(skip_modules) + captured["allow_partial_loading"] = allow_partial_loading + captured["weights"] = weights + + monkeypatch.setattr(DecoderModelForCausalLM, "load_weights", fake_load_weights) + return captured + + +def test_mtp_head_module_names_covers_both_aliases(): + model = _make_one_engine_stub( + MTPDecodingConfig(max_draft_len=3, speculative_model="/path/to/mtp") + ) + assert set(model.mtp_head_module_names()) == { + "model.layers.1", + "draft_model.mtp_layers.0", + } + + +def test_separate_mtp_target_load_skips_heads_without_partial_loading(monkeypatch): + """The target load must never fall back to partial loading. + + ``allow_partial_loading=True`` suppresses ``process_weights_after_loading`` + on every quantized Linear/MoE it touches, which silently leaves the whole + target model's quant scales uninitialized (garbage output). The MTP heads + have to be excluded by module instead. + """ + captured = _capture_parent_load_weights(monkeypatch) + + model = _make_one_engine_stub( + MTPDecodingConfig(max_draft_len=3, speculative_model="/path/to/mtp") + ) + model.load_weights( + weights={ + "backbone.layers.0.norm.weight": torch.ones(4), + "mtp.layers.0.enorm.weight": torch.ones(4), + } + ) + + assert captured["allow_partial_loading"] is False + assert set(captured["skip_modules"]) == { + "draft_model", + "model.layers.1", + "draft_model.mtp_layers.0", + } + assert "mtp.layers.0.enorm.weight" not in captured["weights"] + assert "backbone.layers.0.norm.weight" in captured["weights"] + + +def test_embedded_mtp_target_load_is_unchanged(monkeypatch): + captured = _capture_parent_load_weights(monkeypatch) + + model = _make_one_engine_stub(MTPDecodingConfig(max_draft_len=3)) + model.load_weights(weights={"mtp.layers.0.enorm.weight": torch.ones(4)}) + + assert captured["skip_modules"] == ["draft_model"] + assert captured["allow_partial_loading"] is False + # Embedded heads still load from the target checkpoint. + assert "mtp.layers.0.enorm.weight" in captured["weights"] + + +def test_target_load_keeps_heads_when_speculative_model_is_target(monkeypatch, tmp_path): + captured = _capture_parent_load_weights(monkeypatch) + + target_dir = tmp_path / "target" + target_dir.mkdir() + spec_config = MTPDecodingConfig(max_draft_len=3, speculative_model=str(target_dir)) + resolve_mtp_checkpoint_source(spec_config, str(target_dir)) + + model = _make_one_engine_stub(spec_config) + model.load_weights(weights={"mtp.layers.0.enorm.weight": torch.ones(4)}) + + assert captured["skip_modules"] == ["draft_model"] + assert "mtp.layers.0.enorm.weight" in captured["weights"] + + +def test_nemotron_mapper_remaps_mtp_layers_keys(): + mapper = NemotronHHfWeightMapper() + pretrained = SimpleNamespace( + num_hidden_layers=52, + mamba_head_dim=64, + mamba_num_heads=8, + n_groups=8, + ssm_state_size=128, + num_key_value_heads=2, + tie_word_embeddings=False, + ) + mapping = SimpleNamespace(enable_attention_dp=False, tp_size=1, tp_rank=0) + model_config = SimpleNamespace( + pretrained_config=pretrained, mapping=mapping, moe_backend="TRTLLM" + ) + mapper._config = model_config + mapper._model = SimpleNamespace(model_config=model_config, config=pretrained) + mapper._tp_size = 1 + + weights = { + "mtp.layers.0.enorm.weight": torch.ones(4), + "mtp.layers.1.norm.weight": torch.ones(4), + "backbone.layers.0.norm.weight": torch.ones(4), + } + remapped = mapper.preprocess_weights(weights) + assert "model.layers.52.layers.0.enorm.weight" in remapped + assert "model.layers.52.layers.1.norm.weight" in remapped + assert "model.layers.0.norm.weight" in remapped + assert "mtp.layers.0.enorm.weight" not in remapped + + +def _nemotron_style_mtp_weights(*, include_shared_head: bool) -> dict: + """Minimal remappable mtp.* tensors that satisfy the Nemotron required-suffix check.""" + weights = { + "mtp.layers.0.enorm.weight": torch.ones(4), + "mtp.layers.0.hnorm.weight": torch.ones(4), + "mtp.layers.0.eh_proj.weight": torch.ones(4, 4), + "mtp.layers.1.final_layernorm.weight": torch.ones(4), + } + if include_shared_head: + weights["mtp.shared_head.norm.weight"] = torch.ones(4) + return weights + + +class _PassthroughMtpMapper: + """Nemotron-like remap: ``mtp.layers.{{i}}.*`` -> ``model.layers.{{N}}.layers.{{i}}.*``.""" + + def __init__(self, num_hidden_layers: int): + self._num_hidden_layers = num_hidden_layers + + def preprocess_weights(self, weights: dict) -> dict: + out = {} + for key, value in weights.items(): + if key.startswith("mtp.layers."): + _, _, sublayer_idx, rest = key.split(".", 3) + out[f"model.layers.{self._num_hidden_layers}.layers.{sublayer_idx}.{rest}"] = value + elif key.startswith("mtp."): + out[f"model.layers.{self._num_hidden_layers}.{key[len('mtp.') :]}"] = value + else: + out[key] = value + return out + + +def test_separate_mtp_draft_load_skip_shared_head_scales(monkeypatch): + """Draft load skips shared_head only when the remapped checkpoint omits it.""" + from tensorrt_llm._torch.models import modeling_utils + + captured = {} + + def fake_load_weights_impl_v2( + model, weights, weight_mapper, skip_modules=(), allow_partial_loading=False, **kwargs + ): + captured["skip_modules"] = list(skip_modules) + captured["allow_partial_loading"] = allow_partial_loading + captured["weight_keys"] = set(weights) + + monkeypatch.setattr(modeling_utils, "_load_weights_impl_v2", fake_load_weights_impl_v2) + + spec_config = MTPDecodingConfig(max_draft_len=1, speculative_model="/path/to/mtp") + model = _make_one_engine_stub(spec_config, num_hidden_layers=52) + mapper = _PassthroughMtpMapper(num_hidden_layers=52) + + model.load_draft_weights( + weights=_nemotron_style_mtp_weights(include_shared_head=False), + weight_mapper=mapper, + ) + assert captured["skip_modules"] == ["shared_head"] + assert captured["allow_partial_loading"] is False + assert not any("shared_head" in k for k in captured["weight_keys"]) + + model.load_draft_weights( + weights=_nemotron_style_mtp_weights(include_shared_head=True), + weight_mapper=mapper, + ) + assert captured["skip_modules"] == [] + assert "mtp_layers.0.shared_head.norm.weight" in captured["weight_keys"]