diff --git a/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py b/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py index 6852b2b36..dbcf42da5 100644 --- a/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py @@ -42,6 +42,7 @@ async def run_streaming(inputs, sampling_params_list, model_name, args, output_d model=model_name, deploy_config=args.deploy_config, log_stats=args.log_stats, + quantization=args.quantization, ) # Normalize to a list so batch and single-input share the same code path @@ -194,6 +195,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir model=model_name, log_stats=args.log_stats, deploy_config=args.deploy_config, + quantization=args.quantization, ) if args.profiling_mode: @@ -305,6 +307,12 @@ def parse_args() -> Namespace: default=None, help="CFG alpha for flow-matching guidance (default: use value from stage config, typically 1.2).", ) + parser.add_argument( + "--quantization", + type=str, + default=None, + help="Quantization method (e.g. 'fp8'). Applied to the language model only.", + ) nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/tests/e2e/offline_inference/test_voxtral_tts.py b/tests/e2e/offline_inference/test_voxtral_tts.py index 53ac80b78..de88dc220 100644 --- a/tests/e2e/offline_inference/test_voxtral_tts.py +++ b/tests/e2e/offline_inference/test_voxtral_tts.py @@ -31,7 +31,7 @@ from vllm import SamplingParams from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner -from tests.helpers.stage_config import get_deploy_config_path +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.entrypoints.async_omni import AsyncOmni MODEL = "mistralai/Voxtral-4B-TTS-2603" @@ -42,14 +42,6 @@ MIN_AUDIO_SAMPLES = 1000 VOICE = "casual_female" TEST_TEXT = "Hello, how are you?" -# (model, stage_config_path, extra_omni_kwargs) for indirect parametrize on -# ``omni_runner_function`` (function-scoped: must exit before ``AsyncOmni`` test). -_OMNI_RUNNER_PARAM = ( - MODEL, - STAGE_CONFIG, - {"enforce_eager": True}, -) - def _compose_request(model_name: str, text: str, voice: str) -> dict: """Build the TTS input dict using mistral tokenizer.""" @@ -66,60 +58,83 @@ def _compose_request(model_name: str, text: str, voice: str) -> dict: } -@pytest.mark.advanced_model -@pytest.mark.tts -@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True) -@hardware_test(res={"cuda": "H100"}, num_cards=1) -def test_voxtral_tts_offline_basic(omni_runner_function: OmniRunner) -> None: - """Offline sync path; function-scoped runner so AsyncOmni test does not overlap a live Omni.""" - omni = omni_runner_function.omni - inputs = _compose_request(MODEL, TEST_TEXT, VOICE) - - sampling_params = SamplingParams(max_tokens=2500) - sampling_params_list = [sampling_params, sampling_params] - - outputs = list(omni.generate([inputs], sampling_params_list)) - - assert len(outputs) > 0, "No outputs generated" - - # Find audio output from the final stage - audio_data = None - for o in outputs: - mm = getattr(o, "multimodal_output", None) - if mm and "audio" in mm: - audio_data = mm["audio"] - break - - assert audio_data is not None, "No audio output found in any stage output" - - # Concatenate audio chunks if returned as a list of tensors - if isinstance(audio_data, list): - audio_tensor = torch.cat(audio_data) - elif isinstance(audio_data, torch.Tensor): - audio_tensor = audio_data - else: - audio_tensor = torch.tensor(audio_data) - - audio_array = audio_tensor.float().cpu().numpy() - - assert len(audio_array) > MIN_AUDIO_SAMPLES, ( - f"Audio too short: {len(audio_array)} samples, expected > {MIN_AUDIO_SAMPLES}" - ) - - # Verify audio isn't all zeros / silence - assert np.max(np.abs(audio_array)) > 0.01, "Audio appears to be silence" +def _resolve_stage_config(run_level: str) -> str: + """Resolve stage config: strip load_format for advanced_model (real weights).""" + if run_level == "advanced_model": + return modify_stage_config( + STAGE_CONFIG, + deletes={ + "stages": { + 0: ["load_format"], + 1: ["load_format"], + } + }, + ) + return STAGE_CONFIG @pytest.mark.advanced_model -@pytest.mark.tts +@pytest.mark.omni @hardware_test(res={"cuda": "H100"}, num_cards=1) -def test_voxtral_tts_offline_streaming(): +def test_voxtral_tts_offline_basic(run_level): + """Test basic Voxtral TTS offline inference with a voice preset.""" + stage_config = _resolve_stage_config(run_level) + + with OmniRunner( + MODEL, + stage_configs_path=stage_config, + enforce_eager=True, + ) as runner: + omni = runner.omni + inputs = _compose_request(MODEL, TEST_TEXT, VOICE) + + sampling_params = SamplingParams(max_tokens=2500) + sampling_params_list = [sampling_params, sampling_params] + + outputs = list(omni.generate([inputs], sampling_params_list)) + + assert len(outputs) > 0, "No outputs generated" + + # Find audio output from the final stage + audio_data = None + for o in outputs: + mm = getattr(o, "multimodal_output", None) + if mm and "audio" in mm: + audio_data = mm["audio"] + break + + assert audio_data is not None, "No audio output found in any stage output" + + # Concatenate audio chunks if returned as a list of tensors + if isinstance(audio_data, list): + audio_tensor = torch.cat(audio_data) + elif isinstance(audio_data, torch.Tensor): + audio_tensor = audio_data + else: + audio_tensor = torch.tensor(audio_data) + + audio_array = audio_tensor.float().cpu().numpy() + + assert len(audio_array) > MIN_AUDIO_SAMPLES, ( + f"Audio too short: {len(audio_array)} samples, expected > {MIN_AUDIO_SAMPLES}" + ) + + # Verify audio isn't all zeros / silence + assert np.max(np.abs(audio_array)) > 0.01, "Audio appears to be silence" + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_voxtral_tts_offline_streaming(run_level): """Test AsyncOmni streaming inference for Voxtral TTS.""" async def _run(): + stage_config = _resolve_stage_config(run_level) + async_omni = AsyncOmni( model=MODEL, - stage_configs_path=STAGE_CONFIG, + stage_configs_path=stage_config, stage_init_timeout=300, enforce_eager=True, ) diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py index c78089150..4d2421eea 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py @@ -1,4 +1,5 @@ from collections.abc import Iterable +from dataclasses import replace from functools import cached_property from pathlib import Path @@ -144,6 +145,9 @@ class VoxtralTTSForConditionalGeneration( elif self.model_stage == "audio_tokenizer": self.requires_raw_input_tokens = True self.audio_generation = None + if vllm_config.quant_config is not None: + vllm_config = replace(vllm_config, quant_config=None) + logger.info("Voxtral TTS FP8 routing: stage-1 audio_tokenizer=bf16") self.audio_tokenizer = init_vllm_registered_model( vllm_config=vllm_config, hf_config=config, diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py index 8b7dd7d13..2e50e34f3 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py @@ -3,7 +3,7 @@ import math from collections.abc import Iterable, Mapping, Sequence -from dataclasses import dataclass, fields, is_dataclass +from dataclasses import dataclass, fields, is_dataclass, replace from enum import Enum from functools import cached_property from math import ceil @@ -58,6 +58,8 @@ from vllm.sequence import IntermediateTensors from vllm.tokenizers import cached_tokenizer_from_config from vllm.tokenizers.mistral import MistralTokenizer +from vllm_omni.quantization.component_config import ComponentQuantizationConfig + weight_norm = torch.nn.utils.parametrizations.weight_norm @@ -197,6 +199,10 @@ def from_nested_dict(cls, d): return cls(**kwargs) +def _is_fp8_quant_config(quant_config) -> bool: + return quant_config is not None and quant_config.get_name() == "fp8" + + class FeedForward(nn.Module): def __init__( self, @@ -916,6 +922,28 @@ class VoxtralTTSAudioGenerationForConditionalGeneration(nn.Module, SupportsMulti config = vllm_config.model_config.hf_config self.config = config + quant_config = vllm_config.quant_config + if isinstance(quant_config, ComponentQuantizationConfig): + component_configs = dict(quant_config.component_configs) + component_configs[maybe_prefix(prefix, "audio_tokenizer")] = None + component_configs[maybe_prefix(prefix, "acoustic_transformer")] = None + quant_config = ComponentQuantizationConfig( + component_configs=component_configs, + default_config=None, + ) + vllm_config = replace(vllm_config, quant_config=quant_config) + elif _is_fp8_quant_config(quant_config): + quant_config = ComponentQuantizationConfig( + component_configs={ + maybe_prefix(prefix, "language_model"): quant_config, + maybe_prefix(prefix, "audio_tokenizer"): None, + maybe_prefix(prefix, "acoustic_transformer"): None, + }, + default_config=None, + ) + vllm_config = replace(vllm_config, quant_config=quant_config) + logger.info("Voxtral TTS FP8 routing: language_model=fp8, acoustic_transformer=bf16, audio_tokenizer=bf16") + self.language_model = init_vllm_registered_model( vllm_config=vllm_config, hf_config=config.text_config,