| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import enum |
| 5 | import math |
| 6 | from collections.abc import Iterable, Mapping, Sequence |
| 7 | from contextlib import nullcontext |
| 8 | from typing import Annotated |
| 9 | |
| 10 | import numpy as np |
| 11 | import torch |
| 12 | from torch import nn |
| 13 | from transformers import ( |
| 14 | BatchFeature, |
| 15 | WhisperConfig, |
| 16 | WhisperFeatureExtractor, |
| 17 | ) |
| 18 | from transformers.models.whisper.modeling_whisper import sinusoids |
| 19 | |
| 20 | from vllm.compilation.decorators import support_torch_compile |
| 21 | from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig |
| 22 | from vllm.config.multimodal import MultiModalDummyOptions |
| 23 | from vllm.config.speech_to_text import SpeechToTextParams |
| 24 | from vllm.distributed import get_tensor_model_parallel_world_size |
| 25 | from vllm.inputs import ( |
| 26 | ExplicitEncoderDecoderPrompt, |
| 27 | MultiModalDataDict, |
| 28 | PromptType, |
| 29 | TextPrompt, |
| 30 | ) |
| 31 | from vllm.logger import init_logger |
| 32 | from vllm.model_executor.layers.activation import get_act_fn |
| 33 | from vllm.model_executor.layers.attention import ( |
| 34 | Attention, |
| 35 | CrossAttention, |
| 36 | MMEncoderAttention, |
| 37 | ) |
| 38 | from vllm.model_executor.layers.linear import ( |
| 39 | ColumnParallelLinear, |
| 40 | MergedColumnParallelLinear, |
| 41 | QKVParallelLinear, |
| 42 | RowParallelLinear, |
| 43 | ) |
| 44 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 45 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 46 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 47 | from vllm.model_executor.models.whisper_utils import ( |
| 48 | ISO639_1_SUPPORTED_LANGS, |
| 49 | ) |
| 50 | from vllm.model_executor.offloader import get_offloader |
| 51 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 52 | from vllm.multimodal.inputs import ( |
| 53 | MultiModalFieldConfig, |
| 54 | MultiModalKwargsItems, |
| 55 | ) |
| 56 | from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser |
| 57 | from vllm.multimodal.processing import ( |
| 58 | BaseDummyInputsBuilder, |
| 59 | BaseProcessingInfo, |
| 60 | EncDecMultiModalProcessor, |
| 61 | PromptReplacement, |
| 62 | PromptUpdate, |
| 63 | ) |
| 64 | from vllm.multimodal.processing.processor import HFMultiModalInputs |
| 65 | from vllm.renderers import TokenizeParams |
| 66 | from vllm.tokenizers import TokenizerLike |
| 67 | from vllm.transformers_utils.processor import cached_processor_from_config |
| 68 | from vllm.utils.jsontree import json_map_leaves |
| 69 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 70 | from vllm.utils.torch_utils import set_default_torch_dtype |
| 71 | from vllm.v1.attention.backend import ( |
| 72 | AttentionType, |
| 73 | ) |
| 74 | |
| 75 | from .interfaces import ( |
| 76 | MultiModalEmbeddings, |
| 77 | SupportsLoRA, |
| 78 | SupportsMultiModal, |
| 79 | SupportsTranscription, |
| 80 | ) |
| 81 | from .utils import ( |
| 82 | AutoWeightsLoader, |
| 83 | WeightsMapper, |
| 84 | cast_overflow_tensors, |
| 85 | make_layers, |
| 86 | maybe_prefix, |
| 87 | ) |
| 88 | |
| 89 | logger = init_logger(__name__) |
| 90 | |
| 91 | |
| 92 | class WhisperPosEmbedType(enum.Enum): |
| 93 | SINUSOIDAL = "sinusoidal" |
| 94 | ROPE = "rope" |
| 95 | LEARNED = "learned" |
| 96 | |
| 97 | |
| 98 | class WhisperAudioInputs(TensorSchema): |
| 99 | """Dimensions: |
| 100 | - b: Batch size |
| 101 | - nmb: Number of mel bins |
| 102 | - t: Time frames (M) |
| 103 | """ |
| 104 | |
| 105 | input_features: Annotated[ |
| 106 | list[torch.Tensor] | None, |
| 107 | TensorShape("b", "nmb", "t"), |
| 108 | ] |
| 109 | |
| 110 | |
| 111 | class WhisperEncoderAttention(MMEncoderAttention): |
| 112 | """Multi-headed attention for Whisper encoder with 2D tensor support.""" |
| 113 | |
| 114 | def forward( |
| 115 | self, |
| 116 | query: torch.Tensor, |
| 117 | key: torch.Tensor, |
| 118 | value: torch.Tensor, |
| 119 | cu_seqlens: torch.Tensor | None = None, |
| 120 | max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention |
| 121 | # Only used for FlashInfer CuDNN backend. |
| 122 | sequence_lengths: torch.Tensor | None = None, |
| 123 | ) -> torch.Tensor: |
| 124 | """Input shape: batch_size x seq_len x hidden_size |
| 125 | or seq_len x hidden_size |
| 126 | """ |
| 127 | is_2d = query.dim() == 2 |
| 128 | if is_2d: |
| 129 | query = query.unsqueeze(0) |
| 130 | key = key.unsqueeze(0) |
| 131 | value = value.unsqueeze(0) |
| 132 | |
| 133 | # Call the parent forward method |
| 134 | out = super().forward( |
| 135 | query, |
| 136 | key, |
| 137 | value, |
| 138 | cu_seqlens=cu_seqlens, |
| 139 | max_seqlen=max_seqlen, |
| 140 | sequence_lengths=sequence_lengths, |
| 141 | ) |
| 142 | |
| 143 | if is_2d: |
| 144 | out = out.squeeze(0) |
| 145 | |
| 146 | return out |
| 147 | |
| 148 | |
| 149 | class WhisperPositionalEmbedding(nn.Embedding): |
| 150 | def __init__(self, num_positions: int, embedding_dim: int): |
| 151 | super().__init__(num_positions, embedding_dim) |
| 152 | |
| 153 | def forward(self, position_ids): |
| 154 | return self.weight[position_ids] |
| 155 | |
| 156 | |
| 157 | class WhisperAttention(nn.Module): |
| 158 | def __init__( |
| 159 | self, |
| 160 | embed_dim: int, |
| 161 | num_heads: int, |
| 162 | bias: bool = True, |
| 163 | attn_type: AttentionType = AttentionType.DECODER, |
| 164 | per_layer_sliding_window: int | None = None, |
| 165 | cache_config: CacheConfig | None = None, |
| 166 | quant_config: QuantizationConfig | None = None, |
| 167 | prefix: str = "", |
| 168 | ): |
| 169 | super().__init__() |
| 170 | self.embed_dim = embed_dim |
| 171 | tp_size = get_tensor_model_parallel_world_size() |
| 172 | self.total_num_heads = num_heads |
| 173 | assert self.total_num_heads % tp_size == 0 |
| 174 | self.num_heads = self.total_num_heads // tp_size |
| 175 | if self.total_num_heads >= tp_size: |
| 176 | # Number of heads is greater than TP size, so we partition |
| 177 | # the KV heads across multiple tensor parallel GPUs. |
| 178 | assert self.total_num_heads % tp_size == 0 |
| 179 | else: |
| 180 | # Number of heads is less than TP size, so we replicate |
| 181 | # the KV heads across multiple tensor parallel GPUs. |
| 182 | assert tp_size % self.total_num_heads == 0 |
| 183 | self.num_kv_heads = max(1, self.total_num_heads // tp_size) |
| 184 | self.head_dim = self.embed_dim // self.total_num_heads |
| 185 | self.q_size = self.num_heads * self.head_dim |
| 186 | self.kv_size = self.num_kv_heads * self.head_dim |
| 187 | self.attn_type = attn_type |
| 188 | |
| 189 | if (self.head_dim * num_heads) != self.embed_dim: |
| 190 | raise ValueError( |
| 191 | f"embed_dim must be divisible by num_heads (got `embed_dim`: " |
| 192 | f"{self.embed_dim} and `num_heads`: {num_heads})." |
| 193 | ) |
| 194 | self.scaling = self.head_dim**-0.5 |
| 195 | |
| 196 | self._init_qkv(embed_dim, bias, quant_config, prefix=prefix) |
| 197 | self.out_proj = RowParallelLinear( |
| 198 | input_size=embed_dim, |
| 199 | output_size=embed_dim, |
| 200 | bias=bias, |
| 201 | quant_config=quant_config, |
| 202 | prefix=f"{prefix}.out_proj", |
| 203 | ) |
| 204 | if attn_type == AttentionType.ENCODER: |
| 205 | self.attn = WhisperEncoderAttention( |
| 206 | self.num_heads, |
| 207 | self.head_dim, |
| 208 | self.scaling, |
| 209 | num_kv_heads=self.num_kv_heads, |
| 210 | ) |
| 211 | elif self.attn_type == AttentionType.ENCODER_DECODER: |
| 212 | self.attn = CrossAttention( |
| 213 | self.num_heads, |
| 214 | self.head_dim, |
| 215 | self.scaling, |
| 216 | num_kv_heads=self.num_kv_heads, |
| 217 | cache_config=cache_config, |
| 218 | quant_config=quant_config, |
| 219 | prefix=f"{prefix}.attn", |
| 220 | attn_type=self.attn_type, |
| 221 | ) |
| 222 | else: # AttentionType.DECODER (regular decoder self-attention) |
| 223 | self.attn = Attention( |
| 224 | self.num_heads, |
| 225 | self.head_dim, |
| 226 | self.scaling, |
| 227 | num_kv_heads=self.num_kv_heads, |
| 228 | cache_config=cache_config, |
| 229 | quant_config=quant_config, |
| 230 | prefix=f"{prefix}.attn", |
| 231 | attn_type=self.attn_type, |
| 232 | per_layer_sliding_window=per_layer_sliding_window, |
| 233 | ) |
| 234 | |
| 235 | def _init_qkv( |
| 236 | self, |
| 237 | embed_dim: int, |
| 238 | bias: bool = True, |
| 239 | quant_config: QuantizationConfig | None = None, |
| 240 | prefix: str = "", |
| 241 | ) -> None: |
| 242 | self.qkv_proj = QKVParallelLinear( |
| 243 | hidden_size=embed_dim, |
| 244 | head_size=self.head_dim, |
| 245 | total_num_heads=self.total_num_heads, |
| 246 | total_num_kv_heads=self.total_num_heads, |
| 247 | bias=bias, |
| 248 | quant_config=quant_config, |
| 249 | prefix=f"{prefix}.qkv_proj", |
| 250 | ) |
| 251 | |
| 252 | def forward( |
| 253 | self, |
| 254 | hidden_states: torch.Tensor, |
| 255 | cu_seqlens: torch.Tensor | None = None, |
| 256 | max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention |
| 257 | # Only used for FlashInfer CuDNN backend. |
| 258 | sequence_lengths: torch.Tensor | None = None, |
| 259 | ): |
| 260 | qkv, _ = self.qkv_proj(hidden_states) |
| 261 | q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) |
| 262 | |
| 263 | if cu_seqlens is None: |
| 264 | attn_output = self.attn(q, k, v) |
| 265 | else: |
| 266 | assert self.attn_type == AttentionType.ENCODER, ( |
| 267 | "Variable-length attention metadata is only supported for " |
| 268 | "encoder self-attention." |
| 269 | ) |
| 270 | attn_output = self.attn( |
| 271 | q, |
| 272 | k, |
| 273 | v, |
| 274 | cu_seqlens=cu_seqlens, |
| 275 | max_seqlen=max_seqlen, |
| 276 | sequence_lengths=sequence_lengths, |
| 277 | ) |
| 278 | |
| 279 | output, _ = self.out_proj(attn_output) |
| 280 | |
| 281 | return output |
| 282 | |
| 283 | |
| 284 | class WhisperCrossAttention(WhisperAttention): |
| 285 | def __init__( |
| 286 | self, |
| 287 | embed_dim: int, |
| 288 | num_heads: int, |
| 289 | bias: bool = True, |
| 290 | cache_config: CacheConfig | None = None, |
| 291 | quant_config: QuantizationConfig | None = None, |
| 292 | prefix: str = "", |
| 293 | ): |
| 294 | super().__init__( |
| 295 | embed_dim=embed_dim, |
| 296 | num_heads=num_heads, |
| 297 | bias=bias, |
| 298 | cache_config=cache_config, |
| 299 | quant_config=quant_config, |
| 300 | prefix=prefix, |
| 301 | attn_type=AttentionType.ENCODER_DECODER, |
| 302 | ) |
| 303 | |
| 304 | def _init_qkv( |
| 305 | self, |
| 306 | embed_dim: int, |
| 307 | bias: bool = True, |
| 308 | quant_config: QuantizationConfig | None = None, |
| 309 | prefix: str = "", |
| 310 | ) -> None: |
| 311 | self.q_proj = ColumnParallelLinear( |
| 312 | input_size=embed_dim, |
| 313 | output_size=embed_dim, |
| 314 | bias=bias, |
| 315 | quant_config=quant_config, |
| 316 | prefix=f"{prefix}.q_proj", |
| 317 | ) |
| 318 | # Use MergedColumnParallelLinear for K and V projections. |
| 319 | # This enables LoRA support via MergedColumnParallelLinearWithLoRA |
| 320 | # which handles 2-slice configurations. |
| 321 | self.kv_proj = MergedColumnParallelLinear( |
| 322 | input_size=embed_dim, |
| 323 | output_sizes=[embed_dim, embed_dim], |
| 324 | bias=bias, |
| 325 | quant_config=quant_config, |
| 326 | prefix=f"{prefix}.kv_proj", |
| 327 | ) |
| 328 | |
| 329 | def forward( # type: ignore[override] |
| 330 | self, |
| 331 | hidden_states: torch.Tensor, |
| 332 | encoder_hidden_states: torch.Tensor | None, |
| 333 | ): |
| 334 | q, _ = self.q_proj(hidden_states) |
| 335 | |
| 336 | # Encoder hidden states are only computed once during prefill phase. |
| 337 | # Afterwards, the keys and values should be available in the kv-cache. |
| 338 | if encoder_hidden_states is not None: |
| 339 | kv, _ = self.kv_proj(encoder_hidden_states) |
| 340 | k, v = kv.split([self.kv_size, self.kv_size], dim=-1) |
| 341 | else: |
| 342 | k = v = None |
| 343 | |
| 344 | attn_output = self.attn(q, k, v) |
| 345 | |
| 346 | output, _ = self.out_proj(attn_output) |
| 347 | |
| 348 | return output |
| 349 | |
| 350 | |
| 351 | class WhisperMLP(nn.Module): |
| 352 | def __init__( |
| 353 | self, |
| 354 | embed_dim: int, |
| 355 | ffn_dim: int, |
| 356 | act_fn: str, |
| 357 | quant_config: QuantizationConfig | None = None, |
| 358 | prefix: str = "", |
| 359 | ): |
| 360 | super().__init__() |
| 361 | |
| 362 | self.activation_fn = get_act_fn(act_fn) |
| 363 | self.fc1 = ColumnParallelLinear( |
| 364 | input_size=embed_dim, |
| 365 | output_size=ffn_dim, |
| 366 | quant_config=quant_config, |
| 367 | prefix=f"{prefix}.fc1", |
| 368 | ) |
| 369 | self.fc2 = RowParallelLinear( |
| 370 | input_size=ffn_dim, |
| 371 | output_size=embed_dim, |
| 372 | quant_config=quant_config, |
| 373 | prefix=f"{prefix}.fc2", |
| 374 | ) |
| 375 | |
| 376 | def forward(self, hidden_states: torch.Tensor): |
| 377 | hidden_states, _ = self.fc1(hidden_states) |
| 378 | hidden_states = self.activation_fn(hidden_states) |
| 379 | hidden_states, _ = self.fc2(hidden_states) |
| 380 | return hidden_states |
| 381 | |
| 382 | |
| 383 | class WhisperEncoderLayer(nn.Module): |
| 384 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 385 | super().__init__() |
| 386 | config = vllm_config.model_config.hf_config |
| 387 | sliding_window = getattr(config, "sliding_window", None) |
| 388 | cache_config = vllm_config.cache_config |
| 389 | quant_config = vllm_config.quant_config |
| 390 | |
| 391 | self.embed_dim = config.d_model |
| 392 | self.self_attn = WhisperAttention( |
| 393 | embed_dim=self.embed_dim, |
| 394 | num_heads=config.encoder_attention_heads, |
| 395 | attn_type=AttentionType.ENCODER, |
| 396 | per_layer_sliding_window=sliding_window, |
| 397 | cache_config=cache_config, |
| 398 | quant_config=quant_config, |
| 399 | prefix=f"{prefix}.self_attn", |
| 400 | ) |
| 401 | self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) |
| 402 | self.mlp = WhisperMLP( |
| 403 | embed_dim=config.d_model, |
| 404 | ffn_dim=config.encoder_ffn_dim, |
| 405 | act_fn=config.activation_function, |
| 406 | quant_config=quant_config, |
| 407 | prefix=f"{prefix}.mlp", |
| 408 | ) |
| 409 | self.final_layer_norm = nn.LayerNorm(self.embed_dim) |
| 410 | |
| 411 | def forward( |
| 412 | self, |
| 413 | hidden_states: torch.Tensor, |
| 414 | cu_seqlens: torch.Tensor | None = None, |
| 415 | max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention |
| 416 | # Only used for FlashInfer CuDNN backend. |
| 417 | sequence_lengths: torch.Tensor | None = None, |
| 418 | ): |
| 419 | residual = hidden_states |
| 420 | hidden_states = self.self_attn_layer_norm(hidden_states) |
| 421 | hidden_states = self.self_attn( |
| 422 | hidden_states=hidden_states, |
| 423 | cu_seqlens=cu_seqlens, |
| 424 | max_seqlen=max_seqlen, |
| 425 | sequence_lengths=sequence_lengths, |
| 426 | ) |
| 427 | hidden_states = residual + hidden_states |
| 428 | residual = hidden_states |
| 429 | hidden_states = self.final_layer_norm(hidden_states) |
| 430 | hidden_states = self.mlp(hidden_states) |
| 431 | hidden_states = residual + hidden_states |
| 432 | |
| 433 | hidden_states = cast_overflow_tensors(hidden_states) |
| 434 | |
| 435 | return hidden_states |
| 436 | |
| 437 | |
| 438 | class WhisperDecoderLayer(nn.Module): |
| 439 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 440 | super().__init__() |
| 441 | config = vllm_config.model_config.hf_config |
| 442 | cache_config = vllm_config.cache_config |
| 443 | quant_config = vllm_config.quant_config |
| 444 | |
| 445 | self.self_attn = WhisperAttention( |
| 446 | embed_dim=config.d_model, |
| 447 | num_heads=config.decoder_attention_heads, |
| 448 | attn_type=AttentionType.DECODER, |
| 449 | cache_config=cache_config, |
| 450 | quant_config=quant_config, |
| 451 | prefix=f"{prefix}.self_attn", |
| 452 | ) |
| 453 | self.self_attn_layer_norm = nn.LayerNorm(config.d_model) |
| 454 | self.encoder_attn = WhisperCrossAttention( |
| 455 | embed_dim=config.d_model, |
| 456 | num_heads=config.decoder_attention_heads, |
| 457 | cache_config=cache_config, |
| 458 | quant_config=quant_config, |
| 459 | prefix=f"{prefix}.encoder_attn", |
| 460 | ) |
| 461 | self.encoder_attn_layer_norm = nn.LayerNorm(config.d_model) |
| 462 | self.mlp = WhisperMLP( |
| 463 | embed_dim=config.d_model, |
| 464 | ffn_dim=config.decoder_ffn_dim, |
| 465 | act_fn=config.activation_function, |
| 466 | quant_config=quant_config, |
| 467 | prefix=f"{prefix}.mlp", |
| 468 | ) |
| 469 | self.final_layer_norm = nn.LayerNorm(config.d_model) |
| 470 | |
| 471 | def forward( |
| 472 | self, |
| 473 | hidden_states: torch.Tensor, |
| 474 | encoder_hidden_states: torch.Tensor | None, |
| 475 | ): |
| 476 | residual = hidden_states |
| 477 | hidden_states = self.self_attn_layer_norm(hidden_states) |
| 478 | hidden_states = self.self_attn(hidden_states=hidden_states) |
| 479 | hidden_states = residual + hidden_states |
| 480 | |
| 481 | residual = hidden_states |
| 482 | hidden_states = self.encoder_attn_layer_norm(hidden_states) |
| 483 | hidden_states = self.encoder_attn( |
| 484 | hidden_states=hidden_states, |
| 485 | encoder_hidden_states=encoder_hidden_states, |
| 486 | ) |
| 487 | hidden_states = residual + hidden_states |
| 488 | |
| 489 | residual = hidden_states |
| 490 | hidden_states = self.final_layer_norm(hidden_states) |
| 491 | hidden_states = self.mlp(hidden_states) |
| 492 | hidden_states = residual + hidden_states |
| 493 | |
| 494 | return hidden_states |
| 495 | |
| 496 | |
| 497 | class WhisperEncoder(nn.Module): |
| 498 | def __init__( |
| 499 | self, |
| 500 | *, |
| 501 | vllm_config: VllmConfig, |
| 502 | prefix: str = "", |
| 503 | init_in_fp32: bool = False, |
| 504 | enable_pp: bool = True, |
| 505 | ): |
| 506 | super().__init__() |
| 507 | config = vllm_config.model_config.hf_config |
| 508 | embed_dim = config.d_model |
| 509 | |
| 510 | self.pos_embed_type = WhisperPosEmbedType( |
| 511 | getattr(config, "pos_embed", "sinusoidal") |
| 512 | ) |
| 513 | self.num_mel_bins = config.num_mel_bins |
| 514 | self.max_source_positions = config.max_source_positions |
| 515 | self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0 |
| 516 | |
| 517 | self.conv1 = nn.Conv1d(self.num_mel_bins, embed_dim, kernel_size=3, padding=1) |
| 518 | self.conv2 = nn.Conv1d(embed_dim, embed_dim, stride=2, kernel_size=3, padding=1) |
| 519 | |
| 520 | self.total_stride = self.conv1.stride[0] * self.conv2.stride[0] |
| 521 | |
| 522 | def create_layer(prefix: str) -> nn.Module: |
| 523 | return WhisperEncoderLayer( |
| 524 | vllm_config=vllm_config, prefix=f"{prefix}.layers" |
| 525 | ) |
| 526 | |
| 527 | if enable_pp: |
| 528 | self.start_layer, self.end_layer, self.layers = make_layers( |
| 529 | config.encoder_layers, |
| 530 | create_layer, |
| 531 | prefix=f"{prefix}.layers", |
| 532 | ) |
| 533 | else: |
| 534 | self.start_layer = 0 |
| 535 | self.end_layer = config.encoder_layers |
| 536 | self.layers = nn.ModuleList( |
| 537 | get_offloader().wrap_modules( |
| 538 | create_layer(prefix=f"{prefix}.layers.{idx}") |
| 539 | for idx in range(config.encoder_layers) |
| 540 | ) |
| 541 | ) |
| 542 | self.layer_norm = nn.LayerNorm(config.d_model) |
| 543 | |
| 544 | if self.pos_embed_type not in ( |
| 545 | WhisperPosEmbedType.SINUSOIDAL, |
| 546 | WhisperPosEmbedType.LEARNED, |
| 547 | ): |
| 548 | raise ValueError( |
| 549 | "Only sinusoidal or learned position embeddings are supported " |
| 550 | f"for non-causal models, but got {self.pos_embed_type}" |
| 551 | ) |
| 552 | |
| 553 | maybe_fp32_init_ctx = ( |
| 554 | set_default_torch_dtype(torch.float32) if init_in_fp32 else nullcontext() |
| 555 | ) |
| 556 | |
| 557 | with ( |
| 558 | torch.no_grad(), |
| 559 | maybe_fp32_init_ctx, |
| 560 | ): |
| 561 | self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim) |
| 562 | self.embed_positions.weight.copy_( |
| 563 | sinusoids(*self.embed_positions.weight.shape) |
| 564 | ) |
| 565 | |
| 566 | def forward( |
| 567 | self, input_features: torch.Tensor | list[torch.Tensor] |
| 568 | ) -> torch.Tensor: |
| 569 | hidden_states = [] |
| 570 | input_is_batched = False |
| 571 | for features in input_features: |
| 572 | embeds = nn.functional.gelu(self.conv1(features)) |
| 573 | embeds = nn.functional.gelu(self.conv2(embeds)) |
| 574 | |
| 575 | embeds = embeds.transpose(-1, -2) |
| 576 | embeds = (embeds + self.embed_positions.weight[: embeds.size(-2), :]).to( |
| 577 | embeds.dtype |
| 578 | ) |
| 579 | |
| 580 | hidden_states.append(embeds) |
| 581 | input_is_batched = embeds.ndim > 2 |
| 582 | # Input to MHA must be B x T x D |
| 583 | if input_is_batched: |
| 584 | # Models using WhisperEncoder may handle batching internally. |
| 585 | hidden_states = torch.cat(hidden_states) |
| 586 | else: |
| 587 | hidden_states = torch.stack(hidden_states, dim=0) |
| 588 | |
| 589 | for encoder_layer in self.layers: |
| 590 | hidden_states = encoder_layer(hidden_states) |
| 591 | |
| 592 | hidden_states = self.layer_norm(hidden_states) |
| 593 | return hidden_states |
| 594 | |
| 595 | |
| 596 | @support_torch_compile(dynamic_arg_dims={"input_ids": 0, "positions": -1}) |
| 597 | class WhisperDecoder(nn.Module): |
| 598 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 599 | super().__init__() |
| 600 | config = vllm_config.model_config.hf_config |
| 601 | self.layerdrop = config.decoder_layerdrop |
| 602 | self.padding_idx = config.pad_token_id |
| 603 | self.max_target_positions = config.max_target_positions |
| 604 | self.max_source_positions = config.max_source_positions |
| 605 | self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0 |
| 606 | |
| 607 | self.embed_tokens = nn.Embedding( |
| 608 | config.vocab_size, config.d_model, self.padding_idx |
| 609 | ) |
| 610 | self.embed_positions = WhisperPositionalEmbedding( |
| 611 | self.max_target_positions, config.d_model |
| 612 | ) |
| 613 | self.start_layer, self.end_layer, self.layers = make_layers( |
| 614 | config.decoder_layers, |
| 615 | lambda prefix: WhisperDecoderLayer( |
| 616 | vllm_config=vllm_config, prefix=f"{prefix}.layers" |
| 617 | ), |
| 618 | prefix=f"{prefix}.layers", |
| 619 | ) |
| 620 | self.layer_norm = nn.LayerNorm(config.d_model) |
| 621 | |
| 622 | def forward( |
| 623 | self, |
| 624 | input_ids, |
| 625 | positions: torch.Tensor, |
| 626 | encoder_hidden_states: torch.Tensor | None, |
| 627 | ): |
| 628 | inputs_embeds = self.embed_input_ids(input_ids) |
| 629 | positions = self.embed_positions(positions) |
| 630 | hidden_states = inputs_embeds + positions |
| 631 | |
| 632 | for decoder_layer in self.layers: |
| 633 | hidden_states = decoder_layer( |
| 634 | hidden_states, |
| 635 | encoder_hidden_states=encoder_hidden_states, |
| 636 | ) |
| 637 | |
| 638 | hidden_states = self.layer_norm(hidden_states) |
| 639 | return hidden_states |
| 640 | |
| 641 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 642 | return self.embed_tokens(input_ids) |
| 643 | |
| 644 | |
| 645 | class WhisperModel(nn.Module): |
| 646 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 647 | super().__init__() |
| 648 | self.encoder = WhisperEncoder( |
| 649 | vllm_config=vllm_config, prefix=f"{prefix}.encoder" |
| 650 | ) |
| 651 | self.decoder = WhisperDecoder( |
| 652 | vllm_config=vllm_config, prefix=f"{prefix}.decoder" |
| 653 | ) |
| 654 | |
| 655 | def forward( |
| 656 | self, |
| 657 | input_ids: torch.Tensor | None, |
| 658 | positions: torch.Tensor, |
| 659 | encoder_outputs: list[torch.Tensor], |
| 660 | ) -> torch.Tensor: |
| 661 | enc_states = torch.cat(encoder_outputs, dim=0) if len(encoder_outputs) else None |
| 662 | decoder_outputs = self.decoder( |
| 663 | input_ids=input_ids, |
| 664 | positions=positions, |
| 665 | encoder_hidden_states=enc_states, |
| 666 | ) |
| 667 | return decoder_outputs |
| 668 | |
| 669 | def get_encoder_outputs( |
| 670 | self, |
| 671 | input_features: torch.Tensor | list[torch.Tensor] | None, |
| 672 | ) -> torch.Tensor | None: |
| 673 | if input_features is None: |
| 674 | return None |
| 675 | return self.encoder(input_features) |
| 676 | |
| 677 | |
| 678 | class WhisperProcessingInfo(BaseProcessingInfo): |
| 679 | def get_hf_config(self) -> WhisperConfig: |
| 680 | return self.ctx.get_hf_config(WhisperConfig) |
| 681 | |
| 682 | def get_default_tok_params(self) -> TokenizeParams: |
| 683 | # Special tokens should be provided by the user based on the |
| 684 | # task and language of their request. Also needed to avoid |
| 685 | # appending an EOS token to the prompt which disrupts generation. |
| 686 | return super().get_default_tok_params().with_kwargs(add_special_tokens=False) |
| 687 | |
| 688 | def get_data_parser(self): |
| 689 | feature_extractor = self.get_feature_extractor() |
| 690 | |
| 691 | return MultiModalDataParser( |
| 692 | target_sr=feature_extractor.sampling_rate, |
| 693 | target_channels=self.get_target_channels(), |
| 694 | expected_hidden_size=self._get_expected_hidden_size(), |
| 695 | ) |
| 696 | |
| 697 | @property |
| 698 | def skip_prompt_length_check(self) -> bool: |
| 699 | return True # Because the encoder prompt is padded |
| 700 | |
| 701 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 702 | return {"audio": 1} |
| 703 | |
| 704 | def get_feature_extractor(self, **kwargs: object) -> WhisperFeatureExtractor: |
| 705 | hf_processor = self.get_hf_processor(**kwargs) |
| 706 | feature_extractor = hf_processor.feature_extractor # type: ignore |
| 707 | assert isinstance(feature_extractor, WhisperFeatureExtractor) |
| 708 | return feature_extractor |
| 709 | |
| 710 | def get_target_channels(self) -> int: |
| 711 | """Return target audio channels for Whisper models (mono).""" |
| 712 | return 1 |
| 713 | |
| 714 | def get_num_audio_tokens(self) -> int: |
| 715 | return self.get_hf_config().max_source_positions |
| 716 | |
| 717 | |
| 718 | class WhisperDummyInputsBuilder(BaseDummyInputsBuilder[WhisperProcessingInfo]): |
| 719 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 720 | num_audios = mm_counts.get("audio", 0) |
| 721 | |
| 722 | return "<|startoftranscript|>" * num_audios |
| 723 | |
| 724 | def get_dummy_mm_data( |
| 725 | self, |
| 726 | seq_len: int, |
| 727 | mm_counts: Mapping[str, int], |
| 728 | mm_options: MultiModalDummyOptions, |
| 729 | ) -> MultiModalDataDict: |
| 730 | feature_extractor = self.info.get_feature_extractor() |
| 731 | |
| 732 | sampling_rate = feature_extractor.sampling_rate |
| 733 | audio_len = feature_extractor.chunk_length * sampling_rate |
| 734 | |
| 735 | return { |
| 736 | "audio": self._get_dummy_audios( |
| 737 | length=audio_len, |
| 738 | num_audios=mm_counts.get("audio", 0), |
| 739 | overrides=mm_options.get("audio"), |
| 740 | ) |
| 741 | } |
| 742 | |
| 743 | |
| 744 | class WhisperMultiModalProcessor(EncDecMultiModalProcessor[WhisperProcessingInfo]): |
| 745 | def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: |
| 746 | return self.dummy_inputs.get_dummy_text(mm_counts) |
| 747 | |
| 748 | def create_encoder_prompt( |
| 749 | self, |
| 750 | prompt: list[int], |
| 751 | mm_items: MultiModalDataItems, |
| 752 | ) -> list[int]: |
| 753 | # Strictly speaking, whisper encoder only accept audio features. |
| 754 | # We create a dummy encoder prompt here which will be padded to |
| 755 | # num_audio_tokens. So that we can create dummy data from this |
| 756 | # for encoder profiling. |
| 757 | return [0] |
| 758 | |
| 759 | def _get_hf_mm_inputs( |
| 760 | self, |
| 761 | mm_items: MultiModalDataItems, |
| 762 | hf_kwargs: Mapping[str, object], |
| 763 | ) -> HFMultiModalInputs: |
| 764 | hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) |
| 765 | |
| 766 | feature_extractor = self.info.get_feature_extractor(**hf_kwargs) |
| 767 | return hf_inputs._replace( |
| 768 | hf_kwargs=dict( |
| 769 | hf_inputs.hf_kwargs, |
| 770 | sampling_rate=feature_extractor.sampling_rate, |
| 771 | truncation=True, |
| 772 | ) |
| 773 | ) |
| 774 | |
| 775 | def _get_mm_fields_config( |
| 776 | self, |
| 777 | hf_inputs: BatchFeature, |
| 778 | hf_processor_mm_kwargs: Mapping[str, object], |
| 779 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 780 | return dict(input_features=MultiModalFieldConfig.batched("audio")) |
| 781 | |
| 782 | def _get_prompt_updates( |
| 783 | self, |
| 784 | mm_items: MultiModalDataItems, |
| 785 | hf_processor_mm_kwargs: Mapping[str, object], |
| 786 | out_mm_kwargs: MultiModalKwargsItems, |
| 787 | ) -> Sequence[PromptUpdate]: |
| 788 | num_tokens = self.info.get_num_audio_tokens() |
| 789 | return [ |
| 790 | PromptReplacement( |
| 791 | modality="audio", |
| 792 | target=[0], |
| 793 | replacement=[0] * num_tokens, |
| 794 | ) |
| 795 | ] |
| 796 | |
| 797 | |
| 798 | @MULTIMODAL_REGISTRY.register_processor( |
| 799 | WhisperMultiModalProcessor, |
| 800 | info=WhisperProcessingInfo, |
| 801 | dummy_inputs=WhisperDummyInputsBuilder, |
| 802 | ) |
| 803 | class WhisperForConditionalGeneration( |
| 804 | nn.Module, |
| 805 | SupportsTranscription, |
| 806 | SupportsMultiModal, |
| 807 | SupportsLoRA, |
| 808 | ): |
| 809 | # LoRA-specific attributes |
| 810 | packed_modules_mapping = { |
| 811 | "qkv_proj": ["q_proj", "k_proj", "v_proj"], |
| 812 | "kv_proj": ["k_proj", "v_proj"], |
| 813 | } |
| 814 | |
| 815 | hf_to_vllm_mapper = WeightsMapper( |
| 816 | orig_to_new_substr={".fc1.": ".mlp.fc1.", ".fc2.": ".mlp.fc2."}, |
| 817 | orig_to_new_stacked={ |
| 818 | # weight_name: (param_name, shard_id) |
| 819 | ".self_attn.q_proj": (".self_attn.qkv_proj", "q"), |
| 820 | ".self_attn.k_proj": (".self_attn.qkv_proj", "k"), |
| 821 | ".self_attn.v_proj": (".self_attn.qkv_proj", "v"), |
| 822 | ".encoder_attn.k_proj": (".encoder_attn.kv_proj", 0), |
| 823 | ".encoder_attn.v_proj": (".encoder_attn.kv_proj", 1), |
| 824 | }, |
| 825 | orig_to_new_prefix={"proj_out.": None}, |
| 826 | ) |
| 827 | |
| 828 | # Whisper only supports audio-conditioned generation. |
| 829 | supports_transcription_only = True |
| 830 | supports_segment_timestamp = True |
| 831 | supports_explicit_language_detection = True |
| 832 | supported_languages = ISO639_1_SUPPORTED_LANGS |
| 833 | |
| 834 | @classmethod |
| 835 | def validate_language(cls, language: str | None) -> str | None: |
| 836 | if language is None: |
| 837 | logger.debug( |
| 838 | "No language specified. Language will be auto-detected " |
| 839 | "from audio. To skip detection, pass the `language` field " |
| 840 | "in the TranscriptionRequest." |
| 841 | ) |
| 842 | return None |
| 843 | return super().validate_language(language) |
| 844 | |
| 845 | @classmethod |
| 846 | def get_generation_prompt( |
| 847 | cls, |
| 848 | stt_params: SpeechToTextParams, |
| 849 | ) -> PromptType: |
| 850 | audio = stt_params.audio |
| 851 | stt_config = stt_params.stt_config |
| 852 | language = stt_params.language |
| 853 | task_type = stt_params.task_type |
| 854 | request_prompt = stt_params.request_prompt |
| 855 | |
| 856 | if language is None: |
| 857 | raise ValueError( |
| 858 | "Language must be specified when creating the Whisper prompt" |
| 859 | ) |
| 860 | |
| 861 | decoder_text = ( |
| 862 | f"<|prev|>{request_prompt}" if request_prompt else "" |
| 863 | ) + f"<|startoftranscript|><|{language}|><|{task_type}|><|notimestamps|>" |
| 864 | |
| 865 | return ExplicitEncoderDecoderPrompt( |
| 866 | encoder_prompt=TextPrompt( |
| 867 | prompt="", # Whisper does not support encoder prompt. |
| 868 | multi_modal_data={"audio": (audio, stt_config.sample_rate)}, |
| 869 | ), |
| 870 | decoder_prompt=TextPrompt(prompt=decoder_text), |
| 871 | ) |
| 872 | |
| 873 | @classmethod |
| 874 | def get_language_token_ids( |
| 875 | cls, |
| 876 | tokenizer: TokenizerLike, |
| 877 | ) -> list[int]: |
| 878 | """Return token IDs for all supported language tokens. |
| 879 | |
| 880 | Used with ``SamplingParams.allowed_token_ids`` to constrain |
| 881 | language detection to only produce valid language tokens. |
| 882 | """ |
| 883 | token_ids = [ |
| 884 | tokenizer.convert_tokens_to_ids(f"<|{lang_code}|>") |
| 885 | for lang_code in cls.supported_languages |
| 886 | ] |
| 887 | return token_ids |
| 888 | |
| 889 | @classmethod |
| 890 | def get_language_detection_prompt( |
| 891 | cls, |
| 892 | audio: np.ndarray, |
| 893 | stt_config: SpeechToTextConfig, |
| 894 | ) -> PromptType: |
| 895 | """Return a prompt that elicits a single language token from Whisper. |
| 896 | |
| 897 | Feed only ``<|startoftranscript|>`` as the decoder input so the model |
| 898 | predicts the most likely language token (e.g. ``<|de|>``). |
| 899 | """ |
| 900 | return ExplicitEncoderDecoderPrompt( |
| 901 | encoder_prompt=TextPrompt( |
| 902 | prompt="", |
| 903 | multi_modal_data={"audio": (audio, stt_config.sample_rate)}, |
| 904 | ), |
| 905 | decoder_prompt=TextPrompt(prompt="<|startoftranscript|>"), |
| 906 | ) |
| 907 | |
| 908 | @classmethod |
| 909 | def parse_language_detection_output( |
| 910 | cls, |
| 911 | token_ids: list[int], |
| 912 | tokenizer: TokenizerLike, |
| 913 | ) -> str: |
| 914 | """Parse the language token predicted by Whisper. |
| 915 | |
| 916 | Decodes the first token ID and extracts the language code from the |
| 917 | ``<|xx|>`` format. Expects a valid language token from constrained generation. |
| 918 | """ |
| 919 | decoded = tokenizer.decode( |
| 920 | [token_ids[0]], |
| 921 | skip_special_tokens=False, |
| 922 | ) |
| 923 | # Whisper language tokens have the form <|xx|> |
| 924 | assert decoded.startswith("<|") and decoded.endswith("|>") |
| 925 | lang_code = decoded[2:-2] |
| 926 | assert lang_code in cls.supported_languages |
| 927 | return lang_code |
| 928 | |
| 929 | @classmethod |
| 930 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 931 | if modality.startswith("audio"): |
| 932 | return None |
| 933 | |
| 934 | raise ValueError("Only audio modality is supported") |
| 935 | |
| 936 | @classmethod |
| 937 | def get_speech_to_text_config( |
| 938 | cls, model_config: ModelConfig, task_type: str |
| 939 | ) -> SpeechToTextConfig: |
| 940 | processor = cached_processor_from_config(model_config) |
| 941 | |
| 942 | return SpeechToTextConfig( |
| 943 | max_audio_clip_s=processor.feature_extractor.chunk_length, |
| 944 | sample_rate=processor.feature_extractor.sampling_rate, |
| 945 | ) |
| 946 | |
| 947 | @classmethod |
| 948 | def get_num_audio_tokens( |
| 949 | cls, |
| 950 | audio_duration_s: float, |
| 951 | stt_config: SpeechToTextConfig, |
| 952 | model_config: ModelConfig, |
| 953 | ) -> int | None: |
| 954 | processor = cached_processor_from_config(model_config) |
| 955 | hop_length = processor.feature_extractor.hop_length |
| 956 | assert hop_length is not None |
| 957 | # NOTE(NickLucche) user can't pass encoder |
| 958 | # prompts directly at least not to Whisper. |
| 959 | # One indicator of the encoder amount of processing |
| 960 | # is the log-mel spectogram length. |
| 961 | return math.ceil(audio_duration_s * stt_config.sample_rate / hop_length) |
| 962 | |
| 963 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 964 | super().__init__() |
| 965 | config = vllm_config.model_config.hf_config |
| 966 | quant_config = vllm_config.quant_config |
| 967 | self.config = config |
| 968 | self.dtype = vllm_config.model_config.dtype |
| 969 | |
| 970 | with self._mark_composite_model( |
| 971 | vllm_config, |
| 972 | language_targets=WhisperDecoder, |
| 973 | tower_targets={"audio": WhisperEncoder}, |
| 974 | ): |
| 975 | self.model = WhisperModel(vllm_config=vllm_config, prefix=prefix) |
| 976 | |
| 977 | self.proj_out = ParallelLMHead( |
| 978 | config.vocab_size, |
| 979 | config.d_model, |
| 980 | quant_config=quant_config, |
| 981 | prefix=maybe_prefix(prefix, "proj_out"), |
| 982 | ) |
| 983 | self.proj_out = self.proj_out.tie_weights(self.model.decoder.embed_tokens) |
| 984 | logit_scale = getattr(config, "logit_scale", 1.0) |
| 985 | self.logits_processor = LogitsProcessor(config.vocab_size, scale=logit_scale) |
| 986 | |
| 987 | def forward( |
| 988 | self, |
| 989 | input_ids: torch.Tensor, |
| 990 | positions: torch.Tensor, |
| 991 | encoder_outputs: list[torch.Tensor] | None = None, |
| 992 | **kwargs, |
| 993 | ) -> torch.Tensor: |
| 994 | if encoder_outputs is None: |
| 995 | encoder_outputs = [] |
| 996 | decoder_outputs = self.model( |
| 997 | input_ids=input_ids, |
| 998 | positions=positions, |
| 999 | encoder_outputs=encoder_outputs, |
| 1000 | ) |
| 1001 | return decoder_outputs |
| 1002 | |
| 1003 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 1004 | # Required as part of SupportsMultiModal interface. |
| 1005 | audio_input = self._parse_and_validate_audio_input(**kwargs) |
| 1006 | # Split concatenated encoder outputs into one tensor per audio input |
| 1007 | enc_output = self.model.get_encoder_outputs(audio_input["input_features"]) |
| 1008 | assert enc_output is not None |
| 1009 | # The assumption is we can only process whole mm items (audios) |
| 1010 | return enc_output.unbind(dim=0) |
| 1011 | |
| 1012 | def embed_input_ids( |
| 1013 | self, |
| 1014 | input_ids: torch.Tensor, |
| 1015 | multimodal_embeddings: MultiModalEmbeddings | None = None, |
| 1016 | *, |
| 1017 | is_multimodal: torch.Tensor | None = None, |
| 1018 | ) -> torch.Tensor: |
| 1019 | # This method just returns the decoder sequence embeddings since |
| 1020 | # Whisper does not have encoder text tokens. |
| 1021 | return self.model.decoder.embed_input_ids(input_ids) |
| 1022 | |
| 1023 | def _parse_and_validate_audio_input(self, **kwargs: object) -> WhisperAudioInputs: |
| 1024 | input_features = kwargs.pop("input_features", None) |
| 1025 | |
| 1026 | def to_dtype(value: object) -> torch.Tensor: |
| 1027 | assert isinstance(value, torch.Tensor) |
| 1028 | return value.to(self.dtype) |
| 1029 | |
| 1030 | if input_features is not None: |
| 1031 | input_features = json_map_leaves(to_dtype, input_features) |
| 1032 | |
| 1033 | return WhisperAudioInputs(input_features=input_features) |
| 1034 | |
| 1035 | def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 1036 | logits = self.logits_processor(self.proj_out, hidden_states) |
| 1037 | return logits |
| 1038 | |
| 1039 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1040 | loader = AutoWeightsLoader(self) |
| 1041 | |
| 1042 | # add fake zeros bias for k_proj to state_dict. Covers the |
| 1043 | # ``.k_proj.weight`` and the compressed-tensors packed |
| 1044 | # ``.k_proj.weight_packed`` checkpoints. |
| 1045 | k_proj_weight_suffixes = (".k_proj.weight", ".k_proj.weight_packed") |
| 1046 | for suffix in k_proj_weight_suffixes: |
| 1047 | weights = _create_fake_bias_for_k_proj( |
| 1048 | weights, suffix, out_features=self.config.d_model |
| 1049 | ) |
| 1050 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 1051 | |
| 1052 | |
| 1053 | def _create_fake_bias_for_k_proj( |
| 1054 | weights: Iterable[tuple[str, torch.Tensor]], |
| 1055 | fake_bias_key_name: str, |
| 1056 | out_features: int | None = None, |
| 1057 | ) -> Iterable[tuple[str, torch.Tensor]]: |
| 1058 | """Create full zeros bias for k_proj weight in self-attn and x-attn layers. |
| 1059 | So that the bias for k_proj in qkv_proj or kv_proj can be initialized with |
| 1060 | zeros. |
| 1061 | |
| 1062 | If the checkpoint already provides a real ``.bias`` entry for the given |
| 1063 | weight, it is forwarded as-is and no fake bias is injected for that layer. |
| 1064 | """ |
| 1065 | # Map the weight-name suffix to the corresponding bias-name suffix, e.g. |
| 1066 | # ".k_proj.weight_packed" / ".k_proj.weight" -> ".k_proj.bias", |
| 1067 | # ".wk.weight" -> ".wk.bias". ``.weight_packed`` is stripped first so the |
| 1068 | # longer suffix wins. |
| 1069 | bias_key_name = ( |
| 1070 | fake_bias_key_name.removesuffix(".weight_packed").removesuffix(".weight") |
| 1071 | + ".bias" |
| 1072 | ) |
| 1073 | |
| 1074 | real_bias_names: set[str] = set() |
| 1075 | pending: dict[str, torch.Tensor] = {} |
| 1076 | for name, weight in weights: |
| 1077 | yield name, weight |
| 1078 | |
| 1079 | if name.endswith(bias_key_name): |
| 1080 | real_bias_names.add(name) |
| 1081 | pending.pop(name, None) |
| 1082 | continue |
| 1083 | |
| 1084 | if name.endswith(fake_bias_key_name): |
| 1085 | bias_name = name[: -len(fake_bias_key_name)] + bias_key_name |
| 1086 | if bias_name not in real_bias_names: |
| 1087 | pending[bias_name] = torch.zeros( |
| 1088 | out_features if out_features is not None else weight.size(0) |
| 1089 | ) |
| 1090 | |
| 1091 | for bias_name, bias in pending.items(): |
| 1092 | yield bias_name, bias |