| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inference-only MOSS-Audio model compatible with HuggingFace weights.""" |
| 4 | |
| 5 | import math |
| 6 | from collections.abc import Iterable, Mapping, Sequence |
| 7 | from dataclasses import dataclass, field |
| 8 | from functools import partial |
| 9 | from typing import Annotated, Any |
| 10 | |
| 11 | import numpy as np |
| 12 | import regex as re |
| 13 | import torch |
| 14 | import torch.nn.functional as F |
| 15 | from torch import nn |
| 16 | from transformers import BatchFeature, PreTrainedConfig, Qwen3Config |
| 17 | from transformers.models.whisper import WhisperFeatureExtractor |
| 18 | from typing_extensions import TypedDict |
| 19 | |
| 20 | from vllm.compilation.decorators import support_torch_compile |
| 21 | from vllm.config import VllmConfig |
| 22 | from vllm.config.multimodal import MultiModalDummyOptions |
| 23 | from vllm.distributed import ( |
| 24 | get_pp_group, |
| 25 | get_tensor_model_parallel_world_size, |
| 26 | ) |
| 27 | from vllm.inputs import ModalityData, MultiModalDataDict |
| 28 | from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY, SiluAndMul |
| 29 | from vllm.model_executor.layers.linear import ( |
| 30 | ColumnParallelLinear, |
| 31 | MergedColumnParallelLinear, |
| 32 | ReplicatedLinear, |
| 33 | RowParallelLinear, |
| 34 | ) |
| 35 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 36 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 37 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 38 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 39 | from vllm.multimodal.inputs import ( |
| 40 | AudioItem, |
| 41 | MultiModalFieldConfig, |
| 42 | MultiModalKwargsItems, |
| 43 | ) |
| 44 | from vllm.multimodal.parse import ( |
| 45 | DictEmbeddingItems, |
| 46 | ModalityDataItems, |
| 47 | MultiModalDataItems, |
| 48 | MultiModalDataParser, |
| 49 | ) |
| 50 | from vllm.multimodal.processing import ( |
| 51 | BaseDummyInputsBuilder, |
| 52 | BaseMultiModalProcessor, |
| 53 | BaseProcessingInfo, |
| 54 | InputProcessingContext, |
| 55 | PromptReplacement, |
| 56 | PromptUpdate, |
| 57 | PromptUpdateDetails, |
| 58 | cached_encode, |
| 59 | ) |
| 60 | from vllm.sequence import IntermediateTensors |
| 61 | from vllm.transformers_utils.repo_utils import get_hf_file_to_dict |
| 62 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 63 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 64 | |
| 65 | from .interfaces import ( |
| 66 | MultiModalEmbeddings, |
| 67 | SupportsLoRA, |
| 68 | SupportsMultiModal, |
| 69 | SupportsPP, |
| 70 | _require_is_multimodal, |
| 71 | ) |
| 72 | from .module_mapping import MultiModelKeys |
| 73 | from .qwen3 import Qwen3ForCausalLM, Qwen3Model |
| 74 | from .utils import ( |
| 75 | AutoWeightsLoader, |
| 76 | WeightsMapper, |
| 77 | _merge_multimodal_embeddings, |
| 78 | maybe_prefix, |
| 79 | ) |
| 80 | |
| 81 | MOSS_AUDIO_TOKEN = "<|AUDIO|>" |
| 82 | MOSS_AUDIO_BOS_TOKEN = "<|audio_bos|>" |
| 83 | MOSS_AUDIO_EOS_TOKEN = "<|audio_eos|>" |
| 84 | MOSS_AUDIO_TOKEN_ID = 151654 |
| 85 | MOSS_AUDIO_BOS_TOKEN_ID = 151669 |
| 86 | MOSS_AUDIO_EOS_TOKEN_ID = 151670 |
| 87 | DEFAULT_MAX_AUDIO_SECONDS = 30 |
| 88 | DEFAULT_MOSS_AUDIO_MEL_CONFIG = { |
| 89 | "mel_dim": 128, |
| 90 | "mel_sr": 16000, |
| 91 | "mel_hop_length": 160, |
| 92 | "mel_n_fft": 400, |
| 93 | } |
| 94 | |
| 95 | |
| 96 | class MossAudioProcessorKwargs(TypedDict): |
| 97 | audio_token_id: int |
| 98 | audio_start_id: int |
| 99 | audio_end_id: int |
| 100 | enable_time_marker: bool |
| 101 | mel_config: dict[str, int] |
| 102 | |
| 103 | |
| 104 | MOSS_AUDIO_PLACEHOLDER = ( |
| 105 | f"{MOSS_AUDIO_BOS_TOKEN}{MOSS_AUDIO_TOKEN}{MOSS_AUDIO_EOS_TOKEN}" |
| 106 | ) |
| 107 | MOSS_AUDIO_SPAN_RE = re.compile( |
| 108 | f"{re.escape(MOSS_AUDIO_BOS_TOKEN)}" |
| 109 | f"(?:{re.escape(MOSS_AUDIO_TOKEN)})+" |
| 110 | f"{re.escape(MOSS_AUDIO_EOS_TOKEN)}" |
| 111 | ) |
| 112 | MOSS_AUDIO_PROCESSOR_CONFIG_KEYS = { |
| 113 | "audio_token_id", |
| 114 | "audio_start_id", |
| 115 | "audio_end_id", |
| 116 | "enable_time_marker", |
| 117 | "mel_config", |
| 118 | } |
| 119 | MOSS_AUDIO_PLACEHOLDER_TOKENS = ( |
| 120 | MOSS_AUDIO_BOS_TOKEN, |
| 121 | MOSS_AUDIO_TOKEN, |
| 122 | MOSS_AUDIO_EOS_TOKEN, |
| 123 | ) |
| 124 | |
| 125 | |
| 126 | def _ensure_moss_audio_placeholder_tokens(tokenizer: object) -> None: |
| 127 | """Register the audio placeholder tokens if the tokenizer lacks them. |
| 128 | |
| 129 | The MOSS-Audio hub tokenizer does not define the audio placeholder |
| 130 | tokens (the reference HF processor patches them in at runtime). Left |
| 131 | as plain text, they merge with adjacent tokens under BPE, so prompts |
| 132 | no longer contain a stable token sequence for prompt updates to match. |
| 133 | """ |
| 134 | convert_tokens_to_ids = getattr(tokenizer, "convert_tokens_to_ids", None) |
| 135 | add_tokens = getattr(tokenizer, "add_tokens", None) |
| 136 | if convert_tokens_to_ids is None or add_tokens is None: |
| 137 | return |
| 138 | |
| 139 | unk_token_id = getattr(tokenizer, "unk_token_id", None) |
| 140 | missing_tokens = [ |
| 141 | token |
| 142 | for token in MOSS_AUDIO_PLACEHOLDER_TOKENS |
| 143 | if convert_tokens_to_ids(token) in (None, unk_token_id) |
| 144 | ] |
| 145 | if missing_tokens: |
| 146 | add_tokens(missing_tokens, special_tokens=True) |
| 147 | |
| 148 | |
| 149 | class MossAudioAudioInputs(TensorSchema): |
| 150 | """Dimensions: |
| 151 | - b: Batch size |
| 152 | - nmb: Number of mel bins |
| 153 | - t: Time frames |
| 154 | """ |
| 155 | |
| 156 | audio_data: Annotated[torch.Tensor, TensorShape("b", "nmb", "t")] |
| 157 | audio_data_seqlens: Annotated[torch.Tensor, TensorShape("b")] |
| 158 | |
| 159 | |
| 160 | def _normalize_moss_audio_mel_config( |
| 161 | mel_config: Mapping[str, object] | None = None, |
| 162 | ) -> dict[str, int]: |
| 163 | config = dict(DEFAULT_MOSS_AUDIO_MEL_CONFIG) |
| 164 | config.update(_extract_moss_audio_mel_config(mel_config)) |
| 165 | return config |
| 166 | |
| 167 | |
| 168 | def _extract_moss_audio_mel_config( |
| 169 | mel_config: Mapping[str, object] | None = None, |
| 170 | ) -> dict[str, int]: |
| 171 | config: dict[str, int] = {} |
| 172 | if mel_config is None: |
| 173 | return config |
| 174 | |
| 175 | aliases = { |
| 176 | "mel_dim": ("mel_dim", "feature_size", "n_mels", "num_mel_bins"), |
| 177 | "mel_sr": ("mel_sr", "sampling_rate", "sample_rate"), |
| 178 | "mel_hop_length": ("mel_hop_length", "hop_length"), |
| 179 | "mel_n_fft": ("mel_n_fft", "n_fft"), |
| 180 | } |
| 181 | for target_key, source_keys in aliases.items(): |
| 182 | for source_key in source_keys: |
| 183 | if source_key in mel_config: |
| 184 | value = mel_config[source_key] |
| 185 | assert isinstance(value, int) |
| 186 | config[target_key] = value |
| 187 | break |
| 188 | |
| 189 | return config |
| 190 | |
| 191 | |
| 192 | def _filter_moss_audio_processor_config( |
| 193 | config: Mapping[str, object] | None, |
| 194 | ) -> dict[str, object]: |
| 195 | if not config: |
| 196 | return {} |
| 197 | |
| 198 | return { |
| 199 | key: value |
| 200 | for key, value in config.items() |
| 201 | if key in MOSS_AUDIO_PROCESSOR_CONFIG_KEYS |
| 202 | } |
| 203 | |
| 204 | |
| 205 | def _merge_moss_audio_processor_configs( |
| 206 | *configs: Mapping[str, object] | None, |
| 207 | ) -> dict[str, object]: |
| 208 | merged: dict[str, object] = {} |
| 209 | merged_mel_config: dict[str, int] = {} |
| 210 | for config in configs: |
| 211 | filtered = _filter_moss_audio_processor_config(config) |
| 212 | mel_config = filtered.pop("mel_config", None) |
| 213 | merged.update(filtered) |
| 214 | if isinstance(mel_config, Mapping): |
| 215 | merged_mel_config.update(_extract_moss_audio_mel_config(mel_config)) |
| 216 | |
| 217 | if merged_mel_config: |
| 218 | merged["mel_config"] = merged_mel_config |
| 219 | return merged |
| 220 | |
| 221 | |
| 222 | @dataclass |
| 223 | class MossAudioEncoderConfig: |
| 224 | d_model: int = 1280 |
| 225 | output_dim: int = 1280 |
| 226 | num_mel_bins: int = 128 |
| 227 | encoder_layers: int = 32 |
| 228 | encoder_attention_heads: int = 20 |
| 229 | encoder_ffn_dim: int = 5120 |
| 230 | downsample_rate: int = 8 |
| 231 | downsample_hidden_size: int = 480 |
| 232 | encoder_attention_window_size: int = 100 |
| 233 | max_source_positions: int = 1500 |
| 234 | dropout: float = 0.1 |
| 235 | attention_dropout: float = 0.1 |
| 236 | activation_dropout: float = 0.0 |
| 237 | activation_function: str = "gelu" |
| 238 | layer_norm_eps: float = 1e-5 |
| 239 | _attn_implementation: str = "eager" |
| 240 | pretrained_path: str = "" |
| 241 | n_window: int = 200 |
| 242 | conv_chunksize: int = 64 |
| 243 | deepstack_encoder_layer_indexes: list[int] = field( |
| 244 | default_factory=lambda: [8, 16, 24] |
| 245 | ) |
| 246 | |
| 247 | @classmethod |
| 248 | def from_config(cls, config: object) -> "MossAudioEncoderConfig": |
| 249 | if isinstance(config, cls): |
| 250 | return config |
| 251 | if isinstance(config, Mapping): |
| 252 | values = { |
| 253 | key: value |
| 254 | for key, value in config.items() |
| 255 | if key in cls.__dataclass_fields__ |
| 256 | } |
| 257 | else: |
| 258 | values = { |
| 259 | key: getattr(config, key) |
| 260 | for key in cls.__dataclass_fields__ |
| 261 | if hasattr(config, key) |
| 262 | } |
| 263 | return cls(**values) |
| 264 | |
| 265 | |
| 266 | class MossAudioConfig(PreTrainedConfig): |
| 267 | model_type = "moss_audio" |
| 268 | is_composition = True |
| 269 | |
| 270 | def __init__( |
| 271 | self, |
| 272 | audio_config: Mapping[str, object] | MossAudioEncoderConfig | None = None, |
| 273 | language_config: Mapping[str, object] | Qwen3Config | None = None, |
| 274 | adapter_hidden_size: int = 8192, |
| 275 | ignore_index: int = -100, |
| 276 | deepstack_num_inject_layers: int | None = None, |
| 277 | **kwargs: object, |
| 278 | ) -> None: |
| 279 | self.audio_config = MossAudioEncoderConfig.from_config(audio_config or {}) |
| 280 | if isinstance(language_config, Qwen3Config): |
| 281 | self.language_config = language_config |
| 282 | else: |
| 283 | self.language_config = Qwen3Config(**(language_config or {})) |
| 284 | |
| 285 | self.adapter_hidden_size = adapter_hidden_size |
| 286 | self.ignore_index = ignore_index |
| 287 | self.deepstack_num_inject_layers = deepstack_num_inject_layers |
| 288 | |
| 289 | for key in ("num_hidden_layers", "eos_token_id", "bos_token_id", "vocab_size"): |
| 290 | kwargs.setdefault(key, getattr(self.language_config, key, None)) |
| 291 | kwargs.setdefault("tie_word_embeddings", False) |
| 292 | super().__init__(**kwargs) |
| 293 | |
| 294 | for key in ( |
| 295 | "hidden_size", |
| 296 | "num_attention_heads", |
| 297 | "num_key_value_heads", |
| 298 | "head_dim", |
| 299 | "max_position_embeddings", |
| 300 | "rms_norm_eps", |
| 301 | ): |
| 302 | if hasattr(self.language_config, key): |
| 303 | setattr(self, key, getattr(self.language_config, key)) |
| 304 | |
| 305 | def get_text_config(self, decoder: bool = False) -> Qwen3Config: |
| 306 | return self.language_config |
| 307 | |
| 308 | |
| 309 | class SinusoidsPositionEmbedding(nn.Module): |
| 310 | def __init__(self, num_positions: int, embedding_dim: int) -> None: |
| 311 | super().__init__() |
| 312 | del num_positions # Kept for config compatibility. |
| 313 | max_timescale = 10000.0 |
| 314 | log_timescale_increment = math.log(max_timescale) / (embedding_dim // 2 - 1) |
| 315 | inv_timescales = torch.exp( |
| 316 | -log_timescale_increment * torch.arange(embedding_dim // 2).float() |
| 317 | ) |
| 318 | self.register_buffer("inv_timescales", inv_timescales, persistent=False) |
| 319 | |
| 320 | def forward(self, seq_len: int, device: torch.device) -> torch.Tensor: |
| 321 | scaled_time = ( |
| 322 | torch.arange(seq_len, device=device, dtype=self.inv_timescales.dtype)[ |
| 323 | :, None |
| 324 | ] |
| 325 | * self.inv_timescales[None, :] |
| 326 | ) |
| 327 | return torch.cat([scaled_time.sin(), scaled_time.cos()], dim=1).unsqueeze(0) |
| 328 | |
| 329 | |
| 330 | class MossAudioAttention(nn.Module): |
| 331 | def __init__( |
| 332 | self, |
| 333 | config: MossAudioEncoderConfig, |
| 334 | quant_config: QuantizationConfig | None = None, |
| 335 | prefix: str = "", |
| 336 | ) -> None: |
| 337 | super().__init__() |
| 338 | self.embed_dim = config.d_model |
| 339 | self.num_heads = config.encoder_attention_heads |
| 340 | self.head_dim = self.embed_dim // self.num_heads |
| 341 | if self.head_dim * self.num_heads != self.embed_dim: |
| 342 | raise ValueError( |
| 343 | f"d_model ({self.embed_dim}) must be divisible by " |
| 344 | f"encoder_attention_heads ({self.num_heads})." |
| 345 | ) |
| 346 | |
| 347 | tp_size = get_tensor_model_parallel_world_size() |
| 348 | if self.num_heads % tp_size != 0: |
| 349 | raise ValueError( |
| 350 | "MOSS-Audio audio encoder attention heads must be divisible by " |
| 351 | f"tensor parallel size. Got {self.num_heads=} and {tp_size=}." |
| 352 | ) |
| 353 | self.num_local_heads = self.num_heads // tp_size |
| 354 | # TODO: can use QKVParallelLinear |
| 355 | self.q_proj = ColumnParallelLinear( |
| 356 | input_size=self.embed_dim, |
| 357 | output_size=self.embed_dim, |
| 358 | bias=True, |
| 359 | quant_config=quant_config, |
| 360 | prefix=f"{prefix}.q_proj", |
| 361 | ) |
| 362 | self.k_proj = ColumnParallelLinear( |
| 363 | input_size=self.embed_dim, |
| 364 | output_size=self.embed_dim, |
| 365 | bias=False, |
| 366 | quant_config=quant_config, |
| 367 | prefix=f"{prefix}.k_proj", |
| 368 | ) |
| 369 | self.v_proj = ColumnParallelLinear( |
| 370 | input_size=self.embed_dim, |
| 371 | output_size=self.embed_dim, |
| 372 | bias=True, |
| 373 | quant_config=quant_config, |
| 374 | prefix=f"{prefix}.v_proj", |
| 375 | ) |
| 376 | self.out_proj = RowParallelLinear( |
| 377 | input_size=self.embed_dim, |
| 378 | output_size=self.embed_dim, |
| 379 | bias=True, |
| 380 | quant_config=quant_config, |
| 381 | prefix=f"{prefix}.out_proj", |
| 382 | ) |
| 383 | |
| 384 | def forward( |
| 385 | self, |
| 386 | hidden_states: torch.Tensor, |
| 387 | attention_mask: torch.Tensor, |
| 388 | ) -> torch.Tensor: |
| 389 | batch_size, seq_len, _ = hidden_states.shape |
| 390 | q, _ = self.q_proj(hidden_states) |
| 391 | k, _ = self.k_proj(hidden_states) |
| 392 | v, _ = self.v_proj(hidden_states) |
| 393 | q = q.view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2) |
| 394 | k = k.view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2) |
| 395 | v = v.view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2) |
| 396 | attn_output = F.scaled_dot_product_attention( |
| 397 | q, |
| 398 | k, |
| 399 | v, |
| 400 | attn_mask=attention_mask[:, None, None, :], |
| 401 | dropout_p=0.0, |
| 402 | scale=self.head_dim**-0.5, |
| 403 | ) |
| 404 | output, _ = self.out_proj( |
| 405 | attn_output.transpose(1, 2).reshape( |
| 406 | batch_size, |
| 407 | seq_len, |
| 408 | -1, |
| 409 | ) |
| 410 | ) |
| 411 | return output |
| 412 | |
| 413 | |
| 414 | class MossAudioEncoderLayer(nn.Module): |
| 415 | def __init__( |
| 416 | self, |
| 417 | config: MossAudioEncoderConfig, |
| 418 | quant_config: QuantizationConfig | None = None, |
| 419 | prefix: str = "", |
| 420 | ) -> None: |
| 421 | super().__init__() |
| 422 | self.self_attn = MossAudioAttention( |
| 423 | config, |
| 424 | quant_config=quant_config, |
| 425 | prefix=f"{prefix}.self_attn", |
| 426 | ) |
| 427 | self.self_attn_layer_norm = nn.LayerNorm( |
| 428 | config.d_model, eps=config.layer_norm_eps |
| 429 | ) |
| 430 | self.activation_fn = _ACTIVATION_REGISTRY[config.activation_function] |
| 431 | self.activation_dropout = config.activation_dropout |
| 432 | self.dropout = config.dropout |
| 433 | self.fc1 = ColumnParallelLinear( |
| 434 | config.d_model, |
| 435 | config.encoder_ffn_dim, |
| 436 | bias=True, |
| 437 | quant_config=quant_config, |
| 438 | prefix=f"{prefix}.fc1", |
| 439 | ) |
| 440 | self.fc2 = RowParallelLinear( |
| 441 | config.encoder_ffn_dim, |
| 442 | config.d_model, |
| 443 | bias=True, |
| 444 | quant_config=quant_config, |
| 445 | prefix=f"{prefix}.fc2", |
| 446 | ) |
| 447 | self.final_layer_norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps) |
| 448 | |
| 449 | def forward( |
| 450 | self, |
| 451 | hidden_states: torch.Tensor, |
| 452 | attention_mask: torch.Tensor, |
| 453 | ) -> torch.Tensor: |
| 454 | residual = hidden_states |
| 455 | hidden_states = self.self_attn_layer_norm(hidden_states) |
| 456 | hidden_states = self.self_attn(hidden_states, attention_mask) |
| 457 | hidden_states = residual + F.dropout( |
| 458 | hidden_states, p=self.dropout, training=self.training |
| 459 | ) |
| 460 | |
| 461 | residual = hidden_states |
| 462 | hidden_states = self.final_layer_norm(hidden_states) |
| 463 | hidden_states, _ = self.fc1(hidden_states) |
| 464 | hidden_states = self.activation_fn(hidden_states) |
| 465 | hidden_states = F.dropout( |
| 466 | hidden_states, p=self.activation_dropout, training=self.training |
| 467 | ) |
| 468 | hidden_states, _ = self.fc2(hidden_states) |
| 469 | hidden_states = residual + F.dropout( |
| 470 | hidden_states, p=self.dropout, training=self.training |
| 471 | ) |
| 472 | return hidden_states |
| 473 | |
| 474 | |
| 475 | class MossAudioEncoder(nn.Module): |
| 476 | def __init__( |
| 477 | self, |
| 478 | config: MossAudioEncoderConfig, |
| 479 | quant_config: QuantizationConfig | None = None, |
| 480 | prefix: str = "", |
| 481 | ) -> None: |
| 482 | super().__init__() |
| 483 | self.config = config |
| 484 | self.gelu = nn.GELU() |
| 485 | self.conv1 = nn.Conv2d( |
| 486 | 1, |
| 487 | config.downsample_hidden_size, |
| 488 | kernel_size=(3, 3), |
| 489 | stride=(2, 2), |
| 490 | padding=(1, 1), |
| 491 | ) |
| 492 | self.conv2 = nn.Conv2d( |
| 493 | config.downsample_hidden_size, |
| 494 | config.downsample_hidden_size, |
| 495 | kernel_size=(3, 3), |
| 496 | stride=(2, 2), |
| 497 | padding=(1, 1), |
| 498 | ) |
| 499 | self.conv3 = nn.Conv2d( |
| 500 | config.downsample_hidden_size, |
| 501 | config.downsample_hidden_size, |
| 502 | kernel_size=(3, 3), |
| 503 | stride=(2, 2), |
| 504 | padding=(1, 1), |
| 505 | ) |
| 506 | |
| 507 | conv_freq = self._compute_downsampled_length( |
| 508 | torch.tensor(config.num_mel_bins) |
| 509 | ).item() |
| 510 | self.stem_proj = ReplicatedLinear( |
| 511 | config.downsample_hidden_size * int(conv_freq), |
| 512 | config.d_model, |
| 513 | bias=True, |
| 514 | quant_config=quant_config, |
| 515 | return_bias=False, |
| 516 | prefix=f"{prefix}.stem_proj", |
| 517 | ) |
| 518 | self.embed_positions = SinusoidsPositionEmbedding( |
| 519 | config.max_source_positions, config.d_model |
| 520 | ) |
| 521 | self.layers = nn.ModuleList( |
| 522 | [ |
| 523 | MossAudioEncoderLayer( |
| 524 | config, |
| 525 | quant_config=quant_config, |
| 526 | prefix=f"{prefix}.layers.{layer_idx}", |
| 527 | ) |
| 528 | for layer_idx in range(config.encoder_layers) |
| 529 | ] |
| 530 | ) |
| 531 | self.layer_norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps) |
| 532 | if config.output_dim != config.d_model: |
| 533 | self.out_proj = ReplicatedLinear( |
| 534 | config.d_model, |
| 535 | config.output_dim, |
| 536 | bias=False, |
| 537 | quant_config=quant_config, |
| 538 | return_bias=False, |
| 539 | prefix=f"{prefix}.out_proj", |
| 540 | ) |
| 541 | else: |
| 542 | self.out_proj = nn.Identity() |
| 543 | |
| 544 | self.deepstack_encoder_layer_indexes = list( |
| 545 | config.deepstack_encoder_layer_indexes or [] |
| 546 | ) |
| 547 | self._deepstack_capture_map = { |
| 548 | layer_idx: capture_idx |
| 549 | for capture_idx, layer_idx in enumerate( |
| 550 | self.deepstack_encoder_layer_indexes |
| 551 | ) |
| 552 | } |
| 553 | self.n_window = int(config.n_window) |
| 554 | self.chunk_frames = int(self.n_window * 2) |
| 555 | self.conv_chunksize = int(config.conv_chunksize) |
| 556 | |
| 557 | @property |
| 558 | def dtype(self) -> torch.dtype: |
| 559 | return self.conv1.weight.dtype |
| 560 | |
| 561 | @staticmethod |
| 562 | def _compute_downsampled_length(lengths: torch.Tensor) -> torch.Tensor: |
| 563 | def conv_out_len(length: torch.Tensor) -> torch.Tensor: |
| 564 | return (length - 1) // 2 + 1 |
| 565 | |
| 566 | return conv_out_len(conv_out_len(conv_out_len(lengths))) |
| 567 | |
| 568 | @staticmethod |
| 569 | def compute_num_audio_tokens(raw_mel_len: int) -> int: |
| 570 | lengths = torch.tensor(raw_mel_len, dtype=torch.long) |
| 571 | return int(MossAudioEncoder._compute_downsampled_length(lengths).item()) |
| 572 | |
| 573 | def _encode_chunk_batch( |
| 574 | self, |
| 575 | input_features: torch.Tensor, |
| 576 | seq_lengths: torch.Tensor, |
| 577 | output_deepstack_hidden_states: bool = True, |
| 578 | ) -> tuple[torch.Tensor, list[torch.Tensor]]: |
| 579 | if input_features.dim() == 2: |
| 580 | input_features = input_features.unsqueeze(0) |
| 581 | |
| 582 | downsampled_lengths = self._compute_downsampled_length(seq_lengths) |
| 583 | x = input_features.unsqueeze(1) |
| 584 | x = self.gelu(self.conv1(x)) |
| 585 | x = self.gelu(self.conv2(x)) |
| 586 | x = self.gelu(self.conv3(x)) |
| 587 | x = x.permute(0, 3, 1, 2).contiguous().flatten(2) |
| 588 | x = self.stem_proj(x) |
| 589 | |
| 590 | max_len = int(downsampled_lengths.max().item()) |
| 591 | if x.size(1) > max_len: |
| 592 | x = x[:, :max_len, :] |
| 593 | x = x + self.embed_positions(x.shape[1], x.device).to(x.dtype) |
| 594 | |
| 595 | attention_mask = ( |
| 596 | torch.arange(x.size(1), device=x.device)[None, :] |
| 597 | < downsampled_lengths[:, None] |
| 598 | ) |
| 599 | |
| 600 | deepstack_hidden_states: list[torch.Tensor | None] = [] |
| 601 | if output_deepstack_hidden_states: |
| 602 | deepstack_hidden_states = [None] * len(self.deepstack_encoder_layer_indexes) |
| 603 | for layer_idx, layer in enumerate(self.layers): |
| 604 | x = layer(x, attention_mask) |
| 605 | if output_deepstack_hidden_states: |
| 606 | capture_idx = self._deepstack_capture_map.get(layer_idx) |
| 607 | if capture_idx is not None: |
| 608 | deepstack_hidden_states[capture_idx] = x |
| 609 | |
| 610 | x = self.layer_norm(x) |
| 611 | x = self.out_proj(x) |
| 612 | |
| 613 | if not output_deepstack_hidden_states: |
| 614 | return x, [] |
| 615 | |
| 616 | ordered_deepstack_hidden_states = [ |
| 617 | hidden_states |
| 618 | for hidden_states in deepstack_hidden_states |
| 619 | if hidden_states is not None |
| 620 | ] |
| 621 | ordered_deepstack_hidden_states = [ |
| 622 | self.out_proj(hidden_states) |
| 623 | for hidden_states in ordered_deepstack_hidden_states |
| 624 | ] |
| 625 | return x, ordered_deepstack_hidden_states |
| 626 | |
| 627 | def forward( |
| 628 | self, |
| 629 | input_features: torch.Tensor, |
| 630 | feature_lens: torch.Tensor | None = None, |
| 631 | output_deepstack_hidden_states: bool = True, |
| 632 | ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...] | None]: |
| 633 | if input_features.dim() == 3: |
| 634 | if feature_lens is None: |
| 635 | feature_lens = torch.full( |
| 636 | (input_features.size(0),), |
| 637 | input_features.size(-1), |
| 638 | dtype=torch.long, |
| 639 | device=input_features.device, |
| 640 | ) |
| 641 | else: |
| 642 | feature_lens = feature_lens.to( |
| 643 | device=input_features.device, dtype=torch.long |
| 644 | ) |
| 645 | valid_chunks = [ |
| 646 | input_features[i, :, : int(feature_lens[i].item())] |
| 647 | for i in range(int(input_features.shape[0])) |
| 648 | ] |
| 649 | input_features = torch.cat(valid_chunks, dim=1) |
| 650 | elif input_features.dim() != 2: |
| 651 | raise ValueError( |
| 652 | f"Expected [n_mels, T] or [B, n_mels, T], got " |
| 653 | f"{tuple(input_features.shape)}." |
| 654 | ) |
| 655 | |
| 656 | if feature_lens is None: |
| 657 | feature_lens = torch.tensor( |
| 658 | [int(input_features.shape[1])], |
| 659 | device=input_features.device, |
| 660 | dtype=torch.long, |
| 661 | ) |
| 662 | else: |
| 663 | feature_lens = feature_lens.to( |
| 664 | device=input_features.device, dtype=torch.long |
| 665 | ) |
| 666 | |
| 667 | chunk_num = torch.ceil( |
| 668 | feature_lens.to(torch.float32) / self.chunk_frames |
| 669 | ).long() |
| 670 | chunk_lengths = torch.full( |
| 671 | (int(chunk_num.sum().item()),), |
| 672 | self.chunk_frames, |
| 673 | dtype=torch.long, |
| 674 | device=feature_lens.device, |
| 675 | ) |
| 676 | tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:] |
| 677 | chunk_lengths[tail_chunk_index] = feature_lens % self.chunk_frames |
| 678 | chunk_lengths[chunk_lengths == 0] = self.chunk_frames |
| 679 | |
| 680 | chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0) |
| 681 | padded_feature = nn.utils.rnn.pad_sequence( |
| 682 | chunk_list, batch_first=True |
| 683 | ).transpose(1, 2) |
| 684 | |
| 685 | feature_lens_after_cnn = self._compute_downsampled_length(chunk_lengths) |
| 686 | t_down_max = ( |
| 687 | int(feature_lens_after_cnn.max().item()) |
| 688 | if feature_lens_after_cnn.numel() > 0 |
| 689 | else 0 |
| 690 | ) |
| 691 | indices = torch.arange(t_down_max, device=padded_feature.device) |
| 692 | padded_mask_after_cnn = indices[None, :] < feature_lens_after_cnn[:, None] |
| 693 | |
| 694 | num_deepstack = len(self.deepstack_encoder_layer_indexes) |
| 695 | should_output_deepstack = output_deepstack_hidden_states and num_deepstack > 0 |
| 696 | padded_embeds: list[torch.Tensor] = [] |
| 697 | deepstack_padded_embeds: list[list[torch.Tensor]] = [ |
| 698 | [] for _ in range(num_deepstack if should_output_deepstack else 0) |
| 699 | ] |
| 700 | for feat_chunk, len_chunk in zip( |
| 701 | padded_feature.split(self.conv_chunksize, dim=0), |
| 702 | chunk_lengths.split(self.conv_chunksize, dim=0), |
| 703 | ): |
| 704 | out, deepstack_outs = self._encode_chunk_batch( |
| 705 | feat_chunk, |
| 706 | len_chunk, |
| 707 | output_deepstack_hidden_states=should_output_deepstack, |
| 708 | ) |
| 709 | if out.shape[1] < t_down_max: |
| 710 | out = F.pad(out, (0, 0, 0, t_down_max - out.shape[1])) |
| 711 | padded_embeds.append(out) |
| 712 | |
| 713 | if should_output_deepstack: |
| 714 | if len(deepstack_outs) != num_deepstack: |
| 715 | raise RuntimeError( |
| 716 | "DeepStack output count does not match configured " |
| 717 | "layer indexes." |
| 718 | ) |
| 719 | for capture_idx, ds in enumerate(deepstack_outs): |
| 720 | if ds.shape[1] < t_down_max: |
| 721 | ds = F.pad(ds, (0, 0, 0, t_down_max - ds.shape[1])) |
| 722 | deepstack_padded_embeds[capture_idx].append(ds) |
| 723 | |
| 724 | if padded_embeds: |
| 725 | padded_embed = torch.cat(padded_embeds, dim=0) |
| 726 | else: |
| 727 | padded_embed = torch.empty( |
| 728 | (0, t_down_max, self.config.output_dim), |
| 729 | device=padded_feature.device, |
| 730 | dtype=padded_feature.dtype, |
| 731 | ) |
| 732 | |
| 733 | last_hidden_state = padded_embed[padded_mask_after_cnn].unsqueeze(0) |
| 734 | |
| 735 | deepstack_states: tuple[torch.Tensor, ...] | None = None |
| 736 | if should_output_deepstack: |
| 737 | collected: list[torch.Tensor] = [] |
| 738 | for chunks_list in deepstack_padded_embeds: |
| 739 | if chunks_list: |
| 740 | ds = torch.cat(chunks_list, dim=0) |
| 741 | collected.append(ds[padded_mask_after_cnn].unsqueeze(0)) |
| 742 | else: |
| 743 | collected.append( |
| 744 | torch.empty( |
| 745 | (1, 0, self.config.output_dim), |
| 746 | device=padded_feature.device, |
| 747 | dtype=padded_embed.dtype, |
| 748 | ) |
| 749 | ) |
| 750 | deepstack_states = tuple(collected) |
| 751 | |
| 752 | return last_hidden_state, deepstack_states |
| 753 | |
| 754 | |
| 755 | class GatedMLP(nn.Module): |
| 756 | def __init__( |
| 757 | self, |
| 758 | input_size: int, |
| 759 | hidden_size: int, |
| 760 | output_size: int, |
| 761 | quant_config: QuantizationConfig | None = None, |
| 762 | prefix: str = "", |
| 763 | ) -> None: |
| 764 | super().__init__() |
| 765 | self.gate_up_proj = MergedColumnParallelLinear( |
| 766 | input_size, |
| 767 | [hidden_size, hidden_size], |
| 768 | bias=False, |
| 769 | quant_config=quant_config, |
| 770 | prefix=f"{prefix}.gate_up_proj", |
| 771 | ) |
| 772 | self.down_proj = RowParallelLinear( |
| 773 | hidden_size, |
| 774 | output_size, |
| 775 | bias=False, |
| 776 | input_is_parallel=True, |
| 777 | quant_config=quant_config, |
| 778 | prefix=f"{prefix}.down_proj", |
| 779 | ) |
| 780 | self.act_fn = SiluAndMul() |
| 781 | |
| 782 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 783 | gate_up, _ = self.gate_up_proj(x) |
| 784 | x = self.act_fn(gate_up) |
| 785 | x, _ = self.down_proj(x) |
| 786 | return x |
| 787 | |
| 788 | |
| 789 | @support_torch_compile( |
| 790 | dynamic_arg_dims={ |
| 791 | "input_ids": 0, |
| 792 | "positions": -1, |
| 793 | "intermediate_tensors": 0, |
| 794 | "inputs_embeds": 0, |
| 795 | "deepstack_input_embeds": 0, |
| 796 | } |
| 797 | ) |
| 798 | class MossQwen3Model(Qwen3Model): |
| 799 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 800 | super().__init__(vllm_config=vllm_config, prefix=prefix) |
| 801 | self.deepstack_inject_layer_indices: Iterable[int] = range(0) |
| 802 | |
| 803 | def forward( |
| 804 | self, |
| 805 | input_ids: torch.Tensor | None, |
| 806 | positions: torch.Tensor, |
| 807 | intermediate_tensors: IntermediateTensors | None = None, |
| 808 | inputs_embeds: torch.Tensor | None = None, |
| 809 | deepstack_input_embeds: IntermediateTensors | None = None, |
| 810 | ) -> torch.Tensor | IntermediateTensors: |
| 811 | if get_pp_group().is_first_rank: |
| 812 | if inputs_embeds is not None: |
| 813 | hidden_states = inputs_embeds |
| 814 | else: |
| 815 | hidden_states = self.embed_input_ids(input_ids) |
| 816 | residual = None |
| 817 | else: |
| 818 | assert intermediate_tensors is not None |
| 819 | hidden_states = intermediate_tensors["hidden_states"] |
| 820 | residual = intermediate_tensors["residual"] |
| 821 | |
| 822 | aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) |
| 823 | for layer_idx, layer in enumerate( |
| 824 | self.layers[self.start_layer : self.end_layer], |
| 825 | start=self.start_layer, |
| 826 | ): |
| 827 | hidden_states, residual = layer(positions, hidden_states, residual) |
| 828 | deepstack_key = f"deepstack_input_embeds_{layer_idx}" |
| 829 | if ( |
| 830 | deepstack_input_embeds is not None |
| 831 | and deepstack_key in deepstack_input_embeds.tensors |
| 832 | ): |
| 833 | hidden_states = hidden_states + deepstack_input_embeds[deepstack_key] |
| 834 | self._maybe_add_hidden_state( |
| 835 | aux_hidden_states, |
| 836 | layer_idx - self.start_layer + 1, |
| 837 | hidden_states, |
| 838 | residual, |
| 839 | ) |
| 840 | |
| 841 | if not get_pp_group().is_last_rank: |
| 842 | tensors = {"hidden_states": hidden_states, "residual": residual} |
| 843 | # Keep the DeepStack PP schema config-driven, but only carry |
| 844 | # payloads needed by downstream injection points across this rank. |
| 845 | # Missing downstream payloads are zero-filled below to clear |
| 846 | # receive buffers instead of leaving stale tensors. |
| 847 | for layer_idx in self.deepstack_inject_layer_indices: |
| 848 | if layer_idx < self.end_layer: |
| 849 | continue |
| 850 | deepstack_key = f"deepstack_input_embeds_{layer_idx}" |
| 851 | if ( |
| 852 | deepstack_input_embeds is not None |
| 853 | and deepstack_key in deepstack_input_embeds.tensors |
| 854 | ): |
| 855 | tensors[deepstack_key] = deepstack_input_embeds[deepstack_key] |
| 856 | else: |
| 857 | tensors[deepstack_key] = hidden_states.new_zeros( |
| 858 | hidden_states.shape |
| 859 | ) |
| 860 | return IntermediateTensors(tensors) |
| 861 | |
| 862 | hidden_states, _ = self.norm(hidden_states, residual) |
| 863 | if len(aux_hidden_states) > 0: |
| 864 | return hidden_states, aux_hidden_states |
| 865 | return hidden_states |
| 866 | |
| 867 | |
| 868 | class MossQwen3ForCausalLM(Qwen3ForCausalLM): |
| 869 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 870 | super(Qwen3ForCausalLM, self).__init__() |
| 871 | config = vllm_config.model_config.hf_config |
| 872 | quant_config = vllm_config.quant_config |
| 873 | |
| 874 | self.config = config |
| 875 | self.vllm_config = vllm_config |
| 876 | self.quant_config = quant_config |
| 877 | self.model = MossQwen3Model( |
| 878 | vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") |
| 879 | ) |
| 880 | |
| 881 | if get_pp_group().is_last_rank: |
| 882 | self.lm_head = ParallelLMHead( |
| 883 | config.vocab_size, |
| 884 | config.hidden_size, |
| 885 | quant_config=quant_config, |
| 886 | prefix=maybe_prefix(prefix, "lm_head"), |
| 887 | ) |
| 888 | if config.tie_word_embeddings: |
| 889 | self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) |
| 890 | else: |
| 891 | from .utils import PPMissingLayer |
| 892 | |
| 893 | self.lm_head = PPMissingLayer() |
| 894 | |
| 895 | self.logits_processor = LogitsProcessor(config.vocab_size) |
| 896 | self.deepstack_inject_layer_indices: Iterable[int] = range(0) |
| 897 | |
| 898 | def make_empty_intermediate_tensors( |
| 899 | self, |
| 900 | batch_size: int, |
| 901 | dtype: torch.dtype, |
| 902 | device: torch.device, |
| 903 | ) -> IntermediateTensors: |
| 904 | intermediate_tensors = self.model.make_empty_intermediate_tensors( |
| 905 | batch_size, dtype, device |
| 906 | ) |
| 907 | for layer_idx in self.deepstack_inject_layer_indices: |
| 908 | # Non-first PP ranks only receive DeepStack payloads for layers |
| 909 | # at or after their local start layer. |
| 910 | if layer_idx < self.model.start_layer: |
| 911 | continue |
| 912 | intermediate_tensors[f"deepstack_input_embeds_{layer_idx}"] = torch.zeros( |
| 913 | (batch_size, self.config.hidden_size), |
| 914 | dtype=dtype, |
| 915 | device=device, |
| 916 | ) |
| 917 | return intermediate_tensors |
| 918 | |
| 919 | def forward( |
| 920 | self, |
| 921 | input_ids: torch.Tensor | None, |
| 922 | positions: torch.Tensor, |
| 923 | intermediate_tensors: IntermediateTensors | None = None, |
| 924 | inputs_embeds: torch.Tensor | None = None, |
| 925 | deepstack_input_embeds: IntermediateTensors | None = None, |
| 926 | ) -> torch.Tensor | IntermediateTensors: |
| 927 | return self.model( |
| 928 | input_ids, |
| 929 | positions, |
| 930 | intermediate_tensors, |
| 931 | inputs_embeds, |
| 932 | deepstack_input_embeds=deepstack_input_embeds, |
| 933 | ) |
| 934 | |
| 935 | |
| 936 | def _moss_audio_field_config( |
| 937 | hf_inputs: Mapping[str, torch.Tensor], |
| 938 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 939 | return { |
| 940 | "audio_data": MultiModalFieldConfig.batched("audio"), |
| 941 | "audio_data_seqlens": MultiModalFieldConfig.batched("audio", keep_on_cpu=True), |
| 942 | } |
| 943 | |
| 944 | |
| 945 | class MossAudioMultiModalDataParser(MultiModalDataParser): |
| 946 | def _parse_audio_data( |
| 947 | self, |
| 948 | data: dict[str, torch.Tensor] | ModalityData[AudioItem], |
| 949 | ) -> ModalityDataItems[Any, Any] | None: |
| 950 | if isinstance(data, dict): |
| 951 | return DictEmbeddingItems( |
| 952 | data, |
| 953 | modality="audio", |
| 954 | required_fields={"audio_data", "audio_data_seqlens"}, |
| 955 | fields_factory=_moss_audio_field_config, |
| 956 | ) |
| 957 | |
| 958 | return super()._parse_audio_data(data) |
| 959 | |
| 960 | |
| 961 | class MossAudioProcessor: |
| 962 | model_input_names = [ |
| 963 | "input_ids", |
| 964 | "attention_mask", |
| 965 | "audio_data", |
| 966 | "audio_data_seqlens", |
| 967 | ] |
| 968 | |
| 969 | def __init__( |
| 970 | self, |
| 971 | tokenizer: object, |
| 972 | *, |
| 973 | audio_token_id: int = MOSS_AUDIO_TOKEN_ID, |
| 974 | audio_start_id: int = MOSS_AUDIO_BOS_TOKEN_ID, |
| 975 | audio_end_id: int = MOSS_AUDIO_EOS_TOKEN_ID, |
| 976 | enable_time_marker: bool = False, |
| 977 | mel_config: Mapping[str, object] | None = None, |
| 978 | ) -> None: |
| 979 | assert hasattr(tokenizer, "encode") |
| 980 | assert hasattr(tokenizer, "decode") |
| 981 | assert hasattr(tokenizer, "batch_decode") |
| 982 | self.tokenizer = tokenizer |
| 983 | self.audio_token_id = int(audio_token_id) |
| 984 | self.audio_start_id = int(audio_start_id) |
| 985 | self.audio_end_id = int(audio_end_id) |
| 986 | self.enable_time_marker = bool(enable_time_marker) |
| 987 | self.mel_config = _normalize_moss_audio_mel_config(mel_config) |
| 988 | self.feature_extractor = WhisperFeatureExtractor( |
| 989 | feature_size=self.mel_config["mel_dim"], |
| 990 | sampling_rate=self.mel_config["mel_sr"], |
| 991 | hop_length=self.mel_config["mel_hop_length"], |
| 992 | n_fft=self.mel_config["mel_n_fft"], |
| 993 | ) |
| 994 | self.audio_tokens_per_second = self.mel_config["mel_sr"] / ( |
| 995 | self.mel_config["mel_hop_length"] * 8 |
| 996 | ) |
| 997 | self.time_marker_every_seconds = 2 |
| 998 | self.time_marker_every_audio_tokens = int( |
| 999 | self.audio_tokens_per_second * self.time_marker_every_seconds |
| 1000 | ) |
| 1001 | self._digit_token_ids = { |
| 1002 | "0": 15, |
| 1003 | "1": 16, |
| 1004 | "2": 17, |
| 1005 | "3": 18, |
| 1006 | "4": 19, |
| 1007 | "5": 20, |
| 1008 | "6": 21, |
| 1009 | "7": 22, |
| 1010 | "8": 23, |
| 1011 | "9": 24, |
| 1012 | } |
| 1013 | |
| 1014 | @staticmethod |
| 1015 | def conv3_downsample_len(raw_mel_len: int) -> int: |
| 1016 | return MossAudioEncoder.compute_num_audio_tokens(raw_mel_len) |
| 1017 | |
| 1018 | def _extract_mel(self, audio: np.ndarray | torch.Tensor) -> torch.Tensor: |
| 1019 | if isinstance(audio, torch.Tensor): |
| 1020 | wav = audio.detach().to("cpu", dtype=torch.float32).numpy() |
| 1021 | else: |
| 1022 | wav = np.asarray(audio, dtype=np.float32) |
| 1023 | if wav.size == 0: |
| 1024 | raise ValueError("The audio is too short to be represented.") |
| 1025 | if wav.ndim == 2: |
| 1026 | wav = wav[0] |
| 1027 | feats = self.feature_extractor._np_extract_fbank_features( |
| 1028 | wav[None, ...], device="cpu" |
| 1029 | ) |
| 1030 | return torch.from_numpy(feats[0]) |
| 1031 | |
| 1032 | def _get_default_audio_prompt(self) -> str: |
| 1033 | return MOSS_AUDIO_PLACEHOLDER |
| 1034 | |
| 1035 | def _ensure_audio_placeholders( |
| 1036 | self, |
| 1037 | prompt_text: str, |
| 1038 | num_audios: int, |
| 1039 | ) -> str: |
| 1040 | if num_audios == 0 or MOSS_AUDIO_SPAN_RE.search(prompt_text): |
| 1041 | return prompt_text |
| 1042 | |
| 1043 | audio_prompt = self._get_default_audio_prompt() * num_audios |
| 1044 | if prompt_text: |
| 1045 | return f"{audio_prompt}\n{prompt_text}" |
| 1046 | return audio_prompt |
| 1047 | |
| 1048 | def _build_audio_tokens_with_time_markers(self, audio_seq_len: int) -> list[int]: |
| 1049 | total_duration_seconds = audio_seq_len / self.audio_tokens_per_second |
| 1050 | num_full_seconds = int(total_duration_seconds) |
| 1051 | token_ids: list[int] = [] |
| 1052 | audio_tokens_consumed = 0 |
| 1053 | for second in range( |
| 1054 | self.time_marker_every_seconds, |
| 1055 | num_full_seconds + 1, |
| 1056 | self.time_marker_every_seconds, |
| 1057 | ): |
| 1058 | marker_pos = ( |
| 1059 | second // self.time_marker_every_seconds |
| 1060 | ) * self.time_marker_every_audio_tokens |
| 1061 | audio_segment_len = marker_pos - audio_tokens_consumed |
| 1062 | if audio_segment_len > 0: |
| 1063 | token_ids.extend([self.audio_token_id] * audio_segment_len) |
| 1064 | audio_tokens_consumed += audio_segment_len |
| 1065 | token_ids.extend(self._digit_token_ids[digit] for digit in str(second)) |
| 1066 | |
| 1067 | remaining = audio_seq_len - audio_tokens_consumed |
| 1068 | if remaining > 0: |
| 1069 | token_ids.extend([self.audio_token_id] * remaining) |
| 1070 | return token_ids |
| 1071 | |
| 1072 | def build_audio_placeholder_ids(self, num_audio_tokens: int) -> list[int]: |
| 1073 | if self.enable_time_marker: |
| 1074 | return self._build_audio_tokens_with_time_markers(num_audio_tokens) |
| 1075 | return [self.audio_token_id] * num_audio_tokens |
| 1076 | |
| 1077 | def __call__( |
| 1078 | self, |
| 1079 | text: str | Sequence[str] | None = None, |
| 1080 | audios: Sequence[np.ndarray | torch.Tensor] | None = None, |
| 1081 | audio: Sequence[np.ndarray | torch.Tensor] | None = None, |
| 1082 | return_tensors: str = "pt", |
| 1083 | **kwargs: object, |
| 1084 | ) -> BatchFeature: |
| 1085 | """Build text tokens and audio tensors for one MossAudio prompt. |
| 1086 | |
| 1087 | Example: |
| 1088 | text="Describe this.", audio=[waveform] |
| 1089 | -> input_ids contains audio_start, N audio tokens, audio_end |
| 1090 | -> audio_data has shape [1, mel_dim, max_time] |
| 1091 | -> mel_dim is the number of mel filter-bank bins, 128 by default |
| 1092 | -> audio_data_seqlens stores the unpadded mel length |
| 1093 | |
| 1094 | """ |
| 1095 | del kwargs |
| 1096 | |
| 1097 | # Step 1. Normalize text input; this processor handles one prompt. |
| 1098 | if isinstance(text, (list, tuple)): |
| 1099 | if len(text) != 1: |
| 1100 | raise ValueError(f"Expected text batch size 1, got {len(text)}") |
| 1101 | prompt_text = text[0] |
| 1102 | elif text is None: |
| 1103 | prompt_text = "" |
| 1104 | else: |
| 1105 | prompt_text = text |
| 1106 | |
| 1107 | # Step 2. Accept either `audios` or `audio` and normalize to a list. |
| 1108 | audio_list = audios if audios is not None else audio |
| 1109 | audio_list = [] if audio_list is None else list(audio_list) |
| 1110 | |
| 1111 | # Step 3. Convert waveforms to [mel_dim, time] mel features and token |
| 1112 | # counts. mel_dim is the number of mel filter-bank bins. |
| 1113 | mels: list[torch.Tensor] = [] |
| 1114 | raw_lengths: list[int] = [] |
| 1115 | token_lens: list[int] = [] |
| 1116 | for one_audio in audio_list: |
| 1117 | mel = self._extract_mel(one_audio) |
| 1118 | raw_len = int(mel.shape[-1]) |
| 1119 | num_tokens = self.conv3_downsample_len(raw_len) |
| 1120 | if raw_len <= 0 or num_tokens <= 0: |
| 1121 | raise ValueError("The audio is too short to be represented.") |
| 1122 | mels.append(mel) |
| 1123 | raw_lengths.append(raw_len) |
| 1124 | token_lens.append(num_tokens) |
| 1125 | |
| 1126 | # Step 4. Pad variable-length mel features into a batch tensor. |
| 1127 | if mels: |
| 1128 | max_length = max(raw_lengths) |
| 1129 | audio_batch = torch.zeros( |
| 1130 | (len(mels), self.mel_config["mel_dim"], max_length), |
| 1131 | dtype=torch.float32, |
| 1132 | ) |
| 1133 | for index, mel in enumerate(mels): |
| 1134 | audio_batch[index, :, : mel.shape[-1]] = mel |
| 1135 | audio_data_seqlens = torch.tensor(raw_lengths, dtype=torch.long) |
| 1136 | else: |
| 1137 | audio_batch = None |
| 1138 | audio_data_seqlens = None |
| 1139 | |
| 1140 | # Step 5. Ensure each audio item has a placeholder span in the prompt. |
| 1141 | prompt_text = self._ensure_audio_placeholders(prompt_text, len(audio_list)) |
| 1142 | input_ids = [] |
| 1143 | cursor = 0 |
| 1144 | |
| 1145 | # Step 6. Text-only path: tokenize and preserve placeholder spans. |
| 1146 | if not audio_list: |
| 1147 | for match in MOSS_AUDIO_SPAN_RE.finditer(prompt_text): |
| 1148 | prefix = prompt_text[cursor : match.start()] |
| 1149 | input_ids.extend( |
| 1150 | self.tokenizer.encode(prefix, add_special_tokens=False) |
| 1151 | ) |
| 1152 | input_ids.extend( |
| 1153 | [self.audio_start_id, self.audio_token_id, self.audio_end_id] |
| 1154 | ) |
| 1155 | cursor = match.end() |
| 1156 | suffix = prompt_text[cursor:] |
| 1157 | input_ids.extend(self.tokenizer.encode(suffix, add_special_tokens=False)) |
| 1158 | data: dict[str, torch.Tensor] = { |
| 1159 | "input_ids": torch.tensor([input_ids], dtype=torch.long), |
| 1160 | "attention_mask": torch.ones((1, len(input_ids)), dtype=torch.long), |
| 1161 | } |
| 1162 | return BatchFeature(data=data, tensor_type=return_tensors) |
| 1163 | |
| 1164 | # Step 7. Audio path: expand each placeholder to its audio-token count. |
| 1165 | span_iter = iter(MOSS_AUDIO_SPAN_RE.finditer(prompt_text)) |
| 1166 | for item_idx, _ in enumerate(audio_list): |
| 1167 | match = next(span_iter, None) |
| 1168 | if match is None: |
| 1169 | raise ValueError( |
| 1170 | "Audio placeholder count mismatch: expected one " |
| 1171 | f"{MOSS_AUDIO_PLACEHOLDER!r} span per audio item." |
| 1172 | ) |
| 1173 | prefix = prompt_text[cursor : match.start()] |
| 1174 | input_ids.extend(self.tokenizer.encode(prefix, add_special_tokens=False)) |
| 1175 | input_ids.append(self.audio_start_id) |
| 1176 | input_ids.extend(self.build_audio_placeholder_ids(token_lens[item_idx])) |
| 1177 | input_ids.append(self.audio_end_id) |
| 1178 | cursor = match.end() |
| 1179 | |
| 1180 | # Step 8. Reject extra placeholder spans after all audio items are used. |
| 1181 | suffix = prompt_text[cursor:] |
| 1182 | if MOSS_AUDIO_SPAN_RE.search(suffix): |
| 1183 | raise ValueError( |
| 1184 | "Audio placeholder count mismatch: found more placeholder spans " |
| 1185 | "than audio items." |
| 1186 | ) |
| 1187 | input_ids.extend(self.tokenizer.encode(suffix, add_special_tokens=False)) |
| 1188 | |
| 1189 | # Step 9. Return tokenizer output plus audio tensors for embed_multimodal. |
| 1190 | data = { |
| 1191 | "input_ids": torch.tensor([input_ids], dtype=torch.long), |
| 1192 | "attention_mask": torch.ones((1, len(input_ids)), dtype=torch.long), |
| 1193 | } |
| 1194 | if audio_batch is not None and audio_data_seqlens is not None: |
| 1195 | data["audio_data"] = audio_batch |
| 1196 | data["audio_data_seqlens"] = audio_data_seqlens |
| 1197 | return BatchFeature(data=data, tensor_type=return_tensors) |
| 1198 | |
| 1199 | def decode(self, *args: object, **kwargs: object) -> str: |
| 1200 | return self.tokenizer.decode(*args, **kwargs) |
| 1201 | |
| 1202 | def batch_decode(self, *args: object, **kwargs: object) -> list[str]: |
| 1203 | return self.tokenizer.batch_decode(*args, **kwargs) |
| 1204 | |
| 1205 | |
| 1206 | class MossAudioProcessingInfo(BaseProcessingInfo): |
| 1207 | def __init__(self, ctx: InputProcessingContext) -> None: |
| 1208 | super().__init__(ctx) |
| 1209 | if ctx.tokenizer is not None: |
| 1210 | _ensure_moss_audio_placeholder_tokens(ctx.tokenizer) |
| 1211 | |
| 1212 | def get_hf_config(self) -> MossAudioConfig: |
| 1213 | config = self.ctx.get_hf_config() |
| 1214 | if isinstance(config, MossAudioConfig): |
| 1215 | return config |
| 1216 | return MossAudioConfig( |
| 1217 | audio_config=getattr(config, "audio_config", None), |
| 1218 | language_config=getattr(config, "language_config", None), |
| 1219 | adapter_hidden_size=getattr(config, "adapter_hidden_size", 8192), |
| 1220 | ignore_index=getattr(config, "ignore_index", -100), |
| 1221 | deepstack_num_inject_layers=getattr( |
| 1222 | config, "deepstack_num_inject_layers", None |
| 1223 | ), |
| 1224 | ) |
| 1225 | |
| 1226 | def _get_processor_config_defaults(self) -> dict[str, object]: |
| 1227 | cached_defaults = getattr(self, "_processor_config_defaults", None) |
| 1228 | if cached_defaults is not None: |
| 1229 | return cached_defaults |
| 1230 | |
| 1231 | model_config = self.ctx.model_config |
| 1232 | for file_name in ("processor_config.json", "preprocessor_config.json"): |
| 1233 | config = get_hf_file_to_dict( |
| 1234 | file_name, |
| 1235 | model_config.model, |
| 1236 | model_config.revision, |
| 1237 | ) |
| 1238 | defaults = _filter_moss_audio_processor_config(config) |
| 1239 | if defaults: |
| 1240 | self._processor_config_defaults = defaults |
| 1241 | return defaults |
| 1242 | |
| 1243 | defaults = {} |
| 1244 | self._processor_config_defaults = defaults |
| 1245 | return defaults |
| 1246 | |
| 1247 | @staticmethod |
| 1248 | def _get_processor_cache_key(kwargs: Mapping[str, object]) -> tuple[object, ...]: |
| 1249 | mel_config_arg = kwargs.get("mel_config") |
| 1250 | assert mel_config_arg is None or isinstance(mel_config_arg, Mapping) |
| 1251 | mel_config = _normalize_moss_audio_mel_config(mel_config_arg) |
| 1252 | audio_token_id = kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID) |
| 1253 | audio_start_id = kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID) |
| 1254 | audio_end_id = kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID) |
| 1255 | assert isinstance(audio_token_id, int) |
| 1256 | assert isinstance(audio_start_id, int) |
| 1257 | assert isinstance(audio_end_id, int) |
| 1258 | return ( |
| 1259 | audio_token_id, |
| 1260 | audio_start_id, |
| 1261 | audio_end_id, |
| 1262 | bool(kwargs.get("enable_time_marker", False)), |
| 1263 | tuple(sorted(mel_config.items())), |
| 1264 | ) |
| 1265 | |
| 1266 | def get_hf_processor(self, **kwargs: object) -> MossAudioProcessor: |
| 1267 | merged_kwargs = _merge_moss_audio_processor_configs( |
| 1268 | self._get_processor_config_defaults(), |
| 1269 | self.ctx.get_merged_mm_kwargs({}), |
| 1270 | kwargs, |
| 1271 | ) |
| 1272 | mel_config_arg = merged_kwargs.get("mel_config") |
| 1273 | assert mel_config_arg is None or isinstance(mel_config_arg, Mapping) |
| 1274 | mel_config = _normalize_moss_audio_mel_config(mel_config_arg) |
| 1275 | audio_token_id = merged_kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID) |
| 1276 | audio_start_id = merged_kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID) |
| 1277 | audio_end_id = merged_kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID) |
| 1278 | assert isinstance(audio_token_id, int) |
| 1279 | assert isinstance(audio_start_id, int) |
| 1280 | assert isinstance(audio_end_id, int) |
| 1281 | processor_kwargs: MossAudioProcessorKwargs = { |
| 1282 | "audio_token_id": audio_token_id, |
| 1283 | "audio_start_id": audio_start_id, |
| 1284 | "audio_end_id": audio_end_id, |
| 1285 | "enable_time_marker": bool(merged_kwargs.get("enable_time_marker", False)), |
| 1286 | "mel_config": mel_config, |
| 1287 | } |
| 1288 | |
| 1289 | cache = getattr(self, "_hf_processor_cache", None) |
| 1290 | if cache is None: |
| 1291 | cache = {} |
| 1292 | self._hf_processor_cache = cache |
| 1293 | |
| 1294 | cache_key = self._get_processor_cache_key(processor_kwargs) |
| 1295 | processor = cache.get(cache_key) |
| 1296 | if processor is not None: |
| 1297 | return processor |
| 1298 | |
| 1299 | processor = MossAudioProcessor( |
| 1300 | self.get_tokenizer(), |
| 1301 | **processor_kwargs, |
| 1302 | ) |
| 1303 | cache[cache_key] = processor |
| 1304 | return processor |
| 1305 | |
| 1306 | def get_feature_extractor(self, **kwargs: object) -> WhisperFeatureExtractor: |
| 1307 | return self.get_hf_processor(**kwargs).feature_extractor |
| 1308 | |
| 1309 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 1310 | return {"audio": None} |
| 1311 | |
| 1312 | def get_data_parser(self) -> MultiModalDataParser: |
| 1313 | processor = self.get_hf_processor() |
| 1314 | return MossAudioMultiModalDataParser( |
| 1315 | target_sr=processor.mel_config["mel_sr"], |
| 1316 | target_channels=1, |
| 1317 | expected_hidden_size=self._get_expected_hidden_size(), |
| 1318 | ) |
| 1319 | |
| 1320 | def get_mm_max_tokens_per_item( |
| 1321 | self, |
| 1322 | seq_len: int, |
| 1323 | mm_counts: Mapping[str, int], |
| 1324 | ) -> Mapping[str, int] | None: |
| 1325 | if mm_counts.get("audio", 0) <= 0: |
| 1326 | return {} |
| 1327 | processor = self.get_hf_processor() |
| 1328 | raw_mel_len = math.ceil( |
| 1329 | (processor.mel_config["mel_sr"] * DEFAULT_MAX_AUDIO_SECONDS) |
| 1330 | / processor.mel_config["mel_hop_length"] |
| 1331 | ) |
| 1332 | return {"audio": MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)} |
| 1333 | |
| 1334 | |
| 1335 | class MossAudioDummyInputsBuilder(BaseDummyInputsBuilder[MossAudioProcessingInfo]): |
| 1336 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 1337 | num_audios = mm_counts.get("audio", 0) |
| 1338 | return MOSS_AUDIO_PLACEHOLDER * num_audios |
| 1339 | |
| 1340 | def get_dummy_mm_data( |
| 1341 | self, |
| 1342 | seq_len: int, |
| 1343 | mm_counts: Mapping[str, int], |
| 1344 | mm_options: MultiModalDummyOptions, |
| 1345 | ) -> MultiModalDataDict: |
| 1346 | return { |
| 1347 | "audio": self._get_dummy_audios( |
| 1348 | length=16000, |
| 1349 | num_audios=mm_counts.get("audio", 0), |
| 1350 | overrides=mm_options.get("audio"), |
| 1351 | ) |
| 1352 | } |
| 1353 | |
| 1354 | |
| 1355 | class MossAudioMultiModalProcessor(BaseMultiModalProcessor[MossAudioProcessingInfo]): |
| 1356 | def _call_hf_processor( |
| 1357 | self, |
| 1358 | hf_data: Mapping[str, object], |
| 1359 | hf_kwargs: Mapping[str, object], |
| 1360 | ) -> BatchFeature: |
| 1361 | processor_kwargs = _filter_moss_audio_processor_config(dict(hf_kwargs)) |
| 1362 | return self.info.ctx.call_hf_processor( |
| 1363 | self.info.get_hf_processor(**processor_kwargs), |
| 1364 | hf_data, |
| 1365 | {}, |
| 1366 | ) |
| 1367 | |
| 1368 | def _get_mm_fields_config( |
| 1369 | self, |
| 1370 | hf_inputs: BatchFeature, |
| 1371 | hf_processor_mm_kwargs: Mapping[str, object], |
| 1372 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 1373 | return _moss_audio_field_config(hf_inputs) |
| 1374 | |
| 1375 | def _get_prompt_updates( |
| 1376 | self, |
| 1377 | mm_items: MultiModalDataItems, |
| 1378 | hf_processor_mm_kwargs: Mapping[str, object], |
| 1379 | out_mm_kwargs: MultiModalKwargsItems, |
| 1380 | ) -> Sequence[PromptUpdate]: |
| 1381 | processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 1382 | tokenizer = processor.tokenizer |
| 1383 | |
| 1384 | out_mm_data = out_mm_kwargs.get_data() |
| 1385 | audio_data_seqlens = out_mm_data.get("audio_data_seqlens") |
| 1386 | if audio_data_seqlens is None: |
| 1387 | audio_token_lens: list[int] = [] |
| 1388 | else: |
| 1389 | if isinstance(audio_data_seqlens, torch.Tensor): |
| 1390 | lens = audio_data_seqlens.reshape(-1).tolist() |
| 1391 | else: |
| 1392 | lens = list(audio_data_seqlens) |
| 1393 | audio_token_lens = [ |
| 1394 | MossAudioEncoder.compute_num_audio_tokens(int(length)) |
| 1395 | for length in lens |
| 1396 | ] |
| 1397 | |
| 1398 | def get_replacement( |
| 1399 | item_idx: int, |
| 1400 | suffix_token_ids: list[int] | None = None, |
| 1401 | ) -> PromptUpdateDetails: |
| 1402 | num_tokens = audio_token_lens[item_idx] |
| 1403 | if num_tokens == 0: |
| 1404 | raise ValueError("The audio is too short to be represented.") |
| 1405 | audio_token_ids = processor.build_audio_placeholder_ids(num_tokens) |
| 1406 | suffix_token_ids = suffix_token_ids or [] |
| 1407 | is_embed = torch.tensor( |
| 1408 | [token_id == processor.audio_token_id for token_id in audio_token_ids], |
| 1409 | dtype=torch.bool, |
| 1410 | ) |
| 1411 | return PromptUpdateDetails( |
| 1412 | full=[ |
| 1413 | processor.audio_start_id, |
| 1414 | *audio_token_ids, |
| 1415 | processor.audio_end_id, |
| 1416 | *suffix_token_ids, |
| 1417 | ], |
| 1418 | is_embed=lambda _seq: torch.cat( |
| 1419 | [ |
| 1420 | torch.tensor([False]), |
| 1421 | is_embed, |
| 1422 | torch.tensor([False]), |
| 1423 | torch.zeros(len(suffix_token_ids), dtype=torch.bool), |
| 1424 | ] |
| 1425 | ), |
| 1426 | ) |
| 1427 | |
| 1428 | prompt_update_specs: list[tuple[list[int], list[int]]] = [ |
| 1429 | ( |
| 1430 | [ |
| 1431 | processor.audio_start_id, |
| 1432 | processor.audio_token_id, |
| 1433 | processor.audio_end_id, |
| 1434 | ], |
| 1435 | [], |
| 1436 | ) |
| 1437 | ] |
| 1438 | for suffix in ("", "\n"): |
| 1439 | tokenizer_target = cached_encode( |
| 1440 | tokenizer, |
| 1441 | MOSS_AUDIO_PLACEHOLDER + suffix, |
| 1442 | add_special_tokens=False, |
| 1443 | ) |
| 1444 | suffix_token_ids = cached_encode( |
| 1445 | tokenizer, |
| 1446 | suffix, |
| 1447 | add_special_tokens=False, |
| 1448 | ) |
| 1449 | if any(target == tokenizer_target for target, _ in prompt_update_specs): |
| 1450 | continue |
| 1451 | prompt_update_specs.append((tokenizer_target, suffix_token_ids)) |
| 1452 | |
| 1453 | return [ |
| 1454 | PromptReplacement( |
| 1455 | modality="audio", |
| 1456 | target=target, |
| 1457 | replacement=partial(get_replacement, suffix_token_ids=suffix_token_ids), |
| 1458 | ) |
| 1459 | for target, suffix_token_ids in prompt_update_specs |
| 1460 | ] |
| 1461 | |
| 1462 | |
| 1463 | @MULTIMODAL_REGISTRY.register_processor( |
| 1464 | MossAudioMultiModalProcessor, |
| 1465 | info=MossAudioProcessingInfo, |
| 1466 | dummy_inputs=MossAudioDummyInputsBuilder, |
| 1467 | ) |
| 1468 | class MossAudioModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA): |
| 1469 | packed_modules_mapping = { |
| 1470 | "qkv_proj": [ |
| 1471 | "q_proj", |
| 1472 | "k_proj", |
| 1473 | "v_proj", |
| 1474 | ], |
| 1475 | "gate_up_proj": [ |
| 1476 | "gate_proj", |
| 1477 | "up_proj", |
| 1478 | ], |
| 1479 | } |
| 1480 | |
| 1481 | embedding_modules = { |
| 1482 | "embed_tokens": "input_embeddings", |
| 1483 | "lm_head": "output_embeddings", |
| 1484 | } |
| 1485 | |
| 1486 | hf_to_vllm_mapper = WeightsMapper( |
| 1487 | orig_to_new_prefix={ |
| 1488 | "lm_head.": "language_model.lm_head.", |
| 1489 | "language_model.embed_tokens.": "language_model.model.embed_tokens.", |
| 1490 | "language_model.layers.": "language_model.model.layers.", |
| 1491 | "language_model.norm.": "language_model.model.norm.", |
| 1492 | "audio_encoder.embed_positions": None, |
| 1493 | }, |
| 1494 | orig_to_new_stacked={ |
| 1495 | ".gate_proj": (".gate_up_proj", 0), |
| 1496 | ".up_proj": (".gate_up_proj", 1), |
| 1497 | }, |
| 1498 | ) |
| 1499 | |
| 1500 | def get_mm_mapping(self) -> MultiModelKeys: |
| 1501 | return MultiModelKeys.from_string_field( |
| 1502 | language_model="language_model.", |
| 1503 | ) |
| 1504 | |
| 1505 | @classmethod |
| 1506 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 1507 | if modality.startswith("audio"): |
| 1508 | return MOSS_AUDIO_PLACEHOLDER |
| 1509 | raise ValueError("Only audio modality is supported") |
| 1510 | |
| 1511 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 1512 | super().__init__() |
| 1513 | self.vllm_config = vllm_config |
| 1514 | config = vllm_config.model_config.hf_config |
| 1515 | if not isinstance(config, MossAudioConfig): |
| 1516 | config = MossAudioConfig( |
| 1517 | audio_config=getattr(config, "audio_config", None), |
| 1518 | language_config=getattr(config, "language_config", None), |
| 1519 | adapter_hidden_size=getattr(config, "adapter_hidden_size", 8192), |
| 1520 | ignore_index=getattr(config, "ignore_index", -100), |
| 1521 | deepstack_num_inject_layers=getattr( |
| 1522 | config, "deepstack_num_inject_layers", None |
| 1523 | ), |
| 1524 | ) |
| 1525 | self.config = config |
| 1526 | self.quant_config = vllm_config.quant_config |
| 1527 | self.multimodal_config = vllm_config.model_config.multimodal_config |
| 1528 | |
| 1529 | parallel_config = vllm_config.parallel_config |
| 1530 | tp_size = parallel_config.tensor_parallel_size |
| 1531 | if self.config.adapter_hidden_size % tp_size != 0: |
| 1532 | raise ValueError( |
| 1533 | "MOSS-Audio adapter_hidden_size must be divisible by tensor " |
| 1534 | f"parallel size. Got adapter_hidden_size=" |
| 1535 | f"{self.config.adapter_hidden_size} and tensor_parallel_size=" |
| 1536 | f"{tp_size}." |
| 1537 | ) |
| 1538 | |
| 1539 | audio_config = MossAudioEncoderConfig.from_config(self.config.audio_config) |
| 1540 | if audio_config.encoder_attention_heads % tp_size != 0: |
| 1541 | raise ValueError( |
| 1542 | "MOSS-Audio encoder_attention_heads must be divisible by " |
| 1543 | "tensor parallel size. Got encoder_attention_heads=" |
| 1544 | f"{audio_config.encoder_attention_heads} and " |
| 1545 | f"tensor_parallel_size={tp_size}." |
| 1546 | ) |
| 1547 | language_config = self.config.language_config |
| 1548 | self.audio_token_id = MOSS_AUDIO_TOKEN_ID |
| 1549 | self.deepstack_input_embeds: IntermediateTensors | None = None |
| 1550 | |
| 1551 | with self._mark_tower_model(vllm_config, "audio"): |
| 1552 | self.audio_encoder = MossAudioEncoder( |
| 1553 | audio_config, |
| 1554 | quant_config=self.quant_config, |
| 1555 | prefix=maybe_prefix(prefix, "audio_encoder"), |
| 1556 | ) |
| 1557 | self.audio_adapter = GatedMLP( |
| 1558 | input_size=audio_config.output_dim, |
| 1559 | hidden_size=self.config.adapter_hidden_size, |
| 1560 | output_size=language_config.hidden_size, |
| 1561 | quant_config=self.quant_config, |
| 1562 | prefix=maybe_prefix(prefix, "audio_adapter"), |
| 1563 | ) |
| 1564 | |
| 1565 | deepstack_k = len(audio_config.deepstack_encoder_layer_indexes or []) |
| 1566 | if self.config.deepstack_num_inject_layers is not None: |
| 1567 | deepstack_k = min( |
| 1568 | deepstack_k, |
| 1569 | int(self.config.deepstack_num_inject_layers), |
| 1570 | ) |
| 1571 | self.deepstack_audio_merger_list = nn.ModuleList( |
| 1572 | [ |
| 1573 | GatedMLP( |
| 1574 | input_size=audio_config.output_dim, |
| 1575 | hidden_size=self.config.adapter_hidden_size, |
| 1576 | output_size=language_config.hidden_size, |
| 1577 | quant_config=self.quant_config, |
| 1578 | prefix=maybe_prefix( |
| 1579 | prefix, |
| 1580 | f"deepstack_audio_merger_list.{layer_idx}", |
| 1581 | ), |
| 1582 | ) |
| 1583 | for layer_idx in range(deepstack_k) |
| 1584 | ] |
| 1585 | ) |
| 1586 | |
| 1587 | with self._mark_language_model(vllm_config): |
| 1588 | self.language_model = MossQwen3ForCausalLM( |
| 1589 | vllm_config=vllm_config.with_hf_config( |
| 1590 | language_config, architectures=["Qwen3ForCausalLM"] |
| 1591 | ), |
| 1592 | prefix=maybe_prefix(prefix, "language_model"), |
| 1593 | ) |
| 1594 | self.language_model.deepstack_inject_layer_indices = range(deepstack_k) |
| 1595 | self.language_model.model.deepstack_inject_layer_indices = range( |
| 1596 | deepstack_k |
| 1597 | ) |
| 1598 | |
| 1599 | self.make_empty_intermediate_tensors = ( |
| 1600 | self.language_model.make_empty_intermediate_tensors |
| 1601 | ) |
| 1602 | |
| 1603 | @staticmethod |
| 1604 | def _validate_audio_batch_size( |
| 1605 | audio_batch_size: int, audio_data_seqlens: torch.Tensor |
| 1606 | ) -> None: |
| 1607 | if audio_batch_size != audio_data_seqlens.numel(): |
| 1608 | raise ValueError( |
| 1609 | "audio_data batch size does not match audio_data_seqlens: " |
| 1610 | f"{audio_batch_size} != {audio_data_seqlens.numel()}." |
| 1611 | ) |
| 1612 | |
| 1613 | @staticmethod |
| 1614 | def _pad_audio_data_list( |
| 1615 | audio_data: list[torch.Tensor], |
| 1616 | audio_data_seqlens: torch.Tensor, |
| 1617 | ) -> torch.Tensor: |
| 1618 | if len(audio_data) == 0: |
| 1619 | raise ValueError("audio_data list must not be empty.") |
| 1620 | MossAudioModel._validate_audio_batch_size(len(audio_data), audio_data_seqlens) |
| 1621 | |
| 1622 | # pad_sequence needs every item to share the same trailing feature |
| 1623 | # layout, so validate the mel-major audio tensors before transposing. |
| 1624 | first = audio_data[0] |
| 1625 | if not isinstance(first, torch.Tensor): |
| 1626 | raise TypeError("audio_data list items must be torch.Tensor.") |
| 1627 | if first.ndim != 2: |
| 1628 | raise ValueError("audio_data list items must have shape [mel_dim, time].") |
| 1629 | |
| 1630 | mel_dim = first.shape[0] |
| 1631 | dtype = first.dtype |
| 1632 | device = first.device |
| 1633 | for item in audio_data[1:]: |
| 1634 | if not isinstance(item, torch.Tensor): |
| 1635 | raise TypeError("audio_data list items must be torch.Tensor.") |
| 1636 | if item.ndim != 2: |
| 1637 | raise ValueError( |
| 1638 | "audio_data list items must have shape [mel_dim, time]." |
| 1639 | ) |
| 1640 | if item.shape[0] != mel_dim: |
| 1641 | raise ValueError("audio_data list items must have the same mel_dim.") |
| 1642 | if item.dtype != dtype: |
| 1643 | raise TypeError("audio_data list items must have the same dtype.") |
| 1644 | if item.device != device: |
| 1645 | raise ValueError("audio_data list items must be on the same device.") |
| 1646 | |
| 1647 | # Each item arrives as [mel_dim, time]. pad_sequence pads along dim 1 |
| 1648 | # after converting to [time, mel_dim], then we restore [batch, mel, time]. |
| 1649 | time_major = [item.transpose(0, 1) for item in audio_data] |
| 1650 | padded = torch.nn.utils.rnn.pad_sequence(time_major, batch_first=True) |
| 1651 | return padded.transpose(1, 2).contiguous() |
| 1652 | |
| 1653 | def _parse_and_validate_audio_input( |
| 1654 | self, **kwargs: object |
| 1655 | ) -> MossAudioAudioInputs | None: |
| 1656 | """Normalize and validate model-side audio kwargs. |
| 1657 | |
| 1658 | If audio_data is provided, this checks that audio_data_seqlens is also |
| 1659 | present, flattens sequence lengths to a long tensor, pads list inputs |
| 1660 | to [batch, mel_dim, time], validates batch-size/sequence-length |
| 1661 | agreement, and rejects empty, non-positive, or downsampled-zero audio |
| 1662 | lengths. |
| 1663 | """ |
| 1664 | audio_data = kwargs.pop("audio_data", None) |
| 1665 | audio_data_seqlens = kwargs.pop("audio_data_seqlens", None) |
| 1666 | if audio_data is None: |
| 1667 | return None |
| 1668 | if audio_data_seqlens is None: |
| 1669 | raise ValueError( |
| 1670 | "audio_data_seqlens is required when audio_data is provided." |
| 1671 | ) |
| 1672 | if not isinstance(audio_data_seqlens, torch.Tensor): |
| 1673 | audio_data_seqlens = torch.tensor(audio_data_seqlens, dtype=torch.long) |
| 1674 | audio_data_seqlens = audio_data_seqlens.to(dtype=torch.long).reshape(-1) |
| 1675 | |
| 1676 | if isinstance(audio_data, list): |
| 1677 | audio_data = self._pad_audio_data_list(audio_data, audio_data_seqlens) |
| 1678 | elif isinstance(audio_data, torch.Tensor): |
| 1679 | if audio_data.ndim == 3: |
| 1680 | self._validate_audio_batch_size(audio_data.shape[0], audio_data_seqlens) |
| 1681 | else: |
| 1682 | raise TypeError("audio_data must be a torch.Tensor or list[torch.Tensor].") |
| 1683 | |
| 1684 | audio_token_lens = MossAudioEncoder._compute_downsampled_length( |
| 1685 | audio_data_seqlens |
| 1686 | ) |
| 1687 | if ( |
| 1688 | audio_data_seqlens.numel() == 0 |
| 1689 | or torch.any(audio_data_seqlens <= 0).item() |
| 1690 | or torch.any(audio_token_lens <= 0).item() |
| 1691 | ): |
| 1692 | raise ValueError("The audio is too short to be represented.") |
| 1693 | return MossAudioAudioInputs( |
| 1694 | audio_data=audio_data, |
| 1695 | audio_data_seqlens=audio_data_seqlens, |
| 1696 | ) |
| 1697 | |
| 1698 | def _process_audio_input( |
| 1699 | self, |
| 1700 | audio_input: MossAudioAudioInputs, |
| 1701 | ) -> tuple[torch.Tensor, ...]: |
| 1702 | """Run the audio encoder and return one embedding tensor per audio. |
| 1703 | |
| 1704 | Example: |
| 1705 | audio_data=[2, 128, 1200], audio_data_seqlens=[800, 1200] |
| 1706 | -> returns (audio0_embeds, audio1_embeds), split by token length |
| 1707 | -> DeepStack packs each item as [main, layer0, ...] on dim -1 |
| 1708 | |
| 1709 | """ |
| 1710 | audio_data = audio_input["audio_data"] |
| 1711 | audio_data_seqlens = audio_input["audio_data_seqlens"] |
| 1712 | # The encoder chunks the input by per-audio feature lengths, which |
| 1713 | # needs Python ints for `split`/`pad_sequence`. |
| 1714 | want_deepstack = len(self.deepstack_audio_merger_list) > 0 |
| 1715 | with gpu_sync_allowed(): |
| 1716 | last_hidden_state, deepstack = self.audio_encoder( |
| 1717 | audio_data.to(self.audio_encoder.dtype), |
| 1718 | feature_lens=audio_data_seqlens, |
| 1719 | output_deepstack_hidden_states=want_deepstack, |
| 1720 | ) |
| 1721 | audio_embeds = self.audio_adapter(last_hidden_state) |
| 1722 | audio_lengths = MossAudioEncoder._compute_downsampled_length( |
| 1723 | audio_data_seqlens.to(device=audio_embeds.device, dtype=torch.long) |
| 1724 | ).tolist() |
| 1725 | main_embeddings = tuple(audio_embeds.squeeze(0).split(audio_lengths, dim=0)) |
| 1726 | |
| 1727 | deepstack_embeddings: list[tuple[torch.Tensor, ...]] = [] |
| 1728 | if deepstack is not None: |
| 1729 | if len(deepstack) < len(self.deepstack_audio_merger_list): |
| 1730 | raise RuntimeError( |
| 1731 | "DeepStack output count does not match configured audio " |
| 1732 | "merger count." |
| 1733 | ) |
| 1734 | for idx, hidden_states in enumerate( |
| 1735 | deepstack[: len(self.deepstack_audio_merger_list)] |
| 1736 | ): |
| 1737 | ds_embeds = self.deepstack_audio_merger_list[idx](hidden_states) |
| 1738 | deepstack_embeddings.append( |
| 1739 | tuple(ds_embeds.squeeze(0).split(audio_lengths, dim=0)) |
| 1740 | ) |
| 1741 | |
| 1742 | if not deepstack_embeddings: |
| 1743 | return main_embeddings |
| 1744 | |
| 1745 | return tuple( |
| 1746 | torch.cat( |
| 1747 | [ |
| 1748 | main_embedding, |
| 1749 | *( |
| 1750 | layer_embeddings[item_idx] |
| 1751 | for layer_embeddings in deepstack_embeddings |
| 1752 | ), |
| 1753 | ], |
| 1754 | dim=-1, |
| 1755 | ) |
| 1756 | for item_idx, main_embedding in enumerate(main_embeddings) |
| 1757 | ) |
| 1758 | |
| 1759 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 1760 | audio_input = self._parse_and_validate_audio_input(**kwargs) |
| 1761 | if audio_input is None: |
| 1762 | return () |
| 1763 | return self._process_audio_input(audio_input) |
| 1764 | |
| 1765 | def _split_multimodal_embeddings( |
| 1766 | self, |
| 1767 | multimodal_embeddings: MultiModalEmbeddings, |
| 1768 | hidden_size: int, |
| 1769 | ) -> tuple[tuple[torch.Tensor, ...], tuple[tuple[torch.Tensor, ...], ...]]: |
| 1770 | """Unpack audio embeddings before merging them into token embeddings. |
| 1771 | |
| 1772 | embed_input_ids calls this on the output of embed_multimodal. Plain |
| 1773 | audio embeddings already have width hidden_size and are returned as the |
| 1774 | main embeddings for _merge_multimodal_embeddings. When DeepStack is |
| 1775 | enabled, _process_audio_input packs each audio item as |
| 1776 | [main, layer0, layer1, ...] along the last dimension so the standard |
| 1777 | multimodal path can carry a single embedding object. This method splits |
| 1778 | that packed layout back into main embeddings plus per-layer DeepStack |
| 1779 | embeddings, which _cache_deepstack_input_embeds scatters and forward |
| 1780 | passes into MossQwen3Model for layer injection. |
| 1781 | """ |
| 1782 | if isinstance(multimodal_embeddings, torch.Tensor): |
| 1783 | embeddings = tuple(multimodal_embeddings.unbind(0)) |
| 1784 | else: |
| 1785 | embeddings = tuple(multimodal_embeddings) |
| 1786 | |
| 1787 | if len(embeddings) == 0: |
| 1788 | return (), () |
| 1789 | |
| 1790 | deepstack_count = len(self.deepstack_audio_merger_list) |
| 1791 | if all(embedding.shape[-1] == hidden_size for embedding in embeddings): |
| 1792 | return embeddings, () |
| 1793 | |
| 1794 | packed_hidden_size = hidden_size * (deepstack_count + 1) |
| 1795 | if deepstack_count == 0 or any( |
| 1796 | embedding.shape[-1] != packed_hidden_size for embedding in embeddings |
| 1797 | ): |
| 1798 | got = [int(embedding.shape[-1]) for embedding in embeddings] |
| 1799 | raise ValueError( |
| 1800 | "MOSS-Audio multimodal embedding width mismatch: expected " |
| 1801 | f"{hidden_size} or {packed_hidden_size}, got {got}." |
| 1802 | ) |
| 1803 | |
| 1804 | split_by_item = [ |
| 1805 | torch.split(embedding, hidden_size, dim=-1) for embedding in embeddings |
| 1806 | ] |
| 1807 | main_embeddings = tuple(parts[0] for parts in split_by_item) |
| 1808 | deepstack_embeddings = tuple( |
| 1809 | tuple(parts[layer_idx + 1] for parts in split_by_item) |
| 1810 | for layer_idx in range(deepstack_count) |
| 1811 | ) |
| 1812 | return main_embeddings, deepstack_embeddings |
| 1813 | |
| 1814 | def _cache_deepstack_input_embeds( |
| 1815 | self, |
| 1816 | inputs_embeds: torch.Tensor, |
| 1817 | deepstack_embeddings: tuple[tuple[torch.Tensor, ...], ...], |
| 1818 | is_multimodal: torch.Tensor, |
| 1819 | ) -> None: |
| 1820 | if len(deepstack_embeddings) == 0: |
| 1821 | self.deepstack_input_embeds = None |
| 1822 | return |
| 1823 | flat_by_layer = [ |
| 1824 | torch.cat(layer_embeds, dim=0).to( |
| 1825 | device=inputs_embeds.device, dtype=inputs_embeds.dtype |
| 1826 | ) |
| 1827 | for layer_embeds in deepstack_embeddings |
| 1828 | ] |
| 1829 | num_mm_tokens = int(is_multimodal.sum().item()) |
| 1830 | if any(layer.shape[0] != num_mm_tokens for layer in flat_by_layer): |
| 1831 | got = [int(layer.shape[0]) for layer in flat_by_layer] |
| 1832 | raise ValueError( |
| 1833 | "DeepStack audio token count mismatch: " |
| 1834 | f"expected {num_mm_tokens}, got {got}." |
| 1835 | ) |
| 1836 | data = {} |
| 1837 | for layer_idx, layer_embeds in enumerate(flat_by_layer): |
| 1838 | scattered = inputs_embeds.new_zeros(inputs_embeds.shape) |
| 1839 | scattered[is_multimodal] = layer_embeds |
| 1840 | data[f"deepstack_input_embeds_{layer_idx}"] = scattered |
| 1841 | self.deepstack_input_embeds = IntermediateTensors(data) |
| 1842 | |
| 1843 | def embed_input_ids( |
| 1844 | self, |
| 1845 | input_ids: torch.Tensor, |
| 1846 | multimodal_embeddings: MultiModalEmbeddings | None = None, |
| 1847 | *, |
| 1848 | is_multimodal: torch.Tensor | None = None, |
| 1849 | ) -> torch.Tensor: |
| 1850 | inputs_embeds = self._embed_text_input_ids( |
| 1851 | input_ids, |
| 1852 | self.language_model.embed_input_ids, |
| 1853 | is_multimodal=is_multimodal, |
| 1854 | ) |
| 1855 | |
| 1856 | self.deepstack_input_embeds = None |
| 1857 | if multimodal_embeddings is None or len(multimodal_embeddings) == 0: |
| 1858 | return inputs_embeds |
| 1859 | is_multimodal = _require_is_multimodal(is_multimodal) |
| 1860 | multimodal_embeddings, deepstack_embeddings = self._split_multimodal_embeddings( |
| 1861 | multimodal_embeddings, |
| 1862 | hidden_size=int(inputs_embeds.shape[-1]), |
| 1863 | ) |
| 1864 | |
| 1865 | inputs_embeds = _merge_multimodal_embeddings( |
| 1866 | inputs_embeds=inputs_embeds, |
| 1867 | multimodal_embeddings=multimodal_embeddings, |
| 1868 | is_multimodal=is_multimodal, |
| 1869 | ) |
| 1870 | self._cache_deepstack_input_embeds( |
| 1871 | inputs_embeds, |
| 1872 | deepstack_embeddings, |
| 1873 | is_multimodal, |
| 1874 | ) |
| 1875 | return inputs_embeds |
| 1876 | |
| 1877 | def forward( |
| 1878 | self, |
| 1879 | input_ids: torch.Tensor | None, |
| 1880 | positions: torch.Tensor, |
| 1881 | intermediate_tensors: IntermediateTensors | None = None, |
| 1882 | inputs_embeds: torch.Tensor | None = None, |
| 1883 | **kwargs: object, |
| 1884 | ) -> torch.Tensor | IntermediateTensors: |
| 1885 | if intermediate_tensors is None: |
| 1886 | deepstack_input_embeds = self.deepstack_input_embeds |
| 1887 | else: |
| 1888 | # Non-first PP ranks consume hidden states from intermediate_tensors. |
| 1889 | # The executor may still pass dummy inputs_embeds during profiling. |
| 1890 | inputs_embeds = None |
| 1891 | deepstack_input_embeds = intermediate_tensors |
| 1892 | hidden_states = self.language_model( |
| 1893 | input_ids, |
| 1894 | positions, |
| 1895 | intermediate_tensors=intermediate_tensors, |
| 1896 | inputs_embeds=inputs_embeds, |
| 1897 | deepstack_input_embeds=deepstack_input_embeds, |
| 1898 | ) |
| 1899 | self.deepstack_input_embeds = None |
| 1900 | return hidden_states |
| 1901 | |
| 1902 | def compute_logits( |
| 1903 | self, |
| 1904 | hidden_states: torch.Tensor, |
| 1905 | ) -> torch.Tensor | None: |
| 1906 | return self.language_model.compute_logits(hidden_states) |
| 1907 | |
| 1908 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1909 | loader = AutoWeightsLoader(self) |
| 1910 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |