Add native Transformers inference backend

Add a Transformers 5.14 native Nemotron-H audio adapter and selectable native/legacy Hugging Face inference backends. Document the validated dependency versions. No checkpoint weights or vLLM files are changed.

README.md

@@ -73,7 +73,13 @@ backbone with marginal or no regression. Audex-30B-A3B operates in both **thinki
 We use **vLLM 0.20.0 container image**: [vllm/vllm-openai:v0.20.0-cu129](https://hub.docker.com/layers/vllm/vllm-openai/v0.20.0-cu129/images/sha256-f4ace3494896eeda800dee284d1fc42ca7f5626f31ceae8e24d1383d770567c2)
 
 - **vLLM inference** — text-only reasoning, text-to-speech, text-to-audio, and audio understanding / speech recognition / speech translation: runs on **vLLM 0.20.0**.
-- **Hugging Face / transformers inference** — requires transformers >= 4.53.0 (tested with 4.53.3) and also works with transformers >= 5.0. This additionally needs `mamba-ssm` and `causal-conv1d` (build against your CUDA toolchain, e.g. ```pip install --no-build-isolation causal-conv1d==1.6.2.post1 mamba-ssm==2.3.2.post1```).
+- **Hugging Face / transformers inference** — requires transformers >= 4.53.0 (tested with 4.53.3) and also works with transformers >= 
+5.0. This additionally needs `mamba-ssm` and `causal-conv1d`. Build against your CUDA toolchain:
+
+```bash
+python3 -m pip install transformers==5.14.0 safetensors==0.8.0
+python3 -m pip install --no-build-isolation causal-conv1d==1.6.2.post1 mamba-ssm==2.3.1
+```
 
 **Audio extras:** ```vllm/vllm-openai:v0.20.0``` image does not include audio codecs. This command installs audio-related packages: ```python3 -m pip install "vllm[audio]"```.
 

checkpoint_folder_full/modeling_nemotron_h_audio_native.py

@@ -0,0 +1,299 @@
+# coding=utf-8
+# 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.
+"""Audex audio wrapper for the native Transformers Nemotron-H backbone."""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers.conversion_mapping import get_checkpoint_conversion_mapping, register_checkpoint_conversion_mapping
+from transformers.models.nemotron_h.configuration_nemotron_h import NemotronHConfig
+from transformers.models.nemotron_h.modeling_nemotron_h import NemotronHForCausalLM
+
+
+class NemotronHAudexConfig(NemotronHConfig):
+    """Native Nemotron-H config plus Audex audio metadata."""
+
+    model_type = "nemotron_h_audex"
+
+    def __init__(
+        self,
+        audio_config: Optional[dict[str, Any]] = None,
+        audio_model_type: str = "NV-Whisper",
+        sound_model_type: Optional[str] = None,
+        audio_preprocessor_path: str = "audio_preprocessor",
+        sound_token: str = "<so_embedding>",
+        sound_start_token: str = "<so_start>",
+        sound_end_token: str = "<so_end>",
+        sound_token_id: Optional[int] = None,
+        sound_start_token_id: Optional[int] = None,
+        sound_end_token_id: Optional[int] = None,
+        sound_embedding_size: int = 750,
+        sound_clip_duration: float = 30.0,
+        sound_target_rate: int = 16000,
+        audio_encoder_hidden_size: int = 1280,
+        audio_projector_intermediate_size: int = 4096,
+        audio_projector_activation: str = "relu2",
+        audio_projector_norm_eps: float = 1e-5,
+        **kwargs,
+    ):
+        self.audio_config = audio_config or {
+            "model_type": "qwen2_audio_encoder",
+            "num_mel_bins": 128,
+            "encoder_layers": 32,
+            "encoder_attention_heads": 20,
+            "encoder_ffn_dim": 5120,
+            "d_model": audio_encoder_hidden_size,
+            "activation_function": "gelu",
+            "scale_embedding": False,
+            "max_source_positions": 1500,
+        }
+        self.audio_model_type = audio_model_type
+        self.sound_model_type = sound_model_type
+        self.audio_preprocessor_path = audio_preprocessor_path
+        self.sound_token = sound_token
+        self.sound_start_token = sound_start_token
+        self.sound_end_token = sound_end_token
+        self.sound_token_id = sound_token_id
+        self.sound_start_token_id = sound_start_token_id
+        self.sound_end_token_id = sound_end_token_id
+        self.sound_embedding_size = sound_embedding_size
+        self.sound_clip_duration = sound_clip_duration
+        self.sound_target_rate = sound_target_rate
+        self.audio_encoder_hidden_size = audio_encoder_hidden_size
+        self.audio_projector_intermediate_size = audio_projector_intermediate_size
+        self.audio_projector_activation = audio_projector_activation
+        self.audio_projector_norm_eps = audio_projector_norm_eps
+        super().__init__(**kwargs)
+
+
+class NemotronHAudexRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.float()
+        variance = hidden_states.pow(2).mean(dim=-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        return (self.weight.float() * hidden_states).to(input_dtype)
+
+
+class NemotronHAudexProjector(nn.Module):
+    def __init__(self, config: NemotronHAudexConfig):
+        super().__init__()
+        self.norm = NemotronHAudexRMSNorm(
+            config.audio_encoder_hidden_size,
+            eps=config.audio_projector_norm_eps,
+        )
+        self.fc1 = nn.Linear(
+            config.audio_encoder_hidden_size,
+            config.audio_projector_intermediate_size,
+            bias=False,
+        )
+        self.fc2 = nn.Linear(
+            config.audio_projector_intermediate_size,
+            config.hidden_size,
+            bias=False,
+        )
+        self.activation = config.audio_projector_activation
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.norm(hidden_states)
+        hidden_states = self.fc1(hidden_states)
+        if self.activation == "relu2":
+            hidden_states = F.relu(hidden_states).pow(2)
+        elif self.activation == "gelu":
+            hidden_states = F.gelu(hidden_states)
+        else:
+            raise ValueError(f"Unsupported audio projector activation: {self.activation}")
+        return self.fc2(hidden_states)
+
+
+def _build_qwen2_audio_encoder(audio_config: dict) -> nn.Module:
+    from transformers.models.qwen2_audio.configuration_qwen2_audio import Qwen2AudioEncoderConfig
+    from transformers.models.qwen2_audio.modeling_qwen2_audio import Qwen2AudioEncoder
+
+    return Qwen2AudioEncoder(Qwen2AudioEncoderConfig(**dict(audio_config)))
+
+
+register_checkpoint_conversion_mapping(
+    "NemotronHAudexForConditionalGeneration",
+    get_checkpoint_conversion_mapping("nemotron_h"),
+    overwrite=True,
+)
+
+
+class NemotronHAudexForConditionalGeneration(NemotronHForCausalLM):
+    """Audex audio modules attached to native Transformers Nemotron-H."""
+
+    config_class = NemotronHAudexConfig
+    _tp_plan = None
+    _base_model_tp_plan = None
+    base_model_tp_plan = None
+
+    def __init__(self, config: NemotronHAudexConfig):
+        super().__init__(config)
+        self.audio_encoder = _build_qwen2_audio_encoder(config.audio_config)
+        self.audio_projector = NemotronHAudexProjector(config)
+
+    def encode_audio(self, input_features: torch.Tensor) -> torch.Tensor:
+        encoder_param = next(self.audio_encoder.parameters())
+        input_features = input_features.to(device=encoder_param.device, dtype=encoder_param.dtype)
+        encoder_outputs = self.audio_encoder(input_features=input_features, return_dict=True)
+        audio_hidden = encoder_outputs.last_hidden_state
+        projector_param = next(self.audio_projector.parameters())
+        audio_hidden = audio_hidden.to(device=projector_param.device, dtype=projector_param.dtype)
+        return self.audio_projector(audio_hidden)
+
+    def _audio_embeddings_by_sample(
+        self,
+        input_features: Optional[torch.Tensor],
+        audio_embeddings: Optional[torch.Tensor],
+        batch_size: int,
+    ) -> list[torch.Tensor]:
+        if audio_embeddings is None:
+            if input_features is None:
+                raise ValueError("input_features or audio_embeddings must be provided for audio injection")
+            if input_features.ndim == 3:
+                projected = self.encode_audio(input_features)
+                if batch_size != 1:
+                    raise ValueError(
+                        "3D input_features represent a single sample. "
+                        "Use 4D (batch, clips, mel_bins, frames) features for batched audio."
+                    )
+                return [projected.reshape(-1, projected.shape[-1])]
+            if input_features.ndim == 4:
+                batch, clips, mel_bins, frames = input_features.shape
+                if batch != batch_size:
+                    raise ValueError(f"input_features batch {batch} != input_ids batch {batch_size}")
+                flat_features = input_features.reshape(batch * clips, mel_bins, frames)
+                projected = self.encode_audio(flat_features)
+                projected = projected.reshape(batch, clips * projected.shape[1], projected.shape[-1])
+                return [projected[index] for index in range(batch)]
+            raise ValueError(f"Expected 3D or 4D input_features, got {tuple(input_features.shape)}")
+
+        if audio_embeddings.ndim == 2:
+            if batch_size != 1:
+                raise ValueError("2D audio_embeddings only support batch_size=1")
+            return [audio_embeddings.to(device=self.device)]
+        if audio_embeddings.ndim == 3:
+            if audio_embeddings.shape[0] != batch_size:
+                raise ValueError(f"audio_embeddings batch {audio_embeddings.shape[0]} != input_ids batch {batch_size}")
+            return [audio_embeddings[index].to(device=self.device) for index in range(batch_size)]
+        raise ValueError(f"Expected 2D or 3D audio_embeddings, got {tuple(audio_embeddings.shape)}")
+
+    def prepare_inputs_embeds(
+        self,
+        input_ids: torch.LongTensor,
+        input_features: Optional[torch.Tensor] = None,
+        audio_embeddings: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if input_ids is None:
+            raise ValueError("input_ids are required when injecting audio embeddings")
+        if self.config.sound_token_id is None:
+            raise ValueError("config.sound_token_id is required for audio embedding injection")
+
+        embed_device = self.model.embeddings.weight.device
+        input_ids = input_ids.to(embed_device)
+        inputs_embeds = self.model.embeddings(input_ids).clone()
+        audio_by_sample = self._audio_embeddings_by_sample(
+            input_features=input_features,
+            audio_embeddings=audio_embeddings,
+            batch_size=input_ids.shape[0],
+        )
+
+        for batch_index, audio in enumerate(audio_by_sample):
+            mask = input_ids[batch_index].to(self.device) == self.config.sound_token_id
+            mask = mask.to(embed_device)
+            expected = int(mask.sum().item())
+            if expected != audio.shape[0]:
+                raise ValueError(
+                    "Mismatch between <so_embedding> token count and projected audio tokens: "
+                    f"sample={batch_index} placeholders={expected} audio_tokens={audio.shape[0]}"
+                )
+            inputs_embeds[batch_index, mask] = audio.to(
+                device=embed_device,
+                dtype=inputs_embeds.dtype,
+            )
+        return inputs_embeds
+
+    @staticmethod
+    def _is_prefill(past_key_values) -> bool:
+        if past_key_values is None:
+            return True
+        get_length = getattr(past_key_values, "get_seq_length", None)
+        if callable(get_length):
+            return get_length() == 0
+        return len(past_key_values) == 0
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        input_features=None,
+        audio_embeddings=None,
+        position_ids=None,
+        use_cache=True,
+        is_first_iteration=False,
+        **kwargs,
+    ):
+        prefill = self._is_prefill(past_key_values)
+        if prefill and inputs_embeds is None and (
+            input_features is not None or audio_embeddings is not None
+        ):
+            inputs_embeds = self.prepare_inputs_embeds(
+                input_ids=input_ids,
+                input_features=input_features,
+                audio_embeddings=audio_embeddings,
+            )
+        if prefill:
+            past_key_values = None
+        return super().prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            position_ids=position_ids,
+            use_cache=use_cache,
+            is_first_iteration=is_first_iteration,
+            **kwargs,
+        )
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        input_features: Optional[torch.Tensor] = None,
+        audio_embeddings: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        if inputs_embeds is None and (
+            input_features is not None or audio_embeddings is not None
+        ):
+            inputs_embeds = self.prepare_inputs_embeds(
+                input_ids=input_ids,
+                input_features=input_features,
+                audio_embeddings=audio_embeddings,
+            )
+            input_ids = None
+        return super().forward(input_ids=input_ids, inputs_embeds=inputs_embeds, **kwargs)

inference_scripts_hf/inference_hf.py

@@ -17,14 +17,20 @@
 from __future__ import annotations
 
 import argparse
+import hashlib
+import importlib.util
 import json
 import os
 import shutil
 import sys
+from importlib.metadata import PackageNotFoundError, version
 from pathlib import Path
+from types import ModuleType
 from typing import Optional
 
 import torch
+import transformers
+from packaging.version import Version
 from transformers import AutoConfig, AutoFeatureExtractor, AutoModelForCausalLM, AutoTokenizer
 
 SCRIPT_DIR = Path(__file__).resolve().parent
@@ -44,6 +50,8 @@ from audio_utils import (
     split_thinking,
 )
 
+NATIVE_TRANSFORMERS_MIN_VERSION = Version("5.14.0")
+
 
 def refresh_remote_code_cache(model_path: str) -> None:
     """Drop stale HF dynamic-module cache for this local checkpoint folder."""
@@ -65,6 +73,104 @@ def resolve_device_map(device_map: str, device: str):
     return device_map
 
 
+def resolve_hf_backend(requested_backend: str) -> str:
+    transformers_version = Version(transformers.__version__)
+    if requested_backend == "legacy":
+        return "legacy"
+    if transformers_version >= NATIVE_TRANSFORMERS_MIN_VERSION:
+        return "native"
+    if requested_backend == "native":
+        raise RuntimeError(
+            f"The native backend requires transformers >= {NATIVE_TRANSFORMERS_MIN_VERSION}; "
+            f"found {transformers.__version__}."
+        )
+    print(
+        f"transformers {transformers.__version__} does not include the validated "
+        f"Nemotron-H cache fixes; using the slower legacy backend."
+    )
+    return "legacy"
+
+
+def validate_native_kernels() -> None:
+    try:
+        from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
+        from mamba_ssm.ops.triton.selective_state_update import selective_state_update
+        from mamba_ssm.ops.triton.ssd_combined import (
+            mamba_chunk_scan_combined,
+            mamba_split_conv1d_scan_combined,
+        )
+    except ImportError as exc:
+        raise RuntimeError(
+            "The native backend requires causal-conv1d==1.6.2.post1 and "
+            "mamba-ssm==2.3.1 built against the active PyTorch/CUDA environment."
+        ) from exc
+
+    kernels = (
+        causal_conv1d_fn,
+        causal_conv1d_update,
+        selective_state_update,
+        mamba_chunk_scan_combined,
+        mamba_split_conv1d_scan_combined,
+    )
+    if any(kernel is None for kernel in kernels):
+        raise RuntimeError("The native Nemotron-H fast-path kernels are unavailable.")
+
+    packages = []
+    for package in ("mamba-ssm", "causal-conv1d"):
+        try:
+            packages.append(f"{package}={version(package)}")
+        except PackageNotFoundError:
+            packages.append(f"{package}=unknown")
+    print("native kernels:", ", ".join(packages))
+
+
+def load_native_module(model_path: str) -> ModuleType:
+    module_path = Path(model_path).resolve() / "modeling_nemotron_h_audio_native.py"
+    if not module_path.is_file():
+        raise FileNotFoundError(f"Native Audex model adapter not found: {module_path}")
+
+    path_hash = hashlib.sha256(str(module_path).encode()).hexdigest()[:12]
+    module_name = f"audex_native_{path_hash}"
+    if module_name in sys.modules:
+        return sys.modules[module_name]
+
+    spec = importlib.util.spec_from_file_location(module_name, module_path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Unable to load native Audex model adapter: {module_path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def load_native_model(model_path: str, model_load_kwargs: dict):
+    validate_native_kernels()
+    native_module = load_native_module(model_path)
+    config_class = native_module.NemotronHAudexConfig
+    model_class = native_module.NemotronHAudexForConditionalGeneration
+
+    config = config_class.from_pretrained(model_path)
+    model, loading_info = model_class.from_pretrained(
+        model_path,
+        config=config,
+        output_loading_info=True,
+        **model_load_kwargs,
+    )
+    invalid_keys = {
+        key: loading_info[key]
+        for key in ("missing_keys", "unexpected_keys", "mismatched_keys")
+        if loading_info[key]
+    }
+    if invalid_keys:
+        raise RuntimeError(f"Native checkpoint loading was incomplete: {invalid_keys}")
+
+    from transformers.models.nemotron_h import modeling_nemotron_h
+
+    if not modeling_nemotron_h.is_fast_path_available:
+        raise RuntimeError("Transformers loaded without the native Nemotron-H fast path.")
+    return model, config
+
+
 def load_model(
     model_path: str,
     device_map: str,
@@ -72,26 +178,41 @@ def load_model(
     torch_dtype: str,
     tp_plan: str,
     refresh_code_cache: bool,
+    hf_backend: str,
 ):
     print("loading:", model_path)
-    if refresh_code_cache:
+    backend = resolve_hf_backend(hf_backend)
+    print("hf backend:", backend)
+    if backend == "legacy" and refresh_code_cache:
         refresh_remote_code_cache(model_path)
     dtype = getattr(torch, torch_dtype) if torch_dtype != "auto" else "auto"
     tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
     print("tokenizer loaded")
-    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
-    print("config loaded")
-    feature_extractor = AutoFeatureExtractor.from_pretrained(resolve_audio_preprocessor_path(model_path, config))
-    print("feature_extractor loaded")
+    if backend == "native":
+        config = None
+    else:
+        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+        print("config loaded")
     tp_plan_arg = None if tp_plan == "none" else tp_plan
-    model = AutoModelForCausalLM.from_pretrained(
-        model_path,
-        trust_remote_code=True,
-        torch_dtype=dtype,
-        device_map=resolve_device_map(device_map, device),
-        tp_plan=tp_plan_arg,
-    )
+    model_load_kwargs = {
+        "device_map": resolve_device_map(device_map, device),
+    }
+    if tp_plan_arg is not None:
+        model_load_kwargs["tp_plan"] = tp_plan_arg
+    if backend == "native":
+        model_load_kwargs["dtype"] = dtype
+        model, config = load_native_model(model_path, model_load_kwargs)
+        print("config loaded")
+    else:
+        model_load_kwargs["torch_dtype"] = dtype
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            trust_remote_code=True,
+            **model_load_kwargs,
+        )
     print("model loaded")
+    feature_extractor = AutoFeatureExtractor.from_pretrained(resolve_audio_preprocessor_path(model_path, config))
+    print("feature_extractor loaded")
     model.eval()
     return model, tokenizer, feature_extractor, config
 
@@ -295,6 +416,12 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Do not clear this checkpoint's stale Hugging Face dynamic-module cache before loading.",
     )
+    parser.add_argument(
+        "--hf-backend",
+        default="auto",
+        choices=["auto", "native", "legacy"],
+        help="Use native Transformers when available, or select the reproducibility legacy path.",
+    )
     return parser.parse_args()
 
 
@@ -307,6 +434,7 @@ def main() -> None:
         torch_dtype=args.torch_dtype,
         tp_plan=args.tp_plan,
         refresh_code_cache=not args.no_refresh_remote_code_cache,
+        hf_backend=args.hf_backend,
     )
 
     with open(args.input_json, "r", encoding="utf-8") as f: