A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 19:45 UTC
[Model] Extend Ultravox to accept audio longer than 30s (#13631)
Signed-off-by: Farzad Abdolhosseini <farzad@fixie.ai>
Farzad Abdolhosseini committed
Mar 11, 2025 at 19:27 UTC
80e78d02aceaf55528926c800e89568d1358e1dc
4 files changed
+233
-81
tests/models/decoder_only/audio_language/test_ultravox.py
+1
-1
@@ -15,7 +15,7 @@ from ....conftest import HfRunner, VllmRunner
15
from ....utils import RemoteOpenAIServer
16
from ...utils import check_logprobs_close
17
18
-MODEL_NAME = "fixie-ai/ultravox-v0_4"
18
+MODEL_NAME = "fixie-ai/ultravox-v0_5-llama-3_2-1b"
19
20
AudioTuple = tuple[np.ndarray, int]
21
tests/models/multimodal/processing/test_common.py
+50
-7
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
+import copy
4
from functools import partial
5
+from typing import Optional
6
7
import numpy as np
8
import pytest
@@ -21,6 +23,7 @@ def _test_processing_correctness(
23
hit_rate: float,
24
num_batches: int,
25
simplify_rate: float,
26
+ ignore_mm_keys: Optional[list[str]] = None,
27
):
28
model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)
29
model_info.check_available_online(on_fail="skip")
@@ -123,8 +126,10 @@ def _test_processing_correctness(
126
hf_processor_mm_kwargs={},
127
)
128
126
- assert baseline_result == cached_result, (
127
- f"Failed ({batch_idx=}, {prompt=}, {mm_data=})")
129
+ assert _drop_mm_kwargs_keys(
130
+ baseline_result, ignore_mm_keys) == _drop_mm_kwargs_keys(
131
+ cached_result, ignore_mm_keys), (
132
+ f"Failed ({batch_idx=}, {prompt=}, {mm_data=})")
133
134
baseline_tokenized_result = baseline_processor.apply(
135
tokenizer.encode(prompt, **tokenizer_encode_kwargs),
@@ -132,8 +137,10 @@ def _test_processing_correctness(
137
hf_processor_mm_kwargs={},
138
)
139
135
- assert baseline_result == baseline_tokenized_result, (
136
- f"Failed ({batch_idx=}, {prompt=}, {mm_data=})")
140
+ assert _drop_mm_kwargs_keys(
141
+ baseline_result, ignore_mm_keys) == _drop_mm_kwargs_keys(
142
+ baseline_tokenized_result, ignore_mm_keys), (
143
+ f"Failed ({batch_idx=}, {prompt=}, {mm_data=})")
144
145
cached_tokenized_result = cached_processor.apply(
146
tokenizer.encode(prompt, **tokenizer_encode_kwargs),
@@ -141,8 +148,10 @@ def _test_processing_correctness(
148
hf_processor_mm_kwargs={},
149
)
150
144
- assert cached_result == cached_tokenized_result, (
145
- f"Failed ({batch_idx=}, {prompt=}, {mm_data=})")
151
+ assert _drop_mm_kwargs_keys(
152
+ cached_result, ignore_mm_keys) == _drop_mm_kwargs_keys(
153
+ cached_tokenized_result, ignore_mm_keys), (
154
+ f"Failed ({batch_idx=}, {prompt=}, {mm_data=})")
155
156
157
# yapf: disable
@@ -173,7 +182,7 @@ def _test_processing_correctness(
182
"Qwen/Qwen2-VL-2B-Instruct",
183
"Qwen/Qwen2.5-VL-3B-Instruct",
184
"Qwen/Qwen2-Audio-7B-Instruct",
176
- "fixie-ai/ultravox-v0_4",
185
+ "fixie-ai/ultravox-v0_5-llama-3_2-1b",
186
"openai/whisper-large-v3",
187
"google/paligemma-3b-mix-224",
188
"google/paligemma2-3b-ft-docci-448",
@@ -188,11 +197,19 @@ def test_processing_correctness(
197
num_batches: int,
198
simplify_rate: float,
199
):
200
+ ignore_mm_keys = None
201
+ if 'ultravox' in model_id:
202
+ # In Ultravox, the audio_features can be different depending on padding
203
+ # The slight difference should not be a problem though, since
204
+ # attention_mask lets us ignore the difference.
205
+ ignore_mm_keys = ['audio_features']
206
+
207
_test_processing_correctness(
208
model_id,
209
hit_rate=hit_rate,
210
num_batches=num_batches,
211
simplify_rate=simplify_rate,
212
+ ignore_mm_keys=ignore_mm_keys,
213
)
214
215
@@ -221,3 +238,29 @@ def test_processing_correctness_phi3v(
238
num_batches=num_batches,
239
simplify_rate=simplify_rate,
240
)
241
+
242
+
243
+def _drop_mm_kwargs_keys(result: dict,
244
+ ignore_mm_keys: Optional[list[str]] = None) -> dict:
245
+ """Drop specified keys from result['mm_kwargs'].
246
+
247
+ This is mainly to avoid doing exact match of audio_features in ultravox.
248
+
249
+ Args:
250
+ result: Result to drop keys from
251
+ ignore_mm_keys: List of keys to ignore, e.g. ['audio_features']
252
+ """
253
+ if not ignore_mm_keys:
254
+ return result
255
+
256
+ if 'mm_kwargs' in result:
257
+ result = copy.deepcopy(result)
258
+ mm_kwargs = result['mm_kwargs']
259
+ for key in ignore_mm_keys:
260
+ mm_kwargs.pop(key, None)
261
+ for items in mm_kwargs._items_by_modality.values():
262
+ for item in items:
263
+ for key in ignore_mm_keys:
264
+ item.pop(key, None)
265
+
266
+ return result
tests/models/registry.py
+1
-2
@@ -284,8 +284,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
284
"Qwen2VLForConditionalGeneration": _HfExamplesInfo("Qwen/Qwen2-VL-2B-Instruct"), # noqa: E501
285
"Qwen2_5_VLForConditionalGeneration": _HfExamplesInfo("Qwen/Qwen2.5-VL-3B-Instruct", # noqa: E501
286
min_transformers_version="4.49"), # noqa: E501
287
- "UltravoxModel": _HfExamplesInfo("fixie-ai/ultravox-v0_4",
288
- extras={"v0.5": "fixie-ai/ultravox-v0_5-llama-3_2-1b"}, # noqa: E501
287
+ "UltravoxModel": _HfExamplesInfo("fixie-ai/ultravox-v0_5-llama-3_2-1b", # noqa: E501
288
trust_remote_code=True),
289
# [Encoder-decoder]
290
# Florence-2 uses BartFastTokenizer which can't be loaded from AutoTokenizer
vllm/model_executor/models/ultravox.py
+181
-71
@@ -5,7 +5,7 @@
5
import math
6
from collections.abc import Iterable, Mapping, Sequence
7
from functools import cached_property
8
-from typing import Any, Literal, Optional, Set, Tuple, TypedDict, Union
8
+from typing import Any, List, Literal, Optional, Set, Tuple, TypedDict, Union
9
10
import torch
11
import torch.utils.checkpoint
@@ -44,12 +44,23 @@ from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
44
_AUDIO_PLACEHOLDER_OVERRIDE = "<|reserved_special_token_0|>"
45
_AUDIO_PLACEHOLDER_TOKEN = 128002
46
_AUDIO_TOKENS_PER_SECOND = 6.25
47
+_MAX_ENCODER_BATCH_SIZE = 16
48
49
50
class UltravoxAudioFeatureInputs(TypedDict):
51
type: Literal["audio_features"]
52
data: NestedTensors
52
- """Shape: `(batch_size, num_audios, 80, M)`"""
53
+ """Shape: `(batch_size, num_chunks, 80, M)`"""
54
+ lens: NestedTensors
55
+ """
56
+ Length of the audio frames. Used for attention mask in WhisperEncoder.
57
+ Shape: `(batch_size, num_chunks)`
58
+ """
59
+ token_len: NestedTensors
60
+ """
61
+ Length of the audio tokens. Used for flattening the audio features.
62
+ Shape: `(batch_size, num_chunks)`
63
+ """
64
65
66
class UltravoxAudioEmbeddingInputs(TypedDict):
@@ -78,6 +89,7 @@ class UltravoxProcessingInfo(BaseProcessingInfo):
89
# token, thus we override placeholder with a reserved special
90
# token.
91
hf_processor.audio_token_replacement = _AUDIO_PLACEHOLDER_OVERRIDE
92
+ hf_processor.audio_replacement_token_id = _AUDIO_PLACEHOLDER_TOKEN
93
return hf_processor
94
95
def get_feature_extractor(
@@ -104,7 +116,7 @@ class UltravoxProcessingInfo(BaseProcessingInfo):
116
max_audio_tokens = math.ceil(feature_extractor.chunk_length *
117
_AUDIO_TOKENS_PER_SECOND)
118
107
- return {"audio": max_audio_tokens}
119
+ return {"audio": max_audio_tokens * _MAX_ENCODER_BATCH_SIZE}
120
121
122
class UltravoxDummyInputsBuilder(BaseDummyInputsBuilder[UltravoxProcessingInfo]
@@ -118,7 +130,8 @@ class UltravoxDummyInputsBuilder(BaseDummyInputsBuilder[UltravoxProcessingInfo]
130
feature_extractor = self.info.get_feature_extractor()
131
132
sampling_rate = feature_extractor.sampling_rate
121
- audio_len = feature_extractor.chunk_length * sampling_rate
133
+ audio_len = (feature_extractor.chunk_length * sampling_rate *
134
+ _MAX_ENCODER_BATCH_SIZE)
135
num_audios = mm_counts.get("audio", 0)
136
137
mm_data = {
@@ -160,41 +173,38 @@ class UltravoxMultiModalProcessor(
173
mm_kwargs = dict(
174
**mm_kwargs,
175
sampling_rate=feature_extractor.sampling_rate,
176
+ include_audio_num_chunks=True,
177
)
178
165
- # Ultravox processor doesn't support multiple inputs,
166
- # therefore we need to input text and audio one by one
167
- audio_features, audio_token_len = [], []
168
- shared_outputs = {}
169
- for audio in audios:
170
- # NOTE: Ultravox processor accepts "audio" instead of "audios"
171
- item_processor_data = dict(**mm_data, audio=audio)
172
-
173
- item_outputs = super()._call_hf_processor(
174
- prompt=prompt,
175
- mm_data=item_processor_data,
176
- mm_kwargs=mm_kwargs,
177
- )
178
-
179
- audio_features.append(item_outputs.pop("audio_values")[0])
180
- audio_token_len.append(item_outputs.pop("audio_token_len").item())
181
- shared_outputs = item_outputs
179
+ item_processor_data = dict(**mm_data, audios=audios)
180
183
- combined_outputs = dict(
184
- **shared_outputs,
185
- audio_features=audio_features,
186
- audio_token_len=audio_token_len,
181
+ output = super()._call_hf_processor(
182
+ prompt=prompt,
183
+ mm_data=item_processor_data,
184
+ mm_kwargs=mm_kwargs,
185
)
188
- return BatchFeature(combined_outputs)
186
+ output['audio_features'] = output.pop('audio_values')
187
+
188
+ return output
189
190
def _get_mm_fields_config(
191
self,
192
hf_inputs: BatchFeature,
193
hf_processor_mm_kwargs: Mapping[str, object],
194
) -> Mapping[str, MultiModalFieldConfig]:
195
+ num_chunks = hf_inputs.get('audio_num_chunks', torch.zeros(0))
196
return dict(
196
- audio_features=MultiModalFieldConfig.batched("audio"),
197
- audio_token_len=MultiModalFieldConfig.batched("audio"),
197
+ # to handle longer than 30s audio, each audio might be split
198
+ # into multiple chunks as such, their batch dimension can be
199
+ # higher than the number of audio samples
200
+ audio_features=MultiModalFieldConfig.flat_from_sizes(
201
+ "audio", num_chunks),
202
+ audio_token_len=MultiModalFieldConfig.flat_from_sizes(
203
+ "audio", num_chunks),
204
+ audio_lens=MultiModalFieldConfig.flat_from_sizes(
205
+ "audio", num_chunks),
206
+ # num_chunks can convert audio_chunked to audio batch dimension
207
+ audio_num_chunks=MultiModalFieldConfig.batched("audio"),
208
audio_embeds=MultiModalFieldConfig.batched("audio"),
209
)
210
@@ -205,14 +215,23 @@ class UltravoxMultiModalProcessor(
215
out_mm_kwargs: MultiModalKwargs,
216
) -> Sequence[PromptUpdate]:
217
hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
208
- tokenizer = self.info.get_tokenizer()
209
- vocab = tokenizer.get_vocab()
218
211
- replacement_id = vocab[
212
- hf_processor.audio_token_replacement] # type: ignore
219
+ replacement_id = hf_processor.audio_replacement_token_id # type: ignore
220
+
221
+ # Each audio can be split into multiple chunks.
222
+ # chunks_start_idx[i] indicates the start index of the chunks
223
+ # belonging to the i-th audio.
224
+ num_chunks = out_mm_kwargs.get("audio_num_chunks", torch.zeros(0))
225
+ chunks_start_idx: torch.Tensor = torch.cumsum(num_chunks,
226
+ dim=0,
227
+ dtype=torch.int32)
228
+ chunks_start_idx = torch.cat(
229
+ [torch.tensor([0], dtype=torch.int32), chunks_start_idx])
230
231
def get_replacement_ultravox(item_idx: int):
215
- audio_token_len = out_mm_kwargs["audio_token_len"][item_idx]
232
+ start = chunks_start_idx[item_idx]
233
+ end = chunks_start_idx[item_idx + 1]
234
+ audio_token_len = out_mm_kwargs["audio_token_len"][start:end].sum()
235
return [replacement_id] * int(audio_token_len) # type: ignore
236
237
return [
@@ -304,12 +323,49 @@ class ModifiedWhisperEncoder(WhisperEncoder):
323
324
base_model_prefix = "model.encoder"
325
326
+ def __init__(self, *args, **kwargs):
327
+ super().__init__(*args, **kwargs)
328
+ self.config.is_decoder = False
329
+
330
+ @property
331
+ def max_context_length(self):
332
+ return (self.config.max_source_positions * self.conv1.stride[0] *
333
+ self.conv2.stride[0])
334
+
335
+ def get_attention_mask_by_audio_len(self,
336
+ audio_lens: Optional[torch.Tensor],
337
+ hidden_states: torch.Tensor):
338
+ """
339
+ Create attention mask based on audio lengths to mask out padding tokens
340
+ For each sample in batch:
341
+ - Convert raw audio length to feature length after convolutions
342
+ - Create bool mask: True for valid positions and False for padding
343
+ - Convert to attention mask format expected by transformer layers
344
+ (1.0 for positions to attend to, large negative for positions to ignore)
345
+ This masking ensures consistent behavior between training and inference
346
+ by preventing the model from attending to padding tokens in both cases
347
+ """
348
+ if audio_lens is None:
349
+ return None
350
+
351
+ audio_feature_len = self._get_feat_extract_output_lengths(audio_lens)
352
+ max_seq_len = hidden_states.shape[1]
353
+ attention_mask = torch.arange(max_seq_len,
354
+ device=hidden_states.device)[None, :].lt(
355
+ audio_feature_len.view(-1, 1))
356
+ attention_mask = self.get_extended_attention_mask(
357
+ attention_mask,
358
+ None,
359
+ dtype=hidden_states.dtype,
360
+ )
361
+ return attention_mask
362
+
363
def forward(
364
self,
309
- input_features,
365
+ input_features: torch.Tensor,
366
+ audio_lens: Optional[torch.Tensor] = None,
367
):
311
- expected_seq_length = (self.config.max_source_positions *
312
- self.conv1.stride[0] * self.conv2.stride[0])
368
+ expected_seq_length = self.max_context_length
369
if input_features.shape[-1] > expected_seq_length:
370
raise ValueError(
371
f"Whisper expects the mel input features to be of length "
@@ -328,10 +384,13 @@ class ModifiedWhisperEncoder(WhisperEncoder):
384
p=self.dropout,
385
training=self.training)
386
387
+ attention_mask = self.get_attention_mask_by_audio_len(
388
+ audio_lens, hidden_states)
389
+
390
for encoder_layer in self.layers:
391
layer_outputs = encoder_layer(
392
hidden_states,
334
- None,
393
+ attention_mask,
394
layer_head_mask=None,
395
)
396
@@ -409,17 +468,34 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
468
)
469
470
def _audio_features_to_embeddings(
412
- self, input_features: torch.Tensor) -> torch.Tensor:
413
- audio_input = input_features.to(self.audio_tower.dtype)
414
- audio_features = self.audio_tower(audio_input)
415
- audio_features = audio_features.to(self.audio_tower.dtype)
416
- audio_embeddings = self.multi_modal_projector(audio_features)
471
+ self, input_features: torch.Tensor,
472
+ audio_lens: torch.Tensor) -> torch.Tensor:
473
+ audio_features = input_features.to(self.audio_tower.dtype)
474
+ batch_size = audio_features.size(0)
475
+ audio_embeddings = []
476
+
477
+ # Process audio features in batches to keep memory usage predictable
478
+ for start in range(0, batch_size, _MAX_ENCODER_BATCH_SIZE):
479
+ end = min(start + _MAX_ENCODER_BATCH_SIZE, batch_size)
480
+ # Process through audio tower
481
+ batch_features = self.audio_tower(audio_features[start:end],
482
+ audio_lens[start:end])
483
+ batch_features = batch_features.to(self.audio_tower.dtype)
484
+
485
+ # Process through projector
486
+ batch_embeddings = self.multi_modal_projector(batch_features)
487
+ audio_embeddings.append(batch_embeddings)
488
+
489
+ # Concatenate results
490
+ audio_embeddings = torch.cat(audio_embeddings, dim=0)
491
return audio_embeddings
492
493
def _parse_and_validate_audio_input(
494
self, **kwargs: object) -> Optional[UltravoxAudioInputs]:
495
audio_features = kwargs.pop("audio_features", None)
496
audio_embeds = kwargs.pop("audio_embeds", None)
497
+ audio_lens = kwargs.pop("audio_lens", None)
498
+ audio_token_len = kwargs.pop("audio_token_len", None)
499
500
if audio_features is None and audio_embeds is None:
501
return None
@@ -430,7 +506,9 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
506
f"Got type: {type(audio_features)}")
507
508
return UltravoxAudioFeatureInputs(type="audio_features",
433
- data=audio_features)
509
+ data=audio_features,
510
+ lens=audio_lens,
511
+ token_len=audio_token_len)
512
513
if audio_embeds is not None:
514
if not isinstance(audio_embeds, (torch.Tensor, list)):
@@ -447,34 +525,34 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
525
if audio_input["type"] == "audio_embeds":
526
return audio_input["data"]
527
450
- audio_features = audio_input["data"]
451
- if isinstance(audio_features, torch.Tensor):
452
- # Combine the B and N dimensions for the encoder/projector
453
- flattened = flatten_bn(audio_features)
454
- flattened_embeddings = self._audio_features_to_embeddings(
455
- flattened)
456
-
457
- # Restore the original dimensions
458
- embeddings = flattened_embeddings.unflatten(
459
- 0, audio_features.shape[:2])
460
- return embeddings
461
-
462
- result = []
463
- # TODO: Batch heterogeneous tensors through the encoder/projector
464
- for audio_features_item in audio_features:
465
- if isinstance(audio_features_item, torch.Tensor):
466
- result.append(
467
- self._audio_features_to_embeddings(audio_features_item))
468
- else:
469
- embeddings = [
470
- # Add a batch dimension to embed it, then remove it.
471
- self._audio_features_to_embeddings(tensor.unsqueeze(0)
472
- ).squeeze(0)
473
- for tensor in audio_features_item
474
- ]
475
- result.append(embeddings)
528
+ # Pad and concatenate audio features
529
+ # [[B1, 80, M1], [B2, 80, M2]] -> [B1+B2, 80, max(M1, M2)]
530
+ audio_features = pad_and_concat_to_dim3(audio_input["data"])
531
+
532
+ if isinstance(audio_input['lens'], list):
533
+ # [B1, B2] -> [B1+B2]
534
+ audio_lens = torch.cat(audio_input['lens'])
535
+ audio_token_len = torch.cat(audio_input['token_len'])
536
+ else:
537
+ audio_lens = flatten_bn(audio_input['lens'])
538
+ audio_token_len = flatten_bn(audio_input['token_len'])
539
+
540
+ embeddings = self._audio_features_to_embeddings(
541
+ audio_features, audio_lens)
542
+
543
+ # We should flatten and concatenate embeddings based on token lengths
544
+ # For example, with token_len = [4, 2, 3], flattened_embeddings will be
545
+ # concat(embeddings[0][:4], embeddings[1][:2], embeddings[2][:3])
546
477
- return result
547
+ # Create a mask of valid indices based on token lengths
548
+ max_len = embeddings.shape[1]
549
+ indices = torch.arange(max_len, device=embeddings.device).expand(
550
+ embeddings.shape[0], -1)
551
+ mask = indices < audio_token_len[:, None]
552
+ # Apply mask and flatten
553
+ flattened_embeddings = embeddings[mask]
554
+
555
+ return flattened_embeddings
556
557
def get_multimodal_embeddings(
558
self, **kwargs
@@ -521,7 +599,11 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
599
with the `input_ids`.
600
601
Args:
524
- audio_features: A batch of audio inputs [B, N, 80, M].
602
+ audio_features: A batch of audio input chunks [B, N, 80, M].
603
+ audio_lens: Length of audio frames for each audio chunk [B].
604
+ audio_token_len: Length of audio tokens for each audio chunk [B'].
605
+ Note: batch dim is different from batch dim in audio chunks.
606
+
607
"""
608
609
if intermediate_tensors is not None:
@@ -560,3 +642,31 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
642
loader = AutoWeightsLoader(self,
643
ignore_unexpected_prefixes=["audio_tower."])
644
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
645
+
646
+
647
+def pad_and_concat_to_dim3(
648
+ features: Union[torch.Tensor, List[torch.Tensor], List[List[torch.Tensor]]]
649
+) -> torch.Tensor:
650
+ """
651
+ Pad and concatenate a list of tensors.
652
+
653
+ output:
654
+ Tensor of shape [B, C, M] where M is the maximum length of the input
655
+ tensors, B is the sum of the batch sizes of the input tensors.
656
+ C must be the same for all input tensors.
657
+ """
658
+ if isinstance(features, torch.Tensor):
659
+ if features.ndim > 3:
660
+ # Flatten [B, N, 80, M] -> [B * N, 80, M]
661
+ features = flatten_bn(features)
662
+ return features
663
+
664
+ features = [pad_and_concat_to_dim3(f) for f in features]
665
+
666
+ max_len = max(f.shape[-1] for f in features)
667
+ # Ensure all features have dim=3
668
+ features = [f.view(-1, *f.shape[-2:]) for f in features]
669
+ # Pad and oncatenate:
670
+ # [[B1, 80, M1], [B2, 80, M2]] -> [B1+B2, 80, max(M1, M2)]
671
+ features = [F.pad(f, (0, max_len - f.shape[-1])) for f in features]
672
+ return torch.cat(features)