diff --git a/docs/source/speechlm2/intro.rst b/docs/source/speechlm2/intro.rst index 8aac63f9f12f..9fb805304f97 100644 --- a/docs/source/speechlm2/intro.rst +++ b/docs/source/speechlm2/intro.rst @@ -393,3 +393,4 @@ For more information, see additional sections in the SpeechLM2 docs: datasets configs training_and_scaling + vllm_dflash diff --git a/docs/source/speechlm2/vllm_dflash.rst b/docs/source/speechlm2/vllm_dflash.rst new file mode 100644 index 000000000000..bf56081ce9ab --- /dev/null +++ b/docs/source/speechlm2/vllm_dflash.rst @@ -0,0 +1,95 @@ +DFlash speculative decoding with vLLM +====================================== + +The NeMo SpeechLM vLLM plugin supports checkpoint-backed DFlash and DFlash2 +speculative decoding. Both use intermediate hidden states from the SpeechLM +language tower to condition a separate draft model; generated draft tokens are +verified by the target model, so accepted output remains lossless relative to +the target. + +Requirements +------------ + +* A vLLM-ready NeMo SpeechLM checkpoint whose language backbone is compatible + with the DFlash draft. +* vLLM 0.27.1 or later for the published Nemotron 3.5 Lightning recipe. +* An attention backend that supports the draft model's non-causal attention. + +The following example uses the published NVFP4 DFlash draft for the Nemotron +3.5 Lightning 30B-A3B backbone and proposes six tokens per decoding step: + +.. code-block:: bash + + vllm serve /path/to/vllm-ready-speechlm-checkpoint \ + --trust-remote-code \ + --speculative-config '{ + "method": "dflash", + "model": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4-DFlash", + "num_speculative_tokens": 6 + }' + +The target SpeechLM checkpoint must use +``nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16`` as its language backbone. +The draft checkpoint provides the auxiliary target-layer selection and mask +token configuration consumed by vLLM; no draft weights are bundled with NeMo. +Automodel-trained drafts may retain ``Qwen3DFlashDraftModel`` as their +architecture so they can be reopened for training. The NeMo plugin registers +that name as an alias of vLLM's native ``DFlashDraftModel`` implementation; +serve the original checkpoint without rewriting its ``config.json``. + +DFlash2 inference +----------------- + +DFlash2 adds dynamic convolutions and a candidate-path selector to the draft +model. Its checkpoint must be trained or fine-tuned separately for the target +language backbone; NeMo's vLLM inference plugin does not create or convert +DFlash2 weights. + +At the time of writing, DFlash2 requires vLLM commit +``3406ec1dae9916f920b90f0dbf90dcf54923d042`` from pull request 52816. The +immutable commit pin keeps the DFlash2 runtime reproducible. DFlash2 uses the +same ``method`` value as DFlash; vLLM selects it from the trained draft +checkpoint's architecture: + +.. code-block:: bash + + pip install -U "vllm @ git+https://github.com/vllm-project/vllm.git@3406ec1dae9916f920b90f0dbf90dcf54923d042" + + vllm serve /path/to/vllm-ready-speechlm-checkpoint \ + --trust-remote-code \ + --speculative-config '{ + "method": "dflash", + "model": "/path/to/trained-lightning-dflash2-checkpoint", + "num_speculative_tokens": 6 + }' + +NeMo Automodel training exports ``Qwen3DFlash2DraftModel`` so the checkpoint +can still be reopened by the training stack. The SpeechLM plugin normalizes +that architecture to vLLM's canonical ``DFlash2DraftModel`` before vLLM wraps +the speculative config. This is required for vLLM to force its V2 model runner +and execute the DFlash2 candidate-selector speculator instead of silently +falling back to plain DFlash. Native configs that already declare +``DFlash2DraftModel`` remain supported. + +The draft's ``dflash_config`` must include ``target_layer_ids``, +``conv_group_size``, ``conv_kernel_size``, ``selector_rank``, and +``selector_top_k``. Set +``num_speculative_tokens`` to one less than the convolution block size used to +train the draft: vLLM constructs each runtime block from one anchor plus the +configured number of draft tokens and does not reject a training/inference +block-size mismatch. + +The DFlash2 architecture forces vLLM's V2 model runner, including for hybrid +NemotronH targets that would otherwise use V1. vLLM raises an error for features +it knows are incompatible with V2, but the target and serving configuration +should still be qualified on that runner. The SpeechLM target uses the same +``SupportsEagle3`` hidden-state contract for DFlash and DFlash2, so no draft +weights or training logic are bundled with NeMo. + +Validation +---------- + +Compare greedy generation with and without ``--speculative-config`` using the +same text and audio prompts. The generated token IDs must match. Also inspect +vLLM's speculative-decoding metrics to confirm that draft tokens are proposed +and accepted; matching output alone does not prove that DFlash was active. diff --git a/examples/speechlm2/to_hf.py b/examples/speechlm2/to_hf.py index e0734fbce70e..55996171fc92 100644 --- a/examples/speechlm2/to_hf.py +++ b/examples/speechlm2/to_hf.py @@ -102,6 +102,37 @@ 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) + # Remote-code trust is a runtime security decision. Do not persist a + # training-time opt-in in checkpoints that may be loaded by another user. + config.pop("trust_remote_code", None) + 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 +147,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 +165,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 +192,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 +210,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) @@ -195,15 +232,27 @@ def prepare_for_vllm(output_dir: str, model_cfg: dict) -> None: # 1. Patch config.json (arch, model_type, audio_locator_tag for vLLM plugin). arch_model_cfg = dict(model_cfg) llm_backbone_dir = output_dir / LLM_BACKBONE_DIR - if (llm_backbone_dir / "config.json").exists(): + llm_backbone_config_path = llm_backbone_dir / "config.json" + llm_backbone_config = None + if llm_backbone_config_path.exists(): arch_model_cfg["pretrained_llm"] = str(llm_backbone_dir) - arch = _detect_vllm_architecture(arch_model_cfg) + llm_backbone_config = json.loads(llm_backbone_config_path.read_text()) + 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") + if llm_backbone_config is not None: + # Keep the export portable while making the bundled config authoritative. + # NeMo's HF loader resolves this relative marker to a cached local path; + # the vLLM config consumes the embedded copy without another Hub lookup. + config["pretrained_llm"] = LLM_BACKBONE_DIR + config["llm_config"] = llm_backbone_config + 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 = [ @@ -213,9 +262,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, fix_mistral_regex=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. diff --git a/nemo/collections/asr/modules/parallel_expert_encoder.py b/nemo/collections/asr/modules/parallel_expert_encoder.py index 73e1c488c85d..e626e9dfa751 100644 --- a/nemo/collections/asr/modules/parallel_expert_encoder.py +++ b/nemo/collections/asr/modules/parallel_expert_encoder.py @@ -14,10 +14,10 @@ """Parallel Expert Speech Encoder. -Runs a Sortformer speaker-diarization expert and an ASR Conformer encoder on the +Runs a Sortformer speaker-diarization expert and an ASR encoder on the same mel input, then fuses their outputs (LayerNorm + sinusoidal speaker-kernel + ADD). Expects un-normalised mels; the ASR branch re-applies ``normalize_batch`` -internally. I/O matches :class:`ConformerEncoder` (drop-in). Only self-contained PE +internally. I/O matches the supported ASR encoders (drop-in). Only self-contained PE bundles (inline ``asr_encoder_cfg`` + ``diarization_model_cfg`` in ``model_config.yaml``) are supported. """ @@ -38,6 +38,7 @@ from tqdm import tqdm from nemo.collections.asr.modules.conformer_encoder import ConformerEncoder +from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder from nemo.collections.asr.parts.preprocessing.features import normalize_batch from nemo.core.classes import ModelPT from nemo.core.classes.common import PretrainedModelInfo @@ -284,16 +285,17 @@ def save_to_nemo( @experimental class ParallelExpertEncoder(nn.Module): - """Sortformer-diarizer + ASR Conformer encoder; I/O identical to :class:`ConformerEncoder`. + """Sortformer-diarizer plus a supported ASR encoder with a shared I/O contract. Reconstructed from inline configs in the PE bundle's ``model_config.yaml``. Args: - asr_encoder_cfg (DictConfig): Inline config for the ASR-side :class:`ConformerEncoder`. + asr_encoder_cfg (DictConfig): Inline config for an ASR-side :class:`ConformerEncoder` + or :class:`TransformerEncoder`. diarization_model_cfg (DictConfig): Inline config for the :class:`SortformerEncLabelModel`. asr_normalize_type (str, optional): Normalization replayed on the ASR branch. Defaults to ``per_feature``. freeze_diar (bool): Freeze the Sortformer parameters. Defaults to ``True``. - freeze_asr (bool): Freeze the wrapped ASR ConformerEncoder. Defaults to ``False``. + freeze_asr (bool): Freeze the wrapped ASR encoder. Defaults to ``False``. online_inference_length (int): Online-inference window in encoder output frames (default ``500`` ~= 40s); ``<= 0`` disables it. chunk_left_context (int): Left context (output frames) per online window, shared by @@ -332,10 +334,10 @@ def __init__( ) self.asr_encoder = ConformerEncoder.from_config_dict(_clone_config(asr_encoder_cfg)) - if not isinstance(self.asr_encoder, ConformerEncoder): + if not isinstance(self.asr_encoder, (ConformerEncoder, TransformerEncoder)): raise TypeError( - f"Expected `asr_encoder_cfg._target_` to instantiate a " - f"ConformerEncoder, got {type(self.asr_encoder).__name__} instead." + "Expected `asr_encoder_cfg._target_` to instantiate a ConformerEncoder " + f"or TransformerEncoder, got {type(self.asr_encoder).__name__} instead." ) self.asr_normalize_type = asr_normalize_type or 'per_feature' self._feat_in = self.asr_encoder._feat_in @@ -406,7 +408,7 @@ def train(self, mode: bool = True) -> "ParallelExpertEncoder": self.asr_encoder.eval() return self - # ConformerEncoder-compatible properties (drop-in for SALM perception). + # ASR-encoder-compatible properties (drop-in for SALM perception). @property def d_model(self) -> int: return self.asr_d_model @@ -486,7 +488,7 @@ def _fuse_diar_and_asr(self, asr_encoded: torch.Tensor, spk_targets: torch.Tenso return fused.transpose(1, 2) # (B, D, T) - # Forward — identical signature to ConformerEncoder.forward + # Forward — identical signature across the supported ASR encoders. def forward( self, audio_signal, diff --git a/nemo/collections/speechlm2/vllm/salm/__init__.py b/nemo/collections/speechlm2/vllm/salm/__init__.py index 177e6918aa91..c8365dd9f011 100644 --- a/nemo/collections/speechlm2/vllm/salm/__init__.py +++ b/nemo/collections/speechlm2/vllm/salm/__init__.py @@ -24,6 +24,148 @@ """ _PKG = "nemo.collections.speechlm2.vllm.salm" +_ORIGINAL_VLLM_HF_CONFIG_OVERRIDE = None +_AUTOMODEL_DFLASH2_ARCHITECTURES = frozenset( + { + "Qwen3DFlash2DraftModel", + "DFlashQwen3DFlash2DraftModel", + } +) + + +def _normalize_dflash2_architecture(hf_config): + """Route Automodel DFlash2 exports to vLLM's canonical runtime. + + Automodel keeps its training class in ``config.json`` so the checkpoint can + be reopened for training. The pinned vLLM DFlash2 implementation dispatches + both its V2 runner and candidate-selector speculator only when it sees the + canonical ``DFlash2DraftModel`` architecture. Normalize before vLLM wraps + the draft in ``EAGLEConfig``; otherwise ``method=dflash`` prefixes the + Automodel name and silently selects the plain-DFlash speculator. + """ + architectures = getattr(hf_config, "architectures", None) or [] + if len(architectures) == 1 and architectures[0] in _AUTOMODEL_DFLASH2_ARCHITECTURES: + hf_config.architectures = ["DFlash2DraftModel"] + return hf_config + + +def _nemo_speechlm_mtp_hf_config_override(hf_config): + """Apply SpeechLM speculative-config rewrites, then defer 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. + 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 + hf_config = _ORIGINAL_VLLM_HF_CONFIG_OVERRIDE(hf_config) + return _normalize_dflash2_architecture(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 for SpeechLM drafts. + + Four patches are applied on supported vLLM 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. + + 4. Automodel DFlash2 architecture names are normalized to vLLM's canonical + ``DFlash2DraftModel`` before ``EAGLEConfig`` wrapping, which activates + the V2 model runner and candidate-selector speculator. + """ + 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 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 + _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(): @@ -44,3 +186,18 @@ def register(): "NeMoSpeechLMForConditionalGeneration", f"{_PKG}.model:NeMoSpeechLMForConditionalGeneration", ) + supported_archs = ModelRegistry.get_supported_archs() + if "DFlashDraftModel" in supported_archs: + native_dflash_model = ModelRegistry.models["DFlashDraftModel"] + native_dflash_model_ref = f"{native_dflash_model.module_name}:{native_dflash_model.class_name}" + for automodel_arch in ("Qwen3DFlashDraftModel", "DFlashQwen3DFlashDraftModel"): + if automodel_arch not in supported_archs: + ModelRegistry.register_model(automodel_arch, native_dflash_model_ref) + if "DFlash2DraftModel" in supported_archs: + native_dflash2_model = ModelRegistry.models["DFlash2DraftModel"] + native_dflash2_model_ref = f"{native_dflash2_model.module_name}:{native_dflash2_model.class_name}" + for automodel_arch in _AUTOMODEL_DFLASH2_ARCHITECTURES: + if automodel_arch not in supported_archs: + ModelRegistry.register_model(automodel_arch, native_dflash2_model_ref) + + _patch_vllm_for_nemo_speechlm_mtp() diff --git a/nemo/collections/speechlm2/vllm/salm/audio.py b/nemo/collections/speechlm2/vllm/salm/audio.py index 7ce90e79a2f1..b17a5e8d1540 100644 --- a/nemo/collections/speechlm2/vllm/salm/audio.py +++ b/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}) @@ -105,7 +108,62 @@ def _load_nemo_perception(perception_cfg: dict) -> nn.Module: return perception -def _maybe_mount_pe_encoder(perception: nn.Module, pe_encoder_path: str | None) -> bool: +def _build_pe_encoder_from_config(pe_encoder_config: Mapping[str, object]) -> nn.Module: + """Reconstruct a PE encoder whose config and weights are embedded in an HF export.""" + from omegaconf import OmegaConf + + from nemo.collections.asr.modules.parallel_expert_encoder import ParallelExpertEncoder + + if hasattr(pe_encoder_config, "to_dict"): + pe_encoder_config = pe_encoder_config.to_dict() + if not isinstance(pe_encoder_config, Mapping): + raise TypeError( + "pe_encoder_config must be a mapping containing inline ASR and diarization configs; " + f"got {type(pe_encoder_config).__name__}." + ) + + cfg = OmegaConf.create(dict(pe_encoder_config)) + return ParallelExpertEncoder( + asr_encoder_cfg=cfg.get("asr_encoder_cfg"), + diarization_model_cfg=cfg.get("diarization_model_cfg"), + asr_normalize_type=cfg.get("asr_normalize_type"), + freeze_diar=cfg.get("freeze_diar", True), + freeze_asr=cfg.get("freeze_asr", False), + online_inference_length=cfg.get("online_inference_length", 500), + chunk_left_context=cfg.get("chunk_left_context", 50), + chunk_right_context=cfg.get("chunk_right_context", 50), + diar_fifo_len=cfg.get("diar_fifo_len", 40), + diar_spkcache_update_period=cfg.get("diar_spkcache_update_period", 300), + diar_spkcache_len=cfg.get("diar_spkcache_len", 188), + ) + + +def _attach_pe_encoder(perception: nn.Module, pe_encoder: nn.Module, source: str) -> None: + """Replace the direct perception encoder while preserving placement and preprocessing.""" + if not hasattr(perception, "encoder"): + raise RuntimeError(f"{source} is set but perception has no encoder attribute to replace.") + + existing_d_model = int(getattr(perception.encoder, "d_model", -1)) + if existing_d_model > 0 and int(pe_encoder.d_model) != existing_d_model: + raise ValueError( + f"ParallelExpertEncoder d_model={pe_encoder.d_model} does not match the existing " + f"perception encoder d_model={existing_d_model}." + ) + + ref_param = next(perception.encoder.parameters(), None) + if ref_param is not None: + pe_encoder = pe_encoder.to(device=ref_param.device, dtype=ref_param.dtype) + + try: + perception.preprocessor.featurizer.normalize = None + except AttributeError: + pass + + perception.encoder = pe_encoder + perception.eval() + + +def _mount_pe_encoder_from_path(perception: nn.Module, pe_encoder_path: str | None) -> bool: """Replace ``perception.encoder`` with a ParallelExpertEncoder bundle so PE-trained checkpoints (nested ``asr_encoder.*`` / ``diarization_model.*`` weights) load correctly. @@ -132,45 +190,50 @@ def _maybe_mount_pe_encoder(perception: nn.Module, pe_encoder_path: str | None) """ if pe_encoder_path in (None, "", False): return False - if not hasattr(perception, "encoder"): - raise RuntimeError("pe_encoder_path is set but perception has no `encoder` attribute to replace.") + if not isinstance(pe_encoder_path, str): + raise TypeError(f"pe_encoder_path must be a string, got {type(pe_encoder_path).__name__}.") from nemo.collections.asr.modules.parallel_expert_encoder import ParallelExpertEncoderPT # Only fail-fast for a *local* ``.nemo`` file that is not a PE bundle. A # non-local reference (HF repo id / NGC alias) is resolved offline from the # HuggingFace cache by load_from_nemo -> from_pretrained, so do not reject it. - is_local_nemo_file = ( - isinstance(pe_encoder_path, str) and pe_encoder_path.endswith(".nemo") and os.path.isfile(pe_encoder_path) - ) + is_local_nemo_file = pe_encoder_path.endswith(".nemo") and os.path.isfile(pe_encoder_path) if is_local_nemo_file and not ParallelExpertEncoderPT.is_pe_nemo(pe_encoder_path): raise ValueError(f"pe_encoder_path={pe_encoder_path!r} is not a ParallelExpertEncoderPT .nemo bundle.") pe_encoder = ParallelExpertEncoderPT.load_from_nemo(pe_encoder_path, map_location="cpu", strict=True) + _attach_pe_encoder(perception, pe_encoder, "pe_encoder_path") + return True - existing_d_model = int(getattr(perception.encoder, "d_model", -1)) - if existing_d_model > 0 and int(pe_encoder.d_model) != existing_d_model: - raise ValueError( - f"ParallelExpertEncoder d_model={pe_encoder.d_model} does not match the existing " - f"perception encoder d_model={existing_d_model}." - ) - # load_from_nemo restores onto CPU; copy the replaced encoder's device/dtype to avoid CPU/dtype mismatches. - ref_param = next(perception.encoder.parameters(), None) - if ref_param is not None: - pe_encoder = pe_encoder.to(device=ref_param.device, dtype=ref_param.dtype) +def _maybe_mount_pe_encoder( + perception: nn.Module, + pe_encoder_path: str | None, + pe_encoder_config: Mapping[str, object] | None = None, +) -> bool: + """Mount the PE encoder represented by an exported SpeechLM checkpoint. - # PE encoder consumes un-normalised mels and replays ASR norm internally, so disable preprocessor norm. - try: - perception.preprocessor.featurizer.normalize = None - except AttributeError: - # Preprocessor/featurizer layout varies across backends; if the attribute is - # absent there is no outer normalization to disable, so skipping is correct. - pass + Self-contained HF exports store the constructor inputs in pe_encoder_config + and the nested encoder weights in their own safetensors file. That embedded + form takes precedence over pe_encoder_path so serving does not depend on a + path from the training machine. Legacy path-only checkpoints retain their + existing local-bundle and pretrained-model behavior. - perception.encoder = pe_encoder - perception.eval() - return True + Args: + perception: Perception module whose direct encoder is replaced. + pe_encoder_path: Legacy local bundle path or pretrained model identifier. + pe_encoder_config: Optional self-contained PE constructor config. + + Returns: + True when a PE encoder was mounted, otherwise False. + """ + if pe_encoder_config: + pe_encoder = _build_pe_encoder_from_config(pe_encoder_config) + _attach_pe_encoder(perception, pe_encoder, "pe_encoder_config") + return True + + return _mount_pe_encoder_from_path(perception, pe_encoder_path) def _pad_to_vocab_size(tensor: torch.Tensor, target_vocab: int) -> torch.Tensor: diff --git a/nemo/collections/speechlm2/vllm/salm/config.py b/nemo/collections/speechlm2/vllm/salm/config.py index 6d9f55d1b3fc..1ba78a7bd2dd 100644 --- a/nemo/collections/speechlm2/vllm/salm/config.py +++ b/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 @@ -71,6 +74,7 @@ def __init__( self, perception: dict | None = None, pretrained_llm: str | None = None, + llm_config: dict | None = None, pretrained_asr: str | None = None, audio_locator_tag: str | None = None, prompt_format: str | None = None, @@ -100,6 +104,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) @@ -109,12 +114,14 @@ def __init__( # path inert; real checkpoint loads continue through validation below. self.perception = {} self.pretrained_llm = None + self.llm_config = None self.pretrained_asr = None self.audio_locator_tag = None self.prompt_format = None 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(): @@ -135,6 +142,7 @@ def __init__( ) self.perception = perception or {} self.pretrained_llm = pretrained_llm + self.llm_config = llm_config self.pretrained_asr = pretrained_asr self.audio_locator_tag = audio_locator_tag self.prompt_format = prompt_format @@ -142,7 +150,16 @@ def __init__( self.lora = lora self.encoder_chunk_size_seconds = encoder_chunk_size_seconds - self.text_config = AutoConfig.from_pretrained(pretrained_llm, trust_remote_code=True) + if llm_config is None: + self.text_config = AutoConfig.from_pretrained(pretrained_llm, trust_remote_code=True) + else: + if not isinstance(llm_config, dict): + raise ValueError(f"NeMo SpeechLM llm_config must be a dict, got {type(llm_config).__name__}.") + embedded_config = dict(llm_config) + model_type = embedded_config.pop("model_type", None) + if not model_type: + raise ValueError("NeMo SpeechLM llm_config must declare model_type.") + self.text_config = AutoConfig.for_model(model_type, **embedded_config) raw_archs = getattr(self.text_config, "architectures", []) if len(raw_archs) != 1: @@ -177,6 +194,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 +210,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 +277,7 @@ def __getattr__(self, name): "pretrained_llm", "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 a7160f5f8a32..d2979df36fc7 100644 --- a/nemo/collections/speechlm2/vllm/salm/model.py +++ b/nemo/collections/speechlm2/vllm/salm/model.py @@ -22,8 +22,9 @@ ``backends.py`` and is selected once at ``__init__`` time via ``make_backend(config)``. The class declares ``IsHybrid`` / ``SupportsMambaPrefixCaching`` so vLLM's hybrid KV-cache allocator picks up -NemotronH backbones; for transformer backbones the runtime -``ModelConfig.is_hybrid`` property returns False because ``config.py`` +NemotronH backbones, and ``SupportsEagle3`` so DFlash and DFlash2 can consume +auxiliary hidden states from the language tower. For transformer backbones the +runtime ``ModelConfig.is_hybrid`` property returns False because ``config.py`` populates ``text_config.layer_types`` with all-attention markers (vLLM's granite-4.0-micro escape hatch). @@ -40,6 +41,7 @@ from vllm.model_executor.models.interfaces import ( IsHybrid, MultiModalEmbeddings, + SupportsEagle3, SupportsMambaPrefixCaching, SupportsMultiModal, SupportsPP, @@ -78,6 +80,7 @@ class NeMoSpeechLMForConditionalGeneration( SupportsPP, IsHybrid, SupportsMambaPrefixCaching, + SupportsEagle3, ): """Backbone-agnostic NeMo SpeechLM. Composition with a backend handles per-backbone details.""" @@ -106,12 +109,35 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): with self._mark_tower_model(vllm_config, {"audio"}): self.perception = _load_nemo_perception(config.perception) - _maybe_mount_pe_encoder(self.perception, getattr(config, "pe_encoder_path", None)) + _maybe_mount_pe_encoder( + self.perception, + getattr(config, "pe_encoder_path", None), + getattr(config, "pe_encoder_config", None), + ) self._uses_pe_encoder = isinstance(getattr(self.perception, "encoder", None), ParallelExpertEncoder) self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors + # ── language-model integration ── + + def get_language_model(self) -> nn.Module: + """Return the wrapped decoder used by vLLM speculative decoders. + + DFlash and DFlash2 resolve the target embedding table and LM head + through this hook. Returning the registered vLLM language tower also + lets the ``SupportsEagle3`` interface reach its inner ``EagleModelMixin``. + """ + return self.language_model + + def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + """Select target layers whose hidden states are consumed by DFlash drafters.""" + self.language_model.set_aux_hidden_state_layers(layers) + + def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]: + """Delegate vLLM's fallback auxiliary-layer selection to the decoder.""" + return self.language_model.get_eagle3_default_aux_hidden_state_layers() + # ── audio processing ── def _parse_audio_input( @@ -189,7 +215,7 @@ def forward( intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, **kwargs, - ) -> torch.Tensor | IntermediateTensors: + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: if intermediate_tensors is not None: inputs_embeds = None return self.language_model(input_ids, positions, intermediate_tensors, inputs_embeds) @@ -222,6 +248,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..2d9ab84db1b3 --- /dev/null +++ b/nemo/collections/speechlm2/vllm/salm/mtp.py @@ -0,0 +1,154 @@ +# 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 SpeechLM 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 ``NemotronHMTP`` 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 (including perception weights) are intentionally + omitted because the target model loads them separately. + """ + for name, tensor in items: + 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." + ) + + # 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 NeMo SpeechLM ``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 the serialized + # configuration instead of silently depending on a predictor-internal + # attribute that could drift between vLLM releases. + 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}.") + expected_layer_modules = len(pattern) + return super().load_weights( + _remap_nemo_mtp_weights( + weights, + target_vocab, + expected_layer_modules=expected_layer_modules, + ) + ) diff --git a/tests/collections/asr/test_parallel_expert_encoder.py b/tests/collections/asr/test_parallel_expert_encoder.py index b9c099cc6661..d7a57ec94ddc 100644 --- a/tests/collections/asr/test_parallel_expert_encoder.py +++ b/tests/collections/asr/test_parallel_expert_encoder.py @@ -30,6 +30,7 @@ _default_dtype, _disable_dist_feature_sync, ) +from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder # ``@experimental`` wraps the class in a wrapt proxy, so ``__new__`` (used to build # bare instances that skip the heavy real ``__init__``) must target the underlying @@ -347,6 +348,29 @@ def toy_asr_encoder_cfg() -> DictConfig: ) +def toy_transformer_asr_encoder_cfg() -> DictConfig: + """Tiny production-style TransformerEncoder config for the PE ASR branch.""" + return DictConfig( + { + '_target_': 'nemo.collections.asr.modules.transformer_encoder.TransformerEncoder', + 'feat_in': _MEL_FEATURES, + 'd_model': 64, + 'n_heads': 4, + 'n_layers': 1, + 'drop_rate': 0.0, + 'qkv_bias': False, + 'qk_norm': True, + 'ff_expansion': 2, + 'pre_block_norm': True, + 'subsampling_factor': _SUBSAMPLING_FACTOR, + 'attn_mode': 'full', + 'self_attention_model': 'rope', + 'rope_base': 10000.0, + 'rotary_fraction': 0.5, + } + ) + + def toy_diarization_model_cfg() -> DictConfig: """Tiny SortformerEncLabelModel config the PE encoder mounts as its diar branch.""" model_defaults = {'fc_d_model': _DIAR_FC_D_MODEL, 'tf_d_model': _DIAR_TF_D_MODEL} @@ -462,6 +486,17 @@ def test_pe_encoder_builds_and_wires_both_real_encoders(): assert any(p.requires_grad for p in enc.asr_encoder.parameters()) +@pytest.mark.unit +def test_pe_encoder_accepts_transformer_asr_branch(): + enc = build_toy_pe_encoder(asr_encoder_cfg=toy_transformer_asr_encoder_cfg()) + + assert isinstance(enc.asr_encoder, TransformerEncoder) + assert enc.d_model == 64 + assert enc.subsampling_factor == _SUBSAMPLING_FACTOR + assert enc.pre_encode is enc.asr_encoder.pre_encode + assert enc.diar_kernel.shape == (_N_SPK, 64) + + @pytest.mark.unit @pytest.mark.parametrize( "high_resolution, requested_diar_subsampling_factor, expected_asr_aligned_factor", diff --git a/tests/collections/speechlm2/test_to_hf.py b/tests/collections/speechlm2/test_to_hf.py index 8cd3b08bd7cb..d8d2309188e9 100644 --- a/tests/collections/speechlm2/test_to_hf.py +++ b/tests/collections/speechlm2/test_to_hf.py @@ -20,6 +20,7 @@ import importlib.util import json from pathlib import Path +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -97,6 +98,8 @@ def _seed_output_dir(tmp_path, llm_arch="Qwen2ForCausalLM"): "architectures": [llm_arch], "hidden_size": 2048, "num_hidden_layers": 24, + "audio_token_index": 17, + "image_token_index": 18, } ) ) @@ -130,6 +133,106 @@ class _FakeExportModel: llm = type("_FakeLLM", (), {"config": _FakeLLMConfig()})() +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": "*"}}, + llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern="*E")), + ) + + 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}) + + config = to_hf._hf_export_config(model, "bfloat16") + + assert "trust_remote_code" not in config + 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)), + ) + + 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": "*"}}, + 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": "*"}}, + llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern=pattern)), + ) + + 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": "*"}}, + llm=SimpleNamespace(config=SimpleNamespace(mtp_hybrid_override_pattern="")), + ) + 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", @@ -185,15 +288,41 @@ def test_prepare_for_vllm_missing_audio_locator_tag(tmp_path): to_hf.prepare_for_vllm(str(tmp_path), {"pretrained_llm": "fake-model"}) +def test_detect_vllm_architecture_returns_model_embedding_vocab(): + """Exporter bounds must use the model config, not tokenizer.vocab_size.""" + backbone = SimpleNamespace(architectures=["Qwen2ForCausalLM"], vocab_size=151936) + with patch("transformers.AutoConfig.from_pretrained", return_value=backbone): + architecture, vocab_size = to_hf._detect_vllm_architecture({"pretrained_llm": "fake-model"}) + + assert architecture == "NeMoSpeechLMForConditionalGeneration" + assert vocab_size == 151936 + + +@pytest.mark.parametrize("vocab_size", [None, True, 0, -1]) +def test_detect_vllm_architecture_rejects_invalid_model_vocab(vocab_size): + backbone = SimpleNamespace(architectures=["Qwen2ForCausalLM"], vocab_size=vocab_size) + with ( + patch("transformers.AutoConfig.from_pretrained", return_value=backbone), + pytest.raises(ValueError, match="vocab_size"), + ): + to_hf._detect_vllm_architecture({"pretrained_llm": "fake-model"}) + + # ────────────────────────────────────────────────────────────────────── # Happy paths (mock AutoTokenizer + _detect_vllm_architecture) # ────────────────────────────────────────────────────────────────────── -def _run_prepare(tmp_path, fake_tok, arch="NeMoSpeechLMForConditionalGeneration", llm_arch="Qwen2ForCausalLM"): +def _run_prepare( + tmp_path, + fake_tok, + arch="NeMoSpeechLMForConditionalGeneration", + llm_arch="Qwen2ForCausalLM", + backbone_vocab_size=100, +): output_dir = _seed_output_dir(tmp_path, llm_arch=llm_arch) with ( - patch.object(to_hf, "_detect_vllm_architecture", return_value=arch), + patch.object(to_hf, "_detect_vllm_architecture", return_value=(arch, backbone_vocab_size)), patch("transformers.AutoTokenizer.from_pretrained", return_value=fake_tok), ): to_hf.prepare_for_vllm( @@ -204,16 +333,55 @@ def _run_prepare(tmp_path, fake_tok, arch="NeMoSpeechLMForConditionalGeneration" def test_prepare_for_vllm_patches_config_json(tmp_path): - """config.json gets model_type, architectures, and audio_locator_tag.""" + """config.json gets model metadata without persisting a token-index field.""" output_dir = _run_prepare(tmp_path, _FakeTokenizer()) cfg = json.loads((output_dir / "config.json").read_text()) assert cfg["model_type"] == "nemo_speechlm" 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 +def test_prepare_for_vllm_embeds_bundled_backbone_config(tmp_path): + """The vLLM root config must not reload the stale training backbone reference.""" + output_dir = _seed_output_dir(tmp_path) + backbone_dir = output_dir / "llm_backbone" + backbone_dir.mkdir() + backbone_config = { + "model_type": "qwen2", + "architectures": ["Qwen2ForCausalLM"], + "hidden_size": 2048, + "vocab_size": 100, + } + (backbone_dir / "config.json").write_text(json.dumps(backbone_config)) + + with ( + patch.object( + to_hf, + "_detect_vllm_architecture", + return_value=("NeMoSpeechLMForConditionalGeneration", 100), + ) as detect_architecture, + patch("transformers.AutoTokenizer.from_pretrained", return_value=_FakeTokenizer()), + ): + to_hf.prepare_for_vllm( + str(output_dir), + {"pretrained_llm": "stale-training-model", "audio_locator_tag": AUDIO_TOKEN}, + ) + + cfg = json.loads((output_dir / "config.json").read_text()) + assert cfg["pretrained_llm"] == "llm_backbone" + assert cfg["llm_config"] == backbone_config + detect_architecture.assert_called_once_with( + { + "pretrained_llm": str(backbone_dir), + "audio_locator_tag": AUDIO_TOKEN, + } + ) + + def test_prepare_for_vllm_adds_audio_token_to_vocab(tmp_path): """Audio token is registered via add_special_tokens when not already in vocab.""" fake_tok = _FakeTokenizer(vocab_tokens=["<|im_start|>", "<|im_end|>"]) @@ -229,6 +397,57 @@ def test_prepare_for_vllm_skips_add_if_audio_token_already_in_vocab(tmp_path): assert fake_tok.add_special_tokens_calls == [] +def test_prepare_for_vllm_rejects_invalid_audio_token_id(tmp_path): + fake_tok = _FakeTokenizer(vocab_tokens=[AUDIO_TOKEN]) + fake_tok._vocab[AUDIO_TOKEN] = -1 + + with pytest.raises(ValueError, match="valid ID"): + _run_prepare(tmp_path, fake_tok) + + +def test_prepare_for_vllm_rejects_audio_token_outside_padded_embeddings(tmp_path): + tokens = [f"" for i in range(11)] + [AUDIO_TOKEN] + fake_tok = _FakeTokenizer(vocab_tokens=tokens) + + with pytest.raises(ValueError, match="outside the SpeechLM embedding table"): + _run_prepare(tmp_path, fake_tok, backbone_vocab_size=1) + + +def test_prepare_for_vllm_accepts_audio_token_inside_non_boundary_embedding_row(tmp_path): + """Qwen-style reserved rows may put the audio token below the model-vocab boundary.""" + fake_tok = _FakeTokenizer(vocab_tokens=["", AUDIO_TOKEN, ""]) + + output_dir = _run_prepare(tmp_path, fake_tok, backbone_vocab_size=100) + + 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): + """Conversion must persist the tokenizer that supplied training token IDs.""" + output_dir = _seed_output_dir(tmp_path) + fake_tok = _FakeTokenizer() + with ( + patch.object( + to_hf, + "_detect_vllm_architecture", + return_value=("NeMoSpeechLMForConditionalGeneration", 100), + ), + patch("transformers.AutoTokenizer.from_pretrained", return_value=fake_tok) as load_tokenizer, + ): + to_hf.prepare_for_vllm( + str(output_dir), + { + "pretrained_llm": "fake-model", + "tokenizer_path": "custom-tokenizer", + "audio_locator_tag": AUDIO_TOKEN, + }, + ) + + load_tokenizer.assert_called_once_with("custom-tokenizer", trust_remote_code=True, fix_mistral_regex=True) + + def test_prepare_for_vllm_tokenizer_config_normalized(tmp_path): """tokenizer_config.json has dict-form extra_special_tokens + forced tokenizer_class.""" output_dir = _run_prepare(tmp_path, _FakeTokenizer(tokenizer_class="TokenizersBackend")) diff --git a/tests/collections/speechlm2/test_vllm_plugin.py b/tests/collections/speechlm2/test_vllm_plugin.py index 3c65829d2dd5..8e228f3f5f86 100644 --- a/tests/collections/speechlm2/test_vllm_plugin.py +++ b/tests/collections/speechlm2/test_vllm_plugin.py @@ -25,15 +25,22 @@ import pytest try: - from nemo.collections.speechlm2.vllm.salm import config as _config_module + import nemo.collections.speechlm2.vllm.salm as _salm_module + import nemo.collections.speechlm2.vllm.salm.config as _config_module NeMoSpeechLMConfig = _config_module.NeMoSpeechLMConfig + register = _salm_module.register _HAS_CONFIG = True except (ImportError, RuntimeError): _HAS_CONFIG = False _HAS_VLLM = importlib.util.find_spec("vllm") is not None +if _HAS_VLLM: + import vllm.config.speculative as _spec_mod + + SpeculativeConfig = _spec_mod.SpeculativeConfig + _DEFAULT_CONFIG_KWARGS = { "pretrained_llm": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", "pretrained_asr": "nvidia/canary-1b-v2", @@ -43,6 +50,39 @@ } +def _identity_hf_config_override(hf_config): + """Pickleable target override for vLLM's composed-override API.""" + return hf_config + + +@pytest.mark.skipif(not _HAS_CONFIG, reason="SpeechLM vLLM plugin not available") +@pytest.mark.parametrize( + "architecture", + ("Qwen3DFlash2DraftModel", "DFlashQwen3DFlash2DraftModel"), +) +def test_normalizes_automodel_dflash2_architecture(architecture): + """Normalization does not need vLLM installed or model weights loaded.""" + config = SimpleNamespace(architectures=[architecture]) + + result = _salm_module._normalize_dflash2_architecture(config) + + assert result is config + assert config.architectures == ["DFlash2DraftModel"] + + +@pytest.mark.skipif(not _HAS_CONFIG, reason="SpeechLM vLLM plugin not available") +@pytest.mark.parametrize( + "architectures", + (None, [], ["DFlash2DraftModel"], ["Qwen3DFlashDraftModel"], ["Qwen3DFlash2DraftModel", "Other"]), +) +def test_dflash2_normalization_preserves_unrelated_architectures(architectures): + config = SimpleNamespace(architectures=architectures) + + _salm_module._normalize_dflash2_architecture(config) + + assert config.architectures == architectures + + @pytest.mark.skipif(not _HAS_CONFIG, reason="NeMoSpeechLMConfig not available") class TestNeMoSpeechLMConfig: """Tests for NeMoSpeechLMConfig.""" @@ -78,6 +118,7 @@ def test_default_construction_for_hf_serialization(self): assert cfg.pretrained_llm 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.llm_architectures == [] @@ -90,6 +131,38 @@ def test_loads_text_config(self): assert hasattr(cfg.text_config, "hidden_size") assert cfg.get_text_config() is cfg.text_config + def test_uses_embedded_backbone_config(self, monkeypatch): + """Bundled exports must not reload the stale training backbone reference.""" + embedded_config = { + "model_type": "qwen2", + "architectures": ["Qwen2ForCausalLM"], + "hidden_size": 2048, + "vocab_size": 151936, + "num_hidden_layers": 4, + "rms_norm_eps": 1e-6, + } + text_config = SimpleNamespace(**embedded_config) + from_pretrained = pytest.fail + + monkeypatch.setattr( + _config_module.AutoConfig, + "from_pretrained", + lambda *args, **kwargs: from_pretrained("must not load pretrained_llm when llm_config is embedded"), + ) + monkeypatch.setattr(_config_module.AutoConfig, "for_model", lambda model_type, **kwargs: text_config) + + cfg = NeMoSpeechLMConfig( + **{ + **_DEFAULT_CONFIG_KWARGS, + "pretrained_llm": "llm_backbone", + "llm_config": embedded_config, + } + ) + + assert cfg.pretrained_llm == "llm_backbone" + assert cfg.llm_config == embedded_config + assert cfg.text_config is text_config + def test_hybrid_backbone_aliases_for_vllm(self): cfg = NeMoSpeechLMConfig(**_DEFAULT_CONFIG_KWARGS) assert cfg.is_hybrid is True @@ -109,9 +182,7 @@ def test_hybrid_backbone_aliases_for_vllm(self): ) def test_is_hybrid_backend_helper(self, architectures, expected_is_hybrid): """``_is_hybrid_backend`` should match the documented hybrid allow-list.""" - from nemo.collections.speechlm2.vllm.salm.config import _is_hybrid_backend - - assert _is_hybrid_backend(architectures) is expected_is_hybrid + assert _config_module._is_hybrid_backend(architectures) is expected_is_hybrid @pytest.mark.parametrize( "backbone_archs, expected_is_hybrid", @@ -569,6 +640,60 @@ def test_call_hf_processor_emits_true_audio_lengths(self): assert result["audio_signal"][0].shape[-1] == 12345 assert torch.equal(result["audio_signal_length"], torch.tensor([12345])) + def test_embedded_pe_config_takes_precedence_over_training_path(self, monkeypatch): + from unittest.mock import Mock + + import torch + from torch import nn + + import nemo.collections.asr.modules.parallel_expert_encoder as pe_module + from nemo.collections.speechlm2.vllm.salm.audio import _maybe_mount_pe_encoder + + class _Encoder(nn.Module): + d_model = 4 + + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.ones(1)) + + class _Perception(nn.Module): + def __init__(self): + super().__init__() + self.encoder = _Encoder().to(torch.float64) + self.preprocessor = SimpleNamespace(featurizer=SimpleNamespace(normalize="per_feature")) + + perception = _Perception() + embedded_encoder = _Encoder() + built = {} + + def build_encoder(**kwargs): + built.update(kwargs) + return embedded_encoder + + load_from_nemo = Mock(side_effect=AssertionError("embedded config must avoid external path")) + monkeypatch.setattr(pe_module.ParallelExpertEncoderPT, "load_from_nemo", load_from_nemo) + monkeypatch.setattr(pe_module, "ParallelExpertEncoder", build_encoder) + + mounted = _maybe_mount_pe_encoder( + perception, + "/training/only/encoder.nemo", + { + "target": "nemo.collections.asr.modules.parallel_expert_encoder.ParallelExpertEncoderPT", + "asr_encoder_cfg": {"kind": "transformer"}, + "diarization_model_cfg": {"kind": "sortformer"}, + "asr_normalize_type": "per_feature", + }, + ) + + assert mounted + assert perception.encoder is embedded_encoder + assert perception.encoder.weight.dtype == torch.float64 + assert perception.preprocessor.featurizer.normalize is None + assert not perception.training + assert built["asr_encoder_cfg"]["kind"] == "transformer" + assert built["diarization_model_cfg"]["kind"] == "sortformer" + load_from_nemo.assert_not_called() + def test_perception_forward(self): """A small NeMo perception module should encode dummy audio to embeddings.""" import torch @@ -642,8 +767,6 @@ def test_register_config(self, monkeypatch): """register() should add nemo_speechlm to vLLM's config registry.""" from transformers import AutoConfig - from nemo.collections.speechlm2.vllm.salm import register - monkeypatch.setattr( AutoConfig, "from_pretrained", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError()) ) @@ -662,8 +785,6 @@ def test_register_model(self, monkeypatch): """ from transformers import AutoConfig - from nemo.collections.speechlm2.vllm.salm import register - monkeypatch.setattr( AutoConfig, "from_pretrained", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError()) ) @@ -680,8 +801,6 @@ def test_register_model(self, monkeypatch): def test_register_does_not_patch_fast_tokenizer(self, monkeypatch): from transformers import AutoConfig, PreTrainedTokenizerFast - from nemo.collections.speechlm2.vllm.salm import register - monkeypatch.setattr( AutoConfig, "from_pretrained", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError()) ) @@ -695,8 +814,6 @@ def test_register_does_not_load_backbone_config(self, monkeypatch): from transformers import AutoConfig - from nemo.collections.speechlm2.vllm.salm import register - from_pretrained = Mock(side_effect=AssertionError("register() must not load remote backbone configs")) monkeypatch.setattr(AutoConfig, "from_pretrained", from_pretrained) @@ -705,6 +822,715 @@ def test_register_does_not_load_backbone_config(self, monkeypatch): from_pretrained.assert_not_called() +@pytest.mark.skipif(not _HAS_VLLM, reason="vLLM not installed") +class TestDFlashPlugin: + """Tests for the target-model contract required by DFlash and DFlash2.""" + + def test_registers_automodel_dflash_architecture_alias(self, monkeypatch): + """An untouched Automodel draft config should resolve to vLLM's native DFlash model.""" + from transformers import AutoConfig + from vllm.model_executor.models.registry import ModelRegistry + from vllm.transformers_utils.configs.eagle import EAGLEConfig + + if "DFlashDraftModel" not in ModelRegistry.get_supported_archs(): + pytest.skip("installed vLLM does not provide native DFlash support") + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError()), + ) + + register() + + native_model = ModelRegistry.models["DFlashDraftModel"] + automodel_config = AutoConfig.for_model("qwen3", architectures=["Qwen3DFlashDraftModel"]) + runtime_config = EAGLEConfig(automodel_config, method="dflash", model_type="eagle") + assert runtime_config.architectures == ["DFlashQwen3DFlashDraftModel"] + + for alias in ("Qwen3DFlashDraftModel", *runtime_config.architectures): + alias_model = ModelRegistry.models[alias] + assert (alias_model.module_name, alias_model.class_name) == ( + native_model.module_name, + native_model.class_name, + ) + + @pytest.mark.parametrize( + "automodel_arch", + ("Qwen3DFlash2DraftModel", "DFlashQwen3DFlash2DraftModel"), + ) + def test_routes_automodel_dflash2_to_candidate_selector_runtime(self, monkeypatch, automodel_arch): + """Automodel exports must retain DFlash2 semantics after vLLM config wrapping.""" + from transformers import AutoConfig + from vllm.model_executor.models.registry import ModelRegistry + from vllm.transformers_utils.configs.eagle import EAGLEConfig + + if "DFlash2DraftModel" not in ModelRegistry.get_supported_archs(): + pytest.skip("installed vLLM does not provide the DFlash2 candidate-selector runtime") + + monkeypatch.setattr( + AutoConfig, + "from_pretrained", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError()), + ) + + register() + + native_model = ModelRegistry.models["DFlash2DraftModel"] + automodel_config = AutoConfig.for_model("qwen3", architectures=[automodel_arch]) + normalized_config = SpeculativeConfig.hf_config_override(automodel_config) + runtime_config = EAGLEConfig(normalized_config, method="dflash", model_type="eagle") + + assert normalized_config.architectures == ["DFlash2DraftModel"] + assert runtime_config.architectures == ["DFlash2DraftModel"] + for alias in (automodel_arch, *runtime_config.architectures): + alias_model = ModelRegistry.models[alias] + assert (alias_model.module_name, alias_model.class_name) == ( + native_model.module_name, + native_model.class_name, + ) + + def test_model_advertises_eagle3_support(self): + from vllm.model_executor.models.interfaces import supports_eagle3 + + from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration + + assert supports_eagle3(NeMoSpeechLMForConditionalGeneration) + + def test_get_language_model_exposes_wrapped_decoder(self): + from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration + + model = object.__new__(NeMoSpeechLMForConditionalGeneration) + language_model = object() + object.__setattr__(model, "language_model", language_model) + + assert model.get_language_model() is language_model + + def test_aux_hidden_state_methods_delegate_to_wrapped_decoder(self): + from unittest.mock import Mock + + from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration + + layers = (2, 6, 20, 30, 42, 52) + language_model = Mock() + language_model.get_eagle3_default_aux_hidden_state_layers.return_value = layers + + model = object.__new__(NeMoSpeechLMForConditionalGeneration) + object.__setattr__(model, "language_model", language_model) + + model.set_aux_hidden_state_layers(layers) + + language_model.set_aux_hidden_state_layers.assert_called_once_with(layers) + assert model.get_eagle3_default_aux_hidden_state_layers() == layers + + def test_forward_preserves_auxiliary_hidden_state_output(self): + from unittest.mock import Mock + + from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration + + output = (object(), [object(), object()]) + language_model = Mock(return_value=output) + model = object.__new__(NeMoSpeechLMForConditionalGeneration) + object.__setattr__(model, "language_model", language_model) + + input_ids = object() + positions = object() + assert model.forward(input_ids, positions) is output + language_model.assert_called_once_with(input_ids, positions, None, None) + + +@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.""" + import nemo.collections.speechlm2.vllm.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 + + 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 + + from transformers import AutoConfig + + 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 + + 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 + import nemo.collections.speechlm2.vllm.salm as salm_module + + 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.""" + import nemo.collections.speechlm2.vllm.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.""" + import nemo.collections.speechlm2.vllm.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 + + 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 + + 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 + import nemo.collections.speechlm2.vllm.salm as salm_module + + 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": 0}) + 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 + import nemo.collections.speechlm2.vllm.salm as salm_module + + 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 + import nemo.collections.speechlm2.vllm.salm as salm_module + + 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 + + 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 + + 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 and recurse.""" + from transformers import AutoConfig + import nemo.collections.speechlm2.vllm.salm as salm_module + + 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 multimodal embedding order.""" + import torch + + from nemo.collections.speechlm2.vllm.salm.mtp import NeMoSpeechLMMTP + + m = object.__new__(NeMoSpeechLMMTP) + base_embeds = torch.zeros(5, 2) + m.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 = m.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 + + m = object.__new__(NeMoSpeechLMMTP) + m.model = SimpleNamespace(get_input_embeddings=lambda ids: torch.zeros(2, 2)) + + with pytest.raises(ValueError, match="is_multimodal"): + m.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 unsupported 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 NeMo 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 final NeMo SpeechLM exports, not bare backbone checkpoints.""" + 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 real loader wiring 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, []) + + @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={"num_nextn_predict_layers": 1, "hybrid_override_pattern": "M*M"}, + ) + assert cfg.mtp_hybrid_override_pattern == "M*M" + + @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 0.26 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