voxtral.py 31.5 KB
Newer Older
Patrick von Platen's avatar
Patrick von Platen committed
1
2
3
4
5
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import math
from collections.abc import Iterable, Mapping, Sequence
6
from functools import partial
7
from typing import cast
Patrick von Platen's avatar
Patrick von Platen committed
8
9
10
11
12

import numpy as np
import regex as re
import torch
import torch.nn as nn
13
from mistral_common.audio import Audio, mel_filter_bank
14
15
from mistral_common.protocol.instruct.chunk import AudioChunk, RawAudio, TextChunk
from mistral_common.protocol.instruct.messages import UserMessage
Patrick von Platen's avatar
Patrick von Platen committed
16
17
from mistral_common.protocol.instruct.request import ChatCompletionRequest
from mistral_common.protocol.transcription.request import TranscriptionRequest
18
from transformers import BatchFeature, WhisperConfig
Patrick von Platen's avatar
Patrick von Platen committed
19

20
from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
21
from vllm.config.multimodal import BaseDummyOptions
22
from vllm.config.speech_to_text import SpeechToTextParams
23
from vllm.inputs import MultiModalDataDict, PromptType, TokensPrompt
Patrick von Platen's avatar
Patrick von Platen committed
24
from vllm.logger import init_logger
25
from vllm.model_executor.layers.quantization import QuantizationConfig
Patrick von Platen's avatar
Patrick von Platen committed
26
27
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models import SupportsPP
28
from vllm.model_executor.models.module_mapping import MultiModelKeys
29
30
31
32
33
from vllm.model_executor.models.whisper import (
    WhisperEncoder,
    _create_fake_bias_for_k_proj,
)
from vllm.model_executor.models.whisper_causal import WhisperCausalEncoder
Patrick von Platen's avatar
Patrick von Platen committed
34
from vllm.multimodal import MULTIMODAL_REGISTRY
35
36
37
38
39
40
41
42
43
44
from vllm.multimodal.inputs import (
    MultiModalFieldConfig,
    MultiModalKwargsItems,
    NestedTensors,
)
from vllm.multimodal.parse import (
    AudioProcessorItems,
    MultiModalDataItems,
    MultiModalDataParser,
)
45
from vllm.multimodal.processing import BaseDummyInputsBuilder
46
from vllm.multimodal.processing.processor import (
47
48
49
    BaseMultiModalProcessor,
    BaseProcessingInfo,
    MultiModalProcessingInfo,
50
    PlaceholderFeaturesInfo,
51
    ProcessorInputs,
52
53
    PromptReplacement,
    PromptUpdate,
54
    TimingContext,
55
)
Patrick von Platen's avatar
Patrick von Platen committed
56
from vllm.sequence import IntermediateTensors
57
58
from vllm.tokenizers import cached_tokenizer_from_config
from vllm.tokenizers.mistral import MistralTokenizer
59
60
61
62
from vllm.transformers_utils.processors.voxtral import (
    MistralCommonFeatureExtractor,
    MistralCommonVoxtralProcessor,
)
63
from vllm.utils.collection_utils import is_list_of
Patrick von Platen's avatar
Patrick von Platen committed
64

65
from .interfaces import SupportsLoRA, SupportsMultiModal, SupportsTranscription
66
from .utils import init_vllm_registered_model, maybe_prefix
Patrick von Platen's avatar
Patrick von Platen committed
67
68
69

logger = init_logger(__name__)

70
71
72
73
74
75
76
77
78
79
80
81
ISO639_1_SUPPORTED_LANGS = {
    "ar": "Arabic",
    "nl": "Dutch",
    "en": "English",
    "fr": "French",
    "de": "German",
    "hi": "Hindi",
    "it": "Italian",
    "pt": "Portuguese",
    "es": "Spanish",
}

Patrick von Platen's avatar
Patrick von Platen committed
82
83
84

class VoxtralProcessingInfo(BaseProcessingInfo):
    def get_tokenizer(self) -> MistralTokenizer:
85
        tokenizer = cached_tokenizer_from_config(self.ctx.model_config)
Patrick von Platen's avatar
Patrick von Platen committed
86
87
88
89
90
        if not isinstance(tokenizer, MistralTokenizer):
            raise ValueError("This model requires `--tokenizer-mode mistral`")

        return tokenizer

91
92
93
94
95
    def get_feature_extractor(self) -> MistralCommonFeatureExtractor:
        return MistralCommonFeatureExtractor(
            self.get_tokenizer().instruct.audio_encoder
        )

96
    def get_hf_processor(self, **kwargs) -> MistralCommonVoxtralProcessor:
97
        return MistralCommonVoxtralProcessor(
98
            tokenizer=self.get_tokenizer(),
99
            feature_extractor=self.get_feature_extractor(),
100
        )
Patrick von Platen's avatar
Patrick von Platen committed
101

102
    def get_data_parser(self):
103
        feature_extractor = self.get_feature_extractor()
104

105
        return MultiModalDataParser(
106
            target_sr=feature_extractor.sampling_rate,
107
            target_channels=1,
108
109
110
            expected_hidden_size=self._get_expected_hidden_size(),
        )

111
    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
Patrick von Platen's avatar
Patrick von Platen committed
112
113
114
115
116
117
118
119
120
121
122
123
124
        return {"audio": 5}  # Performance tends to degrade after 5

    def get_mm_max_tokens_per_item(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> Mapping[str, int]:
        return {"audio": self.get_max_audio_tokens()}

    def get_max_audio_tokens(self) -> int:
        return self.ctx.model_config.max_model_len

    def get_max_audio_array_len(self) -> int:
125
        feature_extractor = self.get_feature_extractor()
126

Patrick von Platen's avatar
Patrick von Platen committed
127
        return self.get_max_audio_tokens() * int(
128
            feature_extractor.sampling_rate // feature_extractor.frame_rate
129
        )
Patrick von Platen's avatar
Patrick von Platen committed
130
131
132
133
134
135
136
137
138
139


class VoxtralDummyInputsBuilder(BaseDummyInputsBuilder[VoxtralProcessingInfo]):
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        return ""

    def get_dummy_mm_data(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
140
        mm_options: Mapping[str, BaseDummyOptions],
Patrick von Platen's avatar
Patrick von Platen committed
141
142
143
144
145
    ) -> MultiModalDataDict:
        num_audios = mm_counts.get("audio", 0)

        target_length = self.info.get_max_audio_array_len()

146
        audio_overrides = mm_options.get("audio")
147

Patrick von Platen's avatar
Patrick von Platen committed
148
        return {
149
            "audio": self._get_dummy_audios(
150
151
152
                length=target_length,
                num_audios=num_audios,
                overrides=audio_overrides,
153
            )
Patrick von Platen's avatar
Patrick von Platen committed
154
155
156
157
158
159
        }

    def get_dummy_processor_inputs(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
160
        mm_options: Mapping[str, BaseDummyOptions],
161
        mm_data: MultiModalDataDict | None = None,
Patrick von Platen's avatar
Patrick von Platen committed
162
163
    ) -> ProcessorInputs:
        tokenizer = self.info.get_tokenizer()
164
        feature_extractor = self.info.get_feature_extractor()
Patrick von Platen's avatar
Patrick von Platen committed
165
166

        dummy_text = self.get_dummy_text(mm_counts)
167
168
169
170
171
172
173
174
175
        dummy_mm_data = (
            self.get_dummy_mm_data(seq_len, mm_counts, mm_options)
            if mm_data is None
            else mm_data
        )
        dummy_mm_items = self.info.parse_mm_data(dummy_mm_data)
        dummy_audios = (
            [] if "audio" not in dummy_mm_data else dummy_mm_items["audio"].get_all()
        )
Patrick von Platen's avatar
Patrick von Platen committed
176
177
178
179
180
181

        audio_chunks: list[AudioChunk] = []
        format = "wav"
        for audio in dummy_audios:
            audio_item = Audio(
                audio_array=audio,
182
                sampling_rate=feature_extractor.sampling_rate,
Patrick von Platen's avatar
Patrick von Platen committed
183
184
185
186
187
                format=format,
            )
            chunk = AudioChunk(input_audio=RawAudio.from_audio(audio_item))
            audio_chunks.append(chunk)

188
189
190
191
192
        request = ChatCompletionRequest(
            messages=[
                UserMessage(content=[TextChunk(text=dummy_text), *audio_chunks]),
            ]
        )
Patrick von Platen's avatar
Patrick von Platen committed
193
194
195
        res = tokenizer.mistral.encode_chat_completion(request)
        dummy_tokens = res.tokens

196
        dummy_mm_items = self.info.parse_mm_data(
197
198
199
200
201
            # whixtral tokenizer adds padding to the audio
            # so we need to update the audio arrays
            {**dummy_mm_data, "audio": [a.audio_array for a in res.audios]},
        )

202
        return ProcessorInputs(prompt=dummy_tokens, mm_data_items=dummy_mm_items)
Patrick von Platen's avatar
Patrick von Platen committed
203
204


205
class VoxtralMultiModalProcessor(BaseMultiModalProcessor[VoxtralProcessingInfo]):
Patrick von Platen's avatar
Patrick von Platen committed
206
207
208
209
210
211
212
    def _get_mm_fields_config(
        self,
        hf_inputs: Mapping[str, NestedTensors],
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
        return dict(audio_arrays=MultiModalFieldConfig.batched("audio"))

213
214
215
216
217
218
219
    def _validate_mm_placeholders(
        self,
        mm_placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],
        mm_item_counts: Mapping[str, int],
    ) -> None:
        # mistral_common's tokenizer's does not follow HF's placeholder norms
        # skip validation here
220
        pass
221

222
    def _call_hf_processor(
223
        self,
224
225
226
227
        prompt: str,
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
        tok_kwargs: Mapping[str, object],
228
    ) -> BatchFeature:
229
230
231
232
233
234
235
        mm_data = dict(mm_data)
        audios = mm_data.pop("audios", [])

        if audios:
            # MistralCommonVoxtralProcessor accepts "audio"
            mm_data["audio"] = audios

236
        outputs = super()._call_hf_processor(
237
238
239
            prompt=prompt,
            mm_data=mm_data,
            mm_kwargs=mm_kwargs,
240
241
            # Avoid padding issue
            tok_kwargs={**tok_kwargs, "return_tensors": None},
242
        )
243

244
245
246
247
248
249
        # Missing batch dimension
        if is_list_of(outputs["input_ids"], int):
            outputs["input_ids"] = [outputs["input_ids"]]

        return outputs

Patrick von Platen's avatar
Patrick von Platen committed
250
251
252
253
    def _get_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
254
        out_mm_kwargs: MultiModalKwargsItems,
Patrick von Platen's avatar
Patrick von Platen committed
255
256
    ) -> Sequence[PromptUpdate]:
        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
257
        feature_extractor = processor.feature_extractor
Patrick von Platen's avatar
Patrick von Platen committed
258
259

        audio_id = processor.audio_token_id
260
261
        out_mm_data = out_mm_kwargs.require_data()
        out_audio_items = out_mm_data.get("audio", [])
Patrick von Platen's avatar
Patrick von Platen committed
262
263

        def get_replacement(item_idx: int):
264
265
266
267
268
269
270
271
272
273
274
275
276
277
            if item_idx < len(out_audio_items):
                out_audio_data = out_audio_items[item_idx].get_data()
                audio_arr = out_audio_data["audio_arrays"]
                if isinstance(audio_arr, (torch.Tensor, np.ndarray)):
                    audio_len = len(audio_arr)
                else:
                    raise TypeError(
                        "Unexpected type for audio_arrays in out_mm_kwargs: "
                        f"{type(audio_arr)}"
                    )
            else:
                # Fallback for unexpected processor outputs.
                audios = mm_items.get_items("audio", AudioProcessorItems)
                audio_len = audios.get_audio_length(item_idx)
Patrick von Platen's avatar
Patrick von Platen committed
278

279
            nb_audio_tokens = feature_extractor.get_num_audio_tokens(audio_len)
Patrick von Platen's avatar
Patrick von Platen committed
280
281
282
283
284
285
286
287
288
289
290
291
292

            return [audio_id] * nb_audio_tokens

        return [
            PromptReplacement(
                modality="audio",
                target="",  # Never match the prompt (see below note)
                replacement=get_replacement,
            ),
        ]

    def _cached_apply_hf_processor(
        self,
293
294
        inputs: ProcessorInputs,
        timing_ctx: TimingContext,
295
    ) -> tuple[list[int], MultiModalProcessingInfo, bool]:
296
        prompt_ids, mm_info, _ = super()._cached_apply_hf_processor(inputs, timing_ctx)
Patrick von Platen's avatar
Patrick von Platen committed
297
298

        # NOTE: The tokens are already inserted by the chat template
299
        return prompt_ids, mm_info, True
Patrick von Platen's avatar
Patrick von Platen committed
300
301


302
303
304
305
306
307
308
309
@MULTIMODAL_REGISTRY.register_processor(
    VoxtralMultiModalProcessor,
    info=VoxtralProcessingInfo,
    dummy_inputs=VoxtralDummyInputsBuilder,
)
class VoxtralForConditionalGeneration(
    nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA, SupportsTranscription
):
310
    supported_languages = ISO639_1_SUPPORTED_LANGS
311
312
313
    # transformers' currently has limited support for MistralCommon backend
    # and cached_get_processor. Let's skip until fixed
    skip_warmup_audio_preprocessing = True
Patrick von Platen's avatar
Patrick von Platen committed
314

315
316
    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
317
        "gate_up_proj": ["gate_proj", "up_proj"],
318
319
    }

Patrick von Platen's avatar
Patrick von Platen committed
320
321
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
322
        self.tokenizer = cached_tokenizer_from_config(vllm_config.model_config)
Patrick von Platen's avatar
Patrick von Platen committed
323

324
325
326
327
        # update quant config to so that ignored module and target module names
        # match the vLLM model names
        if hasattr(vllm_config, "quant_config"):
            vllm_config.quant_config = self.maybe_update_quant_config(
328
329
                vllm_config.quant_config
            )
330

Patrick von Platen's avatar
Patrick von Platen committed
331
332
333
334
        config = vllm_config.model_config.hf_config
        self.config = config
        self.downsample_factor = self.config.audio_config.downsample_factor

335
336
337
338
339
340
        with self._mark_language_model(vllm_config):
            self.language_model = init_vllm_registered_model(
                vllm_config=vllm_config,
                hf_config=config.text_config,
                prefix=maybe_prefix(prefix, "language_model"),
            )
Patrick von Platen's avatar
Patrick von Platen committed
341

342
343
344
345
346
347
348
349
350
        with self._mark_tower_model(vllm_config, "audio"):
            self.whisper_encoder = VoxtralEncoderModel(
                vllm_config.with_hf_config(config.audio_config),
                prefix=maybe_prefix(prefix, "whisper_encoder"),
            )
            self.audio_language_adapter = AudioLanguageAdapter(
                hidden_size=config.audio_config.d_model * self.downsample_factor,
                dim=config.text_config.hidden_size,
            )
Patrick von Platen's avatar
Patrick von Platen committed
351

352
353
354
355
356
357
358
359
    def get_mm_mapping(self) -> MultiModelKeys:
        """Get module prefix for multimodal models to filter LoRA modules."""
        return MultiModelKeys.from_string_field(
            language_model="language_model",
            connector="audio_language_adapter",
            tower_model=["whisper_encoder"],
        )

Patrick von Platen's avatar
Patrick von Platen committed
360
361
    def forward(
        self,
362
        input_ids: torch.Tensor | None,
Patrick von Platen's avatar
Patrick von Platen committed
363
        positions: torch.Tensor,
364
365
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
Patrick von Platen's avatar
Patrick von Platen committed
366
        **kwargs: object,
367
    ) -> torch.Tensor | IntermediateTensors:
Patrick von Platen's avatar
Patrick von Platen committed
368
369
370
        if intermediate_tensors is not None:
            inputs_embeds = None

371
372
373
        hidden_states = self.language_model.model(
            input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
        )
Patrick von Platen's avatar
Patrick von Platen committed
374
375
376

        return hidden_states

377
    def embed_multimodal(
Patrick von Platen's avatar
Patrick von Platen committed
378
        self, **kwargs
379
    ) -> list[torch.Tensor] | torch.Tensor | tuple[torch.Tensor, ...] | None:
Patrick von Platen's avatar
Patrick von Platen committed
380
381
382
383
384
385
386
387
388
389
        audio_inputs = self._parse_and_validate_audio_arrays(**kwargs)
        if audio_inputs is None:
            return None

        audio_embeddings = self.whisper_encoder(audio_inputs)

        for i, audio_embedding in enumerate(audio_embeddings):
            seq_len, dim = audio_embedding.shape
            # Pad such that seq_len is divisible by downsample_factor
            target_seq_len = self.downsample_factor * math.ceil(
390
391
                seq_len / self.downsample_factor
            )
Patrick von Platen's avatar
Patrick von Platen committed
392
393
394
395
396
            audio_embedding = torch.nn.functional.pad(
                audio_embedding,
                (0, 0, 0, target_seq_len - seq_len),
            )
            audio_embeddings[i] = audio_embedding.reshape(
397
398
                target_seq_len // self.downsample_factor, dim * self.downsample_factor
            )
Patrick von Platen's avatar
Patrick von Platen committed
399
400
401

        # Concat, project and resplit
        audio_embeddings_packed = torch.cat(audio_embeddings, dim=0)
402
403
404
405
        audio_embeddings_packed = self.audio_language_adapter(audio_embeddings_packed)
        audio_embeddings = torch.split(
            audio_embeddings_packed, [a.shape[0] for a in audio_embeddings], dim=0
        )
Patrick von Platen's avatar
Patrick von Platen committed
406
407
408
409

        return audio_embeddings

    def _parse_and_validate_audio_arrays(
410
        self, **kwargs: object
411
    ) -> list[torch.Tensor] | None:
Patrick von Platen's avatar
Patrick von Platen committed
412
413
414
415
416
        audio_arrays = kwargs.pop("audio_arrays", None)
        if audio_arrays is None:
            return None

        if not isinstance(audio_arrays, (torch.Tensor, list)):
417
418
419
            raise ValueError(
                f"Incorrect type of audio_arrays. Got type: {type(audio_arrays)}"
            )
Patrick von Platen's avatar
Patrick von Platen committed
420
421
422
423
424
425
426
427

        if isinstance(audio_arrays, torch.Tensor):
            audio_arrays = list(audio_arrays.unbind(0))
        return audio_arrays

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
428
    ) -> torch.Tensor | None:
429
        return self.language_model.compute_logits(hidden_states)
Patrick von Platen's avatar
Patrick von Platen committed
430
431

    @classmethod
432
    def get_speech_to_text_config(
433
        cls, model_config: ModelConfig, task_type: str
434
    ) -> SpeechToTextConfig:
435
        tokenizer = cached_tokenizer_from_config(model_config)
Patrick von Platen's avatar
Patrick von Platen committed
436
437
438
439
440
441
442
443
444
445
446
447
        audio_config = tokenizer.instruct.audio_encoder.audio_config
        max_audio_clip_s = audio_config.chunk_length_s
        sample_rate = audio_config.sampling_rate
        return SpeechToTextConfig(
            max_audio_clip_s=max_audio_clip_s,
            sample_rate=sample_rate,
            # mistral_common and whisper encoder take care of chunking
            min_energy_split_window_size=None,
        )

    @classmethod
    # for speech-to-text transcription
448
449
    def get_generation_prompt(
        cls,
450
        stt_params: SpeechToTextParams,
451
    ) -> PromptType:
452
453
454
455
456
        audio = stt_params.audio
        model_config = stt_params.model_config
        stt_config = stt_params.stt_config
        language = stt_params.language

457
        tokenizer = cached_tokenizer_from_config(model_config)
458
459
        audio = Audio(audio, int(stt_config.sample_rate), format="wav")  # lossless
        req = TranscriptionRequest(
460
            model=model_config.model,
461
462
463
            audio=RawAudio.from_audio(audio),
            language=language,
        )
Patrick von Platen's avatar
Patrick von Platen committed
464
465

        tokenized = tokenizer.instruct.encode_transcription(req)
466
467
468
469

        return TokensPrompt(
            prompt_token_ids=tokenized.tokens,
            multi_modal_data={
470
471
472
473
                "audio": [
                    (audio.audio_array, stt_config.sample_rate)
                    for audio in tokenized.audios
                ],
474
475
            },
        )
Patrick von Platen's avatar
Patrick von Platen committed
476
477

    @classmethod
478
479
480
481
    def get_num_audio_tokens(
        cls,
        audio_duration_s: float,
        stt_config: SpeechToTextConfig,
482
        model_config: ModelConfig,
483
    ) -> int | None:
Patrick von Platen's avatar
Patrick von Platen committed
484
        """
485
        Map from audio duration to number of audio tokens produced by the ASR
Patrick von Platen's avatar
Patrick von Platen committed
486
487
488
        model, without running a forward pass.
        This is used for estimating the amount of processing for this audio.
        """
489
        tokenizer = cached_tokenizer_from_config(model_config)
490
491
492
493
        feature_extractor = MistralCommonFeatureExtractor(
            tokenizer.instruct.audio_encoder
        )
        return feature_extractor.get_num_audio_tokens(
494
495
            int(audio_duration_s * stt_config.sample_rate)
        )
Patrick von Platen's avatar
Patrick von Platen committed
496

497
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Patrick von Platen's avatar
Patrick von Platen committed
498
        remapping_rules = [
Patrick von Platen's avatar
Patrick von Platen committed
499
            (r"mm_streams_embeddings.embedding_module\.(.*)", r"\1"),
Patrick von Platen's avatar
Patrick von Platen committed
500
501
            (r"mm_whisper_embeddings\.(.*)", r"\1"),
            (r"audio_language_projection\.(.*)", r"audio_language_adapter.\1"),
502
503
504
505
506
507
508
509
            (
                r"audio_language_adapter\.0\.weight",
                r"audio_language_adapter.w_in.weight",
            ),
            (
                r"audio_language_adapter\.2\.weight",
                r"audio_language_adapter.w_out.weight",
            ),
Patrick von Platen's avatar
Patrick von Platen committed
510
511
512
        ]

        audio_params = dict(
513
514
515
516
517
518
            nn.ModuleDict(
                {
                    "audio_language_adapter": self.audio_language_adapter,
                }
            ).named_parameters()
        )
519
        weights = _create_fake_bias_for_k_proj(weights, ".wk.weight")
Patrick von Platen's avatar
Patrick von Platen committed
520
521
522
523
524
525

        loaded_weights = set()

        def llm_weights_generator():
            nonlocal loaded_weights
            for name, w in weights:
Patrick von Platen's avatar
Patrick von Platen committed
526
527
528
529
530
531
532
533
534
                is_encoder = False
                for k in [
                    "mm_whisper_embeddings",
                    "mm_streams_embeddings.embedding_module",
                ]:
                    is_encoder |= (
                        name.startswith(k)
                        and not name.startswith(f"{k}.tok_embeddings")
                        and not name.startswith(f"{k}.audio_language_projection")
535
                    )
Patrick von Platen's avatar
Patrick von Platen committed
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564

                for pattern, repl in remapping_rules:
                    if re.fullmatch(pattern, name):
                        name = re.sub(pattern, repl, name)

                if is_encoder:
                    name = self.whisper_encoder.load_weight((name, w))
                    loaded_weights.add(f"whisper_encoder.{name}")
                    continue

                if name in audio_params:
                    param = audio_params[name]
                    with torch.no_grad():
                        default_weight_loader(param, w)
                    loaded_weights.add(name)
                else:
                    yield (name, w)

        for name in self.language_model.load_weights(llm_weights_generator()):
            loaded_weights.add(f"language_model.{name}")

        # potentially manually add position embeddings
        sin_key = "whisper_encoder.whisper_encoder.embed_positions.weight"
        if sin_key not in loaded_weights:
            # make sure we don't hit an error here
            loaded_weights.add(sin_key)

        return loaded_weights

565
    def maybe_update_quant_config(
566
567
        self, quant_config: QuantizationConfig
    ) -> QuantizationConfig:
568
569
570
571
572
573
574
575
        """
        Update quant config to so that ignored module and target module names
        match the vLLM model names.
        Right now this is specific for compressed-tensors format and
        load_format mistral.
        """
        remapping_rules = [
            (r"output", r"language_model.lm_head"),
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
            (
                r"layers\.(\d+)\.attention\.wo",
                r"language_model.model.layers.\1.self_attn.out_proj",
            ),
            (
                r"layers\.(\d+)\.attention\.w(.*)",
                r"language_model.model.layers.\1.self_attn.\2_proj",
            ),
            (
                r"layers\.(\d+)\.feed_forward\.w1",
                r"language_model.model.layers.\1.mlp.gate_proj",
            ),
            (
                r"layers\.(\d+)\.feed_forward\.w2",
                r"language_model.model.layers.\1.mlp.down_proj",
            ),
            (
                r"layers\.(\d+)\.feed_forward\.w3",
                r"language_model.model.layers.\1.mlp.up_proj",
            ),
            (
                r"mm_whisper_embeddings\.whisper_encoder\.transformer\.layers\.(\d+)\.attention.w(.*)",
                r"whisper_encoder.whisper_encoder.layers.\1.layers.self_attn.\2_proj",
            ),
            (
                r"mm_whisper_embeddings\.whisper_encoder\.transformer\.layers\.(\d+)\.attention.wo",
                r"whisper_encoder.whisper_encoder.layers.\1.layers.self_attn.out_proj",
            ),
            (
                r"mm_whisper_embeddings\.whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward.w(\d+)",
                r"whisper_encoder.whisper_encoder.layers.\1.layers.mlp.fc\2",
            ),
            (
                r"mm_whisper_embeddings\.whisper_encoder\.conv_layers\.0",
                r"whisper_encoder.whisper_encoder.conv1",
            ),
            (
                r"mm_whisper_embeddings\.whisper_encoder\.conv_layers\.1",
                r"whisper_encoder.whisper_encoder.conv2",
            ),
            (
                r"mm_whisper_embeddings\.audio_language_projection\.0",
                r"audio_language_adapter.w_in",
            ),
            (
                r"mm_whisper_embeddings\.audio_language_projection\.2",
                r"audio_language_adapter.w_out",
            ),
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
        ]

        # Update ignore list
        if hasattr(quant_config, "ignore"):
            mistral_ignore = []
            for name in quant_config.ignore:
                mistral_name = name
                for pattern, repl in remapping_rules:
                    if re.fullmatch(pattern, name):
                        mistral_name = re.sub(pattern, repl, name)
                mistral_ignore.append(mistral_name)
            quant_config.ignore = mistral_ignore

        # Update target list
        if hasattr(quant_config, "config_groups"):
            config_groups = quant_config.config_groups
            for group_name in config_groups:
                if "targets" in config_groups[group_name]:
                    targets = []
                    for name in config_groups[group_name]["targets"]:
                        mistral_name = name
                        for pattern, repl in remapping_rules:
                            if re.fullmatch(pattern, name):
                                mistral_name = re.sub(pattern, repl, name)
                        targets.append(mistral_name)
                config_groups[group_name]["targets"] = targets
            quant_config.config_groups = config_groups

        return quant_config

Patrick von Platen's avatar
Patrick von Platen committed
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669

class AudioLanguageAdapter(nn.Module):
    def __init__(self, hidden_size: int, dim: int) -> None:
        super().__init__()
        self.w_in = nn.Linear(hidden_size, dim, bias=False)
        self.gelu = nn.GELU()
        self.w_out = nn.Linear(dim, dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w_out(self.gelu(self.w_in(x)))


class VoxtralEncoderModel(nn.Module):
    packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}

    mistral_remapping = [
Patrick von Platen's avatar
Patrick von Platen committed
670
        (r"mm_streams_embeddings.embedding_module\.(.*)", r"\1"),
671
672
673
674
675
676
677
678
        (
            r"whisper_encoder\.conv_layers\.0\.(weight|bias)",
            r"whisper_encoder.conv1.\1",
        ),
        (
            r"whisper_encoder\.conv_layers\.1\.(weight|bias)",
            r"whisper_encoder.conv2.\1",
        ),
Patrick von Platen's avatar
Patrick von Platen committed
679
680
681
682
683
684
685
686
        (
            r"whisper_encoder\.conv_layers\.0\.conv\.(weight|bias)",
            r"whisper_encoder.conv1.\1",
        ),  # noqa: E501
        (
            r"whisper_encoder\.conv_layers\.1\.conv\.(weight|bias)",
            r"whisper_encoder.conv2.\1",
        ),  # noqa: E501
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
        (
            r"whisper_encoder\.transformer\.layers\.(\d+)\.attention\.w([qkv])\.(weight|bias)",  # noqa: E501
            r"whisper_encoder.layers.\1.self_attn.\2_proj.\3",
        ),
        (
            r"whisper_encoder\.transformer\.layers\.(\d+)\.attention\.wo\.(weight|bias)",  # noqa: E501
            r"whisper_encoder.layers.\1.self_attn.out_proj.\2",
        ),
        (
            r"whisper_encoder\.transformer\.layers\.(\d+)\.attention_norm\.(weight|bias)",  # noqa: E501
            r"whisper_encoder.layers.\1.self_attn_layer_norm.\2",
        ),
        (
            r"whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward\.w1\.(weight|bias)",  # noqa: E501
            r"whisper_encoder.layers.\1.mlp.fc1.\2",
        ),
        (
            r"whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward\.w2\.(weight|bias)",  # noqa: E501
            r"whisper_encoder.layers.\1.mlp.fc2.\2",
        ),
707
708
709
710
        (
            r"whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward\.w3\.(weight|bias)",
            r"whisper_encoder.layers.\1.mlp.fc3.\2",
        ),  # noqa: E501
711
712
713
714
715
716
717
718
        (
            r"whisper_encoder\.transformer\.layers\.(\d+)\.ffn_norm\.(weight|bias)",
            r"whisper_encoder.layers.\1.final_layer_norm.\2",
        ),
        (
            r"whisper_encoder\.transformer\.norm\.(weight|bias)",
            r"whisper_encoder.layer_norm.\1",
        ),
Patrick von Platen's avatar
Patrick von Platen committed
719
720
721
722
723
724
725
726
727
728
729
    ]

    def __init__(
        self,
        vllm_config: VllmConfig,
        *,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = cast(WhisperConfig, vllm_config.model_config.hf_config)
        self.dtype: torch.dtype = vllm_config.model_config.dtype
730
731
732
733
734
735
736
        self.is_causal = getattr(self.config, "is_causal", False)
        if self.is_causal:
            WhisperEncoderCls = WhisperCausalEncoder
        else:
            WhisperEncoderCls = partial(WhisperEncoder, init_in_fp32=True)

        self.whisper_encoder = WhisperEncoderCls(
737
738
739
            vllm_config=vllm_config,
            prefix=maybe_prefix(prefix, "whisper_encoder"),
        )
Patrick von Platen's avatar
Patrick von Platen committed
740
741
742
743
744
745
746
747
748
749
750
751
752
753
        mel_filters = mel_filter_bank(
            num_frequency_bins=1 + self.config.window_size // 2,
            num_mel_bins=self.config.num_mel_bins,
            min_frequency=0.0,
            max_frequency=8000.0,
            sampling_rate=self.config.sampling_rate,
        )
        self.mel_filters = torch.tensor(mel_filters, dtype=torch.float32)

    def compute_whisper_melspec(
        self,
        audio_waveforms: torch.Tensor,
    ) -> torch.Tensor:
        input_dtype = audio_waveforms.dtype
Andy Lo's avatar
Andy Lo committed
754
755
756
        window = torch.hann_window(
            self.config.window_size, device=audio_waveforms.device
        )
Patrick von Platen's avatar
Patrick von Platen committed
757
758
759
760
761
762
763
        stft = torch.stft(
            audio_waveforms,
            self.config.window_size,
            self.config.hop_length,
            window=window,
            return_complex=True,
        )
764
        magnitudes = stft[..., :-1].abs() ** 2
Patrick von Platen's avatar
Patrick von Platen committed
765
766
        mel_spec = self.mel_filters.T @ magnitudes
        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
767
768
769
770
771
772
773
774
775
776
777
778
779

        if global_log_mel_max := self.config.global_log_mel_max:
            if not isinstance(global_log_mel_max, float):
                raise TypeError(f"{global_log_mel_max=} needs to be of type float.")
            log_spec_max = torch.tensor(
                global_log_mel_max,
                device=log_spec.device,
                dtype=log_spec.dtype,
            )
        else:
            log_spec_max = log_spec.max()

        log_spec = torch.maximum(log_spec, log_spec_max - 8.0)
Patrick von Platen's avatar
Patrick von Platen committed
780
781
782
783
784
        log_spec = (log_spec + 4.0) / 4.0
        return log_spec.to(input_dtype)

    @property
    def downsample_factor(self) -> int:
785
786
787
        return (
            self.whisper_encoder.conv1.stride[0] * self.whisper_encoder.conv2.stride[0]
        )
Patrick von Platen's avatar
Patrick von Platen committed
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814

    @property
    def chunk_size(self) -> int:
        return self.config.max_source_positions * self.downsample_factor

    def prepare_inputs_for_conv(
        self,
        audio_waveforms: list[torch.Tensor],
    ) -> tuple[torch.Tensor, list[int]]:
        assert isinstance(audio_waveforms, list)
        # list[num_mel_bins, seq_len]
        input_features = [
            self.compute_whisper_melspec(audio).to(self.dtype)
            for audio in audio_waveforms
        ]

        chunked_features: list[torch.Tensor] = []
        chunks_per_example: list[int] = []
        for feature in input_features:
            chunks = feature.split(self.chunk_size, dim=-1)
            chunked_features += chunks
            chunks_per_example.append(len(chunks))

        # [total_num_chunks, num_mel_bins, chunk_size]
        return torch.stack(chunked_features), chunks_per_example

    def forward(
815
        self, input_features: torch.Tensor | list[torch.Tensor]
Patrick von Platen's avatar
Patrick von Platen committed
816
817
818
819
820
    ) -> list[torch.Tensor]:
        if not isinstance(input_features, list):
            input_features = [input_features]

        # Split long inputs into chunks
821
        input_embeds, chunks_per_example = self.prepare_inputs_for_conv(input_features)
Patrick von Platen's avatar
Patrick von Platen committed
822
823
824
825
826
827
828
829

        # [total_num_chunks, ceil(chunk_size / downsample_factor), hidden_size]
        out = self.whisper_encoder([input_embeds])

        # Re-concatenate the chunks
        chunk_idx = 0
        results = []
        for n_chunks in chunks_per_example:
830
            result = out[chunk_idx : chunk_idx + n_chunks].flatten(0, 1)
Patrick von Platen's avatar
Patrick von Platen committed
831
832
833
834
835
836
837
838
839
840
841
842
            results.append(result)
            chunk_idx += n_chunks

        return results

    def load_weight(self, weight: tuple[str, torch.Tensor]) -> str:
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
        ]
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
        params_mapping = []

        if self.is_causal:
            # For `WhisperCausalEncoder` we need
            # some more renaming
            stacked_params_mapping.extend(
                [
                    (".mlp.gate_up_proj", ".mlp.fc1", 0),
                    (".mlp.gate_up_proj", ".mlp.fc3", 1),
                ]
            )
            params_mapping.extend(
                [
                    (".mlp.down_proj", ".mlp.fc2"),
                ]
            )
Patrick von Platen's avatar
Patrick von Platen committed
859
860
861
862
863
864
865
        params_dict = dict(self.named_parameters())

        name, loaded_weight = weight
        for pattern, repl in self.mistral_remapping:
            if re.fullmatch(pattern, name):
                name = re.sub(pattern, repl, name)

866
        for param_name, weight_name, shard_id in stacked_params_mapping:
Patrick von Platen's avatar
Patrick von Platen committed
867
868
869
870
871
872
873
874
875
            if weight_name not in name:
                continue
            name = name.replace(weight_name, param_name)

            param = params_dict[name]
            weight_loader = param.weight_loader
            weight_loader(param, loaded_weight, shard_id)
            break
        else:
876
877
878
879
880
            for param_name, weight_name in params_mapping:
                if weight_name not in name:
                    continue
                name = name.replace(weight_name, param_name)

Patrick von Platen's avatar
Patrick von Platen committed
881
            param = params_dict[name]
882
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
Patrick von Platen's avatar
Patrick von Platen committed
883
884
885
            weight_loader(param, loaded_weight)

        return name