Add MTP speculative-decoding support to NeMo SpeechLM vLLM plugin

examples/speechlm2/to_hf.py

@@ -102,6 +102,34 @@ def _canonical_torch_dtype_name(dtype: str | torch.dtype) -> str:
 def _hf_export_config(model: torch.nn.Module, dtype: str | torch.dtype) -> dict[str, Any]:
     """Build the exported root config without mutating the training config."""
     config = OmegaConf.to_container(model.cfg) if isinstance(model.cfg, DictConfig) else deepcopy(model.cfg)
+    mtp_cfg = config.get("mtp")
+    llm_config = getattr(getattr(model, "llm", None), "config", None)
+    actual_mtp_pattern = getattr(llm_config, "mtp_hybrid_override_pattern", None)
+    if isinstance(mtp_cfg, dict) and mtp_cfg.get("enabled", False) and actual_mtp_pattern is not None:
+        if not isinstance(actual_mtp_pattern, str) or not actual_mtp_pattern:
+            raise ValueError(
+                f"Built LLM has invalid mtp_hybrid_override_pattern={actual_mtp_pattern!r}; cannot export it."
+            )
+        # A preserved checkpoint-native MTP head can differ from the recipe's
+        # requested replacement pattern. Persist the pattern of the head that
+        # was actually built so vLLM instantiates the matching physical layers.
+        mtp_cfg["hybrid_override_pattern"] = actual_mtp_pattern
+        actual_mtp_depth = getattr(llm_config, "num_nextn_predict_layers", None)
+        if actual_mtp_depth is not None:
+            if isinstance(actual_mtp_depth, bool) or not isinstance(actual_mtp_depth, int) or actual_mtp_depth <= 0:
+                raise ValueError(
+                    f"Built LLM has invalid num_nextn_predict_layers={actual_mtp_depth!r}; cannot export it."
+                )
+            if mtp_cfg.get("use_repeated_layer", False):
+                if actual_mtp_depth != 1:
+                    raise ValueError(
+                        "A repeated MTP head must serialize exactly one physical layer, but the built LLM "
+                        f"declares num_nextn_predict_layers={actual_mtp_depth}."
+                    )
+            else:
+                # For a preserved native head, the recipe depth is advisory.
+                # Export the physical/logical depth that is actually present.
+                mtp_cfg["num_nextn_predict_layers"] = actual_mtp_depth
     dtype_name = _canonical_torch_dtype_name(dtype)
     config["dtype"] = dtype_name
     config["torch_dtype"] = dtype_name
@@ -116,9 +144,8 @@ def save_hf_checkpoint(model: torch.nn.Module, state_dict: dict, cfg: HfExportCo
     target_dtype = str_to_dtype(cfg.dtype)
     state_dict = {k: v.to(target_dtype) for k, v in state_dict.items()}
 
-    save_file(state_dict, output_dir / "model.safetensors")
-
     config = _hf_export_config(model, cfg.dtype)
+    save_file(state_dict, output_dir / "model.safetensors")
     with open(output_dir / "config.json", "w") as f:
         json.dump(config, f, indent=2)
     save_llm_backbone_config(model, output_dir)
@@ -135,16 +162,18 @@ def save_llm_backbone_config(model: torch.nn.Module, output_dir: str | Path) ->
     llm_config.save_pretrained(str(llm_backbone_dir))
 
 
-def _detect_vllm_architecture(model_cfg: dict) -> str:
-    """Determine the vLLM plugin model class for the checkpoint.
+def _detect_vllm_architecture(model_cfg: dict) -> tuple[str, int]:
+    """Determine the vLLM plugin model class and backbone vocabulary size.
 
     The SALM plugin registers a single architecture name and selects between
     transformer and hybrid backends at instantiation time, so this function
-    just verifies the backbone config is reachable and returns the unified
-    name; the hybrid-vs-transformer split is handled inside the plugin.
+    verifies the backbone config is reachable and returns the unified name
+    plus the embedding-table vocabulary bound. The hybrid-vs-transformer split
+    is handled inside the plugin.
 
     Raises:
-        ValueError: if the HF config can't be loaded or has no 'architectures'.
+        ValueError: If the HF config cannot be loaded, has no architecture, or
+            declares an invalid vocabulary size.
     """
     pretrained_llm = model_cfg.get("pretrained_llm", "")
     try:
@@ -160,8 +189,11 @@ def _detect_vllm_architecture(model_cfg: dict) -> str:
     archs = getattr(llm_cfg, "architectures", [])
     if not archs:
         raise ValueError(f"HF config for {pretrained_llm!r} has empty 'architectures'.")
+    vocab_size = getattr(llm_cfg, "vocab_size", None)
+    if isinstance(vocab_size, bool) or not isinstance(vocab_size, int) or vocab_size <= 0:
+        raise ValueError(f"HF config for {pretrained_llm!r} has invalid 'vocab_size': {vocab_size!r}.")
 
-    return "NeMoSpeechLMForConditionalGeneration"
+    return "NeMoSpeechLMForConditionalGeneration", vocab_size
 
 
 def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
@@ -175,10 +207,12 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
         model_cfg: Model config dict (from experiment YAML).
 
     Raises:
-        ValueError: If ``pretrained_llm`` or ``audio_locator_tag`` is missing.
+        ValueError: If required model metadata is missing, or the tokenizer's
+            audio token does not fit the SpeechLM embedding table.
     """
     from transformers import AutoTokenizer
 
+    from nemo.collections.speechlm2.vllm.salm.config import _SPEECHLM_EMBED_EXTRA_ROWS
     from nemo.utils import logging as LOG
 
     output_dir = Path(output_dir)
@@ -197,13 +231,14 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     llm_backbone_dir = output_dir / LLM_BACKBONE_DIR
     if (llm_backbone_dir / "config.json").exists():
         arch_model_cfg["pretrained_llm"] = str(llm_backbone_dir)
-    arch = _detect_vllm_architecture(arch_model_cfg)
+    arch, base_vocab_size = _detect_vllm_architecture(arch_model_cfg)
     config_path = output_dir / "config.json"
     config = json.loads(config_path.read_text())
     config["model_type"] = "nemo_speechlm"
     config["architectures"] = [arch]
     config["audio_locator_tag"] = audio_token
-    config_path.write_text(json.dumps(config, indent=2) + "\n")
+    config.pop("audio_token_index", None)
+    config.pop("image_token_index", None)
 
     # 2. Save tokenizer (backbone chat_template carries over via save_pretrained)
     existing = [
@@ -213,9 +248,20 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None:
     ]
     if existing:
         LOG.info("Overwriting existing files in %s: %s", output_dir, existing)
-    tok = AutoTokenizer.from_pretrained(pretrained_llm, trust_remote_code=True)
+    tokenizer_src = model_cfg.get("tokenizer_path") or pretrained_llm
+    tok = AutoTokenizer.from_pretrained(tokenizer_src, trust_remote_code=True)
     if audio_token not in tok.get_vocab():
         tok.add_special_tokens({"additional_special_tokens": [audio_token]})
+    audio_token_id = tok.get_vocab().get(audio_token)
+    if isinstance(audio_token_id, bool) or not isinstance(audio_token_id, int) or audio_token_id < 0:
+        raise ValueError(f"Tokenizer did not assign a valid ID to audio token {audio_token!r}.")
+    padded_vocab_size = base_vocab_size + _SPEECHLM_EMBED_EXTRA_ROWS
+    if audio_token_id >= padded_vocab_size:
+        raise ValueError(
+            f"Audio token ID {audio_token_id} is outside the SpeechLM embedding table with "
+            f"{padded_vocab_size} rows. Reduce the tokenizer's added-token count before training/export."
+        )
+    config_path.write_text(json.dumps(config, indent=2) + "\n")
     tok.save_pretrained(str(output_dir))
     # Newer transformers splits long chat_template into a separate
     # ``chat_template.jinja`` file; inline it back and drop the file.

nemo/collections/speechlm2/vllm/salm/__init__.py

@@ -24,6 +24,122 @@
 """
 
 _PKG = "nemo.collections.speechlm2.vllm.salm"
+_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = None
+
+
+def _nemo_speechlm_mtp_hf_config_override(hf_config):
+    """Apply the SpeechLM MTP rewrite, then defer unrelated configs to vLLM.
+
+    This function must remain at module scope: vLLM retains it on the draft
+    ``ModelConfig``, which can cross a spawned process boundary. The original
+    vLLM callable stays in process-local module state because binding the
+    replaced static method inside a ``partial`` also makes that method
+    unresolvable by standard pickle.
+    """
+    if hf_config.model_type == "nemo_speechlm":
+        mtp_cfg = getattr(hf_config, "mtp", None)
+        if not isinstance(mtp_cfg, dict):
+            mtp_cfg = {}
+        # Match SALMAutomodel's training defaults exactly: merely retaining a
+        # recipe depth does not enable MTP, while an enabled block with no
+        # explicit depth constructs one logical head.
+        mtp_enabled = bool(mtp_cfg.get("enabled", False))
+        n_predict = mtp_cfg.get("num_nextn_predict_layers", 1 if mtp_enabled else 0)
+        if mtp_enabled and n_predict > 0:
+            use_repeated_layer = bool(mtp_cfg.get("use_repeated_layer", False))
+            if n_predict > 1 and not use_repeated_layer:
+                raise ValueError(
+                    f"NeMo SpeechLM MTP with {n_predict} distinct head layers is not "
+                    f"supported: vLLM's NemotronHMultiTokenPredictor builds a single "
+                    f"physical MTP layer and reuses it every speculative step. Only "
+                    f"checkpoints trained with mtp.use_repeated_layer=true match that "
+                    f"execution model."
+                )
+            hf_config.model_type = "nemo_speechlm_mtp"
+            hf_config.update(
+                {
+                    # Size of the physical MTP block that vLLM reuses. A
+                    # repeated-layer checkpoint ships one shared head even
+                    # when it was trained for multiple next-token positions,
+                    # so arbitrary inference K values must be multiples of 1.
+                    # Consequently vLLM defaults to K=1 when K is omitted;
+                    # callers should set num_speculative_tokens explicitly.
+                    "n_predict": 1,
+                    # Physical MTP prediction steps to instantiate. Repeated-layer checkpoints
+                    # ship one shared step (one mtp.layers.* module per hybrid-pattern character)
+                    # that is reapplied every speculative iteration, exactly as vLLM drives its
+                    # MTP draft. This also shadows the backbone text_config's
+                    # num_nextn_predict_layers (e.g. 4), which would otherwise trip the
+                    # single-step assert in NemotronHMultiTokenPredictor.
+                    "num_nextn_predict_layers": 1,
+                    "architectures": ["NeMoSpeechLMMTPModel"],
+                }
+            )
+            return hf_config
+
+    global _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE
+    if _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE is None:
+        # A spawn child can import this module while unpickling the function
+        # without running the vLLM plugin hook first. In that case the class
+        # still exposes its native override, which is safe to capture lazily.
+        from vllm.config.speculative import SpeculativeConfig
+
+        current_override = SpeculativeConfig.hf_config_override
+        if current_override is _nemo_speechlm_mtp_hf_config_override:
+            raise RuntimeError("NeMo SpeechLM MTP override was installed without preserving vLLM's original hook.")
+        _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = current_override
+    return _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE(hf_config)
+
+
+_nemo_speechlm_mtp_hf_config_override._nemo_speechlm_mtp_override = True
+
+
+def _patch_vllm_for_nemo_speechlm_mtp() -> None:
+    """Extend vLLM's speculative-decoding framework to support nemo_speechlm MTP.
+
+    Three patches are applied on the supported vLLM 0.19+ releases:
+
+    1. ``MTPModelTypes`` — the Literal type that guards the MTP detection
+       branch in ``SpeculativeConfig.__post_init__`` is extended to include
+       ``"nemo_speechlm_mtp"``.
+
+    2. ``SpeculativeConfig.hf_config_override`` — the static method that
+       rewrites the draft-model HF config is wrapped to detect
+       ``nemo_speechlm`` checkpoints that carry MTP heads (``mtp.enabled``
+       and ``mtp.num_nextn_predict_layers > 0``) and redirect them to the
+       ``NeMoSpeechLMMTPModel`` architecture with the right ``n_predict``.
+
+    3. ``ModelRegistry`` — ``NeMoSpeechLMMTPModel`` is registered so that
+       vLLM can resolve and instantiate it as the draft model.
+    """
+    from typing import Literal, get_args
+
+    import vllm.config.speculative as _spec_mod
+    from vllm.config.speculative import SpeculativeConfig
+
+    # Extend vLLM's recognized MTP model types.
+    old_args = get_args(_spec_mod.MTPModelTypes)
+    if "nemo_speechlm_mtp" not in old_args:
+        _spec_mod.MTPModelTypes = Literal[old_args + ("nemo_speechlm_mtp",)]
+
+    # Route SpeechLM MTP checkpoints through SpeculativeConfig.hf_config_override.
+    current_override = SpeculativeConfig.hf_config_override
+    if not getattr(current_override, "_nemo_speechlm_mtp_override", False):
+        global _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE
+        # Preserve the first native hook for the lifetime of this process.
+        # Replacing it during a later registration could capture a third-party
+        # wrapper that already delegates to us, creating an override cycle.
+        if _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE is None:
+            _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = current_override
+        SpeculativeConfig.hf_config_override = staticmethod(_nemo_speechlm_mtp_hf_config_override)
+
+    # Register the SpeechLM MTP draft architecture with vLLM.
+    from vllm.model_executor.models.registry import ModelRegistry
+
+    ModelRegistry.register_model(
+        "NeMoSpeechLMMTPModel",
+        f"{_PKG}.mtp:NeMoSpeechLMMTP",
+    )
 
 
 def register():
@@ -44,3 +160,5 @@ def register():
         "NeMoSpeechLMForConditionalGeneration",
         f"{_PKG}.model:NeMoSpeechLMForConditionalGeneration",
     )
+
+    _patch_vllm_for_nemo_speechlm_mtp()

nemo/collections/speechlm2/vllm/salm/audio.py

@@ -82,9 +82,12 @@
 
 
 def _ensure_special_tokens(tokenizer: PreTrainedTokenizerBase) -> None:
-    special = [_AUDIO_PLACEHOLDER]
-    existing = set(tokenizer.get_vocab().keys())
-    to_add = [t for t in special if t not in existing]
+    # NOTE: called per request from _call_hf_processor on the API-server event loop.
+    # Use O(1) dict membership; `set(get_vocab().keys())` rebuilt a 131k-entry set
+    # every request (~5-6 ms) purely to check one token. get_vocab() returns vLLM's
+    # cached dict, so membership is O(1).
+    vocab = tokenizer.get_vocab()
+    to_add = [t for t in (_AUDIO_PLACEHOLDER,) if t not in vocab]
     if to_add:
         tokenizer.add_special_tokens({"additional_special_tokens": to_add})
 

nemo/collections/speechlm2/vllm/salm/config.py

@@ -39,9 +39,12 @@
 # silently rendering the wrong placeholder at request time.
 _AUDIO_PLACEHOLDER = "<|audio|>"
 
-# Number of extra embedding rows the SpeechLM adds on top of the backbone's
-# native vocab during training: ``<|audio|>`` locator plus headroom for other
-# special tokens and TensorCore-friendly alignment.
+# Historical serving-time headroom above the backbone vocabulary. vLLM builds
+# the target and draft embedding tables at this padded size, and the weight
+# loader zero-pads the smaller training tensors to match. ``prepare_for_vllm``
+# validates the tokenizer's audio-token ID against this bound during export;
+# the default export flow treats a validation failure as non-fatal and leaves
+# an HF-only checkpoint.
 _SPEECHLM_EMBED_EXTRA_ROWS = 10
 
 
@@ -100,6 +103,7 @@ def __init__(
         # path; real checkpoint loads replace it below after field validation.
         self.text_config = PretrainedConfig()
         self.is_hybrid = False
+        self._pending_image_token_index = None
 
         super().__init__(**kwargs)
 
@@ -115,6 +119,7 @@ def __init__(
             self.pretrained_weights = None
             self.lora = None
             self.encoder_chunk_size_seconds = None
+            self.__dict__.pop("_pending_image_token_index", None)
             return
 
         for name, value in required_fields.items():
@@ -177,6 +182,13 @@ def __init__(
                 self.text_config.layer_types = ["attention"] * num_layers
 
         self.text_config.vocab_size += _SPEECHLM_EMBED_EXTRA_ROWS
+        pending_image_token_index = self.__dict__.pop("_pending_image_token_index", None)
+        if pending_image_token_index is not None and pending_image_token_index != self.image_token_index:
+            raise ValueError(
+                f"image_token_index={pending_image_token_index!r} does not match the backbone vocabulary "
+                f"boundary {self.image_token_index}. Remove this legacy serialized field; SpeechLM derives "
+                f"the vLLM compatibility value at runtime."
+            )
 
     @property
     def llm_architectures(self) -> list[str]:
@@ -186,6 +198,46 @@ def llm_architectures(self) -> list[str]:
     def get_text_config(self, decoder=False) -> PretrainedConfig:
         return self.text_config
 
+    @property
+    def image_token_index(self) -> int | None:
+        """Return the vocabulary-boundary value expected by vLLM's MTP proposer.
+
+        vLLM calls this compatibility field ``image_token_index`` even for an
+        audio multimodal target. Actual audio locations come from vLLM's
+        placeholder ranges; no token-index field is serialized by SpeechLM.
+        """
+        vocab_size = getattr(self.text_config, "vocab_size", None)
+        if vocab_size is None:
+            return None
+        return int(vocab_size) - _SPEECHLM_EMBED_EXTRA_ROWS
+
+    @image_token_index.setter
+    def image_token_index(self, value: int | None) -> None:
+        """Accept vLLM's runtime target-to-draft copy without serializing it."""
+        expected = self.image_token_index
+        if expected is None:
+            # Transformers applies unknown config kwargs in its base-class
+            # constructor, before this wrapper has loaded the real backbone.
+            # Defer validation and discard the temporary value afterwards so
+            # it never becomes serialized state.
+            self._pending_image_token_index = value
+        elif value is not None and value != expected:
+            raise ValueError(
+                f"image_token_index={value!r} does not match the backbone vocabulary boundary {expected}."
+            )
+
+    @property
+    def mtp_hybrid_override_pattern(self) -> str:
+        """Hybrid layer pattern for MTP heads, consumed by NemotronHMultiTokenPredictor.
+
+        Reads from the ``mtp.hybrid_override_pattern`` field in config.json.
+        vLLM supports any sequence of ``"*"`` (attention) and ``"E"`` (MoE),
+        with one physical MTP layer module instantiated per character. Other
+        characters are rejected by vLLM during model construction.
+        """
+        mtp_cfg = self.__dict__.get("mtp") or {}
+        return mtp_cfg.get("hybrid_override_pattern", "*") if isinstance(mtp_cfg, dict) else "*"
+
     _ATTR_ALIASES = {
         "rms_norm_eps": "layer_norm_epsilon",
         "layer_norm_eps": "layer_norm_epsilon",
@@ -213,6 +265,7 @@ def __getattr__(self, name):
             "pretrained_llm",
             "pretrained_asr",
             "audio_locator_tag",
+            "image_token_index",
             "prompt_format",
             "pretrained_weights",
             "text_config",

nemo/collections/speechlm2/vllm/salm/model.py

@@ -222,6 +222,13 @@ def _split_perception_llm(
                 continue
             if name.startswith("perception."):
                 perception[name[len("perception.") :]] = tensor
+            elif name.startswith("llm.mtp."):
+                pass  # MTP draft-head weights; loaded by the speculative draft model, not here
+            elif name.startswith("mtp."):
+                raise ValueError(
+                    f"Unsupported bare MTP tensor {name!r}; NeMo SpeechLM exports must store draft weights "
+                    f"under the 'llm.mtp.*' namespace."
+                )
             else:
                 llm.append((name, tensor))
         return perception, llm