vllm.model_executor.models.ultravox ¶

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

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

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

class UltravoxAudioEmbeddingInputs(TensorSchema):
    """
    Dimensions:
    - b: batch size
    - na: number of audios
    - afs: audio feature size
    - hs: hidden size
    """

    type: Literal["audio_embeds"]
    data: Annotated[
        torch.Tensor | list[torch.Tensor], TensorShape("b", "na", "afs", "hs")
    ]

class UltravoxAudioFeatureInputs(TensorSchema):
    """
    Dimensions:
    - b: batch size
    - n: number of chunks
    - t: Time frames (M)
    - nmb: Number of mel bins
    """

    type: Literal["audio_features"]
    data: Annotated[
        torch.Tensor | list[torch.Tensor] | list[list[torch.Tensor]],
        TensorShape("bn", "nmb", "t"),
    ]
    lens: Annotated[torch.Tensor, TensorShape("bn")]
    """
    Length of the audio frames per chunk. Used for attention mask in WhisperEncoder.
    """
    token_len: Annotated[torch.Tensor, TensorShape("bn")]
    """Length of the audio tokens per chunk. Used for flattening the audio features."""
    num_chunks: Annotated[torch.Tensor, TensorShape("n")]
    """Number of chunks per audio. Used for flattening the audio features."""

@MULTIMODAL_REGISTRY.register_processor(
    UltravoxMultiModalProcessor,
    info=UltravoxProcessingInfo,
    dummy_inputs=UltravoxDummyInputsBuilder,
)
class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"],
    }

    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            "audio_tower.model.encoder.": "audio_tower.",
            # A whisper checkpoint loaded via `audio_model_id` also carries
            # decoder and LM-head weights the tower does not use.
            "audio_tower.model.": None,
            "audio_tower.proj_out.": None,
        }
    )

    supports_tower_connector_lora = True

    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality.startswith("audio"):
            return "<|audio|>"

        raise ValueError("Only audio modality is supported")

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config: UltravoxConfig = vllm_config.model_config.hf_config
        multimodal_config = vllm_config.model_config.multimodal_config
        quant_config = vllm_config.quant_config
        lora_config = vllm_config.lora_config
        self.config = config
        self.multi_modal_config = multimodal_config
        assert self.multi_modal_config

        # LoRA on the tower/connector requires per-item token counts that are
        # predictable from the placeholder count alone (see
        # `get_num_mm_encoder_tokens`), so pad every audio chunk's mel input
        # to the tower's full context instead of the batch's max length.
        self.pad_audio_to_max_context = bool(
            lora_config is not None and lora_config.enable_tower_connector_lora
        )

        self.configure_mm_token_handling(
            self.config.vocab_size,
            [self.config.audio_token_index],
        )

        # The towers live in their own Hub repos, so the revision of the
        # Ultravox repo does not apply to them.
        self.secondary_weights = []
        if config.audio_model_id is not None:
            # this prefix is not for initialization, but for loading weights
            # note the trailing dot
            self.secondary_weights.append(
                DefaultModelLoader.Source(
                    model_or_path=config.audio_model_id,
                    revision=None,
                    prefix="audio_tower.",
                )
            )
        if config.text_model_id is not None:
            # this prefix is not for initialization, but for loading weights
            # note the trailing dot
            self.secondary_weights.append(
                DefaultModelLoader.Source(
                    model_or_path=config.text_model_id,
                    revision=None,
                    prefix="language_model.",
                )
            )

        with self._mark_tower_model(vllm_config, "audio"):
            self.audio_tower = UltravoxWhisperEncoder(
                vllm_config=vllm_config.with_hf_config(
                    config.audio_config, architectures=["UltravoxModel"]
                ),
                prefix=maybe_prefix(prefix, "audio_tower"),
                enable_pp=False,
            )
            if config.num_projector_layers > 0:
                self.multi_modal_projector = UltravoxTransformerProjector(
                    vllm_config,
                    prefix=maybe_prefix(prefix, "multi_modal_projector"),
                )
            else:
                self.multi_modal_projector = UltravoxFeedForwardProjector(
                    config,
                    quant_config=quant_config,
                    prefix=maybe_prefix(prefix, "multi_modal_projector"),
                )

        with self._mark_language_model(vllm_config):
            self.language_model = init_vllm_registered_model(
                vllm_config=vllm_config,
                hf_config=config.wrapped_model_config,
                prefix=maybe_prefix(prefix, "language_model"),
            )

        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors
        )

    def get_mm_mapping(self) -> MultiModelKeys:
        """
        Get the module prefix in multimodal models
        """
        return MultiModelKeys.from_string_field(
            language_model="language_model.",
            connector="multi_modal_projector.",
            tower_model="audio_tower.",
        )

    def _get_max_tokens_per_chunk(self) -> int:
        # Audio is chunked at the tower's full context (30s -> 1500 encoder
        # frames); the projector stacks `stack_factor` frames per LM token, so
        # a full chunk yields ceil(max_source_positions / stack_factor) tokens
        # and only an audio's last chunk can yield fewer.
        return math.ceil(
            self.config.audio_config.max_source_positions / self.config.stack_factor
        )

    def get_num_mm_encoder_tokens(self, num_audio_tokens: int) -> int:
        # With `pad_audio_to_max_context` (enforced when tower/connector LoRA
        # is enabled), the tower's LoRA-wrapped linears always run on
        # `max_source_positions` conv-downsampled tokens per chunk, regardless
        # of the valid frame count.
        num_chunks = math.ceil(num_audio_tokens / self._get_max_tokens_per_chunk())
        return num_chunks * self.config.audio_config.max_source_positions

    def get_num_mm_connector_tokens(self, num_encoder_tokens: int) -> int:
        # The connector runs on the frame-stacked tower output. Stacking pads
        # each chunk to a multiple of `stack_factor`, so this is
        # ceil(max_source_positions / stack_factor) tokens per chunk (188 for
        # whisper's 1500), not `num_encoder_tokens // stack_factor` (187).
        max_source_positions = self.config.audio_config.max_source_positions
        num_chunks = num_encoder_tokens // max_source_positions
        return num_chunks * self._get_max_tokens_per_chunk()

    def _audio_features_to_embeddings(
        self,
        input_features: torch.Tensor,
        audio_lens: torch.Tensor,
        audio_token_len: torch.Tensor,
    ) -> torch.Tensor:
        audio_features = input_features.to(self.audio_tower.dtype)
        batch_size = audio_features.size(0)
        audio_embeddings = []

        # Process audio features in batches to keep memory usage predictable.
        # With tower/connector LoRA, the punica kernels always map the first
        # `x.size(0)` entries of the token-LoRA mapping (set once for the whole
        # scheduled encoder batch) to the rows of each linear's input, so
        # splitting the batch would apply the wrong LoRA ids to every chunk
        # after the first sub-batch; process all chunks in a single pass.
        encoder_batch_size = (
            max(batch_size, 1)
            if self.pad_audio_to_max_context
            else _MAX_ENCODER_BATCH_SIZE
        )
        for start in range(0, batch_size, encoder_batch_size):
            end = min(start + encoder_batch_size, batch_size)
            # Process through audio tower
            batch_features = self.audio_tower(
                audio_features[start:end], audio_lens[start:end]
            )
            batch_features = batch_features.to(self.audio_tower.dtype)

            # Process through projector
            batch_embeddings = self.multi_modal_projector(
                batch_features, audio_token_len[start:end]
            )
            audio_embeddings.append(batch_embeddings)

        # Concatenate results
        audio_embeddings = torch.cat(audio_embeddings, dim=0)
        return audio_embeddings

    def _parse_and_validate_audio_input(
        self, **kwargs: object
    ) -> UltravoxAudioInputs | None:
        audio_features = kwargs.pop("audio_features", None)
        audio_embeds = kwargs.pop("audio_embeds", None)
        audio_lens = kwargs.pop("audio_lens", None)
        audio_token_len = kwargs.pop("audio_token_len", None)
        audio_num_chunks = kwargs.pop("audio_num_chunks", None)

        if audio_features is None and audio_embeds is None:
            return None

        if audio_features is not None:
            return UltravoxAudioFeatureInputs(
                type="audio_features",
                data=audio_features,
                lens=audio_lens,
                token_len=audio_token_len,
                num_chunks=audio_num_chunks,
            )

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

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

    def _process_audio_input(
        self,
        audio_input: UltravoxAudioInputs,
    ) -> NestedTensors | tuple[torch.Tensor, ...]:
        if audio_input["type"] == "audio_embeds":
            return audio_input["data"]

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

        if self.pad_audio_to_max_context:
            # Pad every chunk to the tower's full context so the number of
            # tokens processed by the tower/connector per chunk is constant,
            # keeping the LoRA token mappings computed by
            # `get_num_mm_encoder_tokens`/`get_num_mm_connector_tokens` exact.
            # Attention to the extra padding is masked via `audio_lens`.
            # Over-long inputs are not trimmed here; the tower raises on them.
            pad_len = self.audio_tower.max_context_length - audio_features.shape[-1]
            if pad_len > 0:
                audio_features = F.pad(audio_features, (0, pad_len))

        audio_lens = audio_input["lens"]
        audio_token_len = audio_input["token_len"]

        embeddings = self._audio_features_to_embeddings(
            audio_features, audio_lens, audio_token_len
        )

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

        # Create a mask of valid indices based on token lengths
        max_len = embeddings.shape[1]
        indices = torch.arange(max_len, device=embeddings.device).expand(
            embeddings.shape[0], -1
        )
        mask = indices < audio_token_len[:, None]

        with gpu_sync_allowed():
            # Apply mask and flatten
            flattened_embeddings = embeddings[mask]

            # Return one tensor per input audio
            embed_lens = [
                chunk_lens.sum().item()
                for chunk_lens in audio_token_len.split(
                    audio_input["num_chunks"].tolist()
                )
            ]
            return flattened_embeddings.split(embed_lens)

    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
        audio_input = self._parse_and_validate_audio_input(**kwargs)
        if audio_input is None:
            return []
        audio_embeddings = self._process_audio_input(audio_input)
        return audio_embeddings

    def embed_input_ids(
        self,
        input_ids: torch.Tensor,
        multimodal_embeddings: MultiModalEmbeddings | None = None,
        *,
        is_multimodal: torch.Tensor | None = None,
    ) -> torch.Tensor:
        # This is to satisfy the type checker for each overload
        if multimodal_embeddings is None or is_multimodal is None:
            return super().embed_input_ids(input_ids)

        return super().embed_input_ids(
            input_ids,
            multimodal_embeddings=multimodal_embeddings,
            is_multimodal=is_multimodal,
        )

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor | IntermediateTensors:
        """Run forward pass for Ultravox

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

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

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

        """

        if intermediate_tensors is not None:
            inputs_embeds = None

        language_model = self.language_model
        if hasattr(language_model, "language_model"):
            language_model = language_model.language_model

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

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

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

class UltravoxProcessingInfo(BaseProcessingInfo):
    def get_hf_processor(self, **kwargs: object) -> ProcessorMixin:
        config = self.ctx.model_config.hf_config
        hf_processor = self.ctx.get_hf_processor(**kwargs)

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

        return hf_processor

    def get_feature_extractor(self, **kwargs: object) -> WhisperFeatureExtractor:
        hf_processor = self.get_hf_processor(**kwargs)

        # Changed in https://huggingface.co/fixie-ai/ultravox-v0_5-llama-3_2-1b/commit/9a3c571b8fdaf1e66dd3ea61bbcb6db5c70a438e
        audio_processor = hf_processor.audio_processor  # type: ignore
        if isinstance(audio_processor, WhisperFeatureExtractor):
            return audio_processor

        feature_extractor = audio_processor.feature_extractor  # type: ignore
        assert isinstance(feature_extractor, WhisperFeatureExtractor)
        return feature_extractor

    def get_default_tok_params(self) -> TokenizeParams:
        return super().get_default_tok_params().with_kwargs(add_special_tokens=False)

    def get_data_parser(self):
        feature_extractor = self.get_feature_extractor()

        return MultiModalDataParser(
            target_sr=feature_extractor.sampling_rate,
            target_channels=self.get_target_channels(),
            expected_hidden_size=self._get_expected_hidden_size(),
        )

    def get_target_channels(self) -> int:
        """Return target audio channels for Ultravox models (mono)."""
        return 1

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

class UltravoxWhisperEncoder(WhisperEncoder):
    """Ultravox's `ModifiedWhisperEncoder` on top of vLLM's `WhisperEncoder`.

    Like the original (a modified HF whisper encoder,
    see https://github.com/huggingface/transformers/issues/25744), it accepts
    mel inputs shorter than 30s (positions are sliced to the input length) and
    confines attention to each chunk's valid frames based on `audio_lens`
    (via segmented `cu_seqlens` instead of a dense key-padding mask), so its
    outputs match the HF implementation for valid positions. The linears are
    vLLM-native so the tower can be wrapped for LoRA.
    """

    @property
    def dtype(self) -> torch.dtype:
        return self.conv1.weight.dtype

    @property
    def max_context_length(self) -> int:
        return self.max_source_positions * self.total_stride

    def _get_feat_extract_output_lengths(
        self, input_lengths: torch.Tensor
    ) -> torch.Tensor:
        return (input_lengths - 1) // 2 + 1

    def forward(
        self,
        input_features: torch.Tensor,
        audio_lens: torch.Tensor,
    ) -> torch.Tensor:
        expected_seq_length = self.max_context_length
        if input_features.shape[-1] > expected_seq_length:
            raise ValueError(
                f"Whisper expects the mel input features to be of length "
                f"{expected_seq_length} or less, but found "
                f"{input_features.shape[-1]}. Make sure to pad the input mel "
                f"features to {expected_seq_length}."
            )

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

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

        hidden_states = (inputs_embeds + embed_pos).to(inputs_embeds.dtype)

        batch_size, seq_len, d_model = hidden_states.shape
        attn_metadata = _build_chunk_attn_metadata(
            self.layers[0].self_attn.attn,
            self._get_feat_extract_output_lengths(audio_lens),
            seq_len,
            d_model,
            hidden_states.device,
        )
        hidden_states = hidden_states.reshape(batch_size * seq_len, d_model)
        for encoder_layer in self.layers:
            hidden_states = encoder_layer(hidden_states, **attn_metadata)
        hidden_states = hidden_states.reshape(batch_size, seq_len, d_model)

        hidden_states = self.layer_norm(hidden_states)
        return hidden_states

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        return _load_whisper_layer_weights(self, weights)