From 68071b90b0bb302b9de6c1aff1e3719523d2d7f8 Mon Sep 17 00:00:00 2001 From: BAR HAIM Date: Thu, 30 Jul 2026 15:53:58 +0000 Subject: [PATCH] audio: make ASR dtype configurable; trim comment density Add a first-class asr_dtype config field so the precision the ASR weights load in is tunable per checkpoint instead of hardcoded. Previously ASRTranscriber.load() always requested float16, which crashes any encoder with BatchNorm layers ("Expected weight to have type Float but got Half") because BatchNorm will not promote a float16 weight against float32 features. The only escape was to smuggle torch_dtype through asr_pipeline_kwargs, which is opaque and undocumented. asr_dtype accepts auto/float16/bfloat16/float32 and is validated in GraniteSwitchConfig.__init__ so a typo fails at compose time rather than inside a vLLM worker. None/"auto" preserves today's behavior: float16 on CUDA, float32 elsewhere. asr_pipeline_kwargs still merges last and therefore still wins, so existing checkpoints that use the torch_dtype workaround are unaffected. Also exposed as --asr-dtype on the compose CLI (implies --enable-audio) and documented in docs/AUDIO.md. Separately, reduce comment and docstring volume across the audio modules and their tests. Much of it restated the code or duplicated docs/AUDIO.md; what remains is the comments carrying information the code cannot, such as why kwargs.update must come last and why requires_raw_input_tokens is set. Note tests/vllm/test_audio_processor.py: three transcriber stubs did not accept the new dtype kwarg and failed once processor.py started passing it. Fixed, plus two tests asserting the dtype actually reaches get_transcriber. Signed-off-by: BAR HAIM --- docs/AUDIO.md | 7 + .../composer/compose_granite_switch.py | 17 +- src/granite_switch/config.py | 103 +++++------ src/granite_switch/vllm/audio/__init__.py | 12 +- src/granite_switch/vllm/audio/asr.py | 162 ++++++++---------- src/granite_switch/vllm/audio/chunking.py | 54 ++---- src/granite_switch/vllm/audio/processor.py | 91 +++------- .../vllm/granite_switch_model.py | 11 +- .../test_adapter_routing_audio_enabled.py | 14 +- .../test_answerability_over_audio.py | 33 ++-- tests/integration/test_audio_serving_smoke.py | 65 +++---- tests/unit/test_asr.py | 84 ++++++++- tests/unit/test_config.py | 12 ++ tests/vllm/test_audio_processor.py | 49 +++--- 14 files changed, 342 insertions(+), 372 deletions(-) diff --git a/docs/AUDIO.md b/docs/AUDIO.md index de64363..c05166d 100644 --- a/docs/AUDIO.md +++ b/docs/AUDIO.md @@ -33,6 +33,13 @@ settings into `config.json` so the checkpoint is self-describing: `openai/whisper-small` for multilingual. - `asr_device` — `cpu` (default) keeps vLLM's GPU KV-cache budget clean; set `--asr-device cuda:0` to run transcription on GPU (watch GPU memory). +- `asr_dtype` — precision the ASR weights load in. Unset (default) derives it + from the device: `float16` on CUDA, `float32` on CPU. Half precision halves + the ASR weight footprint and is what the Whisper-family defaults expect, but + it is not universally safe — an encoder with **BatchNorm** layers raises + `Expected weight to have type Float but got Half`, since BatchNorm will not + promote a float16 weight against float32 features. Such a checkpoint needs + `--asr-dtype float32`. Accepted: `auto`, `float16`, `bfloat16`, `float32`. Audio capability is **gated per checkpoint** by `asr_enabled`: a checkpoint built without `--enable-audio` reports no audio modality and never loads the ASR model. diff --git a/src/granite_switch/composer/compose_granite_switch.py b/src/granite_switch/composer/compose_granite_switch.py index 3ec4587..426f5bb 100755 --- a/src/granite_switch/composer/compose_granite_switch.py +++ b/src/granite_switch/composer/compose_granite_switch.py @@ -67,6 +67,7 @@ configure_chat_template, get_alora_first_invocation_token_id, ) +from granite_switch.config import ASR_DTYPES # --------------------------------------------------------------------------- # Utility helpers (kept local — not worth a separate module) @@ -589,6 +590,16 @@ def _compose_argparser(): help="Device the ASR model runs on (default: cpu). Use e.g. cuda:0 to " "run transcription on GPU (watch vLLM's KV-cache memory budget).", ) + parser.add_argument( + "--asr-dtype", + type=str, + default=None, + choices=ASR_DTYPES, + help="Precision the ASR weights load in. Default derives it from " + "--asr-device (float16 on CUDA, float32 on CPU); set float32 for an " + "encoder that cannot run in half precision (e.g. one with BatchNorm). " + "Implies --enable-audio.", + ) parser.add_argument( "--asr-pipeline-kwargs", type=json.loads, @@ -852,6 +863,7 @@ def build(): audio_enabled = ( args.enable_audio or args.asr_model is not None + or args.asr_dtype is not None or args.asr_pipeline_kwargs is not None or args.asr_generate_kwargs is not None or args.asr_max_audio_clips is not None @@ -934,6 +946,7 @@ def build(): model.config.asr_enabled = True model.config.asr_model_id = args.asr_model model.config.asr_device = args.asr_device + model.config.asr_dtype = args.asr_dtype # Optional pipeline-construction extras and default decode kwargs. Only # set when provided so the config stays minimal for the common case. if args.asr_pipeline_kwargs is not None: @@ -953,7 +966,9 @@ def build(): print( f" Audio cascade enabled " f"(asr_model_id={args.asr_model or 'default'}, " - f"asr_device={args.asr_device}, audio_token_id={audio_token_id}, " + f"asr_device={args.asr_device}, " + f"asr_dtype={args.asr_dtype or 'auto'}, " + f"audio_token_id={audio_token_id}, " f"pipeline_kwargs={args.asr_pipeline_kwargs or {}}, " f"generate_kwargs={args.asr_generate_kwargs or {}}, " f"max_audio_clips={args.asr_max_audio_clips or 'default'}, " diff --git a/src/granite_switch/config.py b/src/granite_switch/config.py index 38c20cd..cc788d8 100644 --- a/src/granite_switch/config.py +++ b/src/granite_switch/config.py @@ -3,6 +3,9 @@ from transformers import GraniteMoeHybridConfig +# Accepted asr_dtype values. Keep in sync with vllm.audio.asr._ASR_DTYPE_NAMES. +ASR_DTYPES = ("auto", "float16", "bfloat16", "float32") + class GraniteSwitchConfig(GraniteMoeHybridConfig): """Configuration class for GraniteSwitch model. @@ -37,56 +40,38 @@ class GraniteSwitchConfig(GraniteMoeHybridConfig): Module groups: "qkv_proj", "o_proj", "shared_input_linear", "shared_output_linear". Default: all four groups - Audio (ASR) preprocessing parameters: - asr_enabled (bool): When True, the vLLM backend registers an audio - multimodal preprocessor that transcribes audio inputs to text and - splices the transcript tokens into the prompt before the decoder - runs. The decoder itself is unchanged and only ever sees text - tokens. Default: False. - asr_model_id (Optional[str]): HuggingFace id of the speech-to-text - model the preprocessor loads. When None and asr_enabled is True, - the backend falls back to a small built-in default. This makes the - checkpoint self-describing about which ASR front-end it expects. + Audio (ASR) preprocessing parameters (see docs/AUDIO.md): + asr_enabled (bool): Register the audio preprocessor that transcribes + audio and splices the transcript into the prompt. Default: False. + asr_model_id (Optional[str]): HF id of the speech-to-text model. None + falls back to a small built-in default. asr_device (str): Device the ASR model runs on. Default "cpu" keeps - vLLM's GPU KV-cache budget clean; set to a CUDA device to trade GPU - memory for transcription latency. - asr_pipeline_kwargs (Optional[dict]): Extra keyword arguments merged - into the ``transformers.pipeline(...)`` construction call for the - ASR model (e.g. ``{"chunk_length_s": 15}``). These affect how the - pipeline is built, so they are baked into the transcriber cache - key. None means "no extras". Default: None. - asr_generate_kwargs (Optional[dict]): Default decode-time keyword - arguments passed to the ASR model on every transcription (e.g. - ``{"language": "de", "task": "transcribe"}`` for a multilingual - Whisper). Applied at call time, so a single loaded pipeline can be - reused; per-request values (via ``mm_processor_kwargs``) override - these. Ignored by models that do not generate (e.g. CTC). Default: - None. - asr_max_audio_clips (int): Maximum number of audio clips accepted in a - single request. Each clip's transcript is spliced at its own - ``<|audio|>`` marker. vLLM enforces this as the modality ceiling - (``--limit-mm-per-prompt`` may lower it, not raise it). For the - cascade the clips cost no extra KV (transcripts are ordinary text - tokens bounded by the context window); the ceiling mainly guards - against one request triggering an unbounded number of synchronous - ASR transcriptions, and bounds the startup profiling pass. Default: - 32. - asr_chunk_length_s (float): Window length (seconds) our own long-audio - chunker splits a clip into before transcribing each window. Only - used for backends that do not self-chunk (see asr_self_chunks); the - default Whisper backend chunks internally and ignores this. Default: - 30.0. - asr_chunk_overlap_s (float): Overlap (seconds) between consecutive - chunker windows, so words straddling a boundary are whole in at - least one window; the transcript merge de-duplicates the overlap. - Only used when asr_self_chunks is False. Default: 5.0. - asr_self_chunks (bool): Whether the ASR backend handles long audio - itself. True for the Whisper pipeline (its internal - ``chunk_length_s`` does timestamp-based stitching, higher quality - than our text-level merge), so our chunker is bypassed. Set False - for a backend with a fixed input window (e.g. a future speech - encoder) to route audio through our encoder-agnostic - split/transcribe/merge chunker. Default: True. + vLLM's GPU KV-cache budget clean. + asr_dtype (Optional[str]): Precision the ASR weights load in, one of + ASR_DTYPES. None/"auto" derives it from asr_device (float16 on + CUDA). An encoder with BatchNorm layers must set "float32". + Default: None. + asr_pipeline_kwargs (Optional[dict]): Extra kwargs merged into the + ``transformers.pipeline(...)`` construction, e.g. + ``{"chunk_length_s": 15}``. Baked into the transcriber cache key. + Default: None. + asr_generate_kwargs (Optional[dict]): Default decode-time kwargs, e.g. + ``{"language": "de"}``. Applied per call, so one pipeline is + reused; per-request ``mm_processor_kwargs`` override them. Ignored + by non-generative backends. Default: None. + asr_max_audio_clips (int): Max audio clips per request. Bounds the + synchronous transcriptions one request can trigger and the startup + profiling pass; ``--limit-mm-per-prompt`` may lower it, not raise + it. Default: 32. + asr_chunk_length_s (float): Chunker window length in seconds. Only + used when asr_self_chunks is False. Default: 30.0. + asr_chunk_overlap_s (float): Overlap in seconds between chunker + windows, de-duplicated by the transcript merge. Only used when + asr_self_chunks is False. Default: 5.0. + asr_self_chunks (bool): True when the backend chunks long audio + itself (Whisper's timestamp stitching beats our text-level merge), + bypassing our chunker. False routes audio through the + split/transcribe/merge chunker instead. Default: True. **kwargs: Additional arguments passed to GraniteConfig. """ @@ -109,6 +94,7 @@ def __init__( asr_enabled: bool = False, asr_model_id: str | None = None, asr_device: str = "cpu", + asr_dtype: str | None = None, asr_pipeline_kwargs: dict | None = None, asr_generate_kwargs: dict | None = None, asr_max_audio_clips: int = 32, @@ -195,22 +181,19 @@ def __init__( self.switch_head_dim = switch_head_dim self.fused_add_norm = fused_add_norm - # Audio (ASR) preprocessing parameters. - # The decoder is oblivious to audio: when enabled, the vLLM backend's - # multimodal preprocessor transcribes audio to text and injects the - # transcript tokens into the prompt before embedding. These fields make - # the checkpoint self-describing about its ASR front-end. + # Audio (ASR) preprocessing. The decoder is oblivious to audio; these + # fields make the checkpoint self-describing about its ASR front-end. self.asr_enabled = asr_enabled self.asr_model_id = asr_model_id self.asr_device = asr_device - # Pipeline-construction extras (affect the built pipeline → cache key) - # and default decode kwargs (applied per call, per-request overridable). + # Validated here so a typo fails at compose time, not in a vLLM worker. + if asr_dtype is not None and asr_dtype not in ASR_DTYPES: + raise ValueError( + f"asr_dtype must be one of {ASR_DTYPES} or None, got {asr_dtype!r}" + ) + self.asr_dtype = asr_dtype self.asr_pipeline_kwargs = asr_pipeline_kwargs self.asr_generate_kwargs = asr_generate_kwargs - # Long-audio / multi-clip preprocessing. The transcript is spliced into - # the prompt as ordinary text tokens; a clip that makes the prompt exceed - # the context is rejected by vLLM's standard length check (no truncation). - # The chunker settings only apply to backends that do not self-chunk. if asr_max_audio_clips < 1: raise ValueError( f"asr_max_audio_clips must be >= 1, got {asr_max_audio_clips}" diff --git a/src/granite_switch/vllm/audio/__init__.py b/src/granite_switch/vllm/audio/__init__.py index 486bc25..feab7ce 100644 --- a/src/granite_switch/vllm/audio/__init__.py +++ b/src/granite_switch/vllm/audio/__init__.py @@ -1,16 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 """Audio (ASR) preprocessing for the Granite Switch vLLM backend. -This package implements the **alpha** audio pathway: a speech-to-text *cascade*. -Audio is transcribed to text by a small ASR model and the transcript tokens are -spliced into the prompt before the decoder runs. The Granite Switch decoder is -unchanged and only ever sees text tokens — there is no in-model audio encoder yet -(that is deferred future work: a trained projection from Granite Speech encoder -embeddings into Granite token space). - -The ASR backend (:mod:`asr`) deliberately has no vLLM dependency so it can be -unit-tested on CPU and reused by either the multimodal-processor integration or a -fallback entrypoint wrapper. +Speech-to-text cascade: audio is transcribed and the transcript tokens are +spliced into the prompt, so the decoder only ever sees text. See docs/AUDIO.md. """ from .asr import DEFAULT_ASR_MODEL_ID, ASRTranscriber, transcribe diff --git a/src/granite_switch/vllm/audio/asr.py b/src/granite_switch/vllm/audio/asr.py index d58ea8d..b95085a 100644 --- a/src/granite_switch/vllm/audio/asr.py +++ b/src/granite_switch/vllm/audio/asr.py @@ -1,23 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -"""Speech-to-text backend for the alpha audio cascade. - -A thin wrapper around a HuggingFace ASR model that turns raw audio into a -transcript string. Kept free of any vLLM import so it can be unit-tested on CPU -and reused by both the multimodal-processor integration and a fallback wrapper. - -Design notes: - * The model is loaded **lazily** on first use and cached per (model_id, - device) so engine startup stays cheap and a process loads each ASR model - at most once. - * Default device is CPU — this keeps vLLM's GPU KV-cache budget clean. Set a - CUDA device to trade GPU memory for transcription latency. - * The default model is a small, text-emitting open-source ASR model. The - alpha intentionally uses a complete audio->text model (encoder+decoder); - swapping in the Granite Speech 4.1 encoder is future work and only touches - this module. - * ``transcribe`` accepts the audio item shapes vLLM hands to multimodal - processors: a bare array, an ``(array, sampling_rate)`` tuple, a Python - list of floats, or a torch tensor. +"""Speech-to-text backend for the audio cascade. + +Wraps a HuggingFace ASR pipeline. Free of any vLLM import so it unit-tests on +CPU. The model loads lazily and is cached per (model_id, device, dtype, +pipeline_kwargs), so a process loads each ASR model at most once. + +Device defaults to CPU to keep vLLM's GPU KV-cache budget clean; dtype follows +the device unless a checkpoint sets ``asr_dtype``. See docs/AUDIO.md. """ from __future__ import annotations @@ -32,6 +21,43 @@ # checkpoint does not name its own (config.asr_model_id is None). DEFAULT_ASR_MODEL_ID = "distil-whisper/distil-small.en" +ASR_DTYPE_AUTO = "auto" + +# Keep in sync with config.ASR_DTYPES. +_ASR_DTYPE_NAMES = frozenset({"float16", "bfloat16", "float32"}) +_ASR_DTYPE_ALIASES = { + "fp16": "float16", + "half": "float16", + "bf16": "bfloat16", + "fp32": "float32", + "float": "float32", +} + + +def _resolve_torch_dtype(dtype: str | None, device: str) -> Any: + """Resolve an ``asr_dtype`` name to a ``torch.dtype``. + + None/"auto" derives it from the device: float16 on CUDA, float32 elsewhere + (CPU float16 is slow and partly unimplemented). Name a dtype explicitly for + an encoder that cannot run in half precision — BatchNorm raises on a float16 + weight against float32 features rather than promoting. + """ + import torch + + name = str(dtype or ASR_DTYPE_AUTO).lower() + name = _ASR_DTYPE_ALIASES.get(name, name) + if name == ASR_DTYPE_AUTO: + on_cuda = isinstance(device, str) and device.startswith("cuda") + return torch.float16 if on_cuda else torch.float32 + if name not in _ASR_DTYPE_NAMES: + raise ValueError( + f"Unsupported asr_dtype {dtype!r}. Expected {ASR_DTYPE_AUTO!r} (or " + f"None) to derive it from the device, or one of: " + f"{', '.join(sorted(_ASR_DTYPE_NAMES))}." + ) + return getattr(torch, name) + + _CHUNKING = None @@ -73,23 +99,18 @@ def _load_chunking(): class ASRTranscriber: - """Lazily-loaded ASR model wrapper exposing :meth:`transcribe`. - - Instances are cheap to construct; the underlying model is only materialized - on the first :meth:`transcribe` call (or an explicit :meth:`load`). - """ + """Lazily-loaded ASR model wrapper exposing :meth:`transcribe`.""" def __init__( self, model_id: str = DEFAULT_ASR_MODEL_ID, device: str = "cpu", pipeline_kwargs: Mapping[str, Any] | None = None, + dtype: str | None = None, ) -> None: self.model_id = model_id self.device = device - # Extra kwargs merged into the pipeline() construction. Because they - # change the built pipeline, get_transcriber folds them into the cache - # key so distinct construction options get distinct cached instances. + self.dtype = dtype self.pipeline_kwargs: dict[str, Any] = dict(pipeline_kwargs or {}) self._pipeline = None self._load_lock = threading.Lock() @@ -101,28 +122,17 @@ def load(self) -> None: with self._load_lock: if self._pipeline is not None: return - # Imported lazily so this module stays importable without - # transformers' heavy audio stack until it is actually used. - import torch + # Lazy: keeps this module importable without transformers' audio stack. from transformers import pipeline - torch_dtype = ( - torch.float16 - if isinstance(self.device, str) and self.device.startswith("cuda") - else torch.float32 - ) - # Built-in defaults, then let checkpoint-supplied pipeline_kwargs - # override any of them (e.g. a different chunk_length_s, or model - # kwargs a non-Whisper backend needs). kwargs: dict[str, Any] = { "task": "automatic-speech-recognition", "model": self.model_id, "device": self.device, - "torch_dtype": torch_dtype, - # Enable internal 30s chunking so audio longer than the model's - # native window is handled without us re-implementing it. + "torch_dtype": _resolve_torch_dtype(self.dtype, self.device), "chunk_length_s": 30, } + # pipeline_kwargs last: a checkpoint may override any default above. kwargs.update(self.pipeline_kwargs) self._pipeline = pipeline(**kwargs) @@ -135,32 +145,12 @@ def transcribe( chunk_length_s: float = 30.0, chunk_overlap_s: float = 5.0, ) -> str: - """Transcribe one audio clip to a text string. - - Args: - audio: The audio samples. Accepts a numpy array, a list of floats, a - torch tensor, or an ``(array, sampling_rate)`` tuple (vLLM's - ``AudioItem`` shape). - sampling_rate: Sample rate of ``audio`` in Hz. Required unless - ``audio`` is an ``(array, sampling_rate)`` tuple. Audio is - resampled to 16 kHz when needed. - generate_kwargs: Decode-time kwargs forwarded to the ASR model for - this call (e.g. ``{"language": "fr"}``). Applied per call so the - same loaded pipeline serves many languages. Passed only when - non-empty, so CTC/non-generative backends are unaffected. - self_chunks: True when the backend handles long audio itself (the - Whisper pipeline, via its internal ``chunk_length_s``); the whole - clip is passed in one call. False routes long audio through our - encoder-agnostic chunker (:mod:`.chunking`): split into - overlapping windows, transcribe each, and merge with overlap - de-duplication. - chunk_length_s: Window length for our chunker (only used when - ``self_chunks`` is False). - chunk_overlap_s: Window overlap for our chunker (only used when - ``self_chunks`` is False). - - Returns: - The transcript with surrounding whitespace stripped. + """Transcribe one audio clip, stripped. Resampled to 16 kHz as needed. + + ``sampling_rate`` is required unless ``audio`` is an ``(array, rate)`` + tuple. ``generate_kwargs`` is passed only when non-empty, so CTC backends + are unaffected. ``self_chunks=False`` routes long audio through + :mod:`.chunking` using ``chunk_length_s``/``chunk_overlap_s``. """ samples, sr = _coerce_audio(audio, sampling_rate) samples = _to_mono_float32(samples) @@ -168,12 +158,9 @@ def transcribe( self.load() - # Backend chunks internally (Whisper): one call over the whole clip. if self_chunks: return self._run_pipeline(samples, generate_kwargs) - # Encoder-agnostic path: split into overlapping windows, transcribe each, - # then stitch the per-window transcripts (de-duplicating the overlap). chunking = _load_chunking() segments = chunking.split_waveform( samples, _TARGET_SAMPLE_RATE, chunk_length_s, chunk_overlap_s @@ -187,7 +174,6 @@ def _run_pipeline( generate_kwargs: Mapping[str, Any] | None = None, ) -> str: """Run the loaded pipeline over an already-resampled mono waveform.""" - # Tell the pipeline the rate so it does not attempt its own resampling. call_kwargs: dict[str, Any] = {} if generate_kwargs: call_kwargs["generate_kwargs"] = dict(generate_kwargs) @@ -201,17 +187,14 @@ def _run_pipeline( # ── Module-level cache + convenience function ──────────────────────────────── -# Keyed on (model_id, device, frozen pipeline_kwargs). pipeline_kwargs are in -# the key because they change the constructed pipeline; generate_kwargs are NOT -# — they are applied per transcribe() call, so one cached pipeline serves them -# all (that is what makes per-request language selection cheap). +# Keyed on what changes the constructed pipeline. generate_kwargs are excluded: +# they apply per transcribe() call, so one cached pipeline serves every language. _TRANSCRIBERS: dict[tuple, ASRTranscriber] = {} _CACHE_LOCK = threading.Lock() -# Decode kwargs a client may set per request. Kept to language/task so a client -# cannot inject arbitrary (potentially expensive or unsafe) generation options; -# everything else is fixed by the checkpoint author in config.asr_generate_kwargs. +# Allowlisted so a client cannot inject arbitrary generation options; everything +# else is fixed by the checkpoint in config.asr_generate_kwargs. DEFAULT_ALLOWED_REQUEST_GENERATE_KEYS = frozenset({"language", "task"}) @@ -222,12 +205,9 @@ def resolve_generate_kwargs( ) -> dict[str, Any]: """Merge config-default decode kwargs with allowlisted per-request overrides. - ``config_defaults`` come from ``config.asr_generate_kwargs``. ``request`` is - the per-request mapping (vLLM's ``mm_processor_kwargs``); a top-level - ``language`` or a nested ``asr_generate_kwargs`` object may override the - defaults, but only keys in ``allowed_keys`` are honored. Request values win - so one deployed model can serve many languages. Pure and vLLM-free so it can - be unit-tested on CPU. + ``request`` is vLLM's ``mm_processor_kwargs``: a top-level ``language`` or a + nested ``asr_generate_kwargs`` may override the defaults, but only for keys + in ``allowed_keys``. Request values win, so one model serves many languages. """ merged: dict[str, Any] = dict(config_defaults or {}) if isinstance(request, Mapping): @@ -253,14 +233,15 @@ def get_transcriber( model_id: str | None = None, device: str = "cpu", pipeline_kwargs: Mapping[str, Any] | None = None, + dtype: str | None = None, ) -> ASRTranscriber: """Return a process-wide cached :class:`ASRTranscriber`. - Cached per ``(model_id, device, pipeline_kwargs)``. ``model_id`` of None - resolves to :data:`DEFAULT_ASR_MODEL_ID`. + Cached per ``(model_id, device, dtype, pipeline_kwargs)``. ``model_id`` of + None resolves to :data:`DEFAULT_ASR_MODEL_ID`. """ resolved = model_id or DEFAULT_ASR_MODEL_ID - key = (resolved, device, _freeze(pipeline_kwargs or {})) + key = (resolved, device, dtype, _freeze(pipeline_kwargs or {})) transcriber = _TRANSCRIBERS.get(key) if transcriber is None: with _CACHE_LOCK: @@ -270,6 +251,7 @@ def get_transcriber( model_id=resolved, device=device, pipeline_kwargs=pipeline_kwargs, + dtype=dtype, ) _TRANSCRIBERS[key] = transcriber return transcriber @@ -282,6 +264,7 @@ def transcribe( model_id: str | None = None, device: str = "cpu", pipeline_kwargs: Mapping[str, Any] | None = None, + dtype: str | None = None, generate_kwargs: Mapping[str, Any] | None = None, self_chunks: bool = True, chunk_length_s: float = 30.0, @@ -289,7 +272,10 @@ def transcribe( ) -> str: """Convenience wrapper: transcribe with the cached transcriber for the args.""" return get_transcriber( - model_id=model_id, device=device, pipeline_kwargs=pipeline_kwargs + model_id=model_id, + device=device, + pipeline_kwargs=pipeline_kwargs, + dtype=dtype, ).transcribe( audio, sampling_rate, diff --git a/src/granite_switch/vllm/audio/chunking.py b/src/granite_switch/vllm/audio/chunking.py index bcc8cb5..2dfa7bf 100644 --- a/src/granite_switch/vllm/audio/chunking.py +++ b/src/granite_switch/vllm/audio/chunking.py @@ -1,21 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 """Encoder-agnostic long-audio chunking for the cascade. -A backend with a fixed input window (e.g. a speech encoder that only accepts a -few seconds at a time) cannot ingest a long clip in one shot. This module splits -a long waveform into overlapping fixed-length windows and stitches the resulting -per-window transcripts back into one string — the same job Whisper's internal -``chunk_length_s`` does, but lifted **above** the backend so *any* transcriber -inherits long-audio support. - -Kept pure (numpy + stdlib, no vLLM / torch / transformers import) so it unit- -tests on CPU. The transcriber decides whether to use it via its ``self_chunks`` -flag: backends that already chunk internally (Whisper) bypass this entirely. - -Why overlap + dedup: cutting on hard boundaries can split a word across two -windows, so consecutive windows overlap and each boundary word is whole in at -least one of them. The overlap region is therefore transcribed twice; the merge -finds the repeated span at each seam and keeps one copy. +Splits a waveform into overlapping windows and stitches the per-window +transcripts, so a backend with a fixed input window inherits long-audio support. +Windows overlap because a hard cut can split a word; the overlap is therefore +transcribed twice and the merge keeps one copy. + +Pure numpy + stdlib (no torch/transformers) so it unit-tests on CPU. """ from __future__ import annotations @@ -24,9 +15,7 @@ import numpy as np -# Cap on how many trailing/leading words we search for a seam overlap. Comfortably -# larger than any plausible word count inside a few seconds of overlap, but bounded -# so the merge stays linear in transcript length. +# Bounded seam search so the merge stays linear in transcript length. _MAX_SEAM_WORDS = 60 @@ -36,21 +25,10 @@ def split_waveform( window_s: float, overlap_s: float, ) -> list[np.ndarray]: - """Split ``samples`` into overlapping windows of ``window_s`` seconds. - - Consecutive windows advance by ``window_s - overlap_s`` seconds, so each pair - shares ``overlap_s`` of audio. A clip already shorter than one window is - returned as a single segment (no copying/splitting). - - Args: - samples: 1-D mono waveform. - sr: Sample rate of ``samples`` in Hz. - window_s: Window length in seconds (must be > 0). - overlap_s: Overlap between consecutive windows in seconds (0 <= overlap_s - < window_s). + """Split a 1-D mono waveform into overlapping windows, in order. - Returns: - A list of 1-D numpy views/arrays, in order. + Windows advance by ``window_s - overlap_s``. A clip shorter than one window + is returned as a single segment. """ if window_s <= 0: raise ValueError(f"window_s must be > 0, got {window_s}") @@ -87,9 +65,8 @@ def _norm_word(word: str) -> str: def _seam_overlap(prev: list[str], nxt: list[str]) -> int: """Longest k such that the last k words of ``prev`` match the first k of ``nxt``. - Comparison is punctuation/case-insensitive because ASR often renders the - overlap region slightly differently on each side of the seam. Returns 0 when - there is no matching overlap. + Punctuation/case-insensitive: ASR renders the overlap slightly differently on + each side of a seam. Returns 0 when nothing matches. """ max_k = min(len(prev), len(nxt), _MAX_SEAM_WORDS) for k in range(max_k, 0, -1): @@ -101,12 +78,7 @@ def _seam_overlap(prev: list[str], nxt: list[str]) -> int: def merge_transcripts(transcripts: list[str]) -> str: - """Concatenate per-window transcripts, de-duplicating the overlap at each seam. - - For each new window, find the longest word-level overlap between the tail of - the text so far and the head of the new window, and drop that duplicated span - from the new window before appending. Empty windows are skipped. - """ + """Concatenate per-window transcripts, de-duplicating the overlap at each seam.""" merged: list[str] = [] for text in transcripts: words = text.split() diff --git a/src/granite_switch/vllm/audio/processor.py b/src/granite_switch/vllm/audio/processor.py index 17e39e2..61700c8 100644 --- a/src/granite_switch/vllm/audio/processor.py +++ b/src/granite_switch/vllm/audio/processor.py @@ -1,26 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 -"""vLLM multimodal processor for the alpha audio cascade (Path A). - -Strategy (see also the design notes in :mod:`granite_switch.vllm.audio`): - -* Audio is declared as a multimodal input, so the standard vLLM server accepts - it and developers keep loading a single model. -* In :meth:`_call_hf_processor` we run ASR and tokenize the transcript. -* :meth:`_get_prompt_updates` replaces the single ``<|audio|>`` marker with the - **actual transcript token ids** via ``PromptReplacement``. After this the - transcript tokens are ordinary tokens in ``prompt_token_ids``; the scheduler - sizes KV for the real, runtime-determined length (no fixed audio window). -* The model's ``embed_multimodal`` then returns ``embed_tokens(transcript_ids)`` - for those placeholder positions — identical to what they'd get as normal text. - That redundant-but-consistent step is the exact seam the future projection - model reuses (swap the embedding source for a trained audio encoder). - -Modeled on vLLM 0.19.1's ``ultravox.py`` (the reference audio model), confirmed -against that version's API by the scratch probes. - -ALPHA SCOPE: audio is answered by the base model. We do not place adapter -control tokens for audio requests, so no token-exchange interaction with the -switch is needed here. +"""vLLM multimodal processor for the audio cascade. + +``_call_hf_processor`` runs ASR and tokenizes the transcript; +``_get_prompt_updates`` then replaces the ``<|audio|>`` marker with the real +transcript token ids via ``PromptReplacement``, so the scheduler sizes KV for the +runtime-determined length rather than a fixed audio window. + +Modeled on vLLM 0.19.1's ``ultravox.py``. Audio is answered by the base model — +no adapter control tokens are placed, so the switch is not involved. """ from __future__ import annotations @@ -53,17 +40,10 @@ resolve_generate_kwargs, ) -# The chat-template marker that stands in for an audio clip before replacement. AUDIO_MARKER = "<|audio|>" - -# ASR feature-extractor sample rate the audio is resampled to. _TARGET_SR = 16_000 - -# Fallback context length if the served max_model_len cannot be read (should not -# happen in practice; keeps the transcript budget finite either way). +# Keeps the transcript budget finite if max_model_len cannot be read. _FALLBACK_CONTEXT_LEN = 8192 - -# Dummy clip length (seconds) used during vLLM's profiling run. _DUMMY_AUDIO_SECONDS = 5 @@ -74,14 +54,11 @@ def _asr_enabled(self) -> bool: return bool(getattr(self.get_hf_config(), "asr_enabled", False)) def get_supported_mm_limits(self) -> Mapping[str, int | None]: - # Audio capability is gated per-checkpoint: a non-audio GraniteSwitch - # reports no modalities, so vLLM never profiles audio or loads ASR. + # No modalities on a non-audio checkpoint, so vLLM never loads ASR. if not self._asr_enabled(): return {} - # Configurable ceiling on audio clips per request; each clip's transcript - # is spliced at its own <|audio|> marker. Finite so vLLM can size KV for - # the worst case (clips x per-clip budget). --limit-mm-per-prompt may - # lower this, not raise it above the declared ceiling. + # Finite so vLLM can size KV for the worst case. --limit-mm-per-prompt + # may lower this ceiling, not raise it. return {"audio": self._asr_max_audio_clips()} def get_mm_max_tokens_per_item( @@ -91,16 +68,13 @@ def get_mm_max_tokens_per_item( ) -> Mapping[str, int] | None: if not self._asr_enabled(): return {} - # Worst-case transcript positions one clip can occupy, used only to size - # the encoder cache and profiling pass — NOT to bound requests. A clip - # cannot transcribe to more than the whole context (a longer prompt is - # rejected by vLLM's standard length check), so the per-clip share of the - # context window is the honest upper bound. vLLM calls this with count=1. + # Sizes the encoder cache and profiling pass only — does NOT bound + # requests (vLLM's prompt-length check does). A clip cannot exceed the + # whole context, so the per-clip share of it is the honest upper bound. count = mm_counts.get("audio", 1) or 1 return {"audio": max(1, seq_len // count)} def get_data_parser(self) -> MultiModalDataParser: - # Resample incoming audio to the ASR sample rate. return MultiModalDataParser(target_sr=_TARGET_SR) # --- ASR config resolved from the model's GraniteSwitchConfig --- @@ -113,6 +87,10 @@ def _asr_device(self) -> str: cfg = self.get_hf_config() return getattr(cfg, "asr_device", "cpu") or "cpu" + def _asr_dtype(self) -> str | None: + cfg = self.get_hf_config() + return getattr(cfg, "asr_dtype", None) + def _asr_pipeline_kwargs(self) -> Mapping[str, object]: cfg = self.get_hf_config() return getattr(cfg, "asr_pipeline_kwargs", None) or {} @@ -140,11 +118,8 @@ def _asr_chunk_overlap_s(self) -> float: def _max_model_len(self) -> int: """The served context window, or a safe fallback. - vLLM passes ``seq_len`` into ``get_mm_max_tokens_per_item`` for profiling, - but ``_call_hf_processor`` needs the served context at request time to size - the transcript budget, so read it from the processing context's model - config (falling back to the checkpoint's max_position_embeddings, then a - constant). + ``_call_hf_processor`` needs this at request time to size the transcript + budget; vLLM's profiling ``seq_len`` is not available there. """ model_config = getattr(self.ctx, "model_config", None) max_len = getattr(model_config, "max_model_len", None) @@ -183,19 +158,13 @@ def _transcribe( audio, generate_kwargs: Mapping[str, object] | None = None, ) -> list[int]: - """Transcribe one audio item to a list of token ids (full transcript). - - The transcript is spliced into the prompt as ordinary text tokens; it is - never truncated here. A clip whose transcript makes the prompt exceed the - context is rejected by vLLM's standard prompt-length check. - - Long audio is handled by the backend itself (Whisper) or by our - encoder-agnostic chunker, per the checkpoint's ``asr_self_chunks`` flag. - """ + """Transcribe one audio item to token ids. Never truncated here — an + oversized prompt is rejected by vLLM's own length check.""" transcriber = get_transcriber( model_id=self.info._asr_model_id(), device=self.info._asr_device(), pipeline_kwargs=self.info._asr_pipeline_kwargs(), + dtype=self.info._asr_dtype(), ) # The data parser already resampled to _TARGET_SR. text = transcriber.transcribe( @@ -219,13 +188,11 @@ def _call_hf_processor( tokenizer = self.info.get_tokenizer() audios = mm_data.get("audios", []) or [] - # Text-only request: just tokenize (mirrors Ultravox's text-only branch). if not audios: input_ids = tokenizer.encode(prompt, add_special_tokens=False) return BatchFeature(dict(input_ids=[input_ids]), tensor_type="pt") - # Resolve decode kwargs once (config defaults + allowlisted per-request - # overrides) and apply them to every audio item in this request. + # Resolved once, then applied to every audio item in this request. generate_kwargs = resolve_generate_kwargs( self.info._asr_generate_kwargs(), mm_kwargs, @@ -234,9 +201,7 @@ def _call_hf_processor( input_ids = tokenizer.encode(prompt, add_special_tokens=False) - # Transcribe each audio to token ids; concatenate flat with per-item sizes. - # Transcripts are spliced in full — an oversized request is rejected by - # vLLM's prompt-length check, not silently truncated here. + # Concatenated flat, with per-item sizes to split them back. per_item_ids = [self._transcribe(a, generate_kwargs) for a in audios] sizes = [len(ids) for ids in per_item_ids] flat_ids = [tid for ids in per_item_ids for tid in ids] @@ -257,7 +222,6 @@ def _get_mm_fields_config( ) -> Mapping[str, MultiModalFieldConfig]: num_tokens = hf_inputs.get("audio_num_tokens", torch.zeros(0)) return dict( - # Flat transcript ids, split back per audio by audio_num_tokens. audio_token_ids=MultiModalFieldConfig.flat_from_sizes("audio", num_tokens), audio_num_tokens=MultiModalFieldConfig.batched("audio"), ) @@ -277,7 +241,6 @@ def _get_prompt_updates( def replacement(item_idx: int): s = int(starts[item_idx]) e = int(starts[item_idx + 1]) - # Replace <|audio|> with the real transcript token ids. return [int(t) for t in all_ids[s:e]] return [ diff --git a/src/granite_switch/vllm/granite_switch_model.py b/src/granite_switch/vllm/granite_switch_model.py index 9467846..e24466a 100644 --- a/src/granite_switch/vllm/granite_switch_model.py +++ b/src/granite_switch/vllm/granite_switch_model.py @@ -299,9 +299,8 @@ def forward( adapter_token_ids=self.adapter_token_ids, ) else: - # Either no switch, or the multimodal path where vLLM passes - # pre-merged inputs_embeds and input_ids is None. ALPHA: audio - # requests run on the base model (adapter_id = 0). + # No switch, or the multimodal path (input_ids is None because + # vLLM pre-merged inputs_embeds): run on base, adapter_id 0. if input_ids is not None: num_tokens = input_ids.shape[0] device = input_ids.device @@ -429,10 +428,8 @@ class GraniteSwitchForCausalLM( supports_multimodal = True - # Force vLLM to pass the raw input_ids to forward() on the multimodal path - # (not just precomputed inputs_embeds). The switch needs them to detect - # adapter control tokens, so audio requests route through adapters exactly - # like text requests. + # Without this, the multimodal path passes only inputs_embeds and the switch + # cannot see control tokens — audio requests would bypass adapters. requires_raw_input_tokens = True @classmethod diff --git a/tests/integration/test_adapter_routing_audio_enabled.py b/tests/integration/test_adapter_routing_audio_enabled.py index 131a01f..7bbc63b 100644 --- a/tests/integration/test_adapter_routing_audio_enabled.py +++ b/tests/integration/test_adapter_routing_audio_enabled.py @@ -2,14 +2,10 @@ """Adapters + audio (issue #47): every default adapter routes to its own index on an audio-enabled checkpoint. -Composes the default adapter set (RAG + Core + Guardian) with --enable-audio and -sweeps every adapter control token that resolved, asserting the switch maps -``adapter_token_ids[i]`` to index ``i + 1`` and leaves pre-control positions on -the base (``0``) — i.e. enabling audio does not perturb the token->index map. -For granite-4.1-3b that is 12 adapters: context_relevance ships no 4.1-3b flavor, -so 12 of the 13 defined adapters resolve. Complements the single-adapter sweep in -test_switch_e2e_compose and the serve-time answerability check in -test_answerability_over_audio. +Asserts the switch maps ``adapter_token_ids[i]`` to index ``i + 1`` and leaves +pre-control positions on base (``0``) — enabling audio must not perturb the +token->index map. For granite-4.1-3b, 12 of the 13 defined adapters resolve +(context_relevance ships no 4.1-3b flavor). Markers: slow + requires_model + gpu (opt-in via -m). """ @@ -26,8 +22,6 @@ pytest.skip("requires the HF backend ([hf] extra)", allow_module_level=True) -# The default adapter set is the union of the three granitelib libraries (the -# "12 adapters"): RAG + Core + Guardian. _DEFAULT_ADAPTER_LIBRARIES = [ "ibm-granite/granitelib-rag-r1.0", "ibm-granite/granitelib-core-r1.0", diff --git a/tests/integration/test_answerability_over_audio.py b/tests/integration/test_answerability_over_audio.py index e59f11b..538457b 100644 --- a/tests/integration/test_answerability_over_audio.py +++ b/tests/integration/test_answerability_over_audio.py @@ -2,18 +2,14 @@ """Adapters + audio: the RAG answerability adapter judging a query against a document delivered as audio (issue #47). -Covers the "adapters + audio" box. Where test_audio_serving_smoke.py only proves -an adapter control token doesn't crash the audio path, this checks the switch -reaches the *correct* verdict when the RAG document arrives as speech. - -Document-as-audio: ``documents=[{"text": "<|audio|>"}]`` — the Granite template -renders each document with ``doc | tojson``, so the marker lands in the -```` block where the ASR processor splices the transcript. The -``<|answerability|>`` control token is inserted by the same template (aLoRA -fallback path). Ground truth is one committed speech clip (fixtures/, generated -with SpeechT5 — MIT) driving both classes; the questions key on content that -transcribes cleanly, so this tests the decision, not transcription accuracy -(WER). +Unlike test_audio_serving_smoke.py, which only proves an adapter control token +doesn't crash the audio path, this checks the switch reaches the *correct* +verdict when the document arrives as speech. + +``documents=[{"text": "<|audio|>"}]`` works because the Granite template renders +each document with ``doc | tojson``, so the marker lands in the ```` +block for the ASR processor to splice. Questions key on words that transcribe +cleanly, so this tests the decision, not WER. Markers: slow + requires_model + gpu (opt-in via -m). """ @@ -37,10 +33,7 @@ def _load_experimental_pairs(): - """Extra (base, adapter) pairs from GRANITE_SWITCH_EXPERIMENTAL_MODEL_PAIRS. - - JSON array of {"base": str, "adapter": str}; mirrors the other E2E files. - """ + """Extra (base, adapter) pairs from GRANITE_SWITCH_EXPERIMENTAL_MODEL_PAIRS.""" raw = os.environ.get("GRANITE_SWITCH_EXPERIMENTAL_MODEL_PAIRS", "") if not raw: return [] @@ -134,8 +127,7 @@ def served(audio_rag_checkpoint): def _answerability_adapter_name(config): - # Discovery derives the name from the library layout — match rather than - # hard-code, and skip loudly if this checkpoint has no answerability adapter. + # Discovery derives the name from the library layout, so match, don't hard-code. names = list(getattr(config, "adapter_names", None) or []) for name in names: if "answerab" in name.lower(): @@ -155,8 +147,6 @@ def speech_clip(): def _build_answerability_prompt(tokenizer, adapter_name, question): - # documents=[{"text": "<|audio|>"}] places one audio marker in the template's - # block; adapter_name arms the answerability control-token insert. return tokenizer.apply_chat_template( [{"role": "user", "content": question}], documents=[{"text": _AUDIO_MARKER}], @@ -167,8 +157,7 @@ def _build_answerability_prompt(tokenizer, adapter_name, question): def _parse_label(text): - # Adapter emits the bare enum ("answerable"/"unanswerable"), maybe quoted. - # Match the longer label first so "answerable" doesn't shadow "unanswerable". + # Longer label first, so "answerable" doesn't shadow "unanswerable". norm = text.strip().strip('"').strip().lower() if "unanswerable" in norm: return "unanswerable" diff --git a/tests/integration/test_audio_serving_smoke.py b/tests/integration/test_audio_serving_smoke.py index b1eb3a6..3c0e9c7 100644 --- a/tests/integration/test_audio_serving_smoke.py +++ b/tests/integration/test_audio_serving_smoke.py @@ -1,29 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 """End-to-end vLLM serving smoke for the audio cascade (issue #47). -Boots a real composed, audio-enabled GraniteSwitch checkpoint under vLLM and -drives the three request shapes the #47 checklist calls out as the serving -smoke — **text + adapter + audio x1/x2/x3** — through one live engine. This is -the path the low-level tests deliberately bypass (CUDA graphs, torch.compile, -the ASR processor running inside vLLM's EngineCore subprocess, the multi-clip -splice, and generation), so it is the only test that proves those pieces -compose at serve time. - -Scope is a *smoke*: each request must complete and return well-formed output, -and multi-clip requests must be accepted up to the checkpoint's declared clip -ceiling. It intentionally does NOT assert transcript content — transcription -quality (WER) and adapter-routing correctness are separate #47 boxes covered by -the eval harness and `test_switch_e2e_compose.py` respectively. Synthetic audio -keeps the test asset-free; a silent/tonal clip exercises the serving plumbing -without shipping a speech fixture. - -Model construction goes through the compose CLI (CLAUDE.md gotcha #5): no test -hand-assembles a config. Audio is enabled with `--enable-audio`, which defaults -the ASR front-end to the small built-in model (distil-whisper/distil-small.en) -to keep CI cost down. - -Markers: @pytest.mark.slow + @pytest.mark.requires_model + @pytest.mark.gpu. -CI must opt in explicitly: `pytest -m "slow and requires_model and gpu"`. +Drives text + adapter + audio x1/x2/x3 through one live engine — the serve-time +path the low-level tests bypass (CUDA graphs, the ASR processor inside vLLM's +EngineCore subprocess, the multi-clip splice, generation). + +Deliberately does NOT assert transcript content: WER and adapter-routing +correctness are separate boxes, covered by the eval harness and +test_switch_e2e_compose.py. Synthetic tones keep the test asset-free. + +Opt in explicitly: `pytest -m "slow and requires_model and gpu"`. """ import importlib.util @@ -38,13 +24,8 @@ pytest.skip("requires vLLM installed", allow_module_level=True) -# ---------------------------------------------------------------------------- -# Base-model / adapter-library pairs — kept in lockstep with -# tests/integration/test_switch_e2e_compose.py so the two E2E files exercise -# the same model matrix. A fast CI profile can pin one pair via -k or the -# experimental env var below. -# ---------------------------------------------------------------------------- - +# Kept in lockstep with test_switch_e2e_compose.py so both E2E files exercise the +# same model matrix. _DEFAULT_BASE_MODEL_PAIRS = [ ("ibm-granite/granite-4.0-micro", "ibm-granite/granitelib-core-r1.0"), ("ibm-granite/granite-4.1-3b", "ibm-granite/granitelib-core-r1.0"), @@ -52,11 +33,10 @@ def _load_experimental_pairs(): - """Local/experimental pairings from GRANITE_SWITCH_EXPERIMENTAL_MODEL_PAIRS. + """Local pairings from GRANITE_SWITCH_EXPERIMENTAL_MODEL_PAIRS. - JSON array of {"base": str, "adapter": str}; base/adapter may be HF ids or - local paths. The mechanism is committed; the values are not. Mirrors - test_switch_e2e_compose.py so both E2E files share one extension knob. + JSON array of {"base": str, "adapter": str}, HF ids or local paths. The + mechanism is committed; the values are not. """ raw = os.environ.get("GRANITE_SWITCH_EXPERIMENTAL_MODEL_PAIRS", "") if not raw: @@ -85,9 +65,8 @@ def _load_experimental_pairs(): def audio_checkpoint(request, tmp_path_factory): """Compose one audio-enabled checkpoint per (base, adapter) pair. - Delegates to the compose CLI (same as test_switch_e2e_compose.py) with - `--enable-audio`, then returns the save dir. Module scope amortizes the - download-dominated first run across every smoke case for the pair. + Goes through the compose CLI, never a hand-assembled config (CLAUDE.md + gotcha #5). Module scope amortizes the download across the pair's cases. """ import subprocess import sys @@ -123,8 +102,7 @@ def served(audio_checkpoint): """Boot vLLM once for the checkpoint and share it across smoke cases. Tokenizer init stays ON (unlike the argmax-equivalence test): the ASR - processor needs a tokenizer to encode both the prompt and the transcript, - and the adapter case needs it to tokenize text around the control token. + processor needs it to encode the prompt and the transcript. """ import gc @@ -177,9 +155,8 @@ def test_text_only_serving(served): def test_adapter_control_token_serving(served): """An adapter control token routes through the switch under serving. - Correctness of the routing is `test_switch_e2e_compose.py`'s job; here we - only prove the switch path runs end-to-end in the live engine without - crashing and still generates. + Routing correctness is test_switch_e2e_compose.py's job; the bar here is + that the switch path runs in the live engine and still generates. """ from vllm import SamplingParams from vllm.inputs import TokensPrompt @@ -203,9 +180,7 @@ def test_adapter_control_token_serving(served): def test_audio_clip_serving(served, num_clips): """Audio x1/x2/x3: N markers + N clips transcribe, splice, and generate. - Exercises the multi-clip long-audio serving path through the ASR processor - running inside vLLM's engine subprocess. Content is not asserted (WER is a - separate #47 box); the bar is a completed request with well-formed output. + Content is not asserted; the bar is a completed, well-formed request. """ from vllm import SamplingParams diff --git a/tests/unit/test_asr.py b/tests/unit/test_asr.py index d5e00e0..69187ce 100644 --- a/tests/unit/test_asr.py +++ b/tests/unit/test_asr.py @@ -7,6 +7,7 @@ leaf module directly by file path rather than through the package. """ +import contextlib import importlib.util import pathlib from unittest import mock @@ -282,15 +283,92 @@ def test_config_not_mutated(self): assert cfg == {"language": "de"} +@contextlib.contextmanager +def _patched_pipeline(factory): + """Patch both lookup paths: `from transformers import pipeline` re-resolves to + transformers.pipelines.pipeline, but transformers caches it on the top-level + module after the first access, so a later test would get the real one.""" + with ( + mock.patch("transformers.pipelines.pipeline", factory), + mock.patch("transformers.pipeline", factory), + ): + yield + + +class TestResolveTorchDtype: + """asr_dtype resolution. float16-on-CUDA is the default, but overridable.""" + + def test_auto_on_cuda_is_float16(self): + torch = pytest.importorskip("torch") + assert asr._resolve_torch_dtype(None, "cuda:0") is torch.float16 + assert asr._resolve_torch_dtype("auto", "cuda") is torch.float16 + + def test_auto_on_cpu_is_float32(self): + torch = pytest.importorskip("torch") + assert asr._resolve_torch_dtype(None, "cpu") is torch.float32 + + def test_explicit_float32_overrides_cuda_default(self): + # The BatchNorm-encoder case: CUDA must not force half precision. + torch = pytest.importorskip("torch") + assert asr._resolve_torch_dtype("float32", "cuda:0") is torch.float32 + + def test_explicit_bfloat16(self): + torch = pytest.importorskip("torch") + assert asr._resolve_torch_dtype("bfloat16", "cpu") is torch.bfloat16 + + @pytest.mark.parametrize( + "name,attr", + [ + ("fp16", "float16"), + ("half", "float16"), + ("bf16", "bfloat16"), + ("fp32", "float32"), + ("FLOAT32", "float32"), + ], + ) + def test_aliases_and_case(self, name, attr): + torch = pytest.importorskip("torch") + assert asr._resolve_torch_dtype(name, "cpu") is getattr(torch, attr) + + def test_unknown_name_raises(self): + pytest.importorskip("torch") + with pytest.raises(ValueError, match="Unsupported asr_dtype"): + asr._resolve_torch_dtype("int8", "cpu") + + def test_dtype_is_part_of_cache_key(self): + a = asr.get_transcriber("dt", "cuda:0", dtype="float32") + b = asr.get_transcriber("dt", "cuda:0", dtype="float16") + c = asr.get_transcriber("dt", "cuda:0", dtype="float32") + assert a is not b + assert a is c + + def test_load_passes_resolved_dtype(self): + torch = pytest.importorskip("torch") + factory = mock.Mock(return_value=mock.Mock()) + with _patched_pipeline(factory): + asr.ASRTranscriber(model_id="m", device="cuda:0", dtype="float32").load() + assert factory.call_args.kwargs["torch_dtype"] is torch.float32 + + def test_pipeline_kwargs_torch_dtype_still_wins(self): + torch = pytest.importorskip("torch") + factory = mock.Mock(return_value=mock.Mock()) + with _patched_pipeline(factory): + asr.ASRTranscriber( + model_id="m", + device="cpu", + dtype="float32", + pipeline_kwargs={"torch_dtype": torch.bfloat16}, + ).load() + assert factory.call_args.kwargs["torch_dtype"] is torch.bfloat16 + + class TestLoadMergesPipelineKwargs: def test_pipeline_kwargs_override_defaults(self): # load() must merge config-supplied pipeline_kwargs over the built-in # defaults (e.g. override chunk_length_s, add extra kwargs). pytest.importorskip("torch") fake_pipe_factory = mock.Mock(return_value=mock.Mock()) - # transformers is a lazy module: `from transformers import pipeline` - # re-resolves to transformers.pipelines.pipeline, so patch there. - with mock.patch("transformers.pipelines.pipeline", fake_pipe_factory): + with _patched_pipeline(fake_pipe_factory): t = asr.ASRTranscriber( model_id="m", device="cpu", diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index f25e2a7..6252544 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -106,9 +106,21 @@ def test_asr_defaults_off(self): assert cfg.asr_enabled is False assert cfg.asr_model_id is None assert cfg.asr_device == "cpu" + assert cfg.asr_dtype is None assert cfg.asr_pipeline_kwargs is None assert cfg.asr_generate_kwargs is None + def test_asr_dtype_round_trip(self, tmp_path): + cfg = GraniteSwitchConfig( + num_adapters=0, asr_enabled=True, asr_device="cuda:0", asr_dtype="float32" + ) + cfg.save_pretrained(tmp_path) + assert GraniteSwitchConfig.from_pretrained(tmp_path).asr_dtype == "float32" + + def test_invalid_asr_dtype_raises(self): + with pytest.raises(ValueError, match="asr_dtype"): + GraniteSwitchConfig(num_adapters=0, asr_dtype="fp8") + def test_longaudio_defaults(self): cfg = GraniteSwitchConfig(num_adapters=0) assert cfg.asr_max_audio_clips == 32 diff --git a/tests/vllm/test_audio_processor.py b/tests/vllm/test_audio_processor.py index 2d81f3d..b158b27 100644 --- a/tests/vllm/test_audio_processor.py +++ b/tests/vllm/test_audio_processor.py @@ -1,21 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 """vLLM-tier tests for the audio ASR multimodal processor. -These exercise the code that lives in ``granite_switch.vllm.audio.processor`` -and therefore needs vLLM importable (the base classes come from vLLM). They do -NOT need a GPU or a real ASR model: the transcriber is faked, so what is under -test is our plumbing — - -- per-checkpoint gating of the audio modality on ``asr_enabled``, -- the config accessors for ``asr_model_id`` / ``asr_device`` / - ``asr_pipeline_kwargs`` / ``asr_generate_kwargs``, -- and the Level-2 seam: config-default decode kwargs merged with an allowlisted - per-request override (vLLM's ``mm_processor_kwargs``) reaching the transcriber. - -The pure merge/kwargs logic itself is unit-tested on CPU in -``tests/unit/test_asr.py``; this file checks it is wired through the processor. - -Requires vLLM installed. Skipped otherwise (e.g. CPU-only dev machines). +Needs vLLM importable (the base classes come from it) but no GPU and no real ASR +model — the transcriber is faked, so what is under test is the plumbing: modality +gating on ``asr_enabled``, the config accessors, and decode kwargs reaching the +transcriber. The merge logic itself is unit-tested in tests/unit/test_asr.py. """ import importlib.util @@ -43,9 +32,7 @@ def _make_info(*, max_model_len=131072, **cfg_attrs): """A ProcessingInfo whose get_hf_config() returns a stub config. - Bypasses __init__ (which needs a full vLLM InputProcessingContext) — we only - exercise methods that read the config. A stub ``ctx`` supplies the served - ``max_model_len`` used to size the transcript budget. + Bypasses __init__, which needs a full vLLM InputProcessingContext. """ info = object.__new__(GraniteSwitchASRProcessingInfo) cfg = SimpleNamespace(**cfg_attrs) @@ -167,10 +154,13 @@ def transcribe( capture["chunk_overlap_s"] = chunk_overlap_s return "hello world" - def fake_get_transcriber(model_id=None, device="cpu", pipeline_kwargs=None): + def fake_get_transcriber( + model_id=None, device="cpu", pipeline_kwargs=None, dtype=None + ): capture["model_id"] = model_id capture["device"] = device capture["pipeline_kwargs"] = pipeline_kwargs + capture["dtype"] = dtype return FakeTranscriber() monkeypatch.setattr(proc_mod, "get_transcriber", fake_get_transcriber) @@ -196,6 +186,19 @@ def test_transcribe_forwards_pipeline_and_generate_kwargs(self, monkeypatch): assert capture["generate_kwargs"] == {"language": "fr"} assert capture["sampling_rate"] == proc_mod._TARGET_SR + def test_dtype_forwarded_from_config(self, monkeypatch): + capture = {} + info = _make_info(asr_enabled=True, asr_device="cuda:0", asr_dtype="float32") + proc = _make_processor(info, monkeypatch, capture) + proc._transcribe(np.zeros(1600, dtype=np.float32), {}) + assert capture["dtype"] == "float32" + + def test_dtype_defaults_to_none(self, monkeypatch): + capture = {} + proc = _make_processor(_make_info(asr_enabled=True), monkeypatch, capture) + proc._transcribe(np.zeros(1600, dtype=np.float32), {}) + assert capture["dtype"] is None + def test_empty_generate_kwargs_becomes_none(self, monkeypatch): capture = {} info = _make_info(asr_enabled=True, asr_model_id="w", asr_device="cpu") @@ -340,7 +343,9 @@ def transcribe( monkeypatch.setattr( proc_mod, "get_transcriber", - lambda model_id=None, device="cpu", pipeline_kwargs=None: FakeTranscriber(), + lambda model_id=None, device="cpu", pipeline_kwargs=None, dtype=None: ( + FakeTranscriber() + ), ) return proc @@ -410,7 +415,9 @@ def transcribe(self, audio, **kw): monkeypatch.setattr( proc_mod, "get_transcriber", - lambda model_id=None, device="cpu", pipeline_kwargs=None: FakeTranscriber(), + lambda model_id=None, device="cpu", pipeline_kwargs=None, dtype=None: ( + FakeTranscriber() + ), ) bf = proc._call_hf_processor(