diff --git a/examples/speechlm2/to_hf.py b/examples/speechlm2/to_hf.py index 15e9eab07027..19e989c3e96e 100644 --- a/examples/speechlm2/to_hf.py +++ b/examples/speechlm2/to_hf.py @@ -27,6 +27,7 @@ from safetensors.torch import save_file from nemo.collections.speechlm2.parts.hf_hub import LLM_BACKBONE_DIR +from nemo.collections.speechlm2.vllm.salm.config import _mtp_pattern_from_backbone_config, _resolve_speechlm_mtp_config from nemo.core.classes.common import safe_instantiate from nemo.core.config import hydra_runner from nemo.utils.dtype import str_to_dtype @@ -186,6 +187,77 @@ def _hf_export_config(model: torch.nn.Module, dtype: str | torch.dtype) -> dict[ config["dtype"] = dtype_name config["torch_dtype"] = dtype_name + llm = getattr(model, "llm", None) + text_config = getattr(llm, "config", None) + explicit_mtp = config.get("mtp") + mtp_enabled = ( + bool(explicit_mtp.get("enabled", True)) + if isinstance(explicit_mtp, dict) + else bool(config.get("compute_mtp", False)) + ) + runtime_mtp_config = getattr(llm, "mtp_config", None) + runtime_mtp_depth = getattr(runtime_mtp_config, "num_layers", None) + if runtime_mtp_depth is None and text_config is not None: + runtime_mtp_depth = getattr(text_config, "num_nextn_predict_layers", 0) + if not isinstance(explicit_mtp, dict) and mtp_enabled and not int(runtime_mtp_depth or 0): + # compute_mtp is the legacy switch. Modern SALMAutomodel recipes + # suppress a checkpoint-native head when no explicit mtp block is + # present; do not let a stale flag recreate a serving-only MTP module. + config["compute_mtp"] = False + mtp_enabled = False + + if text_config is not None and mtp_enabled: + use_repeated_layer = getattr( + runtime_mtp_config, + "use_repeated_layer", + bool(explicit_mtp.get("use_repeated_layer", False)) if isinstance(explicit_mtp, dict) else False, + ) + resolved_mtp = dict(explicit_mtp) if isinstance(explicit_mtp, dict) else {} + + raw_mtp_pattern = getattr(text_config, "mtp_hybrid_override_pattern", None) + mtp_block_types = getattr(text_config, "mtp_layers_block_type", None) + if raw_mtp_pattern is not None: + if not isinstance(raw_mtp_pattern, str) or (not raw_mtp_pattern and not mtp_block_types): + raise ValueError( + f"Built LLM has invalid mtp_hybrid_override_pattern={raw_mtp_pattern!r}; " "cannot export it." + ) + actual_mtp_pattern = _mtp_pattern_from_backbone_config(text_config) + if actual_mtp_pattern is not None: + # 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 serving constructs matching + # physical layers. + resolved_mtp["hybrid_override_pattern"] = actual_mtp_pattern + + actual_mtp_depth = getattr(text_config, "num_nextn_predict_layers", None) + logical_mtp_depth = getattr(runtime_mtp_config, "num_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 use_repeated_layer: + 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. + logical_mtp_depth = actual_mtp_depth + + config["mtp"] = _resolve_speechlm_mtp_config( + mtp=resolved_mtp, + compute_mtp=bool(config.get("compute_mtp", False)), + text_config=text_config, + num_nextn_predict_layers=logical_mtp_depth, + use_repeated_layer=use_repeated_layer, + ) + elif mtp_enabled and isinstance(explicit_mtp, dict): + raise ValueError( + "The root mtp config enables MTP, but the instantiated model has no positive-depth MTP head to export." + ) return config @@ -309,6 +381,7 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None: else: config.pop("llm_config", None) config.pop("audio_token_index", None) + config.pop("image_token_index", None) # 2. Save tokenizer (backbone chat_template carries over via save_pretrained) existing = [ diff --git a/nemo/collections/speechlm2/vllm/salm/__init__.py b/nemo/collections/speechlm2/vllm/salm/__init__.py index b6cd7a27c365..387c41f9f988 100644 --- a/nemo/collections/speechlm2/vllm/salm/__init__.py +++ b/nemo/collections/speechlm2/vllm/salm/__init__.py @@ -25,6 +25,113 @@ """ _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 closure 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: 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 = int(mtp_cfg.get("num_nextn_predict_layers", 1 if mtp_enabled else 0) or 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( + { + # vLLM instantiates one physical prediction step and reuses + # it for arbitrary speculative K. Repeated-layer training + # produces exactly that checkpoint layout. + "n_predict": 1, + "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. + import vllm.config.speculative as _spec_mod + + current_override = _spec_mod.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 enabled 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 + + # 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 = _spec_mod.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 later registration could capture a third-party + # wrapper that already delegates to us and create an override cycle. + if _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE is None: + _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = current_override + _spec_mod.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(): @@ -52,6 +159,8 @@ def register(): MODELS_CONFIG_MAP["NeMoSpeechLMForConditionalGeneration"] = NeMoSpeechLMForConditionalGenerationConfig + _patch_vllm_for_nemo_speechlm_mtp() + from nemo.collections.speechlm2.vllm.salm.runtime_compat import install_prompt_contract install_prompt_contract() diff --git a/nemo/collections/speechlm2/vllm/salm/config.py b/nemo/collections/speechlm2/vllm/salm/config.py index 762bb10046f2..64692342e403 100644 --- a/nemo/collections/speechlm2/vllm/salm/config.py +++ b/nemo/collections/speechlm2/vllm/salm/config.py @@ -48,11 +48,117 @@ # an HF-only checkpoint. _SPEECHLM_EMBED_EXTRA_ROWS = 10 +_MTP_BLOCK_TYPE_TO_VLLM_SYMBOL = { + "attention": "*", + "moe": "E", +} + def _is_hybrid_backend(architectures: list[str]) -> bool: return bool(set(architectures) & _HYBRID_ARCHITECTURES) +def _mtp_pattern_from_backbone_config(text_config) -> str | None: + """Return the vLLM Nemotron-H MTP pattern encoded by a backbone config. + + Nemotron 3.5 exports the physical MTP topology in one of two forms: + ``mtp_hybrid_override_pattern`` (symbol string) or + ``mtp_layers_block_type`` (list of block names). vLLM 0.23 supports only + attention (``*``) and MoE (``E``) MTP sublayers, so unsupported topology + must fail before model construction rather than load the wrong draft head. + """ + pattern = getattr(text_config, "mtp_hybrid_override_pattern", None) + if pattern: + unsupported = sorted(set(pattern) - set(_MTP_BLOCK_TYPE_TO_VLLM_SYMBOL.values())) + if unsupported: + raise ValueError( + f"vLLM Nemotron-H MTP does not support pattern symbols {unsupported!r} " + f"in mtp_hybrid_override_pattern={pattern!r}; supported symbols are '*' and 'E'." + ) + return pattern + + block_types = getattr(text_config, "mtp_layers_block_type", None) + if not block_types: + return None + try: + return "".join(_MTP_BLOCK_TYPE_TO_VLLM_SYMBOL[block_type] for block_type in block_types) + except KeyError as error: + raise ValueError( + f"vLLM Nemotron-H MTP does not support block type {error.args[0]!r} in " + f"mtp_layers_block_type={list(block_types)!r}; supported block types are " + f"{sorted(_MTP_BLOCK_TYPE_TO_VLLM_SYMBOL)!r}." + ) from error + + +def _resolve_speechlm_mtp_config( + *, + mtp: dict | None, + compute_mtp: bool, + text_config, + num_nextn_predict_layers: int | None = None, + use_repeated_layer: bool | None = None, +) -> dict | None: + """Normalize the SpeechLM MTP contract consumed by the vLLM plugin. + + New exports carry an explicit root ``mtp`` dictionary. Older SpeechLM + exports, including the first Nemotron 3.5 Lightning checkpoints, only + carry ``compute_mtp`` at the root and keep MTP topology in the saved + backbone config. Derive the missing dictionary for those checkpoints so + they do not need to be re-exported. + """ + explicit_mtp = dict(mtp) if isinstance(mtp, dict) else None + if explicit_mtp is not None: + enabled = bool(explicit_mtp.get("enabled", True)) + else: + enabled = bool(compute_mtp) + + if not enabled: + return None + + if num_nextn_predict_layers is None: + if explicit_mtp is not None and "num_nextn_predict_layers" in explicit_mtp: + num_nextn_predict_layers = explicit_mtp["num_nextn_predict_layers"] + else: + num_nextn_predict_layers = getattr(text_config, "num_nextn_predict_layers", 0) + num_nextn_predict_layers = int(num_nextn_predict_layers or 0) + if num_nextn_predict_layers <= 0: + raise ValueError( + "SpeechLM MTP is enabled but num_nextn_predict_layers is not positive in either " + "the root mtp config or the backbone config." + ) + + explicit_pattern = explicit_mtp.get("hybrid_override_pattern") if explicit_mtp is not None else None + backbone_pattern = _mtp_pattern_from_backbone_config(text_config) + if explicit_pattern and backbone_pattern and explicit_pattern != backbone_pattern: + raise ValueError( + f"Root mtp.hybrid_override_pattern={explicit_pattern!r} disagrees with " + f"backbone MTP topology {backbone_pattern!r}." + ) + pattern = explicit_pattern or backbone_pattern + if not pattern: + raise ValueError( + "SpeechLM MTP is enabled but neither mtp.hybrid_override_pattern nor the backbone's " + "mtp_hybrid_override_pattern/mtp_layers_block_type declares the physical MTP topology." + ) + # Validate explicit patterns as well as backbone-derived patterns. + unsupported = sorted(set(pattern) - set(_MTP_BLOCK_TYPE_TO_VLLM_SYMBOL.values())) + if unsupported: + raise ValueError( + f"vLLM Nemotron-H MTP does not support pattern symbols {unsupported!r} " + f"in hybrid_override_pattern={pattern!r}; supported symbols are '*' and 'E'." + ) + + if use_repeated_layer is None: + use_repeated_layer = bool(explicit_mtp.get("use_repeated_layer", False)) if explicit_mtp else False + + return { + "enabled": True, + "num_nextn_predict_layers": num_nextn_predict_layers, + "use_repeated_layer": bool(use_repeated_layer), + "hybrid_override_pattern": pattern, + } + + class NeMoSpeechLMConfig(PretrainedConfig): """HuggingFace config for NeMo Speech LM multimodal models. @@ -114,6 +220,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) @@ -134,6 +241,7 @@ def __init__( self.pe_encoder_config = None self.pe_encoder_overrides = None self.speaker_encoder = None + self.__dict__.pop("_pending_image_token_index", None) return for name, value in required_fields.items(): @@ -195,6 +303,16 @@ def __init__( self.pe_encoder_overrides = pe_encoder_overrides self.speaker_encoder = speaker_encoder + # Backward compatibility for early Nemotron 3.5 SpeechLM exports: + # they carry ``compute_mtp`` at the root and the MTP topology only in + # llm_backbone/config.json. Normalize that into the explicit contract + # used by the vLLM speculative-config hook. + self.mtp = _resolve_speechlm_mtp_config( + mtp=self.__dict__.get("mtp"), + compute_mtp=bool(self.__dict__.get("compute_mtp", False)), + text_config=self.text_config, + ) + raw_archs = getattr(self.text_config, "architectures", []) if len(raw_archs) != 1: raise ValueError( @@ -228,6 +346,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]: @@ -237,6 +362,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", @@ -265,6 +430,7 @@ def __getattr__(self, name): "llm_config", "pretrained_asr", "audio_locator_tag", + "image_token_index", "prompt_format", "pretrained_weights", "text_config", diff --git a/nemo/collections/speechlm2/vllm/salm/model.py b/nemo/collections/speechlm2/vllm/salm/model.py index 99f34d147c90..3b3d034e5647 100644 --- a/nemo/collections/speechlm2/vllm/salm/model.py +++ b/nemo/collections/speechlm2/vllm/salm/model.py @@ -273,6 +273,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 diff --git a/nemo/collections/speechlm2/vllm/salm/mtp.py b/nemo/collections/speechlm2/vllm/salm/mtp.py new file mode 100644 index 000000000000..9814c5c27155 --- /dev/null +++ b/nemo/collections/speechlm2/vllm/salm/mtp.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MTP speculative-decoding draft model for NeMo SpeechLM checkpoints. + +``NeMoSpeechLMMTP`` wraps vLLM's ``NemotronHMTP`` and adapts the +weight-loading step for the NeMo checkpoint layout where all LLM +weights (including MTP layers) carry an ``llm.`` prefix: + + NeMo checkpoint NemotronHMTP expects + ────────────────────── ──────────────────── + llm.mtp.layers.0.* → mtp.layers.0.* + llm.lm_head.weight → lm_head.weight + llm.model.embed_* → backbone.embeddings.* + +The embedding alias is required by vLLM's Nemotron-H MTP loader even though +the proposer subsequently shares the table with the target model. +""" + +from collections.abc import Iterable + +import torch +from vllm.model_executor.models.nemotron_h_mtp import NemotronHMTP +from vllm.model_executor.models.utils import _merge_multimodal_embeddings + +from nemo.collections.speechlm2.vllm.salm.audio import _pad_to_vocab_size + + +def _remap_nemo_mtp_weights( + items: Iterable[tuple[str, torch.Tensor]], + target_vocab: int | None = None, + expected_layer_modules: int | None = None, +) -> Iterable[tuple[str, torch.Tensor]]: + """Select and map exported SpeechLM draft weights to vLLM aliases. + + ``expected_layer_modules`` is the number of sublayers instantiated for one + EAGLE prediction step (the hybrid pattern length), not the logical draft K. + Rejecting higher checkpoint indices prevents vLLM from silently dropping + weights from a distinct multi-head checkpoint routed as a repeated head. + Non-``llm.`` tensors are omitted because the target model loads them. + """ + for name, tensor in items: + if "._extra_state" in name: + continue + + if not name.startswith("llm."): + continue + name = name[len("llm.") :] + + parts = name.split(".") + if expected_layer_modules is not None and len(parts) > 2 and parts[:2] == ["mtp", "layers"]: + try: + layer_index = int(parts[2]) + except ValueError: + layer_index = None + if layer_index is not None and layer_index >= expected_layer_modules: + raise ValueError( + f"Checkpoint contains {name!r}, but the configured repeated MTP step instantiates " + f"only {expected_layer_modules} hybrid-pattern layer module(s). This looks like a " + f"distinct multi-head checkpoint or stale hybrid-pattern metadata. Verify that the " + f"exported mtp.hybrid_override_pattern describes the built head; vLLM's NemotronH " + f"MTP proposer can reuse only one EAGLE-style prediction step." + ) + + # NeMo stores every expert in one rank-3 tensor, while vLLM's + # NemotronHMTP loader expects one transposed matrix per expert. + if name.endswith(".experts.down_projs"): + prefix = name.removesuffix(".experts.down_projs") + for expert_idx, expert_tensor in enumerate(tensor): + yield f"{prefix}.experts.{expert_idx}.down_proj.weight", expert_tensor.t() + continue + if name.endswith(".experts.gate_and_up_projs"): + prefix = name.removesuffix(".experts.gate_and_up_projs") + for expert_idx, expert_tensor in enumerate(tensor): + yield f"{prefix}.experts.{expert_idx}.up_proj.weight", expert_tensor.t() + continue + + # NemotronHMTP.load_weights only admits embedding names containing + # ``embeddings`` and then maps this backbone alias to + # ``model.embed_tokens``. Passing model.embed_tokens through directly + # is silently skipped and DefaultModelLoader reports it uninitialized. + if name == "model.embed_tokens.weight": + name = "backbone.embeddings.weight" + + if name in {"backbone.embeddings.weight", "lm_head.weight"} and target_vocab is not None: + tensor = _pad_to_vocab_size(tensor, target_vocab) + yield name, tensor + + +class NeMoSpeechLMMTP(NemotronHMTP): + """NemotronH MTP draft model for NeMo SpeechLM checkpoints. + + Extends NemotronHMTP in two ways: + + * ``load_weights`` strips the ``llm.`` prefix from checkpoint names. + * ``embed_input_ids`` fuses audio-feature embeddings into the token + embeddings at placeholder positions, exactly like the target model. + The MTP heads were trained on the same mixed text+audio embedding + stream as the backbone, so the draft must see it too. vLLM probes + ``draft_model.embed_input_ids(ids, multimodal_embeddings=None)`` at + load time (``llm_base_proposer.load_model``); without this method + the probe raises AttributeError and speculative decoding silently + falls back to text-only draft inputs, which collapses acceptance + rates on audio prompts. + """ + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + """Embed token IDs and merge audio embeddings at placeholder positions. + + The target model inherits this fusion from ``SupportsMultiModal``; the + draft must implement it itself because ``NemotronHMTP`` is not + multimodal. The embedding table is shared from the target by vLLM's + MTP framework, so text-token rows stay identical. + """ + inputs_embeds = self.model.get_input_embeddings(input_ids) + + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + if is_multimodal is None: + raise ValueError("is_multimodal is required when multimodal_embeddings are provided.") + + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load only the SALM-prefixed weights required by the reusable draft head.""" + # NeMoSpeechLMConfig delegates this to the already padded text-config + # vocabulary used to construct NemotronHMTP. Reading config directly + # avoids coupling padding to vLLM's internal module names. + target_vocab = int(self.config.vocab_size) + if target_vocab <= 0: + raise ValueError(f"Draft model vocabulary size must be positive, got {target_vocab}.") + + # vLLM instantiates one physical module per character in the configured + # hybrid pattern. Derive the checkpoint bound from serialized metadata + # instead of depending on a predictor-internal implementation detail. + pattern = self.config.mtp_hybrid_override_pattern + if not isinstance(pattern, str) or not pattern: + raise ValueError(f"mtp_hybrid_override_pattern must be a non-empty string, got {pattern!r}.") + return super().load_weights( + _remap_nemo_mtp_weights( + weights, + target_vocab, + expected_layer_modules=len(pattern), + ) + ) diff --git a/nemo/collections/speechlm2/vllm/salm/runtime_compat.py b/nemo/collections/speechlm2/vllm/salm/runtime_compat.py index 34ca0d53e899..a1e5bde7e5d2 100644 --- a/nemo/collections/speechlm2/vllm/salm/runtime_compat.py +++ b/nemo/collections/speechlm2/vllm/salm/runtime_compat.py @@ -36,6 +36,7 @@ def _is_nemo_speechlm_tracker(mm_tracker: Any) -> bool: hf_config = getattr(model_config, "hf_config", None) return getattr(hf_config, "model_type", None) in { "nemo_speechlm", + "nemo_speechlm_mtp", } diff --git a/tests/collections/speechlm2/test_to_hf.py b/tests/collections/speechlm2/test_to_hf.py index 8f8859417743..64c8841e010b 100644 --- a/tests/collections/speechlm2/test_to_hf.py +++ b/tests/collections/speechlm2/test_to_hf.py @@ -110,6 +110,7 @@ def _seed_output_dir(tmp_path, llm_arch="Qwen2ForCausalLM"): "hidden_size": 2048, "num_hidden_layers": 24, "audio_token_index": 17, + "image_token_index": 18, } ) ) @@ -143,6 +144,111 @@ class _FakeExportModel: llm = type("_FakeLLM", (), {"config": _FakeLLMConfig()})() +class _FakeMTPBackboneConfig(_FakeLLMConfig): + num_nextn_predict_layers = 1 + mtp_hybrid_override_pattern = None + mtp_layers_block_type = ["attention", "moe"] + + +class _FakeMTPExportModel: + cfg = { + "pretrained_llm": "fake-model", + "pretrained_asr": "fake-asr", + "pretrained_weights": False, + "compute_mtp": True, + "mtp": None, + "dtype": "bf16", + "torch_dtype": "bf16", + "audio_locator_tag": AUDIO_TOKEN, + } + llm = type( + "_FakeMTPLLM", + (), + { + "config": _FakeMTPBackboneConfig(), + "mtp_config": SimpleNamespace(num_layers=1, use_repeated_layer=False), + }, + )() + + +class _FakeExplicitMTPExportModel: + cfg = { + **_FakeMTPExportModel.cfg, + "compute_mtp": False, + "mtp": { + "enabled": True, + "num_nextn_predict_layers": 1, + "use_repeated_layer": False, + "hybrid_override_pattern": "*E", + }, + } + llm = _FakeMTPExportModel.llm + + +class _FakeStaleLegacyMTPExportModel: + cfg = { + **_FakeMTPExportModel.cfg, + "compute_mtp": True, + "mtp": None, + } + llm = type( + "_FakeDisabledMTPLLM", + (), + { + "config": type( + "_FakeDisabledMTPBackboneConfig", + (_FakeLLMConfig,), + {"num_nextn_predict_layers": 0}, + )(), + "mtp_config": None, + }, + )() + + +def test_hf_export_config_persists_built_mtp_pattern_without_mutating_recipe(): + """A preserved native head's physical pattern must override stale recipe metadata.""" + model = SimpleNamespace( + cfg={ + "mtp": { + "enabled": True, + "hybrid_override_pattern": "*", + "num_nextn_predict_layers": 1, + } + }, + llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern="*E", num_nextn_predict_layers=1)), + ) + + config = to_hf._hf_export_config(model, "bfloat16") + + assert config["mtp"]["hybrid_override_pattern"] == "*E" + assert model.cfg["mtp"]["hybrid_override_pattern"] == "*" + + +def test_hf_export_config_persists_built_list_form_mtp_pattern(): + """A native list-form head must override stale replacement-recipe metadata.""" + model = SimpleNamespace( + cfg={ + "mtp": { + "enabled": True, + "hybrid_override_pattern": "*", + "num_nextn_predict_layers": 1, + } + }, + llm=SimpleNamespace( + config=SimpleNamespace( + mtp_hybrid_override_pattern=None, + mtp_layers_block_type=["attention", "moe"], + num_nextn_predict_layers=1, + ) + ), + ) + + config = to_hf._hf_export_config(model, "bfloat16") + + assert config["mtp"]["hybrid_override_pattern"] == "*E" + assert model.cfg["mtp"]["hybrid_override_pattern"] == "*" + + def test_hf_export_config_does_not_persist_remote_code_trust(): model = SimpleNamespace(cfg={"trust_remote_code": True}) @@ -152,6 +258,105 @@ def test_hf_export_config_does_not_persist_remote_code_trust(): assert model.cfg["trust_remote_code"] is True +def test_hf_export_config_persists_built_non_repeated_mtp_depth(): + """A preserved native head's actual depth must override stale recipe metadata.""" + model = SimpleNamespace( + cfg={ + "mtp": { + "enabled": True, + "hybrid_override_pattern": "*", + "num_nextn_predict_layers": 4, + "use_repeated_layer": False, + } + }, + llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern="*E", num_nextn_predict_layers=1)), + ) + + config = to_hf._hf_export_config(model, "bfloat16") + + assert config["mtp"]["num_nextn_predict_layers"] == 1 + assert model.cfg["mtp"]["num_nextn_predict_layers"] == 4 + + +def test_hf_export_config_keeps_logical_depth_for_repeated_mtp(): + model = SimpleNamespace( + cfg={ + "mtp": { + "enabled": True, + "hybrid_override_pattern": "*", + "num_nextn_predict_layers": 4, + "use_repeated_layer": True, + } + }, + llm=SimpleNamespace( + config=SimpleNamespace(mtp_hybrid_override_pattern="*E", num_nextn_predict_layers=1), + mtp_config=SimpleNamespace(num_layers=4, use_repeated_layer=True), + ), + ) + + config = to_hf._hf_export_config(model, "bfloat16") + + assert config["mtp"]["num_nextn_predict_layers"] == 4 + + +@pytest.mark.parametrize("depth", [False, 0, -1]) +def test_hf_export_config_rejects_invalid_built_mtp_depth(depth): + model = SimpleNamespace( + cfg={ + "mtp": { + "enabled": True, + "hybrid_override_pattern": "*", + "num_nextn_predict_layers": 1, + } + }, + llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern="*", num_nextn_predict_layers=depth)), + ) + + with pytest.raises(ValueError, match="num_nextn_predict_layers"): + to_hf._hf_export_config(model, "bfloat16") + + +@pytest.mark.parametrize("pattern", ["", 3]) +def test_hf_export_config_rejects_invalid_built_mtp_pattern(pattern): + model = SimpleNamespace( + cfg={ + "mtp": { + "enabled": True, + "hybrid_override_pattern": "*", + "num_nextn_predict_layers": 1, + } + }, + llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern=pattern, num_nextn_predict_layers=1)), + ) + + with pytest.raises(ValueError, match="mtp_hybrid_override_pattern"): + to_hf._hf_export_config(model, "bfloat16") + + +def test_save_hf_checkpoint_validates_config_before_writing_weights(tmp_path): + model = SimpleNamespace( + cfg={ + "mtp": { + "enabled": True, + "hybrid_override_pattern": "*", + "num_nextn_predict_layers": 1, + } + }, + llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern="", num_nextn_predict_layers=1)), + ) + cfg = to_hf.HfExportConfig( + class_path="fake.Class", + ckpt_path="fake.ckpt", + ckpt_config="fake.yaml", + output_dir=str(tmp_path), + ) + + with pytest.raises(ValueError, match="mtp_hybrid_override_pattern"): + to_hf.save_hf_checkpoint(model, {"weight": torch.zeros(1)}, cfg) + + assert not (tmp_path / "model.safetensors").exists() + + def test_save_hf_checkpoint_writes_llm_backbone_config(tmp_path): cfg = to_hf.HfExportConfig( class_path="fake.Class", @@ -295,6 +500,56 @@ def test_hf_export_config_embeds_portable_independent_dual_architecture(): assert original_cfg == {"speaker_encoder": {"path": "/models/speaker-transformer", "frozen": True}} +def test_save_hf_checkpoint_writes_explicit_mtp_contract(tmp_path): + """Nemotron 3.5 list-form topology should export as vLLM's *E pattern.""" + cfg = to_hf.HfExportConfig( + class_path="fake.Class", + ckpt_path="fake.ckpt", + ckpt_config="fake.yaml", + output_dir=str(tmp_path), + dtype="bfloat16", + ) + + to_hf.save_hf_checkpoint(_FakeMTPExportModel(), {"weight": torch.zeros(1)}, cfg) + + root_cfg = json.loads((tmp_path / "config.json").read_text()) + assert root_cfg["mtp"] == { + "enabled": True, + "num_nextn_predict_layers": 1, + "use_repeated_layer": False, + "hybrid_override_pattern": "*E", + } + + +def test_save_hf_checkpoint_preserves_explicit_mtp_contract_without_compute_flag( + tmp_path, +): + cfg = to_hf.HfExportConfig( + class_path="fake.Class", + ckpt_path="fake.ckpt", + ckpt_config="fake.yaml", + output_dir=str(tmp_path), + dtype="bfloat16", + ) + + to_hf.save_hf_checkpoint(_FakeExplicitMTPExportModel(), {"weight": torch.zeros(1)}, cfg) + + root_cfg = json.loads((tmp_path / "config.json").read_text()) + assert root_cfg["mtp"] == _FakeExplicitMTPExportModel.cfg["mtp"] + + +def test_hf_export_config_disables_stale_legacy_mtp_flag_without_runtime_head(): + exported = to_hf._hf_export_config(_FakeStaleLegacyMTPExportModel(), "bfloat16") + + assert exported["compute_mtp"] is False + assert exported["mtp"] is None + + +# ────────────────────────────────────────────────────────────────────── +# Error paths (no mocking required — checks run before any HF calls) +# ────────────────────────────────────────────────────────────────────── + + def test_prepare_for_vllm_missing_pretrained_llm(tmp_path): with pytest.raises(ValueError, match="pretrained_llm"): to_hf.prepare_for_vllm(str(tmp_path), {"audio_locator_tag": AUDIO_TOKEN}) @@ -361,6 +616,7 @@ def test_prepare_for_vllm_patches_config_json(tmp_path): assert cfg["architectures"] == ["NeMoSpeechLMForConditionalGeneration"] assert cfg["audio_locator_tag"] == AUDIO_TOKEN assert "audio_token_index" not in cfg + assert "image_token_index" not in cfg # Original LLM fields are preserved. assert cfg["hidden_size"] == 2048 @@ -449,6 +705,7 @@ def test_prepare_for_vllm_accepts_audio_token_inside_non_boundary_embedding_row( cfg = json.loads((output_dir / "config.json").read_text()) assert "audio_token_index" not in cfg + assert "image_token_index" not in cfg def test_prepare_for_vllm_uses_training_tokenizer_path(tmp_path): diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py index 3122a04057fe..c23cdea0fbc2 100644 --- a/tests/collections/speechlm2/test_vllm_plugin.py +++ b/tests/collections/speechlm2/test_vllm_plugin.py @@ -48,6 +48,11 @@ } +def _identity_hf_config_override(hf_config): + """Pickleable target override for vLLM's composed-override API.""" + return hf_config + + def test_full_stack_asr_exports_grouped_encoder_dependencies(): repo_root = Path(__file__).parents[3] code = """ @@ -135,6 +140,7 @@ def test_default_construction_for_hf_serialization(self): assert cfg.llm_config is None assert cfg.pretrained_asr is None assert cfg.audio_locator_tag is None + assert cfg.image_token_index is None assert cfg.prompt_format is None assert cfg.pretrained_weights is None assert cfg.pe_encoder_path is None @@ -285,6 +291,75 @@ def from_pretrained(model_name: str, trust_remote_code: bool = True): cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS) assert cfg.is_hybrid is expected_is_hybrid + def test_derives_multibranch_mtp_contract_from_backbone(self, monkeypatch): + """Legacy compute_mtp exports should derive the exact *E physical head.""" + + def from_pretrained(model_name: str, trust_remote_code: bool = True): + return SimpleNamespace( + architectures=["NemotronHForCausalLM"], + hidden_size=2048, + vocab_size=131072, + num_hidden_layers=4, + num_key_value_heads=2, + layer_norm_epsilon=1e-5, + num_nextn_predict_layers=1, + mtp_hybrid_override_pattern=None, + mtp_layers_block_type=["attention", "moe"], + ) + + monkeypatch.setattr(_config_module.AutoConfig, "from_pretrained", from_pretrained) + + cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS, compute_mtp=True, mtp=None) + + assert cfg.mtp == { + "enabled": True, + "num_nextn_predict_layers": 1, + "use_repeated_layer": False, + "hybrid_override_pattern": "*E", + } + assert cfg.mtp_hybrid_override_pattern == "*E" + + def test_compute_mtp_false_does_not_enable_backbone_head(self, monkeypatch): + """A backbone MTP head must not opt an export into speculative decoding.""" + + def from_pretrained(model_name: str, trust_remote_code: bool = True): + return SimpleNamespace( + architectures=["NemotronHForCausalLM"], + hidden_size=2048, + vocab_size=131072, + num_hidden_layers=4, + num_key_value_heads=2, + layer_norm_epsilon=1e-5, + num_nextn_predict_layers=1, + mtp_layers_block_type=["attention", "moe"], + ) + + monkeypatch.setattr(_config_module.AutoConfig, "from_pretrained", from_pretrained) + + cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS, compute_mtp=False) + + assert cfg.mtp is None + + def test_unsupported_mtp_topology_fails_closed(self, monkeypatch): + """vLLM 0.23 cannot instantiate Mamba or MLP MTP sublayers.""" + + def from_pretrained(model_name: str, trust_remote_code: bool = True): + return SimpleNamespace( + architectures=["NemotronHForCausalLM"], + hidden_size=2048, + vocab_size=131072, + num_hidden_layers=4, + num_key_value_heads=2, + layer_norm_epsilon=1e-5, + num_nextn_predict_layers=1, + mtp_layers_block_type=["mamba", "moe"], + ) + + monkeypatch.setattr(_config_module.AutoConfig, "from_pretrained", from_pretrained) + + with pytest.raises(ValueError, match="does not support block type 'mamba'"): + NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS, compute_mtp=True) + def test_hybrid_backbone_does_not_set_layer_types_shim(self): """Hybrid backbones must NOT have layer_types overridden -- the runtime is_hybrid escape hatch only fires when every layer is 'attention'.""" @@ -1227,6 +1302,765 @@ def test_register_does_not_load_backbone_config(self, monkeypatch): @pytest.mark.skipif(not _HAS_VLLM, reason="vLLM not installed") +class TestMTPPlugin: + """Tests for NeMo SpeechLM MTP speculative-decoding support.""" + + @pytest.fixture(autouse=True) + def restore_original_override(self): + """Keep the process-local fallback hook isolated between tests.""" + from nemo.collections.speechlm2.vllm import salm as salm_module + + original_override = salm_module._ORIGINAL_VLLM_HF_CONFIG_OVERRIDE + yield + salm_module._ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = original_override + + @pytest.fixture(autouse=True) + def mock_backbone_config(self, monkeypatch): + """Keep registration tests independent of Hugging Face network access.""" + if not _HAS_CONFIG: + return + + monkeypatch.setattr( + _config_module.AutoConfig, + "from_pretrained", + lambda *args, **kwargs: SimpleNamespace( + architectures=["NemotronHybridForCausalLM"], + hidden_size=2048, + vocab_size=131072, + num_hidden_layers=4, + num_key_value_heads=2, + layer_norm_epsilon=1e-5, + ), + ) + + class _HFConfigLike: + """Minimal stand-in for a HuggingFace PretrainedConfig.""" + + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + + def update(self, d): + for k, v in d.items(): + setattr(self, k, v) + + def test_mtp_patch_registers_model(self, monkeypatch): + """register() should add NeMoSpeechLMMTPModel to the model registry.""" + from transformers import AutoConfig + from vllm.model_executor.models.registry import ModelRegistry + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()), + ) + + register() + + assert "NeMoSpeechLMMTPModel" in ModelRegistry.get_supported_archs() + + def test_mtp_patch_extends_mtp_model_types(self, monkeypatch): + """register() should add 'nemo_speechlm_mtp' to vLLM's MTPModelTypes Literal.""" + from typing import get_args + + import vllm.config.speculative as _spec_mod + from transformers import AutoConfig + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()), + ) + + register() + + assert "nemo_speechlm_mtp" in get_args(_spec_mod.MTPModelTypes) + + def test_patched_override_routes_nemo_mtp_config(self, monkeypatch): + """hf_config_override should rewrite nemo_speechlm configs with MTP heads.""" + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()), + ) + register() + + hf_cfg = self._HFConfigLike( + model_type="nemo_speechlm", + mtp={ + "enabled": True, + "num_nextn_predict_layers": 1, + "use_repeated_layer": True, + }, + ) + result = SpeculativeConfig.hf_config_override(hf_cfg) + + assert result.model_type == "nemo_speechlm_mtp" + assert result.architectures == ["NeMoSpeechLMMTPModel"] + assert result.n_predict == 1 + assert result.num_nextn_predict_layers == 1 + + def test_patched_override_is_pickleable_for_spawned_engine(self, monkeypatch): + """The callable retained on draft ModelConfig must survive multiprocessing spawn.""" + import pickle + + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm import salm as salm_module + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()), + ) + register() + + restored = pickle.loads(pickle.dumps(SpeculativeConfig.hf_config_override)) + assert restored is salm_module._nemo_speechlm_mtp_hf_config_override + + overrides = [restored] + if hasattr(SpeculativeConfig, "compose_draft_hf_overrides"): + composed = SpeculativeConfig.compose_draft_hf_overrides(_identity_hf_config_override) + assert composed is not SpeculativeConfig.hf_config_override + overrides.append(pickle.loads(pickle.dumps(composed))) + + for override in overrides: + hf_cfg = self._HFConfigLike( + model_type="nemo_speechlm", + mtp={ + "enabled": True, + "num_nextn_predict_layers": 4, + "use_repeated_layer": True, + }, + ) + result = override(hf_cfg) + assert result.model_type == "nemo_speechlm_mtp" + assert result.n_predict == 1 + + def test_patched_override_lazily_captures_native_hook_in_spawn_child(self, monkeypatch): + """A fresh spawn import should delegate unrelated configs to vLLM's native hook.""" + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm import salm as salm_module + + original_calls = [] + + def _recording_native(cfg): + original_calls.append(cfg) + return cfg + + monkeypatch.setattr(salm_module, "_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE", None) + monkeypatch.setattr( + SpeculativeConfig, + "hf_config_override", + staticmethod(_recording_native), + ) + + hf_cfg = self._HFConfigLike(model_type="unrelated") + result = salm_module._nemo_speechlm_mtp_hf_config_override(hf_cfg) + + assert result is hf_cfg + assert original_calls == [hf_cfg] + assert salm_module._ORIGINAL_VLLM_HF_CONFIG_OVERRIDE is _recording_native + + def test_patched_override_rejects_missing_native_hook(self, monkeypatch): + """A corrupted install must fail clearly instead of recursing into our override.""" + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm import salm as salm_module + + monkeypatch.setattr(salm_module, "_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE", None) + monkeypatch.setattr( + SpeculativeConfig, + "hf_config_override", + staticmethod(salm_module._nemo_speechlm_mtp_hf_config_override), + ) + + with pytest.raises(RuntimeError, match="without preserving vLLM's original hook"): + salm_module._nemo_speechlm_mtp_hf_config_override(self._HFConfigLike(model_type="unrelated")) + + def test_patched_override_enabled_mtp_defaults_to_one_head(self, monkeypatch): + """An enabled training block without an explicit depth constructs one head.""" + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()), + ) + register() + + hf_cfg = self._HFConfigLike(model_type="nemo_speechlm", mtp={"enabled": True}) + result = SpeculativeConfig.hf_config_override(hf_cfg) + + assert result.model_type == "nemo_speechlm_mtp" + assert result.n_predict == 1 + assert result.num_nextn_predict_layers == 1 + + def test_patched_override_repeated_layer_exposes_one_reusable_head(self, monkeypatch): + """Repeated-layer training depth must not constrain inference-time K.""" + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()), + ) + register() + + hf_cfg = self._HFConfigLike( + model_type="nemo_speechlm", + mtp={ + "enabled": True, + "num_nextn_predict_layers": 4, + "use_repeated_layer": True, + }, + ) + result = SpeculativeConfig.hf_config_override(hf_cfg) + + assert result.n_predict == 1 + assert result.num_nextn_predict_layers == 1 + + def test_patched_override_no_mtp_falls_through(self, monkeypatch): + """hf_config_override should not alter non-MTP configs.""" + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()), + ) + original_calls = [] + + def _recording_orig(cfg): + original_calls.append(cfg) + return cfg + + from nemo.collections.speechlm2.vllm import salm as salm_module + + monkeypatch.setattr(salm_module, "_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE", None) + monkeypatch.setattr(SpeculativeConfig, "hf_config_override", staticmethod(_recording_orig)) + register() + + hf_cfg = self._HFConfigLike( + model_type="nemo_speechlm", + mtp={"enabled": False, "num_nextn_predict_layers": 1}, + ) + SpeculativeConfig.hf_config_override(hf_cfg) + + assert len(original_calls) == 1 + + def test_patched_override_depth_without_enabled_flag_falls_through(self, monkeypatch): + """A retained recipe depth must not enable a head that training did not construct.""" + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm import salm as salm_module + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()), + ) + original_calls = [] + + def _recording_orig(cfg): + original_calls.append(cfg) + return cfg + + monkeypatch.setattr(salm_module, "_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE", None) + monkeypatch.setattr(SpeculativeConfig, "hf_config_override", staticmethod(_recording_orig)) + register() + + hf_cfg = self._HFConfigLike( + model_type="nemo_speechlm", + mtp={"num_nextn_predict_layers": 4, "use_repeated_layer": True}, + ) + result = SpeculativeConfig.hf_config_override(hf_cfg) + + assert result.model_type == "nemo_speechlm" + assert original_calls == [hf_cfg] + + def test_patched_override_explicitly_disabled_mtp_falls_through(self, monkeypatch): + """An exported recipe depth must not override mtp.enabled=false.""" + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm import salm as salm_module + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()), + ) + original_calls = [] + + def _recording_orig(cfg): + original_calls.append(cfg) + return cfg + + monkeypatch.setattr(salm_module, "_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE", None) + monkeypatch.setattr(SpeculativeConfig, "hf_config_override", staticmethod(_recording_orig)) + register() + + hf_cfg = self._HFConfigLike( + model_type="nemo_speechlm", + mtp={ + "enabled": False, + "num_nextn_predict_layers": 4, + "use_repeated_layer": True, + }, + ) + result = SpeculativeConfig.hf_config_override(hf_cfg) + + assert result.model_type == "nemo_speechlm" + assert original_calls == [hf_cfg] + + def test_patched_override_multi_head_without_repeated_layer_raises(self, monkeypatch): + """hf_config_override should raise for multi-head checkpoints without use_repeated_layer.""" + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()), + ) + register() + + hf_cfg = self._HFConfigLike( + model_type="nemo_speechlm", + mtp={ + "enabled": True, + "num_nextn_predict_layers": 3, + "use_repeated_layer": False, + }, + ) + with pytest.raises(ValueError, match="use_repeated_layer"): + SpeculativeConfig.hf_config_override(hf_cfg) + + def test_mtp_override_registration_is_idempotent(self, monkeypatch): + """register() should not repeatedly wrap the config override.""" + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()), + ) + register() + first_override = SpeculativeConfig.hf_config_override + register() + + assert SpeculativeConfig.hf_config_override is first_override + + def test_mtp_override_reregistration_preserves_first_native_hook(self, monkeypatch): + """A later wrapper that delegates to us must not become our fallback.""" + from transformers import AutoConfig + from vllm.config.speculative import SpeculativeConfig + + from nemo.collections.speechlm2.vllm import salm as salm_module + from nemo.collections.speechlm2.vllm.salm import register + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *a, **kw: (_ for _ in ()).throw(RuntimeError()), + ) + native_calls = [] + + def _recording_native(cfg): + native_calls.append(cfg) + return cfg + + monkeypatch.setattr(salm_module, "_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE", None) + monkeypatch.setattr( + SpeculativeConfig, + "hf_config_override", + staticmethod(_recording_native), + ) + register() + first_override = SpeculativeConfig.hf_config_override + + def _third_party_wrapper(cfg): + return first_override(cfg) + + monkeypatch.setattr( + SpeculativeConfig, + "hf_config_override", + staticmethod(_third_party_wrapper), + ) + register() + + hf_cfg = self._HFConfigLike(model_type="unrelated") + result = SpeculativeConfig.hf_config_override(hf_cfg) + + assert result is hf_cfg + assert native_calls == [hf_cfg] + assert salm_module._ORIGINAL_VLLM_HF_CONFIG_OVERRIDE is _recording_native + + def test_embed_input_ids_text_only(self): + """embed_input_ids with no audio embeddings should return plain text embeddings.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP + + m = object.__new__(NeMoSpeechLMMTP) + base_embeds = torch.arange(6, dtype=torch.float).reshape(3, 2) + m.model = SimpleNamespace(get_input_embeddings=lambda ids: base_embeds.clone()) + + result = m.embed_input_ids(torch.tensor([1, 2, 3]), multimodal_embeddings=None) + + assert result.shape == (3, 2) + assert torch.equal(result, base_embeds) + + def test_embed_input_ids_fuses_audio(self): + """embed_input_ids should replace placeholder positions with audio embeddings.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP + + m = object.__new__(NeMoSpeechLMMTP) + base_embeds = torch.zeros(4, 2) + m.model = SimpleNamespace(get_input_embeddings=lambda ids: base_embeds.clone()) + + audio_feat = torch.ones(2, 2) * 9.0 + is_audio = torch.tensor([False, True, True, False]) + result = m.embed_input_ids( + torch.tensor([0, 1, 2, 3]), + multimodal_embeddings=[audio_feat], + is_multimodal=is_audio, + ) + + assert torch.equal(result[0], torch.zeros(2)) + assert torch.equal(result[1], torch.ones(2) * 9.0) + assert torch.equal(result[2], torch.ones(2) * 9.0) + assert torch.equal(result[3], torch.zeros(2)) + + def test_embed_input_ids_fuses_multiple_audio_chunks(self): + """embed_input_ids should preserve vLLM's nested embedding order.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP + + model = object.__new__(NeMoSpeechLMMTP) + base_embeds = torch.zeros(5, 2) + model.model = SimpleNamespace(get_input_embeddings=lambda ids: base_embeds.clone()) + + audio_feats = [[torch.ones(1, 2) * 3.0], [torch.ones(2, 2) * 7.0]] + is_audio = torch.tensor([True, False, True, True, False]) + result = model.embed_input_ids( + torch.tensor([0, 1, 2, 3, 4]), + multimodal_embeddings=audio_feats, + is_multimodal=is_audio, + ) + + assert torch.equal(result[0], torch.ones(2) * 3.0) + assert torch.equal(result[1], torch.zeros(2)) + assert torch.equal(result[2], torch.ones(2) * 7.0) + assert torch.equal(result[3], torch.ones(2) * 7.0) + assert torch.equal(result[4], torch.zeros(2)) + + def test_embed_input_ids_requires_multimodal_mask(self): + """Audio embeddings without placeholder positions should fail clearly.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP + + model = object.__new__(NeMoSpeechLMMTP) + model.model = SimpleNamespace(get_input_embeddings=lambda ids: torch.zeros(2, 2)) + + with pytest.raises(ValueError, match="is_multimodal"): + model.embed_input_ids( + torch.tensor([0, 1]), + multimodal_embeddings=[torch.ones(1, 2)], + ) + + def test_embed_input_ids_ignores_embeddings_without_placeholder_positions(self): + """vLLM's merge semantics leave embeddings unchanged for an all-text mask.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP + + model = object.__new__(NeMoSpeechLMMTP) + model.model = SimpleNamespace(get_input_embeddings=lambda ids: torch.zeros(2, 2)) + + result = model.embed_input_ids( + torch.tensor([0, 1]), + multimodal_embeddings=[torch.ones(1, 2)], + is_multimodal=torch.tensor([False, False]), + ) + + assert torch.equal(result, torch.zeros(2, 2)) + + def test_target_weight_split_excludes_salm_mtp_and_rejects_bare_mtp(self): + """The target loader should route SALM draft weights and reject native names.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration + + perception_tensor = torch.ones(1) + llm_tensor = torch.ones(2) + perception, llm = NeMoSpeechLMForConditionalGeneration._split_perception_llm( + [ + ("perception.encoder.weight", perception_tensor), + ("llm.mtp.layers.0.weight", torch.ones(4)), + ("llm.model.layers.0.weight", llm_tensor), + ("llm.model.layers.0._extra_state", torch.ones(5)), + ] + ) + + assert perception == {"encoder.weight": perception_tensor} + assert [name for name, _ in llm] == ["llm.model.layers.0.weight"] + assert llm[0][1] is llm_tensor + + with pytest.raises(ValueError, match=r"llm\.mtp\.\*"): + NeMoSpeechLMForConditionalGeneration._split_perception_llm([("mtp.layers.0.weight", torch.ones(3))]) + + def test_mtp_weight_remap_uses_vllm_embedding_alias(self): + """Exported SpeechLM embeddings must pass NemotronHMTP's name filter.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.mtp import _remap_nemo_mtp_weights + + tensor = torch.ones(2, 3) + remapped = dict( + _remap_nemo_mtp_weights( + [ + ("llm.model.embed_tokens.weight", tensor), + ("llm.mtp.layers.0.enorm.weight", tensor), + ("llm.lm_head.weight", tensor), + ] + ) + ) + + assert set(remapped) == { + "backbone.embeddings.weight", + "mtp.layers.0.enorm.weight", + "lm_head.weight", + } + assert remapped["backbone.embeddings.weight"] is tensor + + padded = dict( + _remap_nemo_mtp_weights( + [ + ("llm.model.embed_tokens.weight", tensor), + ("llm.lm_head.weight", tensor), + ], + target_vocab=5, + ) + ) + assert padded["backbone.embeddings.weight"].shape == (5, 3) + assert padded["lm_head.weight"].shape == (5, 3) + + def test_mtp_weight_remap_rejects_distinct_head_layers(self): + """Unsupported distinct heads in a SpeechLM export should fail loudly.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.mtp import _remap_nemo_mtp_weights + + tensor = torch.ones(2, 3) + with pytest.raises(ValueError, match="distinct multi-head"): + list( + _remap_nemo_mtp_weights( + [("llm.mtp.layers.1.enorm.weight", tensor)], + expected_layer_modules=1, + ) + ) + + def test_mtp_weight_remap_ignores_non_salm_names(self): + """The draft loader supports SpeechLM exports, not bare backbone weights.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.mtp import _remap_nemo_mtp_weights + + assert list(_remap_nemo_mtp_weights([("mtp.layers.0.norm.weight", torch.ones(1))])) == [] + + def test_mtp_weight_remap_allows_all_sublayers_in_hybrid_pattern(self): + """Hybrid-pattern sublayers belong to one reusable prediction step.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.mtp import _remap_nemo_mtp_weights + + tensor = torch.ones(2, 3) + remapped = dict( + _remap_nemo_mtp_weights( + [(f"llm.mtp.layers.{i}.norm.weight", tensor) for i in range(3)], + expected_layer_modules=3, + ) + ) + assert set(remapped) == {f"mtp.layers.{i}.norm.weight" for i in range(3)} + + def test_mtp_load_weights_bounds_layers_from_serialized_pattern(self, monkeypatch): + """The loader should derive its physical-layer bound from config.""" + import torch + from vllm.model_executor.models.nemotron_h_mtp import NemotronHMTP + + from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP + + model = object.__new__(NeMoSpeechLMMTP) + object.__setattr__( + model, + "config", + SimpleNamespace(mtp_hybrid_override_pattern="*E", vocab_size=3), + ) + captured = [] + + def _capture_weights(self, weights): + captured.extend(weights) + return set() + + monkeypatch.setattr(NemotronHMTP, "load_weights", _capture_weights) + tensor = torch.ones(1, 2) + NeMoSpeechLMMTP.load_weights( + model, + [ + ("llm.model.embed_tokens.weight", tensor), + ("llm.mtp.layers.0.norm.weight", tensor), + ("llm.mtp.layers.1.norm.weight", tensor), + ("llm.lm_head.weight", tensor), + ], + ) + assert [name for name, _ in captured] == [ + "backbone.embeddings.weight", + "mtp.layers.0.norm.weight", + "mtp.layers.1.norm.weight", + "lm_head.weight", + ] + captured_tensors = dict(captured) + assert captured_tensors["backbone.embeddings.weight"].shape == (3, 2) + assert captured_tensors["lm_head.weight"].shape == (3, 2) + assert captured_tensors["mtp.layers.0.norm.weight"] is tensor + assert captured_tensors["mtp.layers.1.norm.weight"] is tensor + + with pytest.raises(ValueError, match="distinct multi-head"): + NeMoSpeechLMMTP.load_weights(model, [("llm.mtp.layers.2.norm.weight", tensor)]) + + @pytest.mark.parametrize("pattern", [None, "", 3]) + def test_mtp_load_weights_rejects_invalid_serialized_pattern(self, monkeypatch, pattern): + from vllm.model_executor.models.nemotron_h_mtp import NemotronHMTP + + from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP + + model = object.__new__(NeMoSpeechLMMTP) + object.__setattr__( + model, + "config", + SimpleNamespace(mtp_hybrid_override_pattern=pattern, vocab_size=3), + ) + monkeypatch.setattr(NemotronHMTP, "load_weights", lambda self, weights: set()) + + with pytest.raises(ValueError, match="non-empty string"): + NeMoSpeechLMMTP.load_weights(model, []) + + def test_mtp_weight_remap_splits_packed_experts(self): + """Packed Automodel MTP experts must become vLLM per-expert weights.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.mtp import _remap_nemo_mtp_weights + + down_projs = torch.arange(24, dtype=torch.float32).reshape(2, 4, 3) + up_projs = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4) + remapped = dict( + _remap_nemo_mtp_weights( + [ + ("llm.mtp.layers.1.mixer.experts.down_projs", down_projs), + ("llm.mtp.layers.1.mixer.experts.gate_and_up_projs", up_projs), + ("llm.mtp.layers.1.mixer.experts._extra_state", torch.tensor(1)), + ] + ) + ) + + assert set(remapped) == { + "mtp.layers.1.mixer.experts.0.down_proj.weight", + "mtp.layers.1.mixer.experts.1.down_proj.weight", + "mtp.layers.1.mixer.experts.0.up_proj.weight", + "mtp.layers.1.mixer.experts.1.up_proj.weight", + } + for expert_idx in range(2): + down = remapped[f"mtp.layers.1.mixer.experts.{expert_idx}.down_proj.weight"] + up = remapped[f"mtp.layers.1.mixer.experts.{expert_idx}.up_proj.weight"] + assert down.shape == (3, 4) + assert up.shape == (4, 3) + assert down.dtype == down_projs.dtype + assert up.dtype == up_projs.dtype + assert torch.equal(down, down_projs[expert_idx].t()) + assert torch.equal(up, up_projs[expert_idx].t()) + + @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available") + def test_mtp_hybrid_override_pattern_from_config(self): + """mtp_hybrid_override_pattern should read hybrid_override_pattern from mtp config dict.""" + cfg = NeMoSpeechLMConfig( + **_DEFAULT_CONFIG_KWARGS, + mtp={ + "enabled": True, + "num_nextn_predict_layers": 1, + "hybrid_override_pattern": "*E", + }, + ) + assert cfg.mtp_hybrid_override_pattern == "*E" + + @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available") + def test_mtp_hybrid_override_pattern_default_all_attention(self): + """mtp_hybrid_override_pattern should default to '*' (all-attention) when absent.""" + cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS) + assert cfg.mtp_hybrid_override_pattern == "*" + + @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available") + def test_image_token_index_is_unserialized_backbone_vocab_boundary(self): + """vLLM's compatibility property should identify the first padded row.""" + import importlib + + config_mod = importlib.import_module("nemo.collections.speechlm2.vllm.salm.config") + extra_rows = config_mod._SPEECHLM_EMBED_EXTRA_ROWS + + cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS) + base_vocab = cfg.text_config.vocab_size - extra_rows + assert cfg.image_token_index == base_vocab + # vLLM copies the target value onto the draft config. The setter accepts + # that runtime-only assignment without adding serialized state. + cfg.image_token_index = base_vocab + assert cfg.image_token_index == base_vocab + assert "audio_token_index" not in cfg.to_dict() + assert "image_token_index" not in cfg.to_dict() + + @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available") + def test_legacy_image_token_index_is_validated_after_backbone_load(self): + cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS, image_token_index=131072) + assert cfg.image_token_index == 131072 + assert "image_token_index" not in cfg.to_dict() + + with pytest.raises(ValueError, match="backbone vocabulary boundary"): + NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS, image_token_index=42) + + class _FakeTokenizer: def __init__(self): self.added_special_tokens = None