ultravox.py 25.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
# Adapted from https://github.com/fixie-ai/ultravox/blob/ecd58c4041030bae2ad15aa6bcf04ab43199ea02/ultravox/model/ultravox_model.py
"""PyTorch Ultravox model."""
6
from collections.abc import Iterable, Mapping, Sequence
7
from typing import Annotated, Any, Literal, Optional, Union
8
9
10
11

import torch
from torch import nn
from torch.nn import functional as F
12
from transformers import BatchFeature, ProcessorMixin
13
14
15
from transformers.models.whisper import WhisperFeatureExtractor
from transformers.models.whisper.modeling_whisper import WhisperEncoder

16
from vllm.config import VllmConfig
17
from vllm.model_executor.layers.activation import MulAndSilu, get_act_fn
18
from vllm.model_executor.layers.layernorm import RMSNorm
19
from vllm.model_executor.model_loader import DefaultModelLoader
20
from vllm.model_executor.models.module_mapping import MultiModelKeys
21
from vllm.multimodal import MULTIMODAL_REGISTRY
22
from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
23
                                    MultiModalKwargsItems, NestedTensors)
24
from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
25
from vllm.multimodal.processing import (BaseMultiModalProcessor,
26
27
                                        BaseProcessingInfo, PromptReplacement,
                                        PromptUpdate)
28
from vllm.multimodal.profiling import BaseDummyInputsBuilder
29
from vllm.sequence import IntermediateTensors
30
from vllm.transformers_utils.configs.ultravox import UltravoxConfig
31
from vllm.utils.tensor_schema import TensorSchema, TensorShape
32

33
from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
34
                         SupportsMultiModal, SupportsPP)
35
from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
36
                    init_vllm_registered_model, maybe_prefix)
37

38
_AUDIO_PLACEHOLDER_OVERRIDE = "<|audio|>"
39
_MAX_ENCODER_BATCH_SIZE = 16
40
41


42
class UltravoxAudioFeatureInputs(TensorSchema):
43
    """
44
45
46
47
48
    Dimensions:
    - b: batch size
    - n: number of chunks
    - t: Time frames (M)
    - nmb: Number of mel bins
49
    """
50
51
52
53
54
55
56
57
58
59
60
61
62
    type: Literal["audio_features"]
    data: Annotated[Union[torch.Tensor, list[torch.Tensor],
                          list[list[torch.Tensor]]],
                    TensorShape("b", "n", "nmb", "t", dynamic_dims={"n"})]
    lens: Annotated[Union[torch.Tensor, list[torch.Tensor]],
                    TensorShape("b", "n", dynamic_dims={"n"})]
    """Length of the audio frames. Used for attention mask in WhisperEncoder."""
    token_len: Annotated[Union[torch.Tensor, list[torch.Tensor]],
                         TensorShape("b", "n", dynamic_dims={"n"})]
    """Length of the audio tokens. Used for flattening the audio features."""


class UltravoxAudioEmbeddingInputs(TensorSchema):
63
    """
64
65
66
67
68
    Dimensions:
    - b: batch size
    - na: number of audios
    - afs: audio feature size
    - hs: hidden size
69
    """
70
    type: Literal["audio_embeds"]
71
72
    data: Annotated[Union[torch.Tensor, list[torch.Tensor]],
                    TensorShape("b", "na", "afs", "hs")]
73
74
75
76
77
78


UltravoxAudioInputs = Union[UltravoxAudioFeatureInputs,
                            UltravoxAudioEmbeddingInputs]


79
class UltravoxProcessingInfo(BaseProcessingInfo):
80

81
    def get_hf_processor(self, **kwargs: object) -> ProcessorMixin:
82
        config = self.ctx.model_config.hf_config
83
        hf_processor = self.ctx.get_hf_processor(**kwargs)
84
85
86

        # NOTE: Ultravox processing definition uses '<|eot_id|>' as the
        # placeholder that will cause confusion with the actual end of turn
87
        # token, thus we override placeholder with a reserved token.
88
        hf_processor.audio_token_replacement = _AUDIO_PLACEHOLDER_OVERRIDE
89
90
        hf_processor.audio_replacement_token_id = config.audio_token_index

91
        return hf_processor
92

93
94
95
    def get_feature_extractor(self,
                              **kwargs: object) -> WhisperFeatureExtractor:
        hf_processor = self.get_hf_processor(**kwargs)
96
97
98
99
100
        audio_processor = hf_processor.audio_processor  # type: ignore
        feature_extractor = audio_processor.feature_extractor  # type: ignore
        assert isinstance(feature_extractor, WhisperFeatureExtractor)
        return feature_extractor

101
102
    def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
        return {"audio": None}
103

104
105
106
107

class UltravoxDummyInputsBuilder(BaseDummyInputsBuilder[UltravoxProcessingInfo]
                                 ):

108
109
110
111
112
113
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        num_audios = mm_counts.get("audio", 0)

        return "<|audio|>" * num_audios

    def get_dummy_mm_data(
114
        self,
115
116
        seq_len: int,
        mm_counts: Mapping[str, int],
117
    ) -> MultiModalDataDict:
118
        feature_extractor = self.info.get_feature_extractor()
119
120

        sampling_rate = feature_extractor.sampling_rate
121
122
        audio_len = (feature_extractor.chunk_length * sampling_rate *
                     _MAX_ENCODER_BATCH_SIZE)
123
124
        num_audios = mm_counts.get("audio", 0)

125
        return {
126
127
128
129
130
            "audio":
            self._get_dummy_audios(length=audio_len, num_audios=num_audios)
        }


131
132
class UltravoxMultiModalProcessor(
        BaseMultiModalProcessor[UltravoxProcessingInfo]):
133

134
    def _get_data_parser(self) -> MultiModalDataParser:
135
        feature_extractor = self.info.get_feature_extractor()
136
        return MultiModalDataParser(target_sr=feature_extractor.sampling_rate)
137
138

    def _call_hf_processor(
139
140
        self,
        prompt: str,
141
142
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
143
        tok_kwargs: Mapping[str, object],
144
    ) -> BatchFeature:
145
        # Text-only input not supported in composite processor
146
        if not mm_data.get("audios", []):
147
148
            prompt_ids = self.info.get_tokenizer().encode(
                prompt, add_special_tokens=False)
149
            prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
150
151
152
153
            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")

        mm_data = dict(mm_data)
        audios = mm_data.pop("audios", [])
154
        assert isinstance(audios, list)
155

156
        feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
157
158
        mm_kwargs = dict(
            **mm_kwargs,
159
            sampling_rate=feature_extractor.sampling_rate,
160
            include_audio_num_chunks=True,
161
162
        )

163
        item_processor_data = dict(**mm_data, audios=audios)
164

165
166
167
168
        # some tokenizer kwargs are incompatible with UltravoxProcessor
        tok_kwargs.pop("padding", None)
        tok_kwargs.pop("truncation", None)

169
170
171
172
        output = super()._call_hf_processor(
            prompt=prompt,
            mm_data=item_processor_data,
            mm_kwargs=mm_kwargs,
173
            tok_kwargs=tok_kwargs,
174
        )
175
176
177
        output['audio_features'] = output.pop('audio_values')

        return output
178

179
180
181
182
183
    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
184
        num_chunks = hf_inputs.get('audio_num_chunks', torch.zeros(0))
185
        return dict(
186
187
188
189
190
191
192
193
194
195
196
            # to handle longer than 30s audio, each audio might be split
            # into multiple chunks as such, their batch dimension can be
            # higher than the number of audio samples
            audio_features=MultiModalFieldConfig.flat_from_sizes(
                "audio", num_chunks),
            audio_token_len=MultiModalFieldConfig.flat_from_sizes(
                "audio", num_chunks),
            audio_lens=MultiModalFieldConfig.flat_from_sizes(
                "audio", num_chunks),
            # num_chunks can convert audio_chunked to audio batch dimension
            audio_num_chunks=MultiModalFieldConfig.batched("audio"),
197
198
199
            audio_embeds=MultiModalFieldConfig.batched("audio"),
        )

200
    def _get_prompt_updates(
201
202
        self,
        mm_items: MultiModalDataItems,
203
        hf_processor_mm_kwargs: Mapping[str, Any],
204
        out_mm_kwargs: MultiModalKwargsItems,
205
    ) -> Sequence[PromptUpdate]:
206
        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
207

208
209
210
211
212
        replacement_id = hf_processor.audio_replacement_token_id  # type: ignore

        # Each audio can be split into multiple chunks.
        # chunks_start_idx[i] indicates the start index of the chunks
        # belonging to the i-th audio.
213
214
        out_mm_data = out_mm_kwargs.get_data()
        num_chunks = out_mm_data.get("audio_num_chunks", torch.zeros(0))
215
216
217
218
219
        chunks_start_idx: torch.Tensor = torch.cumsum(num_chunks,
                                                      dim=0,
                                                      dtype=torch.int32)
        chunks_start_idx = torch.cat(
            [torch.tensor([0], dtype=torch.int32), chunks_start_idx])
220
221

        def get_replacement_ultravox(item_idx: int):
222
223
            start = chunks_start_idx[item_idx]
            end = chunks_start_idx[item_idx + 1]
224
            audio_token_len = out_mm_data["audio_token_len"][start:end].sum()
225
            return [replacement_id] * int(audio_token_len)  # type: ignore
226
227
228
229

        return [
            PromptReplacement(
                modality="audio",
230
                target="<|audio|>",
231
232
233
                replacement=get_replacement_ultravox,
            )
        ]
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262


class StackAudioFrames(nn.Module):
    """
    Stack the audio embedding frames to reduce the sequence length by a factor
    of `stack_factor`.
    """

    def __init__(self, stack_factor: int = 8):
        super().__init__()
        self.stack_factor = stack_factor

    def forward(self, audio_embeds: torch.Tensor) -> torch.Tensor:
        B, T, C = audio_embeds.shape
        T_pad = (T + self.stack_factor -
                 1) // self.stack_factor * self.stack_factor
        audio_embeds = F.pad(audio_embeds, (0, 0, 0, T_pad - T))
        B, T, C = audio_embeds.shape
        audio_embeds = audio_embeds.view(B, T // self.stack_factor,
                                         C * self.stack_factor)
        return audio_embeds


class UltravoxProjector(nn.Module):

    def __init__(self, config: UltravoxConfig):
        super().__init__()
        self.hidden_dim = config.hidden_size
        self._pad_and_stack = StackAudioFrames(config.stack_factor)
263
264
265
266
        dim_in = config.audio_config.hidden_size * config.stack_factor
        self.ln_pre = RMSNorm(dim_in)
        self.linear_1 = nn.Linear(dim_in, self.hidden_dim, bias=False)
        dim_mid = self.hidden_dim
267
268

        if config.projector_act == "swiglu":
269
            self.act = MulAndSilu()
270
            dim_mid = dim_mid // 2
271
272
273
        else:
            self.act = get_act_fn(config.projector_act)

274
        dim_out = config.text_config.hidden_size
275
276
277
278
279
280
281
282
283
284
        self.linear_2 = nn.Linear(dim_mid, dim_out, bias=False)

        # Ultravox v0.4.1 and below use layer_norm after the second linear layer
        # while v0.5.0 and above uses layer_norm after the first linear layer.
        if config.projector_ln_mid:
            self.ln_mid: nn.Module = RMSNorm(dim_mid)
            self.ln_post = nn.Identity()
        else:
            self.ln_mid = nn.Identity()
            self.ln_post = RMSNorm(dim_out)
285
286
287
288
289
290

    def forward(self, audio_features: torch.Tensor) -> torch.Tensor:
        audio_features = self._pad_and_stack(audio_features)
        audio_features = self.ln_pre(audio_features)
        hidden_states = self.linear_1(audio_features)
        hidden_states = self.act(hidden_states)
291
        hidden_states = self.ln_mid(hidden_states)
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
        hidden_states = self.linear_2(hidden_states)
        hidden_states = self.ln_post(hidden_states)
        return hidden_states


class ModifiedWhisperEncoder(WhisperEncoder):
    """
    Encoder portion of OpenAI's Whisper model.

    This implementation is a slightly modified version of HF Transformers'
    Whisper Encoder, with only a few fixes:
    1. base_model_prefix updated to allow for doing `.from_pretrained`
       directly on the encoder
    2. allow less than 30 second of audio padding to be passed in:
        - relaxed ValueError check for `input_features` length to be less
           than or equal to `expected_seq_length` instead of strictly equal
        - embed_pos is now sliced to match the length of `inputs_embeds`

    Original: https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/modeling_whisper.py
    See commentary: https://github.com/huggingface/transformers/issues/25744
    """

    base_model_prefix = "model.encoder"

316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.config.is_decoder = False

    @property
    def max_context_length(self):
        return (self.config.max_source_positions * self.conv1.stride[0] *
                self.conv2.stride[0])

    def get_attention_mask_by_audio_len(self,
                                        audio_lens: Optional[torch.Tensor],
                                        hidden_states: torch.Tensor):
        """
        Create attention mask based on audio lengths to mask out padding tokens
        For each sample in batch:
        - Convert raw audio length to feature length after convolutions
        - Create bool mask: True for valid positions and False for padding
        - Convert to attention mask format expected by transformer layers
        (1.0 for positions to attend to, large negative for positions to ignore)
        This masking ensures consistent behavior between training and inference
        by preventing the model from attending to padding tokens in both cases
        """
        if audio_lens is None:
            return None

        audio_feature_len = self._get_feat_extract_output_lengths(audio_lens)
        max_seq_len = hidden_states.shape[1]
        attention_mask = torch.arange(max_seq_len,
                                      device=hidden_states.device)[None, :].lt(
                                          audio_feature_len.view(-1, 1))
        attention_mask = self.get_extended_attention_mask(
            attention_mask,
            None,
            dtype=hidden_states.dtype,
        )
        return attention_mask

353
354
    def forward(
        self,
355
356
        input_features: torch.Tensor,
        audio_lens: Optional[torch.Tensor] = None,
357
    ):
358
        expected_seq_length = self.max_context_length
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
        if input_features.shape[-1] > expected_seq_length:
            raise ValueError(
                f"Whisper expects the mel input features to be of length "
                f"{expected_seq_length} or less, but found "
                f"{input_features.shape[-1]}. Make sure to pad the input mel "
                f"features to {expected_seq_length}.")

        inputs_embeds = nn.functional.gelu(self.conv1(input_features))
        inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))

        inputs_embeds = inputs_embeds.permute(0, 2, 1)
        embed_pos = self.embed_positions.weight[:inputs_embeds.size(-2)]

        hidden_states = inputs_embeds + embed_pos
        hidden_states = nn.functional.dropout(hidden_states,
                                              p=self.dropout,
                                              training=self.training)

377
378
379
        attention_mask = self.get_attention_mask_by_audio_len(
            audio_lens, hidden_states)

380
381
382
        for encoder_layer in self.layers:
            layer_outputs = encoder_layer(
                hidden_states,
383
                attention_mask,
384
385
386
387
388
389
390
391
392
                layer_head_mask=None,
            )

            hidden_states = layer_outputs[0]

        hidden_states = self.layer_norm(hidden_states)
        return hidden_states


393
394
395
396
@MULTIMODAL_REGISTRY.register_processor(
    UltravoxMultiModalProcessor,
    info=UltravoxProcessingInfo,
    dummy_inputs=UltravoxDummyInputsBuilder)
397
class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
398
399
400
401
402
403

    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"]
    }

404
405
406
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={"audio_tower.model.encoder.": "audio_tower."})

407
408
409
410
411
412
413
    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
        if modality.startswith("audio"):
            return "<|audio|>"

        raise ValueError("Only audio modality is supported")

414
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
415
        super().__init__()
416
        config: UltravoxConfig = vllm_config.model_config.hf_config
417
        multimodal_config = vllm_config.model_config.multimodal_config
418
419
420
421
        self.config = config
        self.multi_modal_config = multimodal_config
        assert self.multi_modal_config

422
423
        self.secondary_weights = []
        self.audio_tower = ModifiedWhisperEncoder(config.audio_config)
424
        if config.audio_model_id is not None:
425
426
            # this prefix is not for initialization, but for loading weights
            # note the trailing dot
427
428
429
430
431
432
            self.secondary_weights.append(
                DefaultModelLoader.Source(
                    model_or_path=config.audio_model_id,
                    revision=None,
                    prefix="audio_tower.",
                ))
433
434
        self.multi_modal_projector = UltravoxProjector(config)
        self.language_model = init_vllm_registered_model(
435
            vllm_config=vllm_config,
436
            hf_config=config.wrapped_model_config,
437
438
            prefix=maybe_prefix(prefix, "language_model"),
        )
439
        if config.text_model_id is not None:
440
441
            # this prefix is not for initialization, but for loading weights
            # note the trailing dot
442
443
444
445
            self.secondary_weights.append(
                DefaultModelLoader.Source(model_or_path=config.text_model_id,
                                          revision=None,
                                          prefix="language_model."))
446

447
448
449
        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors)

450
451
452
453
454
455
456
457
458
459
    def get_mm_mapping(self) -> MultiModelKeys:
        """
        Get the module prefix in multimodal models
        """
        return MultiModelKeys.from_string_field(
            language_model="language_model.",
            connector="multi_modal_projector.",
            tower_model="audio_tower.",
        )

460
    def _audio_features_to_embeddings(
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
            self, input_features: torch.Tensor,
            audio_lens: torch.Tensor) -> torch.Tensor:
        audio_features = input_features.to(self.audio_tower.dtype)
        batch_size = audio_features.size(0)
        audio_embeddings = []

        # Process audio features in batches to keep memory usage predictable
        for start in range(0, batch_size, _MAX_ENCODER_BATCH_SIZE):
            end = min(start + _MAX_ENCODER_BATCH_SIZE, batch_size)
            # Process through audio tower
            batch_features = self.audio_tower(audio_features[start:end],
                                              audio_lens[start:end])
            batch_features = batch_features.to(self.audio_tower.dtype)

            # Process through projector
            batch_embeddings = self.multi_modal_projector(batch_features)
            audio_embeddings.append(batch_embeddings)

        # Concatenate results
        audio_embeddings = torch.cat(audio_embeddings, dim=0)
481
482
483
484
485
486
        return audio_embeddings

    def _parse_and_validate_audio_input(
            self, **kwargs: object) -> Optional[UltravoxAudioInputs]:
        audio_features = kwargs.pop("audio_features", None)
        audio_embeds = kwargs.pop("audio_embeds", None)
487
488
        audio_lens = kwargs.pop("audio_lens", None)
        audio_token_len = kwargs.pop("audio_token_len", None)
489
490
491
492
493
494

        if audio_features is None and audio_embeds is None:
            return None

        if audio_features is not None:
            return UltravoxAudioFeatureInputs(type="audio_features",
495
496
497
                                              data=audio_features,
                                              lens=audio_lens,
                                              token_len=audio_token_len)
498
499
500
501
502
503
504
505

        if audio_embeds is not None:
            return UltravoxAudioEmbeddingInputs(type="audio_embeds",
                                                data=audio_embeds)

        raise AssertionError("This line should be unreachable.")

    def _process_audio_input(
506
507
508
        self,
        audio_input: UltravoxAudioInputs,
    ) -> Union[NestedTensors, tuple[torch.Tensor, ...]]:
509
510
511
        if audio_input["type"] == "audio_embeds":
            return audio_input["data"]

512
513
514
515
        # Pad and concatenate audio features
        # [[B1, 80, M1], [B2, 80, M2]] -> [B1+B2, 80, max(M1, M2)]
        audio_features = pad_and_concat_to_dim3(audio_input["data"])

516
517
518
        # [B1, B2] -> [B1+B2]
        audio_lens = flatten_bn(audio_input['lens'], concat=True)
        audio_token_len = flatten_bn(audio_input['token_len'], concat=True)
519
520
521
522
523
524
525

        embeddings = self._audio_features_to_embeddings(
            audio_features, audio_lens)

        # We should flatten and concatenate embeddings based on token lengths
        # For example, with token_len = [4, 2, 3], flattened_embeddings will be
        # concat(embeddings[0][:4], embeddings[1][:2], embeddings[2][:3])
526

527
528
529
530
531
532
533
534
        # Create a mask of valid indices based on token lengths
        max_len = embeddings.shape[1]
        indices = torch.arange(max_len, device=embeddings.device).expand(
            embeddings.shape[0], -1)
        mask = indices < audio_token_len[:, None]
        # Apply mask and flatten
        flattened_embeddings = embeddings[mask]

535
536
537
538
539
540
        # Return one tensor per input audio
        embed_lens = [
            token_len_item.sum().item()
            for token_len_item in audio_input['token_len']
        ]
        return flattened_embeddings.split(embed_lens)
541

542
543
544
    def get_language_model(self) -> torch.nn.Module:
        return self.language_model

545
546
    def get_multimodal_embeddings(self,
                                  **kwargs: object) -> MultiModalEmbeddings:
547
548
        audio_input = self._parse_and_validate_audio_input(**kwargs)
        if audio_input is None:
549
            return []
550
551
552
553
554
555
        audio_embeddings = self._process_audio_input(audio_input)
        return audio_embeddings

    def get_input_embeddings(
        self,
        input_ids: torch.Tensor,
556
        multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
557
558
559
560
        *,
        is_multimodal: Optional[torch.Tensor] = None,
        # Multi-modal token ID may exceed vocab size
        handle_oov_mm_token: bool = True,
561
    ) -> torch.Tensor:
562
563
564
565
566
567
568
569
570
571
        # This is to satisfy the type checker for each overload
        if multimodal_embeddings is None or is_multimodal is None:
            return super().get_input_embeddings(input_ids)

        return super().get_input_embeddings(
            input_ids,
            multimodal_embeddings=multimodal_embeddings,
            is_multimodal=is_multimodal,
            handle_oov_mm_token=handle_oov_mm_token,
        )
572
573
574
575
576
577

    def forward(self,
                input_ids: torch.Tensor,
                positions: torch.Tensor,
                intermediate_tensors: Optional[torch.Tensor] = None,
                inputs_embeds: Optional[torch.Tensor] = None,
578
                **kwargs) -> Union[torch.Tensor, IntermediateTensors]:
579
580
581
582
583
584
585
586
587
588
        """Run forward pass for Ultravox

        One key thing to understand is the `input_ids` already accounts for the
        positions of the to-be-inserted audio embeddings. The to-be-inserted
        audio has a size that is essentially 6.25 tokens per second of audio.

        This way, the `positions` and `attn_metadata` are consistent
        with the `input_ids`.

        Args:
589
590
591
592
593
            input_ids: Flattened (concatenated) input_ids corresponding to a
                batch.
            positions: Position indices for the input tokens.
            intermediate_tensors: Intermediate tensors from prior forward pass.
            inputs_embeds: Optional tensor of input embeddings.
594

595
        """
596

597
        if intermediate_tensors is not None:
598
            inputs_embeds = None
599

600
601
602
603
604
605
606
607
        language_model = self.language_model
        if hasattr(language_model, "language_model"):
            language_model = language_model.language_model

        hidden_states = language_model.model(input_ids,
                                             positions,
                                             intermediate_tensors,
                                             inputs_embeds=inputs_embeds)
608
609
        return hidden_states

610
611
    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.language_model.compute_logits(hidden_states)
612

613
614
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
615
616
617

        loader = AutoWeightsLoader(self,
                                   ignore_unexpected_prefixes=["audio_tower."])
618
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
619
620
621


def pad_and_concat_to_dim3(
622
    features: Union[torch.Tensor, list[torch.Tensor], list[list[torch.Tensor]]]
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
) -> torch.Tensor:
    """
    Pad and concatenate a list of tensors.

    output:
        Tensor of shape [B, C, M] where M is the maximum length of the input
        tensors, B is the sum of the batch sizes of the input tensors.
        C must be the same for all input tensors.
    """
    if isinstance(features, torch.Tensor):
        if features.ndim > 3:
            # Flatten [B, N, 80, M] -> [B * N, 80, M]
            features = flatten_bn(features)
        return features

    features = [pad_and_concat_to_dim3(f) for f in features]

    max_len = max(f.shape[-1] for f in features)
    # Ensure all features have dim=3
    features = [f.view(-1, *f.shape[-2:]) for f in features]
643
    # Pad and concatenate:
644
645
646
    # [[B1, 80, M1], [B2, 80, M2]] -> [B1+B2, 80, max(M1, M2)]
    features = [F.pad(f, (0, max_len - f.shape[-1])) for f in features]
    return torch.cat(features)