diff --git a/docs/models/pooling_models/specific_models.md b/docs/models/pooling_models/specific_models.md index 87fed54139aa..1af03848e30d 100644 --- a/docs/models/pooling_models/specific_models.md +++ b/docs/models/pooling_models/specific_models.md @@ -484,3 +484,46 @@ curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{ `embed&token_classify` task is an internal execution contract for this plugin, not a generic Pooling API response format. Without the plugin, select one of the three concrete tasks above. + +## NVIDIA Nemotron-3-Diarization + +`nvidia/Nemotron-3-Diarization` is a streaming Sortformer speaker diarization +model. For every 10 ms frame of a recording it predicts the activity of up to 8 +speakers, numbered in order of first arrival. + +| Architecture | Inputs | Example HF Models | +| - | - | - | +| `Nemotron3DiarizationForAudioFrameClassification` | A | `nvidia/Nemotron-3-Diarization` | + +vLLM runs it offline, one whole recording per request, through the `plugin` +pooling task. The model takes no text, so start the engine with +`skip_tokenizer_init=True`. vLLM runs it in eager mode on the V1 model runner, +and its feature extractor needs the `librosa` package. Use `dtype="float32"`, +the checkpoint's dtype, for logits that match Transformers; half precision moves +a few segment boundaries. Each output is the `[num_frames, 8]` logits; turn +them into speech segments with the Transformers processor: + +```python +from transformers import AutoProcessor +from vllm import LLM + +model = "nvidia/Nemotron-3-Diarization" +llm = LLM(model=model, runner="pooling", skip_tokenizer_init=True, dtype="float32") +processor = AutoProcessor.from_pretrained(model) + +(output,) = llm.encode( + [{"prompt_token_ids": [0], "multi_modal_data": {"audio": (audio, sample_rate)}}], + pooling_task="plugin", +) +segments = processor.extract_speaker_dict(output.outputs.data[None])[0] +# [{'Start': , 'End': , 'Speaker': }, ...] +``` + +A step runs up to `max_num_seqs` recordings, and the features of all of them sit +on the GPU together: about 200 MB per hour of audio in fp32. For long recordings, +lower `max_num_seqs` to bound this memory. + +Streaming sessions, which carry the speaker cache from one request to the next, +and the online Pooling API are not supported yet. See +[examples/pooling/plugin/nemotron3_diarization_offline.py](../../../examples/pooling/plugin/nemotron3_diarization_offline.py) +for a complete example. diff --git a/examples/pooling/plugin/nemotron3_diarization_offline.py b/examples/pooling/plugin/nemotron3_diarization_offline.py new file mode 100644 index 000000000000..4482b964325d --- /dev/null +++ b/examples/pooling/plugin/nemotron3_diarization_offline.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Offline speaker diarization with NVIDIA Nemotron-3-Diarization. + +Each request is one whole recording. The model returns the logits of up to 8 +speakers for every 10 ms frame, `[num_frames, 8]`, and the Transformers processor +turns them into speech segments. + +Usage:: + + python nemotron3_diarization_offline.py --audio meeting.wav +""" + +from argparse import Namespace + +from transformers import AutoProcessor + +from vllm import LLM, EngineArgs +from vllm.assets.audio import AudioAsset +from vllm.multimodal.media.audio import load_audio +from vllm.utils.argparse_utils import FlexibleArgumentParser + + +def parse_args(): + parser = FlexibleArgumentParser() + parser = EngineArgs.add_cli_args(parser) + parser.add_argument( + "--audio", + type=str, + default=None, + help="Path to the recording to diarize. Defaults to a vLLM audio asset.", + ) + parser.set_defaults( + model="nvidia/Nemotron-3-Diarization", + runner="pooling", + skip_tokenizer_init=True, + dtype="float32", + ) + return parser.parse_args() + + +def main(args: Namespace): + audio_path = args.__dict__.pop("audio") + if audio_path is None: + audio, sample_rate = AudioAsset("winning_call").audio_and_sample_rate + else: + audio, sample_rate = load_audio(audio_path, sr=None) + + llm = LLM(**vars(args)) + processor = AutoProcessor.from_pretrained(args.model) + + # The model reads the audio only; the prompt is a placeholder token. + outputs = llm.encode( + [ + { + "prompt_token_ids": [0], + "multi_modal_data": {"audio": (audio, sample_rate)}, + } + ], + pooling_task="plugin", + ) + + for output in outputs: + logits = output.outputs.data # [num_frames, num_speakers] + for segment in processor.extract_speaker_dict(logits[None])[0]: + print( + f"speaker {segment['Speaker']}: " + f"{segment['Start']:.2f}s - {segment['End']:.2f}s" + ) + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/tests/entrypoints/llm/test_prompt_validation.py b/tests/entrypoints/llm/test_prompt_validation.py index 40e5e7b14bdd..fe60c3ffc4b1 100644 --- a/tests/entrypoints/llm/test_prompt_validation.py +++ b/tests/entrypoints/llm/test_prompt_validation.py @@ -37,6 +37,20 @@ def test_out_of_vocab_token_skip_tokenizer_init(vllm_runner): runner.llm.generate({"prompt_token_ids": [999999]}) +def test_raw_input_only_model_requires_mm_data(vllm_runner): + # A prompt without multimodal data would give such a model nothing to run on, + # so it must be rejected before it reaches the engine. + with ( + vllm_runner( + "nvidia/Nemotron-3-Diarization", + runner="pooling", + skip_tokenizer_init=True, + ) as runner, + pytest.raises(VLLMValidationError, match="must include `multi_modal_data`"), + ): + runner.llm.encode({"prompt_token_ids": [0]}, pooling_task="plugin") + + def test_require_mm_embeds(vllm_runner): with ( vllm_runner( diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index 0a0fb3e45c52..71f8a3d68b4e 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -761,6 +761,7 @@ def test_v2_model_runner_accepts_routed_experts(monkeypatch): use_mla=False, logits_processors=None, enable_prompt_embeds=False, + is_multimodal_raw_input_only_model=False, ), aux_output_config=SimpleNamespace(enable_return_routed_experts=True), speculative_config=None, diff --git a/tests/models/registry.py b/tests/models/registry.py index 317db44a57c8..06d376b8a5ef 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -664,6 +664,11 @@ def check_available_online( "nvidia/llama-nemotron-embed-vl-1b-v2", trust_remote_code=True ), "LlavaNextForConditionalGeneration": _HfExamplesInfo("royokong/e5-v"), + "Nemotron3DiarizationForAudioFrameClassification": _HfExamplesInfo( + "nvidia/Nemotron-3-Diarization", + min_transformers_version="5.18.0", + require_embed_inputs=True, + ), "Phi3VForCausalLM": _HfExamplesInfo( "TIGER-Lab/VLM2Vec-Full", trust_remote_code=True ), diff --git a/tests/test_config.py b/tests/test_config.py index 72e25126753c..6dc9589fb28b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1121,12 +1121,31 @@ def test_v2_model_runner_heterogeneous_vocab_draft_falls_back_to_v1( def test_v2_model_runner_supports_custom_logits_processors(): config = VllmConfig() config.model_config = cast( - ModelConfig, SimpleNamespace(logits_processors=["a.b:C"]) + ModelConfig, + SimpleNamespace( + logits_processors=["a.b:C"], is_multimodal_raw_input_only_model=False + ), ) assert config._get_v2_model_runner_unsupported_features() == [] +@pytest.mark.parametrize("raw_input_only", [False, True]) +def test_v2_model_runner_raw_input_only_model_falls_back_to_v1(raw_input_only): + """Model Runner V2 does not pass raw multimodal inputs to the model, so + raw-input-only models (e.g. Terratorch) must stay on V1.""" + config = VllmConfig() + config.model_config = cast( + ModelConfig, + SimpleNamespace(is_multimodal_raw_input_only_model=raw_input_only), + ) + + expected = ["raw-input-only multimodal models"] + assert config._get_v2_model_runner_unsupported_features() == ( + expected if raw_input_only else [] + ) + + @pytest.mark.parametrize( "architecture", ["DFlash2DraftModel", "DFlash2BeamDraftModel", "LiLiCorrDraftModel"], diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index c871b225b621..9e4aedeb122e 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3129,6 +3129,12 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: if self.parallel_config.use_ubatching: unsupported.extend(self._get_dbo_unsupported_features()) + if ( + self.model_config is not None + and self.model_config.is_multimodal_raw_input_only_model + ): + unsupported.append("raw-input-only multimodal models") + return unsupported def _get_standalone_cpu_model_runner_unsupported_features(self) -> list[str]: diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 57411849e2be..3ef84319893c 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -930,6 +930,16 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: ) +class Nemotron3DiarizationForAudioFrameClassificationConfig(VerifyAndUpdateConfig): + @staticmethod + def verify_and_update_model_config(model_config: "ModelConfig") -> None: + # The forward syncs with the host at every chunk of a recording, so it + # cannot be captured in a CUDA graph. + if not model_config.enforce_eager: + logger.info("Nemotron-3-Diarization runs in eager mode.") + model_config.enforce_eager = True + + class NomicBertModelConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -1325,6 +1335,9 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "Nemotron3_5TranscribeForConditionalGeneration": ( Nemotron3_5TranscribeForConditionalGenerationConfig ), + "Nemotron3DiarizationForAudioFrameClassification": ( + Nemotron3DiarizationForAudioFrameClassificationConfig + ), "NomicBertModel": NomicBertModelConfig, "Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig, "Qwen2ForRewardModel": Qwen2ForRewardModelConfig, diff --git a/vllm/model_executor/models/nemotron3_diarization.py b/vllm/model_executor/models/nemotron3_diarization.py new file mode 100644 index 000000000000..b4296d558978 --- /dev/null +++ b/vllm/model_executor/models/nemotron3_diarization.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Wrapper around the Transformers implementation of NVIDIA Nemotron-3-Diarization. + +The model is a streaming Sortformer speaker diarizer: for every 10 ms spectrogram +frame of a recording it returns the logits of up to 8 speakers, numbered in order +of first arrival. vLLM serves it offline, one whole recording per request, through +the `plugin` pooling task. The Transformers forward chunks the recording and keeps +the speaker cache internally, so no state is carried between requests. +""" + +from collections.abc import Iterable, Mapping, Sequence + +import torch +import torch.nn as nn +from transformers import BatchFeature, Nemotron3DiarizationConfig +from transformers import ( + Nemotron3DiarizationForAudioFrameClassification as HFNemotron3Diarization, +) + +from vllm.config import VllmConfig +from vllm.config.multimodal import MultiModalDummyOptions +from vllm.inputs import MultiModalDataDict, MultiModalInput, mm_input +from vllm.model_executor.layers.pooler import IdentityPooler +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import ( + MultiModalFieldConfig, + MultiModalKwargsItems, + PlaceholderRange, +) +from vllm.multimodal.parse import ( + AudioProcessorItems, + MultiModalDataItems, + MultiModalDataParser, +) +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + ProcessorInputs, + PromptUpdate, + TimingContext, +) +from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.v1.outputs import PoolerOutput +from vllm.v1.pool.metadata import PoolingMetadata + +from .interfaces import IsAttentionFree, MultiModalEmbeddings, SupportsMultiModal +from .interfaces_base import attn_type + +# Length of the dummy recording used for profiling. The model has no KV cache to +# size, so a short clip that exercises the forward is enough. +_DUMMY_AUDIO_SECONDS = 5 + + +class Nemotron3DiarizationProcessingInfo(BaseProcessingInfo): + def get_hf_config(self) -> Nemotron3DiarizationConfig: + return self.ctx.get_hf_config(Nemotron3DiarizationConfig) + + def get_feature_extractor(self, **kwargs: object): + return self.get_hf_processor(**kwargs).feature_extractor + + def get_data_parser(self): + return MultiModalDataParser( + target_sr=self.get_feature_extractor().sampling_rate, + target_channels=1, + expected_hidden_size=self._get_expected_hidden_size(), + ) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"audio": 1} + + +class Nemotron3DiarizationDummyInputsBuilder( + BaseDummyInputsBuilder[Nemotron3DiarizationProcessingInfo] +): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return "" + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: MultiModalDummyOptions, + ) -> MultiModalDataDict: + sampling_rate = self.info.get_feature_extractor().sampling_rate + return { + "audio": self._get_dummy_audios( + length=_DUMMY_AUDIO_SECONDS * sampling_rate, + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), + ) + } + + +class Nemotron3DiarizationMultiModalProcessor( + BaseMultiModalProcessor[Nemotron3DiarizationProcessingInfo] +): + # The model consumes a spectrogram, not tokens. + requires_tokenizer = False + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return dict( + input_features=MultiModalFieldConfig.batched("audio"), + attention_mask=MultiModalFieldConfig.batched("audio"), + ) + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + return [] + + def apply( + self, + inputs: ProcessorInputs, + timing_ctx: TimingContext, + ) -> MultiModalInput: + mm_items = inputs.mm_data_items + hf_processor_mm_kwargs = inputs.hf_processor_mm_kwargs + + with timing_ctx.record("apply_hf_processor"): + audios = mm_items.get_items("audio", AudioProcessorItems) + feature_extractor = self.info.get_feature_extractor( + **hf_processor_mm_kwargs + ) + # Same features as `Nemotron3DiarizationProcessor` in offline mode + mm_processed_data = feature_extractor( + [audios.get(i) for i in range(len(audios))], + sampling_rate=feature_extractor.sampling_rate, + return_attention_mask=True, + return_tensors="pt", + ) + + mm_kwargs = MultiModalKwargsItems.from_hf_inputs( + mm_processed_data, + self._get_mm_fields_config(mm_processed_data, hf_processor_mm_kwargs), + ) + + with timing_ctx.record("get_mm_hashes"): + mm_hashes = inputs.get_mm_hashes( + self.info.model_id, + self.info.ctx.get_mm_config().mm_hasher_algorithm, + ) + + return mm_input( + prompt_token_ids=[0], + mm_kwargs=mm_kwargs, + mm_hashes=mm_hashes, + mm_placeholders={"audio": [PlaceholderRange(offset=0, length=0)]}, + ) + + +class Nemotron3DiarizationPooler(IdentityPooler): + """Trims the padded logits of a step back to the frames of each recording.""" + + def __init__(self) -> None: + super().__init__() + self.num_frames: list[int] | None = None + + def forward( + self, + hidden_states: torch.Tensor, + pooling_metadata: PoolingMetadata, + ) -> PoolerOutput: + num_frames, self.num_frames = self.num_frames, None + if num_frames is None: + return hidden_states + return [logits[:n] for logits, n in zip(hidden_states, num_frames)] + + +@attn_type("attention_free") +@MULTIMODAL_REGISTRY.register_processor( + Nemotron3DiarizationMultiModalProcessor, + info=Nemotron3DiarizationProcessingInfo, + dummy_inputs=Nemotron3DiarizationDummyInputsBuilder, +) +class Nemotron3DiarizationForAudioFrameClassification( + nn.Module, IsAttentionFree, SupportsMultiModal +): + supports_multimodal_raw_input_only = True + is_pooling_model = True + + # The Transformers model is wrapped as is, under `self.model` + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"": "model."}) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("audio"): + return None + + raise ValueError("Only audio modality is supported") + + def __init__(self, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config: Nemotron3DiarizationConfig = vllm_config.model_config.hf_config + self.model = HFNemotron3Diarization(config) + self.num_speakers = config.head_config.num_speakers + + # Transformers computes the RoPE frequencies on the CPU. Computing them on + # the GPU, where the model is built, shifts the logits by up to ~2e-3. + rotary_emb = self.model.model.audio_tower.rotary_emb + if rotary_emb.rope_type == "default": + with torch.device("cpu"): + inv_freq, _ = rotary_emb.compute_default_rope_parameters( + rotary_emb.config + ) + rotary_emb.inv_freq.copy_(inv_freq) + rotary_emb.original_inv_freq.copy_(inv_freq) + + self.pooler = Nemotron3DiarizationPooler() + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + # The prompt is a single placeholder token that the model does not use. + return torch.empty((input_ids.shape[0], 0)) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + input_features: torch.Tensor | list[torch.Tensor] | None = None, + attention_mask: torch.Tensor | list[torch.Tensor] | None = None, + **kwargs: object, + ) -> torch.Tensor: + if input_features is None: + # No recording in this step: the output still needs a batch dimension + self.pooler.num_frames = None + return positions.new_zeros( + (positions.shape[0], 0, self.num_speakers), dtype=self.model.dtype + ) + assert attention_mask is not None + + logits = list[torch.Tensor]() + # One recording per forward: the offline forward chunks it and updates the + # speaker cache, which syncs with the host. + with gpu_sync_allowed(): + for features, mask in zip(input_features, attention_mask): + output = self.model( + input_features=features[None].to(self.model.dtype), + attention_mask=mask[None], + ) + logits.append(output.logits[0]) + + # The runner pools a single tensor, so recordings of different lengths + # are padded here and trimmed back by the pooler. + self.pooler.num_frames = [x.shape[0] for x in logits] + return nn.utils.rnn.pad_sequence(logits, batch_first=True) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index da9d0ce9737e..8e9eabfacd1b 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -261,6 +261,11 @@ "llava_next", "LlavaNextForConditionalGeneration", ), + # Like Terratorch, a raw-input model served through the "plugin" task + "Nemotron3DiarizationForAudioFrameClassification": ( + "nemotron3_diarization", + "Nemotron3DiarizationForAudioFrameClassification", + ), "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"), "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"), "SiglipModel": ("siglip", "SiglipEmbeddingModel"), diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 55c88c52cab3..1bcc3b92d726 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -499,6 +499,12 @@ def get_total_num_kv_heads(self) -> int: return 0 +class Nemotron3DiarizationModelArchConfigConvertor(TerratorchModelArchConfigConvertor): + def get_vocab_size(self) -> int: + # The model reads audio only; its prompt is the placeholder token [0]. + return 1 + + class Zamba2ModelArchConfigConvertor(ModelArchConfigConvertorBase): def get_head_size(self) -> int: return getattr(self.hf_text_config, "attention_head_dim", 0) @@ -874,6 +880,7 @@ def derive_max_model_len_and_key(self) -> tuple[float, str | None]: "mimo_v2_omni_mtp": MimoV2MTPModelArchConfigConvertor, "moss_audio": MossAudioModelArchConfigConvertor, "nemotron-nas": NemotronNasModelArchConfigConvertor, + "nemotron3_diarization": Nemotron3DiarizationModelArchConfigConvertor, "bailing_hybrid_v3_mtp": BailingHybridV3MTPModelArchConfigConvertor, "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor, "qwen3_5": Qwen3_5ModelArchConfigConvertor, diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 2ddfe5457f93..f403d59ba4fe 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -670,4 +670,14 @@ def _validate_model_inputs( if encoder_input is not None: self._validate_model_input(encoder_input, prompt_type="encoder") + if ( + self.model_config.is_multimodal_raw_input_only_model + and decoder_input["type"] != "multimodal" + ): + raise VLLMValidationError( + "This model reads only its multimodal input, so the prompt " + "must include `multi_modal_data`.", + parameter="multi_modal_data", + ) + self._validate_model_input(decoder_input, prompt_type="decoder")