From be0f32f972b407b6be8658f2d6ccc360ea8bb1d1 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Fri, 21 Aug 2026 09:56:01 +0800 Subject: [PATCH 01/28] [model, rollout] feat: add Qwen3-TTS talker pipeline Signed-off-by: dongbo910220 <1275604947@qq.com> Co-authored-by: Michael-Zzq <52185141+Michael-Zzq@users.noreply.github.com> --- pyproject.toml | 10 + tests/pipelines/test_qwen3_tts_on_cpu.py | 206 ++++++++++++++++ ...st_qwen3_tts_transformers_compat_on_cpu.py | 97 ++++++++ .../test_qwen3_tts_rollout_on_cpu.py | 232 ++++++++++++++++++ tests/workers/test_omni_fsdp_engine_on_cpu.py | 63 +++++ verl_omni/pipelines/__init__.py | 3 + verl_omni/pipelines/model_base.py | 31 +++ verl_omni/pipelines/qwen3_tts/__init__.py | 19 ++ verl_omni/pipelines/qwen3_tts/agent_loop.py | 61 +++++ .../qwen3_tts/omni_rollout_adapter.py | 195 +++++++++++++++ .../pipelines/qwen3_tts/rollout_utils.py | 147 +++++++++++ .../pipelines/qwen3_tts/talker_forward.py | 229 +++++++++++++++++ .../qwen3_tts/talker_training_adapter.py | 228 +++++++++++++++++ .../qwen3_tts/transformers_compat.py | 74 ++++++ .../workers/rollout/vllm_rollout/utils.py | 26 +- 15 files changed, 1617 insertions(+), 4 deletions(-) create mode 100644 tests/pipelines/test_qwen3_tts_on_cpu.py create mode 100644 tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py create mode 100644 tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py create mode 100644 verl_omni/pipelines/qwen3_tts/__init__.py create mode 100644 verl_omni/pipelines/qwen3_tts/agent_loop.py create mode 100644 verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py create mode 100644 verl_omni/pipelines/qwen3_tts/rollout_utils.py create mode 100644 verl_omni/pipelines/qwen3_tts/talker_forward.py create mode 100644 verl_omni/pipelines/qwen3_tts/talker_training_adapter.py create mode 100644 verl_omni/pipelines/qwen3_tts/transformers_compat.py diff --git a/pyproject.toml b/pyproject.toml index 387c7f9dd..dc094060d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,6 +58,16 @@ vllm-omni = [ audio = [ "qwen-omni-utils>=0.0.9", ] +# Install qwen-tts separately with --no-deps because qwen-tts 0.1.1 pins +# Transformers 4.57.3 while vLLM 0.27 requires Transformers 5.x. +tts = [ + "einops>=0.8.0", + "librosa>=0.10.2", + "onnxruntime>=1.20.0", + "soundfile>=0.12.1", + "sox>=1.5.0", + "torchaudio", +] # CUDA rollout backend (vllm) + actor FA3 (kernels) + liger-kernel. Install in step 1 on GPU. gpu = [ "vllm==0.27.0", diff --git a/tests/pipelines/test_qwen3_tts_on_cpu.py b/tests/pipelines/test_qwen3_tts_on_cpu.py new file mode 100644 index 000000000..633ea7113 --- /dev/null +++ b/tests/pipelines/test_qwen3_tts_on_cpu.py @@ -0,0 +1,206 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Dependency-light Qwen3-TTS actor and rollout contract tests.""" + +import importlib.util +import sys +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +ROOT = Path(__file__).parents[2] + + +def _load(name: str, relative_path: str): + spec = importlib.util.spec_from_file_location(name, ROOT / relative_path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +forward = _load("qwen3_tts_forward_test", "verl_omni/pipelines/qwen3_tts/talker_forward.py") +rollout = _load("qwen3_tts_rollout_test", "verl_omni/pipelines/qwen3_tts/rollout_utils.py") + +TOKENS = forward.TalkerTokens(900, 901, 902, 4196, 4197, 4198, 4203, 4204, 4205) + + +def test_talker_batch_matches_auto_language_teacher_forcing_layout(): + text_ids = torch.tensor([1, 2, 3, 4, 5, 6]) + codes = torch.arange(3 * 16, dtype=torch.long).reshape(3, 16) % 128 + codes[:, 0] = torch.tensor([10, 11, 12]) + + batch = forward.build_talker_batch([text_ids], [codes], TOKENS, sub_codebook_vocab=2048) + + speaker_slot = 6 + codec_start = 8 + text_ids.numel() - 1 + assert batch.input_ids[0, 3:speaker_slot, 1].tolist() == [4203, 4204, 4205] + assert not batch.codec_embedding_mask[0, speaker_slot] + assert batch.input_ids[0, speaker_slot + 1, 1] == TOKENS.codec_pad + torch.testing.assert_close(batch.codec_ids[0, codec_start : codec_start + 3], codes) + assert batch.logit_start == [codec_start - 1] + assert batch.codec_lens == [3] + + +def test_only_validated_auto_language_layout_is_accepted(): + assert forward.require_auto_language("auto") == "Auto" + with pytest.raises(ValueError, match="supports only tts_language=Auto"): + forward.require_auto_language("Chinese") + + +def test_actor_logits_align_to_effective_codec0_response(monkeypatch): + class CodePredictor: + @staticmethod + def get_input_embeddings(): + return [torch.nn.Embedding(2048, 4)] + + talker = SimpleNamespace(code_predictor=CodePredictor()) + model = SimpleNamespace( + talker=talker, + config=SimpleNamespace( + tts_pad_token_id=TOKENS.tts_pad, + tts_bos_token_id=TOKENS.tts_bos, + tts_eos_token_id=TOKENS.tts_eos, + talker_config=SimpleNamespace( + codec_pad_id=TOKENS.codec_pad, + codec_bos_id=TOKENS.codec_bos, + codec_eos_token_id=TOKENS.codec_eos, + codec_nothink_id=TOKENS.codec_nothink, + codec_think_bos_id=TOKENS.codec_think_bos, + codec_think_eos_id=TOKENS.codec_think_eos, + ), + ), + ) + + def fake_logits(_talker, batch, _speaker): + vocab = torch.arange(1, 4301, dtype=torch.float32).reshape(1, 1, -1) + return vocab.expand(1, batch.input_ids.shape[1] - 1, -1) + + monkeypatch.setattr(forward, "codec0_logits", fake_logits) + input_ids = torch.zeros((1, 9), dtype=torch.long) + input_ids[0, -3:] = torch.tensor([10, 11, 12]) + codes = torch.zeros((1, 6, 16), dtype=torch.long) + codes[0, :3, 0] = torch.tensor([10, 11, 12]) + + logits = forward.tts_actor_logits( + model, + input_ids, + torch.ones_like(input_ids), + torch.tensor([[1, 2, 3, 4, 5, 6]]), + codes, + torch.tensor([3]), + torch.tensor([6]), + torch.zeros((1, 4)), + ) + + assert torch.nonzero(logits.abs().sum(dim=-1)[0], as_tuple=False).reshape(-1).tolist() == [5, 6, 7] + torch.testing.assert_close(logits[0, 5], torch.arange(1, 4301, dtype=torch.float32)) + + +@pytest.mark.parametrize("response_length", [2, 14, 15, 16, 17, 32]) +def test_codec_alignment_recovers_exact_prefix_without_final_residual_row(response_length): + token_ids = list(range(100, 100 + response_length - 1)) + [2150] + generated = torch.arange((response_length - 1) * 16, dtype=torch.long).reshape(response_length - 1, 16) + 1 + generated[:, 0] = torch.tensor(token_ids[:-1]) + raw = torch.cat((torch.zeros(12, 16, dtype=torch.long), generated)) + + aligned = rollout.align_audio_codes(raw, token_ids) + + assert aligned[:, 0].tolist() == token_ids + torch.testing.assert_close(aligned[:-1, 1:], generated[:, 1:]) + assert not aligned[-1, 1:].any() + + +def test_codec_alignment_preserves_final_row_and_rejects_heuristic_match(): + token_ids = [101, 102, 103, 2150] + generated = torch.arange(4 * 16, dtype=torch.long).reshape(4, 16) + 1 + generated[:, 0] = torch.tensor(token_ids) + aligned = rollout.align_audio_codes(torch.cat((torch.zeros(12, 16, dtype=torch.long), generated)), token_ids) + torch.testing.assert_close(aligned, generated) + + malformed = torch.zeros(15, 16, dtype=torch.long) + malformed[12:, 0] = torch.tensor([101, 999, 103]) + with pytest.raises(RuntimeError, match="Could not exactly align"): + rollout.align_audio_codes(malformed, token_ids) + + +def test_rollout_chunk_accumulator_handles_cumulative_and_delta_outputs(): + first = torch.zeros(12, 16, dtype=torch.long) + cumulative = torch.cat((first, torch.ones(1, 16, dtype=torch.long))) + assert torch.equal(rollout.append_tensor_chunk(first, cumulative), cumulative) + + delta = torch.full((1, 16), 2, dtype=torch.long) + assert torch.equal(rollout.append_tensor_chunk(cumulative, delta), torch.cat((cumulative, delta))) + with pytest.raises(RuntimeError, match="changed shape"): + rollout.append_tensor_chunk(cumulative, torch.zeros(1, 15, dtype=torch.long)) + + +def test_validation_seed_covers_both_codec_samplers_and_candidates_without_mutation(): + original = {"temperature": 0.8, "extra_args": {"existing": "kept"}} + seeded = rollout.with_rollout_generation_seed( + original, + {"split": "validation", "generation_seed": np.int64(42017)}, + session_id=3, + global_steps=100, + require_session_id=True, + ) + + assert seeded == { + "temperature": 0.8, + "seed": 42020, + "extra_args": {"existing": "kept", "tts_local_seed": 42020}, + } + assert seeded == rollout.with_rollout_generation_seed( + original, + {"split": "validation", "generation_seed": np.int64(42017)}, + session_id=3, + global_steps=0, + require_session_id=True, + ) + assert original == {"temperature": 0.8, "extra_args": {"existing": "kept"}} + + gate_first = rollout.with_rollout_generation_seed( + {}, {"split": "gate", "id": "gate-7"}, session_id=0, global_steps=0, base_seed=42 + ) + gate_second = rollout.with_rollout_generation_seed( + {}, {"split": "gate", "id": "gate-7"}, session_id=1, global_steps=100, base_seed=42 + ) + assert gate_first["seed"] != gate_second["seed"] + assert gate_first == rollout.with_rollout_generation_seed( + {}, {"split": "gate", "id": "gate-7"}, session_id=0, global_steps=100, base_seed=42 + ) + + +def test_training_seeds_are_reproducible_and_group_diverse(): + kwargs = { + "extra_info": {"split": "train", "id": "sample-7"}, + "global_steps": 12, + "uid": "uid-7", + "base_seed": 42, + "require_session_id": True, + } + first = rollout.with_rollout_generation_seed({}, session_id=0, **kwargs) + repeated = rollout.with_rollout_generation_seed({}, session_id=0, **kwargs) + second = rollout.with_rollout_generation_seed({}, session_id=1, **kwargs) + next_step = rollout.with_rollout_generation_seed({}, session_id=0, **{**kwargs, "global_steps": 13}) + + assert first == repeated + assert len({first["seed"], second["seed"], next_step["seed"]}) == 3 + assert first["seed"] == first["extra_args"]["tts_local_seed"] + with pytest.raises(RuntimeError, match="session_id"): + rollout.with_rollout_generation_seed({}, session_id=None, **kwargs) diff --git a/tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py b/tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py new file mode 100644 index 000000000..6e3da9759 --- /dev/null +++ b/tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py @@ -0,0 +1,97 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Optional real-package check for qwen-tts on the repository TF5 stack.""" + +import importlib +import importlib.util +from pathlib import Path + +import pytest +import torch +from packaging.version import Version + +ROOT = Path(__file__).parents[2] + + +def _load_compat_module(): + path = ROOT / "verl_omni/pipelines/qwen3_tts/transformers_compat.py" + spec = importlib.util.spec_from_file_location("qwen3_tts_transformers_compat_under_test", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_qwen_tts_tiny_model_constructs_and_forwards_without_source_patch(): + transformers = pytest.importorskip("transformers") + if Version(transformers.__version__).major < 5: + pytest.skip("Transformers 5.x compatibility test") + if importlib.util.find_spec("qwen_tts") is None: + pytest.skip("qwen-tts is an optional dependency") + + compat = _load_compat_module() + with compat.qwen3_tts_import_context(): + config_module = importlib.import_module("qwen_tts.core.models.configuration_qwen3_tts") + model_module = importlib.import_module("qwen_tts.core.models.modeling_qwen3_tts") + + predictor = { + "vocab_size": 32, + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "max_position_embeddings": 64, + "num_code_groups": 4, + "layer_types": ["full_attention"], + "pad_token_id": None, + } + talker = { + "code_predictor_config": predictor, + "vocab_size": 64, + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "max_position_embeddings": 64, + "num_code_groups": 4, + "text_hidden_size": 8, + "text_vocab_size": 80, + "spk_id": {}, + "codec_language_id": {}, + "pad_token_id": None, + "rope_scaling": { + "rope_type": "default", + "type": "default", + "mrope_section": [1, 1, 0], + "interleaved": True, + }, + } + config = config_module.Qwen3TTSConfig( + talker_config=talker, + speaker_encoder_config={}, + tts_model_type="custom", + tokenizer_type="12hz", + ) + model = model_module.Qwen3TTSForConditionalGeneration(config) + output = model.talker( + inputs_embeds=torch.randn(2, 5, 8), + attention_mask=torch.ones(2, 5, dtype=torch.long), + use_cache=False, + output_hidden_states=False, + ) + + assert output.logits.shape == (2, 5, 64) diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py new file mode 100644 index 000000000..61512b20d --- /dev/null +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -0,0 +1,232 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU contracts for Qwen3-TTS's multi-stage rollout integration.""" + +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("verl") +pytest.importorskip("vllm_omni") + +from vllm import SamplingParams + +from verl_omni.pipelines.qwen3_tts.omni_rollout_adapter import Qwen3TTSRolloutAdapter +from verl_omni.pipelines.qwen3_tts.talker_training_adapter import Qwen3TTSTalkerAdapter +from verl_omni.workers.rollout.vllm_rollout.utils import _receive_model_weight_buckets +from verl_omni.workers.rollout.vllm_rollout.vllm_omni_async_server import vLLMOmniHttpServer + + +class _Tokenizer: + def decode(self, token_ids, **kwargs): + return "first text" if token_ids == [1] else "other text" + + def __call__(self, text, **kwargs): + return {"input_ids": list(range(len(text)))} + + +def test_rollout_adapter_builds_unique_prompt_and_scopes_weight_sync(tmp_path): + speaker = tmp_path / "speaker.json" + speaker.write_text("[0.0, 1.0]") + model_config = SimpleNamespace( + tokenizer=_Tokenizer(), + override_config={"tts_spk_embed_path": str(speaker), "tts_language": "Auto"}, + hf_config=SimpleNamespace(talker_config=SimpleNamespace(codec_eos_token_id=2150)), + ) + + first = Qwen3TTSRolloutAdapter.prepare_engine_prompt([1], model_config, {}) + second = Qwen3TTSRolloutAdapter.prepare_engine_prompt([2], model_config, {}) + + assert first["additional_information"]["text"] == ["first text"] + assert first["cache_salt"] != second["cache_salt"] + assert Qwen3TTSRolloutAdapter.weight_sync_stage_ids("full") == [0] + assert not Qwen3TTSRolloutAdapter.supports_cache_engine_sleep("full") + assert Qwen3TTSRolloutAdapter.get_output_modalities("full") == ["latent", "audio"] + + +def test_rollout_adapter_requires_speaker_embedding(): + model_config = SimpleNamespace( + tokenizer=_Tokenizer(), + override_config={"tts_language": "Auto"}, + hf_config=SimpleNamespace(talker_config=SimpleNamespace(codec_eos_token_id=2150)), + ) + + with pytest.raises(ValueError, match="requires tts_spk_embed_path"): + Qwen3TTSRolloutAdapter.prepare_engine_prompt([1], model_config, {}) + + +def test_talker_adapter_pads_exact_rollout_fields_for_actor_forward(): + model_inputs = {"input_ids": torch.zeros(2, 6, dtype=torch.long)} + micro_batch = { + "extra_fields": [ + {"tts_text_ids": [1, 2], "tts_audio_codes": torch.ones(3, 16, dtype=torch.long)}, + {"tts_text_ids": [3, 4, 5], "tts_audio_codes": torch.full((2, 16), 2, dtype=torch.long)}, + ] + } + + prepared = Qwen3TTSTalkerAdapter.prepare_model_inputs(model_inputs, micro_batch, None) + + assert prepared["tts_text_ids"].shape == (2, 3) + assert prepared["tts_audio_codes"].shape == (2, 3, 16) + assert prepared["text_len"].tolist() == [2, 3] + assert prepared["response_len"].tolist() == [3, 2] + assert not prepared["tts_audio_codes"][1, 2].any() + + +def test_talker_adapter_requires_v1_agent_loop_extra_fields(): + model_inputs = {"input_ids": torch.zeros(1, 4, dtype=torch.long)} + + with pytest.raises(RuntimeError, match="V1 agent-loop trainer"): + Qwen3TTSTalkerAdapter.prepare_model_inputs(model_inputs, {}, None) + + +def test_rollout_adapter_combines_policy_codes_and_waveform(): + token_ids = [101, 102, 2150] + generated = torch.arange(3 * 16, dtype=torch.long).reshape(3, 16) + 1 + generated[:, 0] = torch.tensor(token_ids) + policy = SimpleNamespace( + stage_id=0, + request_output=SimpleNamespace(outputs=[SimpleNamespace(token_ids=token_ids)]), + multimodal_output={"codes": {"audio": torch.cat((torch.zeros(12, 16), generated))}}, + ) + decoder = SimpleNamespace( + stage_id=1, + request_output=None, + multimodal_output={"audio": torch.ones(2400), "sr": 24_000}, + ) + prompt = {"additional_information": {"text": ["first text"]}} + + selected, fields = Qwen3TTSRolloutAdapter.combine_engine_outputs([policy, decoder], prompt) + + assert selected is policy + torch.testing.assert_close(fields["tts_audio_codes"], generated.long()) + torch.testing.assert_close(fields["audio"], torch.ones(2400)) + assert fields["audio_sample_rate"] == 24_000 + assert fields["tts_text"] == "first text" + + +def test_bucketed_weight_sync_rebuilds_derived_codec_table_once(monkeypatch): + class Model: + def __init__(self): + self.rebuilds = 0 + self.loads = 0 + self._stacked_codec_embed = object() + + def _build_stacked_codec_embed(self): + self.rebuilds += 1 + self._stacked_codec_embed = object() + + def load_weights(self, weights): + self.loads += 1 + self._build_stacked_codec_embed() + + class Receiver: + @staticmethod + def receive_weights(on_bucket_received): + on_bucket_received({"first": torch.tensor(1)}) + on_bucket_received({"second": torch.tensor(2)}) + + model = Model() + empty_cache_calls = [] + monkeypatch.setattr( + "verl_omni.workers.rollout.vllm_rollout.utils.get_torch_device", + lambda: SimpleNamespace(empty_cache=lambda: empty_cache_calls.append(True)), + ) + + _receive_model_weight_buckets(Receiver(), model) + + assert model.loads == 2 + assert model.rebuilds == 1 + assert empty_cache_calls == [True] + model._build_stacked_codec_embed() + assert model.rebuilds == 2 + + +def test_server_prepares_stage_specific_sampling_params(): + class Adapter: + @staticmethod + def prepare_engine_prompt(**kwargs): + return { + "prompt_token_ids": [1, 1, 1, 1], + "additional_information": {"text": ["hello"]}, + } + + server = object.__new__(vLLMOmniHttpServer) + server._ar_mode = True + server._omni_rollout_adapter = Adapter + server._stage_sampling_constraints = {0: {}} + server.model_config = SimpleNamespace() + server.config = SimpleNamespace( + max_model_len=64, + prompt_length=16, + response_length=8, + repetition_penalty=1.0, + ) + server.engine = SimpleNamespace(default_sampling_params_list=[SamplingParams(), SimpleNamespace(stage="decoder")]) + + prompt, params = server._preprocess_input( + [5, 6], + {"temperature": 0.8, "logprobs": True}, + {}, + None, + None, + ) + + assert prompt["additional_information"]["max_new_tokens"] == [8] + assert len(params) == 2 + assert params[0].max_tokens == 8 + assert params[0].temperature == pytest.approx(0.8) + assert params[0].logprobs == 0 + assert params[1].stage == "decoder" + + +@pytest.mark.asyncio +async def test_server_retains_requested_stage_outputs_and_targets_weight_sync(): + policy = SimpleNamespace(request_output=SimpleNamespace(outputs=[])) + + class Engine: + def __init__(self): + self.generate_kwargs = None + self.rpc_kwargs = None + + async def generate(self, **kwargs): + self.generate_kwargs = kwargs + yield policy + + async def collective_rpc(self, **kwargs): + self.rpc_kwargs = kwargs + return "rpc-result" + + class Adapter: + @staticmethod + def combine_engine_outputs(outputs, prompt): + assert outputs == [policy] + return policy, {"audio_sample_rate": 24_000} + + server = object.__new__(vLLMOmniHttpServer) + server._ar_mode = True + server.engine = Engine() + server._rollout_output_modalities = ["latent", "audio"] + server._omni_rollout_adapter = Adapter + server._weight_sync_stage_ids = [0] + + result = await server._run_generation({"prompt_token_ids": [1]}, SamplingParams(), "request-0", None, 0) + rpc_result = await server.collective_rpc("update_weights_from_ipc", kwargs={"base_sync_done": True}) + + assert result is policy + assert result._verl_omni_rollout_fields == {"audio_sample_rate": 24_000} + assert server.engine.generate_kwargs["output_modalities"] == ["latent", "audio"] + assert server.engine.rpc_kwargs["stage_ids"] == [0] + assert rpc_result == "rpc-result" diff --git a/tests/workers/test_omni_fsdp_engine_on_cpu.py b/tests/workers/test_omni_fsdp_engine_on_cpu.py index 3cd9137d1..713bf6fee 100644 --- a/tests/workers/test_omni_fsdp_engine_on_cpu.py +++ b/tests/workers/test_omni_fsdp_engine_on_cpu.py @@ -354,6 +354,69 @@ def test_collect_lora_params_import_not_from_verl(): ) +# --------------------------------------------------------------------------- +# Qwen3-TTS reference-policy manual offload +# --------------------------------------------------------------------------- + + +def _make_manual_ref_engine(omni_impl, *, enabled=True): + engine = object.__new__(omni_impl.OmniFSDPEngine) + engine.engine_config = types.SimpleNamespace(forward_only=True, param_offload=False) + engine.model_adapter_cls = types.SimpleNamespace(requires_manual_ref_offload=enabled) + return engine + + +def test_tts_ref_build_disables_forced_fsdp_cpu_offload_temporarily(): + omni_impl = _get_omni_impl_module() + engine = _make_manual_ref_engine(omni_impl) + observed = [] + + def fake_build(module): + observed.append(engine.engine_config.forward_only) + return module + + module = object() + with patch.object(omni_impl.FSDPEngineWithLMHead, "_build_fsdp_module", side_effect=fake_build): + assert engine._build_fsdp_module(module) is module + + assert observed == [False] + assert engine.engine_config.forward_only is True + + +def test_tts_ref_to_uses_manual_load_path_and_restores_on_error(): + omni_impl = _get_omni_impl_module() + engine = _make_manual_ref_engine(omni_impl) + observed = [] + + def fake_to(*args, **kwargs): + observed.append((engine.engine_config.forward_only, args, kwargs)) + raise RuntimeError("test failure") + + with ( + patch.object(omni_impl.FSDPEngineWithLMHead, "to", side_effect=fake_to), + pytest.raises(RuntimeError, match="test failure"), + ): + engine.to("cuda", model=True, optimizer=False, grad=False) + + assert observed == [(False, ("cuda",), {"model": True, "optimizer": False, "grad": False})] + assert engine.engine_config.forward_only is True + + +def test_non_tts_ref_keeps_standard_forward_only_path(): + omni_impl = _get_omni_impl_module() + engine = _make_manual_ref_engine(omni_impl, enabled=False) + observed = [] + + def fake_build(module): + observed.append(engine.engine_config.forward_only) + return module + + with patch.object(omni_impl.FSDPEngineWithLMHead, "_build_fsdp_module", side_effect=fake_build): + engine._build_fsdp_module(object()) + + assert observed == [True] + + # --------------------------------------------------------------------------- # ``_build_module`` calls adapter ``configure_model`` # --------------------------------------------------------------------------- diff --git a/verl_omni/pipelines/__init__.py b/verl_omni/pipelines/__init__.py index 73f555380..24ae44f7f 100644 --- a/verl_omni/pipelines/__init__.py +++ b/verl_omni/pipelines/__init__.py @@ -19,6 +19,7 @@ minimax_h3_diffusion_nft, minimax_h3_flow_grpo, qwen3_omni, + qwen3_tts, qwen_image_diffusion_nft, qwen_image_dpo, qwen_image_dual_grpo, @@ -35,6 +36,7 @@ from .minimax_h3_diffusion_nft import * # noqa: F401, F403 from .minimax_h3_flow_grpo import * # noqa: F401, F403 from .qwen3_omni import * # noqa: F401, F403 +from .qwen3_tts import * # noqa: F401, F403 from .qwen_image_diffusion_nft import * # noqa: F401, F403 from .qwen_image_dpo import * # noqa: F401, F403 from .qwen_image_dual_grpo import * # noqa: F401, F403 @@ -46,6 +48,7 @@ from .wan22_dance_grpo import * # noqa: F401, F403 __all__ = list(qwen3_omni.__all__) +__all__ += list(qwen3_tts.__all__) __all__ += list(qwen_image_flow_grpo.__all__) __all__ += list(qwen_image_diffusion_nft.__all__) __all__ += list(qwen_image_mix_grpo.__all__) diff --git a/verl_omni/pipelines/model_base.py b/verl_omni/pipelines/model_base.py index 245007f45..c8e1451c8 100644 --- a/verl_omni/pipelines/model_base.py +++ b/verl_omni/pipelines/model_base.py @@ -758,6 +758,16 @@ def rollout_flags(cls, pipeline_mode="thinker_only") -> dict[int, dict]: """ return {} + @classmethod + def weight_sync_stage_ids(cls, pipeline_mode="thinker_only") -> list[int] | None: + """Return stages that receive actor weights, or all stages by default.""" + return None + + @classmethod + def supports_cache_engine_sleep(cls, pipeline_mode="thinker_only") -> bool: + """Return whether the pipeline supports rollout cache sleep and wake.""" + return True + @classmethod def get_pipeline_id(cls, pipeline_mode: str = "thinker_only") -> str: """Return the vLLM-Omni pipeline model_type for *pipeline_mode*. @@ -812,3 +822,24 @@ def get_stage_engine_extras(cls, stage_id: int, pipeline_mode: str = "thinker_on dict: Extra key-value pairs merged into the stage's engine args. """ return {} + + @classmethod + def prepare_engine_prompt( + cls, + prompt_ids: list[int], + model_config, + multi_modal_data: dict, + mm_processor_kwargs: Optional[dict] = None, + ) -> dict | None: + """Build an architecture-specific rollout prompt when required.""" + return None + + @classmethod + def get_output_modalities(cls, pipeline_mode: str = "thinker_only") -> list[str] | None: + """Return intermediate modalities that must be retained by the engine.""" + return None + + @classmethod + def combine_engine_outputs(cls, outputs: list, prompt: dict) -> tuple[Any, dict[str, Any]]: + """Select the policy output and collect architecture-specific fields.""" + return (outputs[-1] if outputs else None), {} diff --git a/verl_omni/pipelines/qwen3_tts/__init__.py b/verl_omni/pipelines/qwen3_tts/__init__.py new file mode 100644 index 000000000..47fca4762 --- /dev/null +++ b/verl_omni/pipelines/qwen3_tts/__init__.py @@ -0,0 +1,19 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .agent_loop import Qwen3TTSSingleTurnAgentLoop +from .omni_rollout_adapter import Qwen3TTSRolloutAdapter +from .talker_training_adapter import Qwen3TTSTalkerAdapter + +__all__ = ["Qwen3TTSSingleTurnAgentLoop", "Qwen3TTSTalkerAdapter", "Qwen3TTSRolloutAdapter"] diff --git a/verl_omni/pipelines/qwen3_tts/agent_loop.py b/verl_omni/pipelines/qwen3_tts/agent_loop.py new file mode 100644 index 000000000..b8a8bca18 --- /dev/null +++ b/verl_omni/pipelines/qwen3_tts/agent_loop.py @@ -0,0 +1,61 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Single-turn agent loop that makes codec-0 the policy sequence.""" + +import torch +from verl.experimental.agent_loop.agent_loop import register +from verl.experimental.agent_loop.single_turn_agent_loop import SingleTurnAgentLoop + +from verl_omni.pipelines.qwen3_tts.rollout_utils import is_evaluation_split, with_rollout_generation_seed +from verl_omni.pipelines.qwen3_tts.talker_forward import TEXT_PROMPT_TRAILER_TOKENS, build_assistant_text + + +@register("qwen3_tts_single_turn") +class Qwen3TTSSingleTurnAgentLoop(SingleTurnAgentLoop): + async def run(self, sampling_params, **kwargs): + extra_info = kwargs.get("extra_info") + evaluation = is_evaluation_split(extra_info) + candidate_count = self.rollout_config.val_kwargs.n if evaluation else self.rollout_config.n + sampling_params = with_rollout_generation_seed( + sampling_params, + extra_info, + session_id=kwargs.get("session_id"), + global_steps=kwargs.get("global_steps"), + uid=kwargs.get("uid"), + base_seed=int(self.config.data.get("seed", 0)), + require_session_id=int(candidate_count) > 1, + ) + output = await super().run(sampling_params, **kwargs) + extra = output.extra_fields + codes, text = extra.get("tts_audio_codes"), extra.get("tts_text") + if codes is None or text is None: + raise RuntimeError("Qwen3-TTS rollout did not return codec codes and text.") + codes = torch.as_tensor(codes, dtype=torch.long) + if codes.ndim != 2 or codes.shape[-1] != 16: + raise ValueError("Qwen3-TTS codec codes must have shape (frames, 16).") + codes = codes[: self.response_length] + policy_ids = codes[:, 0].tolist() + if not policy_ids: + raise RuntimeError("Qwen3-TTS rollout returned an empty codec trajectory.") + if output.response_logprobs is not None: + if len(output.response_logprobs) < len(policy_ids): + raise RuntimeError("Qwen3-TTS rollout logprobs are shorter than the policy trajectory.") + output.response_logprobs = output.response_logprobs[: len(policy_ids)] + text_ids = self.tokenizer(build_assistant_text(str(text)), return_tensors="pt", padding=False)["input_ids"] + extra["tts_text_ids"] = text_ids[:, :-TEXT_PROMPT_TRAILER_TOKENS].reshape(-1).tolist() + extra["tts_audio_codes"] = codes + output.prompt_ids = [0] + output.response_ids = policy_ids + output.response_mask = [1] * len(policy_ids) + return output diff --git a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py new file mode 100644 index 000000000..1d0f3bfef --- /dev/null +++ b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py @@ -0,0 +1,195 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Qwen3-TTS two-stage rollout adapter.""" + +import copy +import hashlib +from collections.abc import Mapping +from dataclasses import replace +from functools import lru_cache + +import torch +from vllm_omni.config.pipeline_registry import register_pipeline +from vllm_omni.config.stage_config import PipelineConfig +from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE + +from verl_omni.pipelines.model_base import OmniRolloutPipelineBase +from verl_omni.pipelines.qwen3_tts.rollout_utils import align_audio_codes, append_tensor_chunk +from verl_omni.pipelines.qwen3_tts.talker_forward import ( + build_assistant_text, + load_speaker_xvector, + require_auto_language, +) + +_PIPELINE_ID = "qwen3_tts_rl" +_SYNC_PROCESSOR = "verl_omni.pipelines.qwen3_tts.omni_rollout_adapter.talker2code2wav_token_only" +QWEN3_TTS_RL_PIPELINE = PipelineConfig( + model_type=_PIPELINE_ID, + model_arch=QWEN3_TTS_PIPELINE.model_arch, + stages=( + replace(QWEN3_TTS_PIPELINE.stages[0], final_output=True, final_output_type="latent"), + replace(QWEN3_TTS_PIPELINE.stages[1], sync_process_input_func=_SYNC_PROCESSOR), + ), +) + + +@lru_cache(maxsize=4) +def _load_speaker_vector(path: str) -> list[float]: + return load_speaker_xvector(path).reshape(-1).tolist() + + +def _completion(output): + request_output = getattr(output, "request_output", None) + completions = getattr(request_output, "outputs", None) if request_output is not None else None + return completions[0] if completions else None + + +def _materialize(value): + if isinstance(value, Mapping): + return {key: _materialize(item) for key, item in value.items()} + if isinstance(value, list): + return [_materialize(item) for item in value] + return value + + +def talker2code2wav_token_only(source_outputs, prompt=None, _requires_multimodal_data=False): + """Materialize shared-memory Mapping payloads before the upstream processor mutates them.""" + from vllm_omni.model_executor.stage_input_processors.qwen3_tts import ( + talker2code2wav_token_only as upstream_processor, + ) + + converted = [] + for source_output in source_outputs: + source_copy = copy.copy(source_output) + source_copy.outputs = [] + for completion in getattr(source_output, "outputs", []): + completion_copy = copy.copy(completion) + multimodal = getattr(completion, "multimodal_output", None) + if isinstance(multimodal, Mapping): + completion_copy.multimodal_output = _materialize(multimodal) + source_copy.outputs.append(completion_copy) + converted.append(source_copy) + return upstream_processor(converted, prompt, _requires_multimodal_data) + + +@OmniRolloutPipelineBase.register(_PIPELINE_ID) +class Qwen3TTSRolloutAdapter(OmniRolloutPipelineBase): + @classmethod + def _check_mode(cls, pipeline_mode): + if pipeline_mode != "full": + raise ValueError("Qwen3-TTS RL supports only pipeline_mode='full'.") + + @classmethod + def build_stage_configs(cls, pipeline_mode="full"): + cls._check_mode(pipeline_mode) + return list(QWEN3_TTS_RL_PIPELINE.stages) + + @classmethod + def get_pipeline_id(cls, pipeline_mode="full"): + cls._check_mode(pipeline_mode) + return _PIPELINE_ID + + @classmethod + def ensure_pipeline_registered(cls, pipeline_mode="full"): + cls._check_mode(pipeline_mode) + register_pipeline(QWEN3_TTS_RL_PIPELINE) + + @classmethod + def get_output_modalities(cls, pipeline_mode="full"): + cls._check_mode(pipeline_mode) + return ["latent", "audio"] + + @classmethod + def weight_sync_stage_ids(cls, pipeline_mode="full"): + cls._check_mode(pipeline_mode) + return [0] + + @classmethod + def supports_cache_engine_sleep(cls, pipeline_mode="full"): + cls._check_mode(pipeline_mode) + return False + + @classmethod + def get_stage_engine_extras(cls, stage_id, pipeline_mode="full"): + cls._check_mode(pipeline_mode) + return {"max_model_len": 65536, "max_num_batched_tokens": 65536} if stage_id == 1 else {} + + @classmethod + def prepare_engine_prompt(cls, prompt_ids, model_config, multi_modal_data, mm_processor_kwargs=None): + text = model_config.tokenizer.decode(prompt_ids, skip_special_tokens=True).strip() + if not text: + raise ValueError("Qwen3-TTS received an empty text prompt.") + speaker_path = model_config.override_config.get("tts_spk_embed_path") + if not speaker_path: + raise ValueError("Qwen3-TTS GRPO requires tts_spk_embed_path for the validated non-streaming replay.") + language = require_auto_language(model_config.override_config.get("tts_language", "Auto")) + additional_information = { + "task_type": ["Base"], + "text": [text], + "language": [language], + "non_streaming_mode": [True], + "x_vector_only_mode": [True], + "voice_clone_prompt": [{"ref_spk_embedding": _load_speaker_vector(speaker_path)}], + } + assistant_ids = model_config.tokenizer(build_assistant_text(text), padding=False)["input_ids"] + assistant_ids = torch.as_tensor(assistant_ids).reshape(-1).tolist() + prompt_length = len(assistant_ids) + 2 + identity = "\0".join((text, str(language), str(speaker_path))).encode() + return { + "prompt_token_ids": [1] * prompt_length, + "additional_information": additional_information, + "cache_salt": hashlib.sha256(identity).hexdigest(), + } + + @classmethod + def combine_engine_outputs(cls, outputs, prompt): + policy_output = None + policy_length = -1 + audio_codes = waveform = None + sample_rate = None + diagnostics = [] + for output in outputs: + completion = _completion(output) + if getattr(output, "stage_id", None) == 0 and completion is not None: + length = len(getattr(completion, "token_ids", None) or []) + if length >= policy_length: + policy_output, policy_length = output, length + multimodal = getattr(output, "multimodal_output", None) + diagnostics.append((getattr(output, "stage_id", None), type(multimodal).__name__)) + if not isinstance(multimodal, Mapping): + continue + codes = multimodal.get("codes") + if isinstance(codes, Mapping): + audio_codes = append_tensor_chunk(audio_codes, codes.get("audio")) + if getattr(output, "stage_id", None) == 1: + waveform = append_tensor_chunk( + waveform, multimodal.get("audio", multimodal.get("model_outputs")), flatten=True + ) + sample_rate = multimodal.get("sr", multimodal.get("audio_sample_rate", sample_rate)) + if policy_output is None: + raise RuntimeError("Qwen3-TTS rollout produced no stage-0 policy output.") + token_ids = list(getattr(_completion(policy_output), "token_ids", None) or []) + if audio_codes is None: + raise RuntimeError(f"Qwen3-TTS rollout produced no codec trajectory: {diagnostics}") + fields = { + "tts_audio_codes": align_audio_codes(audio_codes, token_ids), + "tts_text": prompt["additional_information"]["text"][0], + } + if waveform is not None: + fields["audio"] = waveform.float().reshape(-1) + if sample_rate is not None: + if isinstance(sample_rate, list | tuple): + sample_rate = sample_rate[-1] + fields["audio_sample_rate"] = int(sample_rate.item() if hasattr(sample_rate, "item") else sample_rate) + return policy_output, fields diff --git a/verl_omni/pipelines/qwen3_tts/rollout_utils.py b/verl_omni/pipelines/qwen3_tts/rollout_utils.py new file mode 100644 index 000000000..78add32dc --- /dev/null +++ b/verl_omni/pipelines/qwen3_tts/rollout_utils.py @@ -0,0 +1,147 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Helpers for collecting Qwen3-TTS rollout outputs.""" + +import hashlib +from collections.abc import Mapping +from typing import Any + +import torch + +_EVALUATION_SPLITS = {"gate", "official_test", "test", "val", "validation"} +_MAX_SAMPLING_SEED = 2**31 - 1 + + +def _scalar(value): + if hasattr(value, "item"): + try: + value = value.item() + except ValueError: + return None + return value + + +def _extra_info_mapping(extra_info) -> dict: + extra_info = _scalar(extra_info) + return dict(extra_info.items()) if isinstance(extra_info, Mapping) else {} + + +def is_evaluation_split(extra_info) -> bool: + info = _extra_info_mapping(extra_info) + return str(_scalar(info.get("split")) or "").lower() in _EVALUATION_SPLITS + + +def validation_generation_seed(extra_info) -> int | None: + info = _extra_info_mapping(extra_info) + if not is_evaluation_split(info): + return None + seed = _scalar(info.get("generation_seed")) + return None if seed is None else int(seed) + + +def rollout_generation_seed( + extra_info, + *, + session_id=None, + global_steps=None, + uid=None, + base_seed=0, + require_session_id=False, +) -> int: + """Derive a stable, group-diverse seed for both Qwen3-TTS samplers.""" + if require_session_id and session_id is None: + raise RuntimeError("Qwen3-TTS group sampling requires a per-candidate session_id when rollout.n > 1.") + candidate = int(_scalar(session_id) or 0) + explicit_seed = validation_generation_seed(extra_info) + if explicit_seed is not None: + return (explicit_seed + candidate) % _MAX_SAMPLING_SEED + + info = _extra_info_mapping(extra_info) + sample_id = _scalar(info.get("id", info.get("index", uid))) + step = "evaluation" if is_evaluation_split(info) else _scalar(global_steps) + payload = "\0".join(map(str, (int(base_seed), step, sample_id, candidate))).encode() + return int.from_bytes(hashlib.blake2b(payload, digest_size=8).digest(), "big") % _MAX_SAMPLING_SEED + + +def with_rollout_generation_seed(sampling_params, extra_info, **seed_kwargs): + """Seed codec-0 and residual-codebook sampling without mutating input.""" + seed = rollout_generation_seed(extra_info, **seed_kwargs) + seeded = dict(sampling_params) + residual_args = dict(seeded.get("extra_args") or {}) + residual_args["tts_local_seed"] = seed + seeded.update(seed=seed, extra_args=residual_args) + return seeded + + +def append_tensor_chunk(accumulated: torch.Tensor | None, value: Any, *, flatten=False): + if isinstance(value, list | tuple): + value = value[-1] if value else None + if value is None: + return accumulated + chunk = torch.as_tensor(value).detach().cpu() + if flatten: + chunk = chunk.reshape(-1) + if not chunk.numel(): + return accumulated + if accumulated is None: + return chunk + if chunk.shape[1:] != accumulated.shape[1:]: + raise RuntimeError(f"Rollout chunks changed shape from {tuple(accumulated.shape)} to {tuple(chunk.shape)}.") + if chunk.shape[0] >= accumulated.shape[0] and torch.equal(chunk[: accumulated.shape[0]], accumulated): + return chunk + if accumulated.shape[0] >= chunk.shape[0] and torch.equal(accumulated[: chunk.shape[0]], chunk): + return accumulated + return torch.cat((accumulated, chunk), dim=0) + + +def align_audio_codes(audio_codes: torch.Tensor, token_ids: list[int]) -> torch.Tensor: + """Recover residual codebooks using codec-0 policy tokens as an exact invariant. + + The engine can prepend placeholder rows and need not emit residual codes + for the last sampled token, because that token is never consumed as the + context for another policy token. Thus ``token_ids[:-1]`` must match + exactly; a final residual row is retained only when the engine provides it. + """ + if audio_codes.ndim != 2 or audio_codes.shape[-1] != 16: + raise ValueError("Qwen3-TTS codec codes must have shape (frames, 16).") + if not token_ids: + return audio_codes[:0].long() + raw_codes = audio_codes.long() + policy_ids = torch.as_tensor(token_ids, dtype=torch.long, device=raw_codes.device) + required = len(token_ids) - 1 + candidates: list[tuple[int, int]] = [] + if required == 0: + candidates.append((0, 0)) + else: + for start in range(raw_codes.shape[0] - required + 1): + if torch.equal(raw_codes[start : start + required, 0], policy_ids[:required]): + has_final = int( + raw_codes.shape[0] - start >= len(token_ids) + and raw_codes[start + required, 0] == policy_ids[required] + ) + candidates.append((required + has_final, start)) + if not candidates: + raise RuntimeError( + "Could not exactly align Qwen3-TTS residual codebooks with the sampled codec-0 policy: " + f"response_length={len(token_ids)}, raw_codec_rows={raw_codes.shape[0]}." + ) + + copy_length, start = max(candidates) + codes = raw_codes.new_zeros((len(token_ids), 16)) + if copy_length: + codes[:copy_length] = raw_codes[start : start + copy_length] + codes[:, 0] = policy_ids + if required and not torch.equal(codes[:required, 0], policy_ids[:required]): + raise RuntimeError("Qwen3-TTS codec alignment invariant failed after recovery.") + return codes diff --git a/verl_omni/pipelines/qwen3_tts/talker_forward.py b/verl_omni/pipelines/qwen3_tts/talker_forward.py new file mode 100644 index 000000000..cdb5a7b48 --- /dev/null +++ b/verl_omni/pipelines/qwen3_tts/talker_forward.py @@ -0,0 +1,229 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Teacher-forced codec-0 forward math for Qwen3-TTS.""" + +import json +from dataclasses import dataclass + +import torch + +NUM_CODEBOOKS = 16 +TEXT_PROMPT_TRAILER_TOKENS = 5 + + +def build_assistant_text(text: str) -> str: + return f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n" + + +def load_speaker_xvector(path: str) -> torch.Tensor: + with open(path) as file: + vector = json.load(file) + return torch.tensor(vector, dtype=torch.float32).reshape(1, -1) + + +@dataclass +class TalkerTokens: + tts_pad: int + tts_bos: int + tts_eos: int + codec_pad: int + codec_bos: int + codec_eos: int + codec_nothink: int + codec_think_bos: int + codec_think_eos: int + + @classmethod + def from_config(cls, config): + talker = config.talker_config + return cls( + int(config.tts_pad_token_id), + int(config.tts_bos_token_id), + int(config.tts_eos_token_id), + int(talker.codec_pad_id), + int(talker.codec_bos_id), + int(talker.codec_eos_token_id), + int(talker.codec_nothink_id), + int(talker.codec_think_bos_id), + int(talker.codec_think_eos_id), + ) + + +@dataclass +class TalkerBatch: + input_ids: torch.Tensor + codec_ids: torch.Tensor + text_embedding_mask: torch.Tensor + codec_embedding_mask: torch.Tensor + codec_mask: torch.Tensor + attention_mask: torch.Tensor + codec_lens: list[int] + logit_start: list[int] + speaker_slots: list[int] + + +def build_talker_batch( + text_ids, + audio_codes, + tokens, + *, + device=None, + sub_codebook_vocab=None, +) -> TalkerBatch: + text_lens = [int(ids.reshape(-1).shape[0]) for ids in text_ids] + codec_lens = [int(codes.shape[0]) for codes in audio_codes] + speaker_slot = 6 + text_start = 8 + seq_len = max(t + c for t, c in zip(text_lens, codec_lens, strict=True)) + 8 + batch_size = len(text_ids) + input_ids = torch.zeros((batch_size, seq_len, 2), dtype=torch.long, device=device) + codec_ids = torch.zeros((batch_size, seq_len, NUM_CODEBOOKS), dtype=torch.long, device=device) + text_mask = torch.zeros((batch_size, seq_len), dtype=torch.bool, device=device) + codec_embedding_mask = torch.zeros_like(text_mask) + codec_mask = torch.zeros_like(text_mask) + attention_mask = torch.zeros((batch_size, seq_len), dtype=torch.long, device=device) + + for index, (sample_text, sample_codes) in enumerate(zip(text_ids, audio_codes, strict=True)): + ids = sample_text.reshape(-1).to(device=device, dtype=torch.long) + codes = sample_codes.to(device=device, dtype=torch.long) + if sub_codebook_vocab is not None: + codes = codes.clone() + codes[:, 1:].clamp_(0, sub_codebook_vocab - 1) + text_len, codec_len = text_lens[index], codec_lens[index] + + input_ids[index, :3, 0] = ids[:3] + input_ids[index, 3:speaker_slot, 0] = tokens.tts_pad + input_ids[index, speaker_slot, 0] = tokens.tts_pad + input_ids[index, speaker_slot + 1, 0] = tokens.tts_bos + input_ids[index, text_start : text_start + text_len - 3, 0] = ids[3:] + input_ids[index, text_start + text_len - 3, 0] = tokens.tts_eos + input_ids[index, text_start + text_len - 2 : text_start + text_len + codec_len, 0] = tokens.tts_pad + text_mask[index, : text_start + text_len + codec_len] = True + + input_ids[index, 3:speaker_slot, 1] = torch.tensor( + [tokens.codec_nothink, tokens.codec_think_bos, tokens.codec_think_eos], device=device + ) + input_ids[index, speaker_slot + 1, 1] = tokens.codec_pad + input_ids[index, text_start : text_start + text_len - 2, 1] = tokens.codec_pad + input_ids[index, text_start + text_len - 2, 1] = tokens.codec_bos + start = text_start + text_len - 1 + input_ids[index, start : start + codec_len, 1] = codes[:, 0] + input_ids[index, start + codec_len, 1] = tokens.codec_eos + codec_ids[index, start : start + codec_len] = codes + codec_embedding_mask[index, 3 : text_start + text_len + codec_len] = True + codec_embedding_mask[index, speaker_slot] = False + codec_mask[index, start : start + codec_len] = True + attention_mask[index, : text_start + text_len + codec_len] = True + + return TalkerBatch( + input_ids, + codec_ids, + text_mask.unsqueeze(-1), + codec_embedding_mask.unsqueeze(-1), + codec_mask, + attention_mask, + codec_lens, + [text_start + text_len - 2 for text_len in text_lens], + [speaker_slot] * batch_size, + ) + + +def require_auto_language(language) -> str: + """Limit RL training to the prompt layout validated by the actor forward.""" + normalized = str(language or "Auto").strip() + if normalized.lower() != "auto": + raise ValueError( + "Qwen3-TTS RL currently supports only tts_language=Auto; " + "language-specific codec prefixes require a separately validated actor forward." + ) + return "Auto" + + +def codec0_input_embeddings(talker, batch: TalkerBatch, speaker_embedding: torch.Tensor) -> torch.Tensor: + ids = batch.input_ids + text_embeddings = talker.model.text_embedding(ids[:, :, 0]) + if getattr(talker, "text_projection", None) is not None: + text_embeddings = talker.text_projection(text_embeddings) + codec_embeddings = talker.model.codec_embedding(ids[:, :, 1]) * batch.codec_embedding_mask + codec_embeddings = codec_embeddings.clone() + sample_indices = torch.arange(codec_embeddings.shape[0], device=codec_embeddings.device) + codec_embeddings[sample_indices, batch.speaker_slots] = speaker_embedding.to(codec_embeddings.dtype) + embeddings = text_embeddings * batch.text_embedding_mask + codec_embeddings + sub_embeddings = talker.code_predictor.get_input_embeddings() + codec_mask = batch.codec_mask.unsqueeze(-1) + for codebook in range(1, NUM_CODEBOOKS): + embeddings += sub_embeddings[codebook - 1](batch.codec_ids[:, :, codebook]) * codec_mask + return embeddings + + +def codec0_logits(talker, batch: TalkerBatch, speaker_embedding: torch.Tensor) -> torch.Tensor: + embeddings = codec0_input_embeddings(talker, batch, speaker_embedding) + output = talker( + inputs_embeds=embeddings[:, :-1], + attention_mask=batch.attention_mask[:, :-1], + use_cache=False, + output_hidden_states=False, + ) + return output.logits + + +def tts_actor_logits( + model, + input_ids, + attention_mask, + tts_text_ids, + tts_audio_codes, + response_len, + text_len, + speaker_embedding, +) -> torch.Tensor: + batch_size, output_len = input_ids.shape + texts, codes, response_starts = [], [], [] + for index in range(batch_size): + response_size, text_size = int(response_len[index]), int(text_len[index]) + response_start = int(attention_mask[index].sum()) - response_size + if response_start < 1 or response_start + response_size > output_len: + raise ValueError("Invalid Qwen3-TTS response alignment.") + policy_ids = input_ids[index, response_start : response_start + response_size].long() + codec0_ids = tts_audio_codes[index, :response_size, 0].long() + if not torch.equal(policy_ids, codec0_ids): + mismatch = int(torch.nonzero(policy_ids != codec0_ids, as_tuple=False)[0]) + raise RuntimeError( + f"Qwen3-TTS codec trajectory is not aligned with actor labels: sample={index}, frame={mismatch}." + ) + texts.append(tts_text_ids[index, :text_size].long()) + codes.append(tts_audio_codes[index, :response_size].long()) + response_starts.append(response_start) + + talker = model.talker + sub_vocab = int(talker.code_predictor.get_input_embeddings()[0].num_embeddings) + tokens = TalkerTokens.from_config(model.config) + batch = build_talker_batch( + texts, + codes, + tokens, + device=input_ids.device, + sub_codebook_vocab=sub_vocab, + ) + logits = codec0_logits(talker, batch, speaker_embedding) + output_vocab = max(logits.shape[-1], int(input_ids.max()) + 1) + aligned = logits.new_zeros((batch_size, output_len, output_vocab)) + for index, response_start in enumerate(response_starts): + codec_len = batch.codec_lens[index] + target = slice(response_start - 1, response_start - 1 + codec_len) + if output_vocab > logits.shape[-1]: + aligned[index, target, logits.shape[-1] :] = -1e4 + source = batch.logit_start[index] + aligned[index, target, : logits.shape[-1]] = logits[index, source : source + codec_len] + return aligned diff --git a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py new file mode 100644 index 000000000..f3dcc5cd4 --- /dev/null +++ b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py @@ -0,0 +1,228 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Qwen3-TTS talker actor adapter.""" + +import logging +import types +from typing import Any + +import torch + +from verl_omni.pipelines.model_base import OmniModelBase +from verl_omni.pipelines.qwen3_tts.talker_forward import ( + load_speaker_xvector, + require_auto_language, + tts_actor_logits, +) +from verl_omni.pipelines.qwen3_tts.transformers_compat import qwen3_tts_import_context + +logger = logging.getLogger(__name__) +_PASSTHROUGH_TEMPLATE = "{% for message in messages %}{{ message['content'] }}{% endfor %}" +_TRAINABLE_PREFIXES = ("talker.model.", "talker.codec_head.") + + +def _prepare_config_for_checkpoint(config) -> None: + speaker_config = getattr(config, "speaker_encoder_config", None) + if speaker_config is not None: + speaker_config.__dict__.pop("dtype", None) + speaker_config.__dict__.pop("_dtype", None) + + +def register_qwen3_tts_automodel() -> None: + config_cls = model_cls = None + with qwen3_tts_import_context(): + for module_name in ("transformers", "qwen_tts.core.models.configuration_qwen3_tts"): + try: + config_cls = getattr(__import__(module_name, fromlist=["Qwen3TTSConfig"]), "Qwen3TTSConfig", None) + except ImportError: + continue + if config_cls is not None: + break + for module_name in ("transformers", "qwen_tts.core.models.modeling_qwen3_tts"): + try: + model_cls = getattr( + __import__(module_name, fromlist=["Qwen3TTSForConditionalGeneration"]), + "Qwen3TTSForConditionalGeneration", + None, + ) + except ImportError: + continue + if model_cls is not None: + break + if config_cls is None or model_cls is None: + return + + from transformers import AutoConfig, AutoModelForMultimodalLM + + try: + AutoConfig.register(getattr(config_cls, "model_type", "qwen3_tts"), config_cls) + except ValueError: + pass + try: + AutoModelForMultimodalLM.register(config_cls, model_cls) + except ValueError: + pass + + +def _speaker_embedding(model, batch_size, device, dtype): + cached = getattr(model, "_verl_tts_speaker_embedding", None) + if cached is None: + path = getattr(model.config, "tts_spk_embed_path", None) + if not path: + return None + cached = load_speaker_xvector(path) + model._verl_tts_speaker_embedding = cached + return cached.to(device=device, dtype=dtype).expand(batch_size, -1) + + +def _reinitialize_rope_buffers(model): + for submodule in model.modules(): + rope_init = getattr(submodule, "rope_init_fn", None) + inv_freq = getattr(submodule, "inv_freq", None) + if rope_init is None or not torch.is_tensor(inv_freq): + continue + new_inv_freq, scaling = rope_init(submodule.config, device=inv_freq.device) + submodule.inv_freq.data.copy_(new_inv_freq.to(device=inv_freq.device, dtype=inv_freq.dtype)) + submodule.attention_scaling = scaling + + +def _qwen3_tts_forward( + self, + input_ids=None, + attention_mask=None, + tts_text_ids=None, + tts_audio_codes=None, + response_len=None, + text_len=None, + **kwargs, +): + from transformers.modeling_outputs import CausalLMOutputWithPast + + if any(value is None for value in (tts_text_ids, tts_audio_codes, response_len, text_len)): + raise RuntimeError("Qwen3-TTS forward is missing exact rollout codec fields.") + if attention_mask is None: + attention_mask = torch.ones_like(input_ids) + if not getattr(self, "_verl_tts_rope_initialized", False): + _reinitialize_rope_buffers(self) + self._verl_tts_rope_initialized = True + speaker = _speaker_embedding(self, input_ids.shape[0], input_ids.device, next(self.talker.parameters()).dtype) + return CausalLMOutputWithPast( + logits=tts_actor_logits( + self, + input_ids, + attention_mask, + tts_text_ids, + tts_audio_codes, + response_len, + text_len, + speaker, + ) + ) + + +def _get_input_embeddings(self): + return self.talker.model.codec_embedding + + +def _set_input_embeddings(self, value): + self.talker.model.codec_embedding = value + + +@OmniModelBase.register("Qwen3TTSForConditionalGeneration", stage="talker") +class Qwen3TTSTalkerAdapter(OmniModelBase): + requires_manual_ref_offload = True + + @classmethod + def get_strip_modules(cls, model_config): + return ["speaker_encoder", "speech_tokenizer", "code2wav"] + + @classmethod + def configure_model(cls, module, model_config): + module = super().configure_model(module, model_config) + _prepare_config_for_checkpoint(module.config) + module.config.tts_spk_embed_path = model_config.override_config.get("tts_spk_embed_path") + module.config.tts_language = require_auto_language(model_config.override_config.get("tts_language", "Auto")) + if not module.config.tts_spk_embed_path: + raise ValueError("Qwen3-TTS GRPO requires tts_spk_embed_path for the validated non-streaming replay.") + module.forward = types.MethodType(_qwen3_tts_forward, module) + module.get_input_embeddings = types.MethodType(_get_input_embeddings, module) + module.set_input_embeddings = types.MethodType(_set_input_embeddings, module) + module._no_split_modules = ["Qwen3TTSTalkerDecoderLayer", "Qwen3TTSDecoderLayer"] + trainable = 0 + for name, parameter in module.named_parameters(): + parameter.requires_grad_(name.startswith(_TRAINABLE_PREFIXES)) + trainable += int(parameter.requires_grad) + logger.info("Qwen3-TTS talker adapter enabled %d trainable parameter tensors", trainable) + return module + + @classmethod + def configure_processor(cls, model_path: str, model_config) -> Any: + return None + + @classmethod + def configure_tokenizer(cls, model_path: str, model_config): + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=model_config.trust_remote_code) + tokenizer.chat_template = _PASSTHROUGH_TEMPLATE + if tokenizer.pad_token_id is None: + tokenizer.pad_token_id = tokenizer.eos_token_id or 0 + talker_config = getattr(model_config.hf_config, "talker_config", None) + if talker_config is not None: + talker_config.tie_word_embeddings = False + return tokenizer + + @classmethod + def prepare_model_inputs(cls, model_inputs, micro_batch, model_config): + del model_config + fields = micro_batch.get("extra_fields") + if fields is None: + raise RuntimeError( + "Qwen3-TTS actor inputs require AgentLoopOutput.extra_fields; use the V1 agent-loop trainer path." + ) + if hasattr(fields, "tolist"): + fields = fields.tolist() + if isinstance(fields, dict): + fields = [fields] + fields = [getattr(item, "data", item) for item in fields] + if len(fields) != model_inputs["input_ids"].shape[0]: + raise RuntimeError("Qwen3-TTS actor extra_fields do not match its batch size.") + + texts = [torch.as_tensor(item["tts_text_ids"], dtype=torch.long).reshape(-1) for item in fields] + codes = [torch.as_tensor(item["tts_audio_codes"], dtype=torch.long) for item in fields] + if any(item.ndim != 2 or item.shape[-1] != 16 for item in codes): + raise ValueError("Qwen3-TTS codec codes must have shape (frames, 16).") + device = model_inputs["input_ids"].device + text_buffer = torch.zeros((len(fields), max(item.numel() for item in texts)), dtype=torch.long, device=device) + code_buffer = torch.zeros( + (len(fields), max(item.shape[0] for item in codes), 16), dtype=torch.long, device=device + ) + text_lens = torch.empty(len(fields), dtype=torch.long, device=device) + response_lens = torch.empty_like(text_lens) + for index, (text, code) in enumerate(zip(texts, codes, strict=True)): + text_buffer[index, : text.numel()] = text.to(device) + code_buffer[index, : code.shape[0]] = code.to(device) + text_lens[index], response_lens[index] = text.numel(), code.shape[0] + model_inputs.update( + { + "tts_text_ids": text_buffer, + "tts_audio_codes": code_buffer, + "text_len": text_lens, + "response_len": response_lens, + } + ) + return model_inputs + + +register_qwen3_tts_automodel() diff --git a/verl_omni/pipelines/qwen3_tts/transformers_compat.py b/verl_omni/pipelines/qwen3_tts/transformers_compat.py new file mode 100644 index 000000000..3784386ba --- /dev/null +++ b/verl_omni/pipelines/qwen3_tts/transformers_compat.py @@ -0,0 +1,74 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Import-time compatibility for qwen-tts 0.1.1 on Transformers 5.x.""" + +from contextlib import contextmanager +from functools import wraps + +import torch +from packaging.version import Version + + +def _default_rope_init(config, device=None, **kwargs): + del kwargs + base = getattr(config, "rope_theta", 10_000.0) or 10_000.0 + dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + positions = torch.arange(0, dim, 2, dtype=torch.int64, device=device).float() / dim + return 1.0 / (base**positions), 1.0 + + +def _compatible_mask(original): + @wraps(original) + def wrapper(*args, input_embeds=None, inputs_embeds=None, cache_position=None, **kwargs): + del cache_position + embeddings = inputs_embeds if inputs_embeds is not None else input_embeds + return original(*args, inputs_embeds=embeddings, **kwargs) + + return wrapper + + +@contextmanager +def qwen3_tts_import_context(): + """Expose the TF5 APIs expected while qwen-tts binds its imports. + + qwen-tts modules retain the mask wrappers they import. The global + Transformers functions are restored immediately afterwards. + """ + import transformers + + if Version(transformers.__version__).major < 5: + yield + return + + import transformers.masking_utils as masking_utils + import transformers.utils.generic as generic_utils + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + + original_check = generic_utils.check_model_inputs + original_causal_mask = masking_utils.create_causal_mask + original_sliding_mask = masking_utils.create_sliding_window_causal_mask + + def compatible_check_model_inputs(func=None): + return original_check if func is None else original_check(func) + + generic_utils.check_model_inputs = compatible_check_model_inputs + masking_utils.create_causal_mask = _compatible_mask(original_causal_mask) + masking_utils.create_sliding_window_causal_mask = _compatible_mask(original_sliding_mask) + ROPE_INIT_FUNCTIONS.setdefault("default", _default_rope_init) + try: + yield + finally: + generic_utils.check_model_inputs = original_check + masking_utils.create_causal_mask = original_causal_mask + masking_utils.create_sliding_window_causal_mask = original_sliding_mask diff --git a/verl_omni/workers/rollout/vllm_rollout/utils.py b/verl_omni/workers/rollout/vllm_rollout/utils.py index b03205af0..c389fc9a0 100644 --- a/verl_omni/workers/rollout/vllm_rollout/utils.py +++ b/verl_omni/workers/rollout/vllm_rollout/utils.py @@ -16,7 +16,7 @@ import time import torch -from verl.utils.device import get_visible_devices_keyword +from verl.utils.device import get_torch_device, get_visible_devices_keyword from verl.workers.rollout.vllm_rollout.utils import VLLM_LORA_INT_ID, VLLM_LORA_NAME, VLLM_LORA_PATH, set_death_signal from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension @@ -32,6 +32,26 @@ def _split_visible_devices(value: str) -> list[str]: return [entry.strip() for entry in value.split(",") if entry.strip()] +def _receive_model_weight_buckets(receiver, model) -> None: + """Stream model weights and rebuild optional derived tensors once.""" + rebuild_derived_weights = getattr(model, "_build_stacked_codec_embed", None) + if callable(rebuild_derived_weights): + model._build_stacked_codec_embed = lambda: None + try: + receiver.receive_weights( + on_bucket_received=lambda weights, *args, **kwargs: model.load_weights(weights) + ) + finally: + if callable(rebuild_derived_weights): + model._build_stacked_codec_embed = rebuild_derived_weights + if callable(rebuild_derived_weights): + old_derived_weights = getattr(model, "_stacked_codec_embed", None) + model._stacked_codec_embed = None + del old_derived_weights + get_torch_device().empty_cache() + rebuild_derived_weights() + + class vLLMOmniColocateWorkerExtension(CustomPipelineWorkerExtension): """ The class for vLLM-Omni's worker to inherit from, in the colocate setting. @@ -191,9 +211,7 @@ def update_weights_from_ipc( from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader patch_vllm_moe_model_weight_loader(model) - receiver.receive_weights( - on_bucket_received=lambda weights, *args, **kwargs: model.load_weights(weights) - ) + _receive_model_weight_buckets(receiver, model) from vllm.model_executor.model_loader.utils import process_weights_after_loading process_weights_after_loading(model, model_config, self.device) From 78ea2e2ac4b7fc8edd2a398b9c43bb52dbc56f6c Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Fri, 21 Aug 2026 09:56:09 +0800 Subject: [PATCH 02/28] [reward, cfg] feat: add Qwen3-TTS GRPO audio reward recipe Signed-off-by: dongbo910220 <1275604947@qq.com> Co-authored-by: Michael-Zzq <52185141+Michael-Zzq@users.noreply.github.com> --- docs/api/reward.rst | 20 +- docs/start/http_scorer.md | 7 + docs/start/models.md | 18 +- examples/grpo_trainer/qwen3_tts/README.md | 112 ++++++++++ .../qwen3_tts/run_qwen3_tts_grpo.sh | 152 ++++++++++++++ .../test_audio_reward_manager_on_cpu.py | 191 ++++++++++++++++++ .../test_audio_http_scorer_client_on_cpu.py | 158 +++++++++++++++ .../reward_loop/reward_manager/__init__.py | 3 +- verl_omni/reward_loop/reward_manager/audio.py | 121 +++++++++++ .../reward_score/audio_http_scorer_client.py | 175 ++++++++++++++++ 10 files changed, 953 insertions(+), 4 deletions(-) create mode 100644 examples/grpo_trainer/qwen3_tts/README.md create mode 100755 examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh create mode 100644 tests/reward_loop/test_audio_reward_manager_on_cpu.py create mode 100644 tests/utils/reward_score/test_audio_http_scorer_client_on_cpu.py create mode 100644 verl_omni/reward_loop/reward_manager/audio.py create mode 100644 verl_omni/utils/reward_score/audio_http_scorer_client.py diff --git a/docs/api/reward.rst b/docs/api/reward.rst index 1fe28f3ca..1a0a7a23e 100644 --- a/docs/api/reward.rst +++ b/docs/api/reward.rst @@ -6,8 +6,9 @@ Last updated: |today| (API docstrings are auto-generated). VeRL-Omni reward pipelines support both rule-based scoring (e.g. JPEG compressibility) and model-based generative reward models (e.g. OCR via a vision-language model served behind an OpenAI-compatible router). Reward -computation is dispatched per sample by the -:class:`~verl_omni.reward_loop.reward_manager.VisualRewardManager`, which +computation is dispatched per sample by modality-specific reward managers, +including :class:`~verl_omni.reward_loop.reward_manager.VisualRewardManager` +and :class:`~verl_omni.reward_loop.reward_manager.AudioRewardManager`, which plugs into :class:`~verl_omni.reward_loop.reward_loop.OmniRewardLoopManager` — verl's :class:`~verl.experimental.reward_loop.RewardLoopManager` extended with profiler control over the reward-model rollout servers. @@ -17,8 +18,10 @@ profiler control over the reward-model rollout servers. verl_omni.reward_loop.reward_loop.OmniRewardLoopManager verl_omni.reward_loop.reward_manager.VisualRewardManager + verl_omni.reward_loop.reward_manager.AudioRewardManager verl_omni.utils.reward_score.default_compute_score_image verl_omni.utils.reward_score.http_scorer_client.compute_score + verl_omni.utils.reward_score.audio_http_scorer_client.compute_score verl_omni.utils.reward_score.unified_reward.compute_score_unified_reward Reward Loop Manager @@ -33,6 +36,13 @@ Reward Manager .. autoclass:: verl_omni.reward_loop.reward_manager.VisualRewardManager :members: __init__, run_single +.. autoclass:: verl_omni.reward_loop.reward_manager.AudioRewardManager + :members: __init__, run_single + +``AudioRewardManager`` reads ``audio`` and ``audio_sample_rate`` from rollout +``extra_info``, validates a finite CPU float waveform, and calls a synchronous +or asynchronous custom scorer with ``solution_audio=(waveform, sample_rate)``. + Default Score Dispatcher ~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -60,6 +70,12 @@ HTTP Scorer Client .. automodule:: verl_omni.utils.reward_score.http_scorer_client :members: compute_score +Audio HTTP Scorer Client +^^^^^^^^^^^^^^^^^^^^^^^^ + +.. automodule:: verl_omni.utils.reward_score.audio_http_scorer_client + :members: compute_score + UnifiedReward Scorer ^^^^^^^^^^^^^^^^^^^^^ diff --git a/docs/start/http_scorer.md b/docs/start/http_scorer.md index d978ae32e..cae6d9da2 100644 --- a/docs/start/http_scorer.md +++ b/docs/start/http_scorer.md @@ -5,6 +5,13 @@ Last updated: 08/09/2026 VeRL-Omni ships a generic HTTP reward client (`verl_omni.utils.reward_score.http_scorer_client`) that sends generated images to an external scorer service over HTTP and returns the score. This is useful when your reward model is too large to co-locate with training, needs a different runtime (e.g., a separate GPU pool), or is shared across multiple experiments. +Audio rollouts use `verl_omni.utils.reward_score.audio_http_scorer_client`. +That client sends JSON containing a base64-encoded float32 waveform, +`sample_rate`, target `prompt`, and scalar metadata. The service returns a JSON +object with a finite `score` and optional diagnostics. See the +[Qwen3-TTS GRPO example](../../examples/grpo_trainer/qwen3_tts/README.md) for +the complete audio protocol and configuration. + ## How it works ```text diff --git a/docs/start/models.md b/docs/start/models.md index d249000c1..e3a83fde3 100644 --- a/docs/start/models.md +++ b/docs/start/models.md @@ -223,6 +223,21 @@ rather than a separate per-stage YAML file. --- +### Qwen3-TTS-12Hz-0.6B Base + +| Property | Detail | +|----------|--------| +| **Hugging Face ID** | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | +| **Trainable component** | Talker codec-0 policy, full-parameter example | +| **Rollout** | Two-stage vLLM-Omni Talker + code2wav pipeline | +| **Algorithm** | Stock GRPO, vanilla PPO loss, optional direct KL | +| **Reward** | Generic decoded-audio reward; SpeechJudge-BTRM external scorer example | + +The example uses two training GPUs and an independently deployed audio scorer. +See [Qwen3-TTS GRPO with an audio reward](../../examples/grpo_trainer/qwen3_tts/README.md). + +--- + ## Model Architecture Summary | Model | Architecture | Text encoder | @@ -245,7 +260,8 @@ rather than a separate per-stage YAML file. | Qwen2.5-VL-3B-Instruct | `Qwen/Qwen2.5-VL-3B-Instruct` | Vision-Language | SD3.5 (Flow-GRPO) | vLLM, TP=1, dedicated pool | | PickScore | `yuvalkirstain/PickScore_v1` | Vision (preference) | Qwen-Image-Edit (Flow-GRPO), BAGEL (PickScore recipe) | Local CLIP load, async workers | | HPSv3 | Local `.safetensors` | Vision (aesthetic) | Wan2.2 (DanceGRPO) | Local safetensors load | -| HTTP scorer | External HTTP service | Any | Any model | Gunicorn/Flask, pickle protocol | +| HTTP scorer | External HTTP service | Image/audio | Any model | Pickle image or JSON audio protocol | +| SpeechJudge-BTRM | `RMSnow/SpeechJudge-BTRM` | Audio quality | Qwen3-TTS example | External service; CC-BY-NC-4.0 | | JPEG incompressibility | Rule-based | Image stats | Any diffusion model | No model process needed | For end-to-end instructions on setting up each reward, see the respective diff --git a/examples/grpo_trainer/qwen3_tts/README.md b/examples/grpo_trainer/qwen3_tts/README.md new file mode 100644 index 000000000..78cc5cfe5 --- /dev/null +++ b/examples/grpo_trainer/qwen3_tts/README.md @@ -0,0 +1,112 @@ +# Qwen3-TTS GRPO with an audio reward + +This example full-parameter tunes the codec-0 policy of +`Qwen/Qwen3-TTS-12Hz-0.6B-Base`. It uses verl's stock GRPO advantage, +vanilla PPO policy loss, and optional direct reference-model KL. The other +15 codec codebooks and code2wav stage remain frozen but are retained so every +candidate can be decoded and scored as audio. + +The launcher follows the V1 omni-model integration guide: it calls +`verl_omni.trainer.main_omni` and expresses the recipe as CLI overrides on the +standard `omni_trainer` config, without a model-specific Trainer or config tree. + +SpeechJudge-BTRM is one possible scorer. It runs behind the generic audio HTTP +client because its published Transformers environment conflicts with the +Transformers 5.x vLLM stack. The Trainer and `AudioRewardManager` do not contain +SpeechJudge-specific branches, candidate masks, ASR gates, or custom losses. + +## Install + +Install the engine before the training stack: + +```bash +uv pip install -e ".[gpu]" --torch-backend=auto +uv pip install "vllm-omni @ git+https://github.com/vllm-project/vllm-omni.git@$(cat .github/vllm_omni_pin.txt)" +uv pip install -e ".[tts,train,dev]" +uv pip install qwen-tts==0.1.1 --no-deps +``` + +The Qwen3-TTS adapter contains a bounded import compatibility layer for +`qwen-tts==0.1.1` on the repository's Transformers 5.x stack. It does not edit +site-packages. The system `sox` executable is also required by qwen-tts. + +## Data + +Training and validation parquet rows use the normal verl format: + +```python +{ + "data_source": "tts", + "prompt": [{"role": "user", "content": "Text to synthesize"}], + "reward_model": {"style": "model", "ground_truth": "Text to synthesize"}, + "extra_info": {"id": "stable-id", "split": "train"}, +} +``` + +Use disjoint prompts. The default recipe evaluates the same complete 100-row +validation parquet at step 0 and every 20 updates. Give validation rows fixed +`extra_info.generation_seed` values to keep candidate sampling paired across +checkpoints. + +The concatenated replay layout also requires one fixed speaker embedding JSON. +Generate it once with the official Qwen3-TTS Base model's +`extract_speaker_embedding` API from a 24 kHz reference recording, then reuse +the same file for the entire run. + +## Audio scorer protocol + +The configured endpoint receives one JSON request per candidate: + +```json +{ + "protocol_version": "1", + "waveform_f32_base64": "...", + "num_samples": 24000, + "sample_rate": 24000, + "prompt": "Text to synthesize", + "metadata": {"id": "stable-id"} +} +``` + +It must return `{"score": 1.25}` and may include additional scalar metrics. +The client retries only transient network, timeout, HTTP 408/429, and 5xx +failures. Missing, malformed, or non-finite results stop the run instead of +being converted to a valid zero reward. + +For SpeechJudge-BTRM, deploy the official +[`AmphionTeam/SpeechJudge`](https://github.com/AmphionTeam/SpeechJudge) code and +[`RMSnow/SpeechJudge-BTRM`](https://huggingface.co/RMSnow/SpeechJudge-BTRM) +checkpoint in a separate environment, then expose its pointwise score through +this protocol. Pin the SpeechJudge source revision and runtime versions in the +service deployment. SpeechJudge-BTRM is licensed CC-BY-NC-4.0. + +## Train + +```bash +MODEL_PATH=/path/to/Qwen3-TTS-12Hz-0.6B-Base \ +TRAIN_FILE=/path/to/train.parquet \ +VAL_FILE=/path/to/fixed-validation-100.parquet \ +SPK_EMBED_PATH=/path/to/speaker.json \ +SCORER_URL=http://scorer-host:18080/score \ +OUTPUT_DIR=/path/to/output \ +bash examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh +``` + +The example defaults are `B=4`, `G=8`, `lr=2e-7`, direct `low_var_kl` with +coefficient `0.12`, two GPUs, and 500 updates. These are recipe values, not +algorithm requirements. `norm_adv_by_std_in_grpo` remains at the upstream +default. + +For a two-update implementation smoke test: + +```bash +TOTAL_TRAINING_STEPS=2 TEST_FREQ=-1 SAVE_FREQ=1 RESUME_MODE=disable \ +OUTPUT_DIR=outputs/qwen3_tts_grpo_smoke \ +bash examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh \ + trainer.val_before_train=false trainer.log_val_generations=0 +``` + +This smoke proves rollout, finite audio reward, optimizer update, weight sync, +and checkpoint wiring only. It is not evidence that GRPO improves held-out +speech quality; that requires the complete fixed-validation curve and paired +human listening evaluation. diff --git a/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh b/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh new file mode 100755 index 000000000..d390c12cf --- /dev/null +++ b/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh @@ -0,0 +1,152 @@ +#!/usr/bin/env bash +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0 + +set -euo pipefail + +MODEL_PATH="${MODEL_PATH:?Set MODEL_PATH to a local Qwen3-TTS Base directory}" +TRAIN_FILE="${TRAIN_FILE:?Set TRAIN_FILE to the training parquet}" +VAL_FILE="${VAL_FILE:?Set VAL_FILE to the fixed validation parquet}" +SPK_EMBED_PATH="${SPK_EMBED_PATH:?Set SPK_EMBED_PATH to a speaker embedding JSON file}" +SCORER_URL="${SCORER_URL:?Set SCORER_URL to an audio scorer /score endpoint}" +OUTPUT_DIR="${OUTPUT_DIR:-outputs/qwen3_tts_grpo}" +NUM_GPUS="${NUM_GPUS:-2}" +SEED="${SEED:-42}" +TOTAL_TRAINING_STEPS="${TOTAL_TRAINING_STEPS:-500}" +TEST_FREQ="${TEST_FREQ:-20}" +SAVE_FREQ="${SAVE_FREQ:-20}" +RESUME_MODE="${RESUME_MODE:-auto}" +PYTHON_BIN="${PYTHON_BIN:-python3}" +for path in "${MODEL_PATH}" "${TRAIN_FILE}" "${VAL_FILE}" "${SPK_EMBED_PATH}"; do + [[ -e "${path}" ]] || { printf 'Missing required path: %s\n' "${path}" >&2; exit 2; } +done + +mkdir -p "${OUTPUT_DIR}" +export PYTHONHASHSEED="${SEED}" +export TOKENIZERS_PARALLELISM=false +export VERL_USE_EXTERNAL_MODULES=verl_omni +export VLLM_USE_FLASHINFER_SAMPLER=0 + +"${PYTHON_BIN}" -m verl_omni.trainer.main_omni \ + "data.train_files=${TRAIN_FILE}" \ + "data.val_files=${VAL_FILE}" \ + "data.seed=${SEED}" \ + data.train_batch_size=4 \ + data.dataloader_num_workers=0 \ + data.max_prompt_length=512 \ + data.max_response_length=360 \ + data.val_max_samples=100 \ + data.validation_shuffle=false \ + data.filter_overlong_prompts=true \ + data.truncation=error \ + "actor_rollout_ref.model.path=${MODEL_PATH}" \ + actor_rollout_ref.model.model_stage=talker \ + actor_rollout_ref.model.hf_config_name=talker_config \ + actor_rollout_ref.model.trust_remote_code=false \ + actor_rollout_ref.model.use_remove_padding=false \ + actor_rollout_ref.model.enable_gradient_checkpointing=true \ + +actor_rollout_ref.model.override_config.attn_implementation=eager \ + +actor_rollout_ref.model.override_config.tts_language=Auto \ + "+actor_rollout_ref.model.override_config.tts_spk_embed_path=${SPK_EMBED_PATH}" \ + actor_rollout_ref.actor.strategy=fsdp \ + actor_rollout_ref.actor.ppo_mini_batch_size=4 \ + actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.actor.use_dynamic_bsz=false \ + actor_rollout_ref.actor.use_kl_loss=true \ + actor_rollout_ref.actor.kl_loss_coef=0.12 \ + actor_rollout_ref.actor.kl_loss_type=low_var_kl \ + actor_rollout_ref.actor.ppo_epochs=1 \ + actor_rollout_ref.actor.shuffle=false \ + actor_rollout_ref.actor.clip_ratio_low=0.2 \ + actor_rollout_ref.actor.clip_ratio_high=0.2 \ + actor_rollout_ref.actor.entropy_coeff=0.0 \ + actor_rollout_ref.actor.loss_agg_mode=seq-mean-token-mean \ + actor_rollout_ref.actor.use_torch_compile=false \ + actor_rollout_ref.actor.policy_loss.loss_mode=vanilla \ + actor_rollout_ref.actor.optim.lr=2.0e-7 \ + actor_rollout_ref.actor.optim.lr_warmup_steps=0 \ + actor_rollout_ref.actor.optim.weight_decay=0.0 \ + actor_rollout_ref.actor.optim.clip_grad=1.0 \ + "actor_rollout_ref.actor.data_loader_seed=${SEED}" \ + "actor_rollout_ref.actor.fsdp_config.fsdp_size=${NUM_GPUS}" \ + "actor_rollout_ref.actor.fsdp_config.seed=${SEED}" \ + actor_rollout_ref.actor.fsdp_config.model_dtype=float32 \ + actor_rollout_ref.actor.fsdp_config.dtype=float32 \ + actor_rollout_ref.actor.fsdp_config.param_offload=false \ + actor_rollout_ref.actor.fsdp_config.optimizer_offload=false \ + actor_rollout_ref.actor.fsdp_config.use_orig_params=true \ + actor_rollout_ref.actor.fsdp_config.use_torch_compile=false \ + actor_rollout_ref.actor.fsdp_config.wrap_policy.min_num_params=0 \ + actor_rollout_ref.rollout.name=vllm_omni \ + actor_rollout_ref.rollout.mode=async \ + actor_rollout_ref.rollout.n=8 \ + actor_rollout_ref.rollout.temperature=1.0 \ + actor_rollout_ref.rollout.top_p=1.0 \ + actor_rollout_ref.rollout.top_k=-1 \ + actor_rollout_ref.rollout.dtype=float32 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.gpu_memory_utilization=0.20 \ + actor_rollout_ref.rollout.max_num_seqs=8 \ + actor_rollout_ref.rollout.max_num_batched_tokens=1024 \ + actor_rollout_ref.rollout.free_cache_engine=false \ + actor_rollout_ref.rollout.calculate_log_probs=true \ + actor_rollout_ref.rollout.logprobs_mode=processed_logprobs \ + actor_rollout_ref.rollout.load_format=safetensors \ + actor_rollout_ref.rollout.layered_summon=true \ + actor_rollout_ref.rollout.enable_prefix_caching=false \ + actor_rollout_ref.rollout.enforce_eager=true \ + actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.rollout.agent.num_workers=8 \ + actor_rollout_ref.rollout.agent.default_agent_loop=qwen3_tts_single_turn \ + "+actor_rollout_ref.rollout.engine_kwargs.vllm_omni.seed=${SEED}" \ + +actor_rollout_ref.rollout.engine_kwargs.vllm_omni.output_mode=ar \ + +actor_rollout_ref.rollout.engine_kwargs.vllm_omni.pipeline_name=qwen3_tts_rl \ + +actor_rollout_ref.rollout.engine_kwargs.vllm_omni.pipeline_mode=full \ + +actor_rollout_ref.rollout.engine_kwargs.vllm_omni.async_chunk=false \ + +actor_rollout_ref.rollout.engine_kwargs.vllm_omni.attention_config.backend=TRITON_ATTN \ + actor_rollout_ref.rollout.val_kwargs.n=1 \ + actor_rollout_ref.rollout.val_kwargs.do_sample=true \ + actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \ + actor_rollout_ref.rollout.val_kwargs.top_p=1.0 \ + actor_rollout_ref.rollout.val_kwargs.top_k=-1 \ + actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \ + actor_rollout_ref.ref.strategy=fsdp \ + "actor_rollout_ref.ref.fsdp_config.fsdp_size=${NUM_GPUS}" \ + "actor_rollout_ref.ref.fsdp_config.seed=${SEED}" \ + actor_rollout_ref.ref.fsdp_config.model_dtype=float32 \ + actor_rollout_ref.ref.fsdp_config.dtype=float32 \ + actor_rollout_ref.ref.fsdp_config.param_offload=false \ + actor_rollout_ref.ref.fsdp_config.use_orig_params=true \ + actor_rollout_ref.ref.fsdp_config.wrap_policy.min_num_params=0 \ + algorithm.adv_estimator=grpo \ + algorithm.norm_adv_by_std_in_grpo=true \ + algorithm.use_kl_in_reward=false \ + reward.num_workers=1 \ + reward.custom_reward_function.path=pkg://verl_omni.utils.reward_score.audio_http_scorer_client \ + reward.custom_reward_function.name=compute_score \ + "+reward.custom_reward_function.reward_kwargs.server_url=${SCORER_URL}" \ + +reward.custom_reward_function.reward_kwargs.timeout_s=120.0 \ + +reward.custom_reward_function.reward_kwargs.max_retries=2 \ + +reward.custom_reward_function.reward_kwargs.retry_backoff_s=0.5 \ + reward.reward_manager.source=importlib \ + reward.reward_manager.name=AudioRewardManager \ + reward.reward_manager.module.path=pkg://verl_omni.reward_loop.reward_manager \ + trainer.val_before_train=true \ + trainer.balance_batch=true \ + trainer.critic_warmup=0 \ + trainer.logger='["console"]' \ + trainer.project_name=qwen3_tts_grpo \ + trainer.experiment_name=qwen3_tts_0_6b_audio_reward \ + "trainer.n_gpus_per_node=${NUM_GPUS}" \ + trainer.nnodes=1 \ + trainer.total_epochs=100 \ + trainer.log_val_generations=100 \ + "trainer.total_training_steps=${TOTAL_TRAINING_STEPS}" \ + "trainer.test_freq=${TEST_FREQ}" \ + "trainer.save_freq=${SAVE_FREQ}" \ + "trainer.resume_mode=${RESUME_MODE}" \ + trainer.max_actor_ckpt_to_keep=26 \ + "trainer.default_local_dir=${OUTPUT_DIR}/checkpoints" \ + "trainer.validation_data_dir=${OUTPUT_DIR}/validation" \ + "trainer.rollout_data_dir=${OUTPUT_DIR}/rollout" \ + "$@" 2>&1 | tee "${OUTPUT_DIR}/train.log" diff --git a/tests/reward_loop/test_audio_reward_manager_on_cpu.py b/tests/reward_loop/test_audio_reward_manager_on_cpu.py new file mode 100644 index 000000000..547b1289d --- /dev/null +++ b/tests/reward_loop/test_audio_reward_manager_on_cpu.py @@ -0,0 +1,191 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU tests for generic waveform reward routing.""" + +import importlib.util +from pathlib import Path +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch +from omegaconf import OmegaConf +from verl import DataProto + + +def _load_audio_reward_manager(): + path = Path(__file__).parents[2] / "verl_omni/reward_loop/reward_manager/audio.py" + spec = importlib.util.spec_from_file_location("audio_reward_manager_under_test", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module.AudioRewardManager + + +AudioRewardManager = _load_audio_reward_manager() + + +def _config(): + return OmegaConf.create({"reward": {}}) + + +def _manager(compute_score): + return AudioRewardManager(_config(), MagicMock(), compute_score=compute_score) + + +def _data(audio=None, sample_rate=24_000): + fields = {} + if audio is not None: + fields = {"audio": audio, "audio_sample_rate": sample_rate} + return DataProto.from_dict( + tensors={"responses": torch.zeros(1, 4, dtype=torch.long)}, + non_tensors={ + "data_source": ["tts_reward"], + "reward_model": [{"ground_truth": "ni3 hao3"}], + "extra_info": [{"id": "sample-0"}], + "tool_extra_fields": [fields], + }, + ) + + +def test_assemble_scores_preserves_one_reward_per_sample(): + data = DataProto.from_dict( + tensors={ + "prompts": torch.zeros(3, 2, dtype=torch.long), + "responses": torch.zeros(3, 4, dtype=torch.long), + "attention_mask": torch.tensor( + [ + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 0, 0], + [1, 1, 1, 1, 1, 0], + ] + ), + } + ) + + scores = AudioRewardManager.assemble_rm_scores(data, [0.1, -1.0, 2.5]) + + assert scores.shape == (3, 4) + assert scores.dtype == torch.float32 + torch.testing.assert_close( + scores, + torch.tensor( + [ + [0.0, 0.0, 0.0, 0.1], + [0.0, -1.0, 0.0, 0.0], + [0.0, 0.0, 2.5, 0.0], + ] + ), + ) + + +def test_run_single_rejects_multi_sample_batch(): + data = DataProto.from_dict( + tensors={"responses": torch.zeros(2, 4, dtype=torch.long)}, + non_tensors={ + "data_source": ["tts_reward", "tts_reward"], + "reward_model": [{"ground_truth": "first"}, {"ground_truth": "second"}], + }, + ) + manager = _manager(lambda **kwargs: 0.0) + + with pytest.raises(ValueError, match="batch size 2"): + manager.loop.run_until_complete(manager.run_single(data)) + + +def test_run_single_passes_waveform_and_returns_diagnostics(): + def compute_score(data_source, solution_audio, ground_truth, extra_info): + waveform, sample_rate = solution_audio + assert data_source == "tts_reward" + assert waveform.dtype == np.float32 + assert waveform.shape == (24_000,) + assert sample_rate == 24_000 + assert ground_truth == "ni3 hao3" + assert extra_info["id"] == "sample-0" + return {"score": 0.75, "pinyin_error_rate": 0.1} + + manager = _manager(compute_score) + result = manager.loop.run_until_complete(manager.run_single(_data(np.ones(24_000, dtype=np.float32)))) + + assert result == { + "reward_score": 0.75, + "reward_extra_info": {"pinyin_error_rate": 0.1}, + } + + +@pytest.mark.parametrize( + ("data", "message"), + [ + (_data(), "requires extra_info\\['audio'\\]"), + (_data([], 24_000), "empty waveform"), + (_data([0.0, float("nan")], 24_000), "NaN or infinity"), + (_data([0.0], 0), "positive integer"), + (_data([0.0], 24_000.5), "positive integer"), + ], +) +def test_invalid_audio_fails_closed(data, message): + manager = _manager(lambda **kwargs: 0.0) + + with pytest.raises((KeyError, ValueError), match=message): + manager.loop.run_until_complete(manager.run_single(data)) + + +def test_missing_sample_rate_fails_closed(): + data = _data([0.0]) + data.non_tensor_batch["tool_extra_fields"][0].pop("audio_sample_rate") + manager = _manager(lambda **kwargs: 0.0) + + with pytest.raises(KeyError, match="audio_sample_rate"): + manager.loop.run_until_complete(manager.run_single(data)) + + +def test_list_of_chunks_is_concatenated(): + def compute_score(solution_audio, **kwargs): + waveform, sample_rate = solution_audio + np.testing.assert_array_equal(waveform, np.array([0.1, 0.2, 0.3], dtype=np.float32)) + assert sample_rate == 16_000 + return 0.5 + + manager = _manager(compute_score) + result = manager.loop.run_until_complete( + manager.run_single(_data([torch.tensor([0.1, 0.2]), torch.tensor([0.3])], 16_000)) + ) + + assert result == {"reward_score": 0.5, "reward_extra_info": {"acc": 0.5}} + + +@pytest.mark.asyncio +async def test_async_score_function_is_supported(): + async def compute_score(solution_audio, **kwargs): + assert solution_audio[1] == 24_000 + return {"score": -0.25, "judge_margin": 3.0} + + result = await _manager(compute_score).run_single(_data(torch.ones(32))) + + assert result == {"reward_score": -0.25, "reward_extra_info": {"judge_margin": 3.0}} + + +@pytest.mark.parametrize("score", [float("nan"), float("inf"), -float("inf")]) +def test_non_finite_reward_fails_closed(score): + manager = _manager(lambda **kwargs: score) + + with pytest.raises(ValueError, match="must be finite"): + manager.loop.run_until_complete(manager.run_single(_data([0.0]))) + + +def test_reward_dictionary_requires_score(): + manager = _manager(lambda **kwargs: {"metric": 1.0}) + + with pytest.raises(ValueError, match="missing 'score'"): + manager.loop.run_until_complete(manager.run_single(_data([0.0]))) diff --git a/tests/utils/reward_score/test_audio_http_scorer_client_on_cpu.py b/tests/utils/reward_score/test_audio_http_scorer_client_on_cpu.py new file mode 100644 index 000000000..87c25c24a --- /dev/null +++ b/tests/utils/reward_score/test_audio_http_scorer_client_on_cpu.py @@ -0,0 +1,158 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU contracts for the external audio reward client.""" + +import asyncio +import base64 +import importlib.util +import socket +from pathlib import Path + +import numpy as np +import pytest +from aiohttp import web + + +def _load_client_module(): + path = Path(__file__).parents[3] / "verl_omni/utils/reward_score/audio_http_scorer_client.py" + spec = importlib.util.spec_from_file_location("audio_http_scorer_client_under_test", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +client = _load_client_module() + + +def test_request_serialization_preserves_waveform_prompt_and_scalar_metadata(): + waveform = np.array([0.0, -0.25, 0.5], dtype=np.float32) + payload = client._serialize_request( + (waveform, 24_000), + "target text", + {"id": "sample-1", "seed": np.array(7), "audio": waveform, "ignored": [1, 2]}, + ) + + decoded = np.frombuffer(base64.b64decode(payload["waveform_f32_base64"]), dtype=" 1 and waveform.shape[0] == 1: + waveform = waveform[0] + if waveform.ndim == 2: + waveform = waveform.mean(dim=0) + elif waveform.ndim != 1: + raise ValueError(f"Expected audio shape (T,) or (C,T), got {tuple(waveform.shape)}.") + if waveform.numel() == 0: + raise ValueError("Audio reward received an empty waveform.") + if not torch.isfinite(waveform).all(): + raise ValueError("Audio reward received a waveform containing NaN or infinity.") + + if isinstance(sample_rate, list | tuple): + if len(sample_rate) != 1: + raise ValueError("Audio reward requires exactly one sample rate per waveform.") + sample_rate = sample_rate[0] + if hasattr(sample_rate, "item"): + try: + sample_rate = sample_rate.item() + except (RuntimeError, ValueError) as exc: + raise ValueError("Audio reward requires one scalar sample rate per waveform.") from exc + if isinstance(sample_rate, bool) or not isinstance(sample_rate, int | float): + raise TypeError(f"Audio sample rate must be numeric, got {type(sample_rate).__name__}.") + if not math.isfinite(float(sample_rate)) or float(sample_rate) <= 0 or float(sample_rate) != int(sample_rate): + raise ValueError(f"Audio sample rate must be a positive integer, got {sample_rate!r}.") + return waveform.numpy().astype(np.float32, copy=False), int(sample_rate) + + async def run_single(self, data: DataProto) -> dict: + if len(data) != 1: + raise ValueError(f"AudioRewardManager scores one sample at a time, got batch size {len(data)}.") + item = data[0] + batch = item.non_tensor_batch + extra_info = self._mapping(batch.get("extra_info", {})) + extra_info.update(self._mapping(batch.get("tool_extra_fields"))) + extra_info["num_turns"] = batch.get("__num_turns__", extra_info.get("num_turns")) + extra_info["global_steps"] = batch.get("global_steps", extra_info.get("global_steps", 0)) + ground_truth = batch["reward_model"]["ground_truth"] + audio = self._extract_audio(item, extra_info) + kwargs = { + "data_source": batch["data_source"], + "solution_audio": audio, + "ground_truth": ground_truth, + "extra_info": extra_info, + } + if self.is_async_reward_score: + result = await self.compute_score(**kwargs) + else: + result = await self.loop.run_in_executor(None, lambda: self.compute_score(**kwargs)) + if isinstance(result, dict): + if "score" not in result: + raise ValueError("Audio reward result dictionary is missing 'score'.") + score = float(result["score"]) + reward_extra_info = {key: value for key, value in result.items() if key != "score"} + else: + score = float(result) + reward_extra_info = {"acc": score} + if not math.isfinite(score): + raise ValueError(f"Audio reward must be finite, got {score!r}.") + return {"reward_score": score, "reward_extra_info": reward_extra_info} diff --git a/verl_omni/utils/reward_score/audio_http_scorer_client.py b/verl_omni/utils/reward_score/audio_http_scorer_client.py new file mode 100644 index 000000000..29ff3d3d9 --- /dev/null +++ b/verl_omni/utils/reward_score/audio_http_scorer_client.py @@ -0,0 +1,175 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""JSON HTTP client for audio reward services in a separate runtime.""" + +import asyncio +import base64 +import math +from typing import Any + +import aiohttp +import numpy as np + +PROTOCOL_VERSION = "1" + + +class _RetryableHTTPError(RuntimeError): + pass + + +def _scalar_metadata(extra_info: dict | None) -> dict[str, str | int | float | bool | None]: + metadata = {} + for key, value in (extra_info or {}).items(): + if key in {"audio", "audio_sample_rate"}: + continue + if isinstance(value, np.ndarray) and value.shape == (): + value = value.item() + elif hasattr(value, "item"): + try: + value = value.item() + except (RuntimeError, ValueError): + continue + if value is None or isinstance(value, str | int | bool): + metadata[str(key)] = value + elif isinstance(value, float) and math.isfinite(value): + metadata[str(key)] = value + return metadata + + +def _serialize_request(solution_audio, ground_truth: str, extra_info: dict | None) -> dict[str, Any]: + if not isinstance(solution_audio, tuple) or len(solution_audio) != 2: + raise TypeError("Audio HTTP scorer expects solution_audio=(waveform, sample_rate).") + waveform, sample_rate = solution_audio + waveform = np.asarray(waveform, dtype=" dict: + if not isinstance(payload, dict): + raise RuntimeError("Audio scorer response must be a JSON object.") + if "error" in payload: + raise RuntimeError(f"Audio scorer error: {payload['error']}") + if "score" not in payload: + raise RuntimeError("Audio scorer response is missing 'score'.") + try: + score = float(payload["score"]) + except (TypeError, ValueError) as exc: + raise RuntimeError(f"Audio scorer returned an invalid score: {payload['score']!r}.") from exc + if not math.isfinite(score): + raise RuntimeError(f"Audio scorer returned a non-finite score: {score!r}.") + + diagnostics = {} + for key, value in payload.items(): + if key == "score": + continue + if value is None or isinstance(value, str | int | bool): + diagnostics[str(key)] = value + elif isinstance(value, float) and math.isfinite(value): + diagnostics[str(key)] = value + else: + raise RuntimeError(f"Audio scorer diagnostic {key!r} must be a finite JSON scalar, got {value!r}.") + return {"score": score, **diagnostics} + + +async def _session() -> aiohttp.ClientSession: + loop = asyncio.get_running_loop() + session = getattr(compute_score, "_session", None) + session_loop = getattr(compute_score, "_session_loop", None) + if session is None or session.closed or session_loop is not loop: + if session is not None and not session.closed: + await session.close() + session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=None)) + compute_score._session = session + compute_score._session_loop = loop + return session + + +async def _request_score(server_url: str, payload: dict, timeout_s: float) -> dict: + session = await _session() + try: + async with session.post( + server_url, + json=payload, + timeout=aiohttp.ClientTimeout(total=timeout_s), + ) as response: + if response.status != 200: + detail = await response.text() + error = f"Audio scorer returned HTTP {response.status}: {detail}" + if response.status in {408, 429} or 500 <= response.status < 600: + raise _RetryableHTTPError(error) + raise RuntimeError(error) + try: + result = await response.json(content_type=None) + except (aiohttp.ContentTypeError, ValueError) as exc: + raise RuntimeError("Audio scorer returned malformed JSON.") from exc + except asyncio.TimeoutError as exc: + raise _RetryableHTTPError(f"Audio scorer timed out after {timeout_s} seconds.") from exc + return _validate_response(result) + + +async def compute_score( + solution_audio, + ground_truth: str, + extra_info: dict | None = None, + *, + server_url: str, + timeout_s: float = 120.0, + max_retries: int = 2, + retry_backoff_s: float = 0.5, + **kwargs, +) -> dict: + """Send one waveform to an external scorer and return its finite score.""" + del kwargs + if isinstance(timeout_s, bool) or not isinstance(timeout_s, int | float) or not math.isfinite(float(timeout_s)): + raise ValueError("timeout_s must be a finite number.") + if timeout_s <= 0: + raise ValueError("timeout_s must be positive.") + if isinstance(max_retries, bool) or not isinstance(max_retries, int): + raise ValueError("max_retries must be an integer.") + if max_retries < 0: + raise ValueError("max_retries must be non-negative.") + if ( + isinstance(retry_backoff_s, bool) + or not isinstance(retry_backoff_s, int | float) + or not math.isfinite(float(retry_backoff_s)) + ): + raise ValueError("retry_backoff_s must be a finite number.") + if retry_backoff_s < 0: + raise ValueError("retry_backoff_s must be non-negative.") + payload = _serialize_request(solution_audio, ground_truth, extra_info) + + last_error = None + for attempt in range(max_retries + 1): + try: + return await _request_score(server_url, payload, timeout_s) + except (_RetryableHTTPError, aiohttp.ClientConnectionError, aiohttp.ClientPayloadError) as exc: + last_error = exc + if attempt < max_retries: + await asyncio.sleep(retry_backoff_s * (2**attempt)) + raise RuntimeError(f"Audio scoring failed after {max_retries + 1} attempts: {last_error}") from last_error From eb66ec756f85b3809567f7263451666d6eb4c2cf Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Fri, 21 Aug 2026 10:57:16 +0800 Subject: [PATCH 03/28] [rollout] fix: align Qwen3-TTS prompt embedding dtype Signed-off-by: dongbo910220 <1275604947@qq.com> Co-authored-by: Michael-Zzq <52185141+Michael-Zzq@users.noreply.github.com> --- .../test_qwen3_tts_rollout_on_cpu.py | 40 +++++++++++++++++++ .../qwen3_tts/omni_rollout_adapter.py | 6 ++- .../pipelines/qwen3_tts/rollout_model.py | 33 +++++++++++++++ 3 files changed, 78 insertions(+), 1 deletion(-) create mode 100644 verl_omni/pipelines/qwen3_tts/rollout_model.py diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index 61512b20d..56649784c 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -23,7 +23,9 @@ from vllm import SamplingParams +from verl_omni.pipelines.qwen3_tts import omni_rollout_adapter from verl_omni.pipelines.qwen3_tts.omni_rollout_adapter import Qwen3TTSRolloutAdapter +from verl_omni.pipelines.qwen3_tts.rollout_model import _align_prompt_embedding_dtype from verl_omni.pipelines.qwen3_tts.talker_training_adapter import Qwen3TTSTalkerAdapter from verl_omni.workers.rollout.vllm_rollout.utils import _receive_model_weight_buckets from verl_omni.workers.rollout.vllm_rollout.vllm_omni_async_server import vLLMOmniHttpServer @@ -37,6 +39,44 @@ def __call__(self, text, **kwargs): return {"input_ids": list(range(len(text)))} +def test_rollout_pipeline_registers_dtype_aligned_talker(monkeypatch): + registered_models = [] + registered_pipelines = [] + monkeypatch.setattr( + omni_rollout_adapter.ModelRegistry, + "register_model", + lambda architecture, model_class: registered_models.append((architecture, model_class)), + ) + monkeypatch.setattr( + omni_rollout_adapter, + "register_pipeline", + lambda pipeline: registered_pipelines.append(pipeline), + ) + + Qwen3TTSRolloutAdapter.ensure_pipeline_registered() + + assert registered_models == [ + ( + "Qwen3TTSDtypeAlignedTalkerForConditionalGeneration", + "verl_omni.pipelines.qwen3_tts.rollout_model:Qwen3TTSDtypeAlignedTalkerForConditionalGeneration", + ) + ] + assert registered_pipelines == [omni_rollout_adapter.QWEN3_TTS_RL_PIPELINE] + assert registered_pipelines[0].model_arch == registered_models[0][0] + + +def test_rollout_model_aligns_talker_and_prompt_builder_embedding_dtype(): + model = SimpleNamespace( + _embedding_dtype=torch.bfloat16, + _prompt_builder=SimpleNamespace(_embedding_dtype=torch.bfloat16), + ) + + _align_prompt_embedding_dtype(model, torch.float32) + + assert model._embedding_dtype == torch.float32 + assert model._prompt_builder._embedding_dtype == torch.float32 + + def test_rollout_adapter_builds_unique_prompt_and_scopes_weight_sync(tmp_path): speaker = tmp_path / "speaker.json" speaker.write_text("[0.0, 1.0]") diff --git a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py index 1d0f3bfef..b2b94f4f7 100644 --- a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py @@ -20,6 +20,7 @@ from functools import lru_cache import torch +from vllm.model_executor.models import ModelRegistry from vllm_omni.config.pipeline_registry import register_pipeline from vllm_omni.config.stage_config import PipelineConfig from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE @@ -33,10 +34,12 @@ ) _PIPELINE_ID = "qwen3_tts_rl" +_ROLLOUT_MODEL_ARCH = "Qwen3TTSDtypeAlignedTalkerForConditionalGeneration" +_ROLLOUT_MODEL_CLASS = "verl_omni.pipelines.qwen3_tts.rollout_model:Qwen3TTSDtypeAlignedTalkerForConditionalGeneration" _SYNC_PROCESSOR = "verl_omni.pipelines.qwen3_tts.omni_rollout_adapter.talker2code2wav_token_only" QWEN3_TTS_RL_PIPELINE = PipelineConfig( model_type=_PIPELINE_ID, - model_arch=QWEN3_TTS_PIPELINE.model_arch, + model_arch=_ROLLOUT_MODEL_ARCH, stages=( replace(QWEN3_TTS_PIPELINE.stages[0], final_output=True, final_output_type="latent"), replace(QWEN3_TTS_PIPELINE.stages[1], sync_process_input_func=_SYNC_PROCESSOR), @@ -103,6 +106,7 @@ def get_pipeline_id(cls, pipeline_mode="full"): @classmethod def ensure_pipeline_registered(cls, pipeline_mode="full"): cls._check_mode(pipeline_mode) + ModelRegistry.register_model(_ROLLOUT_MODEL_ARCH, _ROLLOUT_MODEL_CLASS) register_pipeline(QWEN3_TTS_RL_PIPELINE) @classmethod diff --git a/verl_omni/pipelines/qwen3_tts/rollout_model.py b/verl_omni/pipelines/qwen3_tts/rollout_model.py new file mode 100644 index 000000000..00d14b8ef --- /dev/null +++ b/verl_omni/pipelines/qwen3_tts/rollout_model.py @@ -0,0 +1,33 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Qwen3-TTS rollout model extensions.""" + +import torch +from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( + Qwen3TTSTalkerForConditionalGeneration, +) + + +def _align_prompt_embedding_dtype(model, dtype: torch.dtype) -> None: + """Keep request embeddings compatible with the vLLM model input buffer.""" + model._embedding_dtype = dtype + model._prompt_builder._embedding_dtype = dtype + + +class Qwen3TTSDtypeAlignedTalkerForConditionalGeneration(Qwen3TTSTalkerForConditionalGeneration): + """Make Qwen3-TTS prompt embeddings follow the configured rollout dtype.""" + + def __init__(self, *, vllm_config, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + _align_prompt_embedding_dtype(self, vllm_config.model_config.dtype) From 498d7a883575dae389557faa0b9ee4df181550c0 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Fri, 21 Aug 2026 11:09:16 +0800 Subject: [PATCH 04/28] [rollout] fix: register Qwen3-TTS model in vLLM workers Signed-off-by: dongbo910220 <1275604947@qq.com> Co-authored-by: Michael-Zzq <52185141+Michael-Zzq@users.noreply.github.com> --- pyproject.toml | 3 +++ .../test_qwen3_tts_rollout_on_cpu.py | 4 ++-- .../qwen3_tts/omni_rollout_adapter.py | 8 +++---- verl_omni/pipelines/qwen3_tts/vllm_plugin.py | 24 +++++++++++++++++++ 4 files changed, 32 insertions(+), 7 deletions(-) create mode 100644 verl_omni/pipelines/qwen3_tts/vllm_plugin.py diff --git a/pyproject.toml b/pyproject.toml index dc094060d..bdbec51c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,9 @@ dependencies = [ "TransferQueue==0.1.8", ] +[project.entry-points."vllm.general_plugins"] +verl_omni_qwen3_tts_rollout = "verl_omni.pipelines.qwen3_tts.vllm_plugin:register_qwen3_tts_rollout_model" + [project.optional-dependencies] # Rollout engine (install in step 2, after gpu step 1). vllm-omni = [ diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index 56649784c..c7becf9dc 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -23,7 +23,7 @@ from vllm import SamplingParams -from verl_omni.pipelines.qwen3_tts import omni_rollout_adapter +from verl_omni.pipelines.qwen3_tts import omni_rollout_adapter, vllm_plugin from verl_omni.pipelines.qwen3_tts.omni_rollout_adapter import Qwen3TTSRolloutAdapter from verl_omni.pipelines.qwen3_tts.rollout_model import _align_prompt_embedding_dtype from verl_omni.pipelines.qwen3_tts.talker_training_adapter import Qwen3TTSTalkerAdapter @@ -43,7 +43,7 @@ def test_rollout_pipeline_registers_dtype_aligned_talker(monkeypatch): registered_models = [] registered_pipelines = [] monkeypatch.setattr( - omni_rollout_adapter.ModelRegistry, + vllm_plugin.ModelRegistry, "register_model", lambda architecture, model_class: registered_models.append((architecture, model_class)), ) diff --git a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py index b2b94f4f7..3bca2c854 100644 --- a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py @@ -20,7 +20,6 @@ from functools import lru_cache import torch -from vllm.model_executor.models import ModelRegistry from vllm_omni.config.pipeline_registry import register_pipeline from vllm_omni.config.stage_config import PipelineConfig from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE @@ -32,14 +31,13 @@ load_speaker_xvector, require_auto_language, ) +from verl_omni.pipelines.qwen3_tts.vllm_plugin import ROLLOUT_MODEL_ARCH, register_qwen3_tts_rollout_model _PIPELINE_ID = "qwen3_tts_rl" -_ROLLOUT_MODEL_ARCH = "Qwen3TTSDtypeAlignedTalkerForConditionalGeneration" -_ROLLOUT_MODEL_CLASS = "verl_omni.pipelines.qwen3_tts.rollout_model:Qwen3TTSDtypeAlignedTalkerForConditionalGeneration" _SYNC_PROCESSOR = "verl_omni.pipelines.qwen3_tts.omni_rollout_adapter.talker2code2wav_token_only" QWEN3_TTS_RL_PIPELINE = PipelineConfig( model_type=_PIPELINE_ID, - model_arch=_ROLLOUT_MODEL_ARCH, + model_arch=ROLLOUT_MODEL_ARCH, stages=( replace(QWEN3_TTS_PIPELINE.stages[0], final_output=True, final_output_type="latent"), replace(QWEN3_TTS_PIPELINE.stages[1], sync_process_input_func=_SYNC_PROCESSOR), @@ -106,7 +104,7 @@ def get_pipeline_id(cls, pipeline_mode="full"): @classmethod def ensure_pipeline_registered(cls, pipeline_mode="full"): cls._check_mode(pipeline_mode) - ModelRegistry.register_model(_ROLLOUT_MODEL_ARCH, _ROLLOUT_MODEL_CLASS) + register_qwen3_tts_rollout_model() register_pipeline(QWEN3_TTS_RL_PIPELINE) @classmethod diff --git a/verl_omni/pipelines/qwen3_tts/vllm_plugin.py b/verl_omni/pipelines/qwen3_tts/vllm_plugin.py new file mode 100644 index 000000000..6582e974c --- /dev/null +++ b/verl_omni/pipelines/qwen3_tts/vllm_plugin.py @@ -0,0 +1,24 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""vLLM model registration for the Qwen3-TTS rollout extension.""" + +from vllm.model_executor.models import ModelRegistry + +ROLLOUT_MODEL_ARCH = "Qwen3TTSDtypeAlignedTalkerForConditionalGeneration" +ROLLOUT_MODEL_CLASS = "verl_omni.pipelines.qwen3_tts.rollout_model:Qwen3TTSDtypeAlignedTalkerForConditionalGeneration" + + +def register_qwen3_tts_rollout_model() -> None: + """Register lazily so every vLLM engine and worker process can resolve it.""" + ModelRegistry.register_model(ROLLOUT_MODEL_ARCH, ROLLOUT_MODEL_CLASS) From 7205e4cd4db356427679110966fc29b9782446d7 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Fri, 21 Aug 2026 11:20:43 +0800 Subject: [PATCH 05/28] [rollout] fix: initialize Qwen3-TTS rollout workers Signed-off-by: dongbo910220 <1275604947@qq.com> Co-authored-by: Michael-Zzq <52185141+Michael-Zzq@users.noreply.github.com> --- pyproject.toml | 3 - .../test_qwen3_tts_rollout_on_cpu.py | 66 ++++++++++++++----- verl_omni/pipelines/model_base.py | 10 +++ .../qwen3_tts/omni_rollout_adapter.py | 21 ++++-- verl_omni/pipelines/qwen3_tts/vllm_plugin.py | 24 ------- .../{rollout_model.py => worker_extension.py} | 20 +++--- .../vllm_rollout/vllm_omni_async_server.py | 5 ++ 7 files changed, 90 insertions(+), 59 deletions(-) delete mode 100644 verl_omni/pipelines/qwen3_tts/vllm_plugin.py rename verl_omni/pipelines/qwen3_tts/{rollout_model.py => worker_extension.py} (58%) diff --git a/pyproject.toml b/pyproject.toml index bdbec51c4..dc094060d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,9 +49,6 @@ dependencies = [ "TransferQueue==0.1.8", ] -[project.entry-points."vllm.general_plugins"] -verl_omni_qwen3_tts_rollout = "verl_omni.pipelines.qwen3_tts.vllm_plugin:register_qwen3_tts_rollout_model" - [project.optional-dependencies] # Rollout engine (install in step 2, after gpu step 1). vllm-omni = [ diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index c7becf9dc..57b8e303a 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -23,10 +23,13 @@ from vllm import SamplingParams -from verl_omni.pipelines.qwen3_tts import omni_rollout_adapter, vllm_plugin +from verl_omni.pipelines.qwen3_tts import omni_rollout_adapter from verl_omni.pipelines.qwen3_tts.omni_rollout_adapter import Qwen3TTSRolloutAdapter -from verl_omni.pipelines.qwen3_tts.rollout_model import _align_prompt_embedding_dtype from verl_omni.pipelines.qwen3_tts.talker_training_adapter import Qwen3TTSTalkerAdapter +from verl_omni.pipelines.qwen3_tts.worker_extension import ( + Qwen3TTSColocateWorkerExtension, + _align_prompt_embedding_dtype, +) from verl_omni.workers.rollout.vllm_rollout.utils import _receive_model_weight_buckets from verl_omni.workers.rollout.vllm_rollout.vllm_omni_async_server import vLLMOmniHttpServer @@ -39,14 +42,8 @@ def __call__(self, text, **kwargs): return {"input_ids": list(range(len(text)))} -def test_rollout_pipeline_registers_dtype_aligned_talker(monkeypatch): - registered_models = [] +def test_rollout_pipeline_registers_upstream_talker(monkeypatch): registered_pipelines = [] - monkeypatch.setattr( - vllm_plugin.ModelRegistry, - "register_model", - lambda architecture, model_class: registered_models.append((architecture, model_class)), - ) monkeypatch.setattr( omni_rollout_adapter, "register_pipeline", @@ -55,14 +52,8 @@ def test_rollout_pipeline_registers_dtype_aligned_talker(monkeypatch): Qwen3TTSRolloutAdapter.ensure_pipeline_registered() - assert registered_models == [ - ( - "Qwen3TTSDtypeAlignedTalkerForConditionalGeneration", - "verl_omni.pipelines.qwen3_tts.rollout_model:Qwen3TTSDtypeAlignedTalkerForConditionalGeneration", - ) - ] assert registered_pipelines == [omni_rollout_adapter.QWEN3_TTS_RL_PIPELINE] - assert registered_pipelines[0].model_arch == registered_models[0][0] + assert registered_pipelines[0].model_arch == omni_rollout_adapter.QWEN3_TTS_PIPELINE.model_arch def test_rollout_model_aligns_talker_and_prompt_builder_embedding_dtype(): @@ -77,6 +68,34 @@ def test_rollout_model_aligns_talker_and_prompt_builder_embedding_dtype(): assert model._prompt_builder._embedding_dtype == torch.float32 +def test_worker_extension_aligns_loaded_model_dtype(): + model = SimpleNamespace( + _embedding_dtype=torch.bfloat16, + _prompt_builder=SimpleNamespace(_embedding_dtype=torch.bfloat16), + ) + worker = SimpleNamespace( + _get_standard_weight_model_and_config=lambda: (model, SimpleNamespace(dtype=torch.float32)) + ) + + Qwen3TTSColocateWorkerExtension.align_qwen3_tts_prompt_embedding_dtype(worker) + + assert model._embedding_dtype == torch.float32 + assert model._prompt_builder._embedding_dtype == torch.float32 + + +@pytest.mark.asyncio +async def test_rollout_adapter_initializes_stage_zero_workers(): + calls = [] + + class Engine: + async def collective_rpc(self, **kwargs): + calls.append(kwargs) + + await Qwen3TTSRolloutAdapter.initialize_rollout_workers(Engine(), "full") + + assert calls == [{"method": "align_qwen3_tts_prompt_embedding_dtype", "stage_ids": [0]}] + + def test_rollout_adapter_builds_unique_prompt_and_scopes_weight_sync(tmp_path): speaker = tmp_path / "speaker.json" speaker.write_text("[0.0, 1.0]") @@ -94,6 +113,17 @@ def test_rollout_adapter_builds_unique_prompt_and_scopes_weight_sync(tmp_path): assert Qwen3TTSRolloutAdapter.weight_sync_stage_ids("full") == [0] assert not Qwen3TTSRolloutAdapter.supports_cache_engine_sleep("full") assert Qwen3TTSRolloutAdapter.get_output_modalities("full") == ["latent", "audio"] + assert Qwen3TTSRolloutAdapter.get_worker_extension_cls("full") == ( + "verl_omni.pipelines.qwen3_tts.worker_extension.Qwen3TTSColocateWorkerExtension" + ) + + +def test_server_selects_qwen3_tts_worker_extension(): + server = object.__new__(vLLMOmniHttpServer) + server._omni_rollout_adapter = Qwen3TTSRolloutAdapter + server._omni_pipeline_mode = "full" + + assert server._get_worker_extension_cls() == Qwen3TTSRolloutAdapter.get_worker_extension_cls("full") def test_rollout_adapter_requires_speaker_embedding(): @@ -138,12 +168,12 @@ def test_rollout_adapter_combines_policy_codes_and_waveform(): generated[:, 0] = torch.tensor(token_ids) policy = SimpleNamespace( stage_id=0, - request_output=SimpleNamespace(outputs=[SimpleNamespace(token_ids=token_ids)]), + outputs=[SimpleNamespace(token_ids=token_ids)], multimodal_output={"codes": {"audio": torch.cat((torch.zeros(12, 16), generated))}}, ) decoder = SimpleNamespace( stage_id=1, - request_output=None, + outputs=[], multimodal_output={"audio": torch.ones(2400), "sr": 24_000}, ) prompt = {"additional_information": {"text": ["first text"]}} diff --git a/verl_omni/pipelines/model_base.py b/verl_omni/pipelines/model_base.py index c8e1451c8..d3b995f96 100644 --- a/verl_omni/pipelines/model_base.py +++ b/verl_omni/pipelines/model_base.py @@ -810,6 +810,16 @@ def get_engine_hf_overrides(cls, pipeline_mode: str = "thinker_only") -> dict: """ return {} + @classmethod + def get_worker_extension_cls(cls, pipeline_mode: str = "thinker_only") -> str | None: + """Return a model-specific vLLM worker extension, if required.""" + return None + + @classmethod + async def initialize_rollout_workers(cls, engine, pipeline_mode: str = "thinker_only") -> None: + """Run model-specific setup after all rollout workers are ready.""" + return + @classmethod def get_stage_engine_extras(cls, stage_id: int, pipeline_mode: str = "thinker_only") -> dict: """Return per-stage ``engine_extras`` to inject into the deploy config. diff --git a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py index 3bca2c854..4c1bc2d3a 100644 --- a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py @@ -31,13 +31,12 @@ load_speaker_xvector, require_auto_language, ) -from verl_omni.pipelines.qwen3_tts.vllm_plugin import ROLLOUT_MODEL_ARCH, register_qwen3_tts_rollout_model _PIPELINE_ID = "qwen3_tts_rl" _SYNC_PROCESSOR = "verl_omni.pipelines.qwen3_tts.omni_rollout_adapter.talker2code2wav_token_only" QWEN3_TTS_RL_PIPELINE = PipelineConfig( model_type=_PIPELINE_ID, - model_arch=ROLLOUT_MODEL_ARCH, + model_arch=QWEN3_TTS_PIPELINE.model_arch, stages=( replace(QWEN3_TTS_PIPELINE.stages[0], final_output=True, final_output_type="latent"), replace(QWEN3_TTS_PIPELINE.stages[1], sync_process_input_func=_SYNC_PROCESSOR), @@ -51,8 +50,8 @@ def _load_speaker_vector(path: str) -> list[float]: def _completion(output): - request_output = getattr(output, "request_output", None) - completions = getattr(request_output, "outputs", None) if request_output is not None else None + request_output = getattr(output, "request_output", None) or output + completions = getattr(request_output, "outputs", None) return completions[0] if completions else None @@ -104,7 +103,6 @@ def get_pipeline_id(cls, pipeline_mode="full"): @classmethod def ensure_pipeline_registered(cls, pipeline_mode="full"): cls._check_mode(pipeline_mode) - register_qwen3_tts_rollout_model() register_pipeline(QWEN3_TTS_RL_PIPELINE) @classmethod @@ -122,6 +120,19 @@ def supports_cache_engine_sleep(cls, pipeline_mode="full"): cls._check_mode(pipeline_mode) return False + @classmethod + def get_worker_extension_cls(cls, pipeline_mode="full"): + cls._check_mode(pipeline_mode) + return "verl_omni.pipelines.qwen3_tts.worker_extension.Qwen3TTSColocateWorkerExtension" + + @classmethod + async def initialize_rollout_workers(cls, engine, pipeline_mode="full"): + cls._check_mode(pipeline_mode) + await engine.collective_rpc( + method="align_qwen3_tts_prompt_embedding_dtype", + stage_ids=[0], + ) + @classmethod def get_stage_engine_extras(cls, stage_id, pipeline_mode="full"): cls._check_mode(pipeline_mode) diff --git a/verl_omni/pipelines/qwen3_tts/vllm_plugin.py b/verl_omni/pipelines/qwen3_tts/vllm_plugin.py deleted file mode 100644 index 6582e974c..000000000 --- a/verl_omni/pipelines/qwen3_tts/vllm_plugin.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright 2026 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""vLLM model registration for the Qwen3-TTS rollout extension.""" - -from vllm.model_executor.models import ModelRegistry - -ROLLOUT_MODEL_ARCH = "Qwen3TTSDtypeAlignedTalkerForConditionalGeneration" -ROLLOUT_MODEL_CLASS = "verl_omni.pipelines.qwen3_tts.rollout_model:Qwen3TTSDtypeAlignedTalkerForConditionalGeneration" - - -def register_qwen3_tts_rollout_model() -> None: - """Register lazily so every vLLM engine and worker process can resolve it.""" - ModelRegistry.register_model(ROLLOUT_MODEL_ARCH, ROLLOUT_MODEL_CLASS) diff --git a/verl_omni/pipelines/qwen3_tts/rollout_model.py b/verl_omni/pipelines/qwen3_tts/worker_extension.py similarity index 58% rename from verl_omni/pipelines/qwen3_tts/rollout_model.py rename to verl_omni/pipelines/qwen3_tts/worker_extension.py index 00d14b8ef..380eedb1a 100644 --- a/verl_omni/pipelines/qwen3_tts/rollout_model.py +++ b/verl_omni/pipelines/qwen3_tts/worker_extension.py @@ -11,12 +11,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Qwen3-TTS rollout model extensions.""" +"""Qwen3-TTS vLLM worker initialization.""" import torch -from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( - Qwen3TTSTalkerForConditionalGeneration, -) + +from verl_omni.workers.rollout.vllm_rollout.utils import vLLMOmniColocateWorkerExtension def _align_prompt_embedding_dtype(model, dtype: torch.dtype) -> None: @@ -25,9 +24,12 @@ def _align_prompt_embedding_dtype(model, dtype: torch.dtype) -> None: model._prompt_builder._embedding_dtype = dtype -class Qwen3TTSDtypeAlignedTalkerForConditionalGeneration(Qwen3TTSTalkerForConditionalGeneration): - """Make Qwen3-TTS prompt embeddings follow the configured rollout dtype.""" +class Qwen3TTSColocateWorkerExtension(vLLMOmniColocateWorkerExtension): + """Apply Qwen3-TTS setup after the upstream model is loaded.""" - def __init__(self, *, vllm_config, prefix: str = ""): - super().__init__(vllm_config=vllm_config, prefix=prefix) - _align_prompt_embedding_dtype(self, vllm_config.model_config.dtype) + def align_qwen3_tts_prompt_embedding_dtype(self) -> None: + standard = self._get_standard_weight_model_and_config() + if standard is None: + raise RuntimeError("Qwen3-TTS rollout worker has no loaded AR model") + model, model_config = standard + _align_prompt_embedding_dtype(model, model_config.dtype) diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py index 256d7657c..5bb28a7c9 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py @@ -178,6 +178,11 @@ async def run_server(self, args: argparse.Namespace): engine_client = AsyncOmni(**engine_args) app = build_app(args) await omni_init_app_state(engine_client, app.state, args) + if self._omni_rollout_adapter is not None: + await self._omni_rollout_adapter.initialize_rollout_workers( + engine_client, + self._omni_pipeline_mode, + ) # Deploy config YAML is consumed by AsyncOmni above; clean up the temp dir. if getattr(self, "_temp_deploy_ctx", None) is not None: From 183a19a769939194246f19ca9a2eabe8e86f1692 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Fri, 21 Aug 2026 12:46:05 +0800 Subject: [PATCH 06/28] [model, worker] fix: align Qwen3-TTS rollout probabilities Signed-off-by: dongbo910220 <1275604947@qq.com> Co-authored-by: Michael-Zzq <52185141+Michael-Zzq@users.noreply.github.com> --- examples/grpo_trainer/qwen3_tts/README.md | 5 +- .../qwen3_tts/run_qwen3_tts_grpo.sh | 2 + tests/pipelines/test_qwen3_tts_on_cpu.py | 16 +++- tests/workers/test_omni_fsdp_engine_on_cpu.py | 22 +++++ .../pipelines/qwen3_tts/talker_forward.py | 15 +++- verl_omni/workers/engine/fsdp/omni_impl.py | 87 ++++++++++++++----- verl_omni/workers/engine_workers.py | 19 +++- 7 files changed, 139 insertions(+), 27 deletions(-) diff --git a/examples/grpo_trainer/qwen3_tts/README.md b/examples/grpo_trainer/qwen3_tts/README.md index 78cc5cfe5..2d94e3ef9 100644 --- a/examples/grpo_trainer/qwen3_tts/README.md +++ b/examples/grpo_trainer/qwen3_tts/README.md @@ -95,7 +95,10 @@ bash examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh The example defaults are `B=4`, `G=8`, `lr=2e-7`, direct `low_var_kl` with coefficient `0.12`, two GPUs, and 500 updates. These are recipe values, not algorithm requirements. `norm_adv_by_std_in_grpo` remains at the upstream -default. +default. Actor, reference, rollout, and weight synchronization all use FP32: +the codec policy sums 16 codebook embeddings autoregressively, so a nominal +FP32 rollout with BF16 FSDP forward or BF16 synchronization does not reproduce +the same selected-token probabilities. For a two-update implementation smoke test: diff --git a/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh b/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh index d390c12cf..6facbd7fc 100755 --- a/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh +++ b/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh @@ -72,6 +72,7 @@ export VLLM_USE_FLASHINFER_SAMPLER=0 "actor_rollout_ref.actor.fsdp_config.seed=${SEED}" \ actor_rollout_ref.actor.fsdp_config.model_dtype=float32 \ actor_rollout_ref.actor.fsdp_config.dtype=float32 \ + '+actor_rollout_ref.actor.fsdp_config.mixed_precision={param_dtype:fp32,reduce_dtype:fp32,buffer_dtype:fp32}' \ actor_rollout_ref.actor.fsdp_config.param_offload=false \ actor_rollout_ref.actor.fsdp_config.optimizer_offload=false \ actor_rollout_ref.actor.fsdp_config.use_orig_params=true \ @@ -115,6 +116,7 @@ export VLLM_USE_FLASHINFER_SAMPLER=0 "actor_rollout_ref.ref.fsdp_config.seed=${SEED}" \ actor_rollout_ref.ref.fsdp_config.model_dtype=float32 \ actor_rollout_ref.ref.fsdp_config.dtype=float32 \ + '+actor_rollout_ref.ref.fsdp_config.mixed_precision={param_dtype:fp32,reduce_dtype:fp32,buffer_dtype:fp32}' \ actor_rollout_ref.ref.fsdp_config.param_offload=false \ actor_rollout_ref.ref.fsdp_config.use_orig_params=true \ actor_rollout_ref.ref.fsdp_config.wrap_policy.min_num_params=0 \ diff --git a/tests/pipelines/test_qwen3_tts_on_cpu.py b/tests/pipelines/test_qwen3_tts_on_cpu.py index 633ea7113..597b326a0 100644 --- a/tests/pipelines/test_qwen3_tts_on_cpu.py +++ b/tests/pipelines/test_qwen3_tts_on_cpu.py @@ -57,6 +57,16 @@ def test_talker_batch_matches_auto_language_teacher_forcing_layout(): assert batch.codec_lens == [3] +def test_codec0_mask_matches_rollout_vocabulary(): + masked = forward.mask_codec0_logits(torch.zeros((1, 2, 4300)), 2048, TOKENS.codec_eos) + + assert (masked[..., 0] < -1e3).all() + assert torch.isfinite(masked[..., 1:2048]).all() + assert (masked[..., 2048 : TOKENS.codec_eos] < -1e3).all() + assert torch.isfinite(masked[..., TOKENS.codec_eos]).all() + assert (masked[..., TOKENS.codec_eos + 1 :] < -1e3).all() + + def test_only_validated_auto_language_layout_is_accepted(): assert forward.require_auto_language("auto") == "Auto" with pytest.raises(ValueError, match="supports only tts_language=Auto"): @@ -109,7 +119,11 @@ def fake_logits(_talker, batch, _speaker): ) assert torch.nonzero(logits.abs().sum(dim=-1)[0], as_tuple=False).reshape(-1).tolist() == [5, 6, 7] - torch.testing.assert_close(logits[0, 5], torch.arange(1, 4301, dtype=torch.float32)) + assert logits[0, 5, 0] == -1e4 + assert logits[0, 5, 1] == 2 + assert logits[0, 5, 2047] == 2048 + assert logits[0, 5, 2048] == -1e4 + assert logits[0, 5, TOKENS.codec_eos] == TOKENS.codec_eos + 1 @pytest.mark.parametrize("response_length", [2, 14, 15, 16, 17, 32]) diff --git a/tests/workers/test_omni_fsdp_engine_on_cpu.py b/tests/workers/test_omni_fsdp_engine_on_cpu.py index 713bf6fee..b9da529ed 100644 --- a/tests/workers/test_omni_fsdp_engine_on_cpu.py +++ b/tests/workers/test_omni_fsdp_engine_on_cpu.py @@ -328,6 +328,28 @@ def prepare_model_inputs(cls, model_inputs, replay_batch, model_config): assert output_args == {"base": True} +def test_weight_sync_dtype_honors_float32_rollout(): + omni_impl = _get_omni_impl_module() + tensor = torch.tensor([1.25], dtype=torch.bfloat16) + + dtype = omni_impl.OmniFSDPEngine._resolve_weight_sync_dtype("float32") + synced = omni_impl.OmniFSDPEngine._cast_weight_for_sync(tensor, dtype) + + assert dtype is torch.float32 + assert synced.dtype is torch.float32 + assert synced.item() == pytest.approx(1.25) + + +def test_weight_sync_dtype_keeps_integer_buffers(): + omni_impl = _get_omni_impl_module() + tensor = torch.tensor([1, 2], dtype=torch.int64) + + synced = omni_impl.OmniFSDPEngine._cast_weight_for_sync(tensor, torch.float32) + + assert synced is tensor + assert synced.dtype is torch.int64 + + # --------------------------------------------------------------------------- # ``collect_lora_params`` import source # --------------------------------------------------------------------------- diff --git a/verl_omni/pipelines/qwen3_tts/talker_forward.py b/verl_omni/pipelines/qwen3_tts/talker_forward.py index cdb5a7b48..924dec9a0 100644 --- a/verl_omni/pipelines/qwen3_tts/talker_forward.py +++ b/verl_omni/pipelines/qwen3_tts/talker_forward.py @@ -178,6 +178,15 @@ def codec0_logits(talker, batch: TalkerBatch, speaker_embedding: torch.Tensor) - return output.logits +def mask_codec0_logits(logits: torch.Tensor, codebook_vocab: int, codec_eos_token_id: int) -> torch.Tensor: + """Match the codec-token vocabulary exposed by the rollout model.""" + valid = torch.zeros(logits.shape[-1], dtype=torch.bool, device=logits.device) + valid[1 : min(codebook_vocab, logits.shape[-1])] = True + if 0 <= codec_eos_token_id < logits.shape[-1]: + valid[codec_eos_token_id] = True + return logits.masked_fill(~valid, -1e4) + + def tts_actor_logits( model, input_ids, @@ -216,7 +225,11 @@ def tts_actor_logits( device=input_ids.device, sub_codebook_vocab=sub_vocab, ) - logits = codec0_logits(talker, batch, speaker_embedding) + logits = mask_codec0_logits( + codec0_logits(talker, batch, speaker_embedding), + sub_vocab, + int(model.config.talker_config.codec_eos_token_id), + ) output_vocab = max(logits.shape[-1], int(input_ids.max()) + 1) aligned = logits.new_zeros((batch_size, output_len, output_vocab)) for index, response_start in enumerate(response_starts): diff --git a/verl_omni/workers/engine/fsdp/omni_impl.py b/verl_omni/workers/engine/fsdp/omni_impl.py index 97d62fb43..e2875b51c 100644 --- a/verl_omni/workers/engine/fsdp/omni_impl.py +++ b/verl_omni/workers/engine/fsdp/omni_impl.py @@ -43,6 +43,53 @@ class OmniFSDPEngine(FSDPEngineWithLMHead): """FSDP engine for omni models""" + @staticmethod + def _resolve_weight_sync_dtype(weight_sync_dtype): + if weight_sync_dtype is None or isinstance(weight_sync_dtype, torch.dtype): + return weight_sync_dtype + + from verl.utils.torch_dtypes import PrecisionType + + return PrecisionType.to_dtype(weight_sync_dtype) + + @staticmethod + def _cast_weight_for_sync(tensor: torch.Tensor, dtype: torch.dtype | None) -> torch.Tensor: + if dtype is not None and tensor.is_floating_point() and tensor.dtype != dtype: + return tensor.to(dtype=dtype, non_blocking=True) + return tensor + + def _materialize_weight_for_sync(self, param, device, dtype: torch.dtype | None) -> torch.Tensor: + if isinstance(param, DTensor): + tensor = param.to(device, non_blocking=True).full_tensor() + dtype = torch.bfloat16 if dtype is None else dtype + else: + tensor = param + return self._cast_weight_for_sync(tensor, dtype) + + def _run_with_manual_ref_offload(self, call): + adapter_cls = getattr(self, "model_adapter_cls", None) + manual_ref_offload = ( + getattr(adapter_cls, "requires_manual_ref_offload", False) + and getattr(self.engine_config, "forward_only", False) + and not getattr(self.engine_config, "param_offload", True) + ) + if not manual_ref_offload: + return call() + + self.engine_config.forward_only = False + try: + return call() + finally: + self.engine_config.forward_only = True + + def _build_fsdp_module(self, module): + parent_build = super()._build_fsdp_module + return self._run_with_manual_ref_offload(lambda: parent_build(module)) + + def to(self, device, model=True, optimizer=True, grad=True): + parent_to = super().to + return self._run_with_manual_ref_offload(lambda: parent_to(device, model=model, optimizer=optimizer, grad=grad)) + def prepare_model_inputs(self, micro_batch): """Prepare standard LM inputs, then add model-native replay fields.""" model_inputs, output_args = super().prepare_model_inputs(micro_batch) @@ -55,8 +102,15 @@ def prepare_model_inputs(self, micro_batch): ) return model_inputs, output_args - def get_per_tensor_param(self, layered_summon=False, base_sync_done=False, **kwargs): + def get_per_tensor_param( + self, + layered_summon=False, + base_sync_done=False, + weight_sync_dtype=None, + **kwargs, + ): log_gpu_memory_usage("Before load_fsdp_model_to_gpu", logger=logger) + sync_dtype = self._resolve_weight_sync_dtype(weight_sync_dtype) # FSDP2 CPUOffloadPolicy owns CPU<->GPU placement; calling model.to(device) here # leaves the module half-moved and crashes state_dict() below (verl#5995). The @@ -82,7 +136,7 @@ def get_per_tensor_param(self, layered_summon=False, base_sync_done=False, **kwa if not base_sync_done: params = {replace_lora_wrapper(k, peft_config): v for k, v in params.items()} else: # merge lora - return self._merged_lora_per_tensor_param(), None + return self._merged_lora_per_tensor_param(sync_dtype), None else: params = self.module.state_dict() @@ -93,20 +147,10 @@ def get_per_tensor_param(self, layered_summon=False, base_sync_done=False, **kwa offload_fsdp_model_to_cpu(self.module) log_gpu_memory_usage("After offload_fsdp_model_to_cpu", logger=logger) - if peft_config is not None and base_sync_done: - per_tensor_param = params.items() - else: - device = get_device_id() # used when fsdp2 set cpu_offload_policy - # TODO: cast fp32 to bf16 to reduce weight sync overhead, need more fine-grained control, e.g MoE gate - per_tensor_param = ( - ( - name, - param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True) - if isinstance(param, DTensor) - else param, - ) - for name, param in params.items() - ) + device = get_device_id() # used when fsdp2 set cpu_offload_policy + per_tensor_param = ( + (name, self._materialize_weight_for_sync(param, device, sync_dtype)) for name, param in params.items() + ) if self._qat_enabled: from verl.utils.qat.quantizer import QATQuantizer @@ -134,20 +178,17 @@ def get_per_tensor_param(self, layered_summon=False, base_sync_done=False, **kwa return per_tensor_param, peft_config_dict - def _merged_lora_per_tensor_param(self): + def _merged_lora_per_tensor_param(self, weight_sync_dtype=None): """Stream materialized merged weights before restoring the actor.""" device = get_device_id() + sync_dtype = self._resolve_weight_sync_dtype(weight_sync_dtype) try: with merged_lora_context(self.module, backup_adapters=True): params = normalize_peft_param_name(self.module.state_dict()) params = convert_weight_keys(params, getattr(self.module, "_fsdp_wrapped_module", self.module)) for name, param in params.items(): - yield ( - name, - param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True) - if isinstance(param, DTensor) - else param.detach().clone(), - ) + materialized = self._materialize_weight_for_sync(param, device, sync_dtype) + yield name, materialized.detach().clone() finally: log_gpu_memory_usage("Before offload_fsdp_model_to_cpu", logger=logger) if self._is_offload_param: diff --git a/verl_omni/workers/engine_workers.py b/verl_omni/workers/engine_workers.py index ce1979343..93a2978d8 100644 --- a/verl_omni/workers/engine_workers.py +++ b/verl_omni/workers/engine_workers.py @@ -792,6 +792,7 @@ def init_model(self): # 3. build rollout engine if "rollout" in self.role: rollout_config: RolloutConfig = omega_conf_to_dataclass(self.config.rollout) + self._rollout_weight_sync_dtype = rollout_config.dtype # TODO: move rollout_device_mesh into ServerAdapter # 3.1 build rollout device mesh (sglang need only) @@ -966,6 +967,17 @@ def _offload_actor_and_empty_cache(self, timings: Optional[dict] = None): if timings is not None: timings["offload_actor_to_cpu"] = time.perf_counter() - start + def _get_rollout_weight_sync_dtype(self): + """Use the dtype of the rollout instance that owns the target weights.""" + rollout_dtype = getattr(self, "_rollout_weight_sync_dtype", None) + if rollout_dtype is not None: + return rollout_dtype + rollout_config = getattr(getattr(self, "rollout", None), "config", None) + rollout_dtype = getattr(rollout_config, "dtype", None) + if rollout_dtype is not None: + return rollout_dtype + return self.config.rollout.get("dtype", None) + def _gather_lora_weights(self, timings: Optional[dict] = None): """Gather LoRA adapter params into a CPU dict, without offloading the actor. @@ -980,6 +992,7 @@ def _gather_lora_weights(self, timings: Optional[dict] = None): layered_summon=self.layered_summon, base_sync_done=True, adapter_name=self.config.rollout.rollout_adapter, + weight_sync_dtype=self._get_rollout_weight_sync_dtype(), ) lora_weights = {name: tensor for name, tensor in per_tensor_param} if timings is not None: @@ -1039,12 +1052,14 @@ async def update_weights(self, global_steps: int = None, mode: str = "auto"): per_tensor_param, _ = self.actor.engine.get_per_tensor_param( base_sync_done=True, adapter_name=self.config.rollout.rollout_adapter, + weight_sync_dtype=self._get_rollout_weight_sync_dtype(), ) await self.checkpoint_engine.send_weights(per_tensor_param) return per_tensor_param, _ = self.actor.engine.get_per_tensor_param( - adapter_name=self.config.rollout.rollout_adapter + adapter_name=self.config.rollout.rollout_adapter, + weight_sync_dtype=self._get_rollout_weight_sync_dtype(), ) await self.checkpoint_engine.send_weights(per_tensor_param) return @@ -1141,6 +1156,7 @@ async def update_weights(self, global_steps: int = None, mode: str = "auto"): layered_summon=self.layered_summon, base_sync_done=True, adapter_name=self.config.rollout.rollout_adapter, + weight_sync_dtype=self._get_rollout_weight_sync_dtype(), ) do_lora_base_sync = False @@ -1154,6 +1170,7 @@ async def update_weights(self, global_steps: int = None, mode: str = "auto"): layered_summon=self.layered_summon, base_sync_done=False, adapter_name=self.config.rollout.rollout_adapter, + weight_sync_dtype=self._get_rollout_weight_sync_dtype(), ) await self.rollout.update_weights( per_tensor_param_base, peft_config=peft_config, base_sync_done=False, global_steps=global_steps From 24d3e5817c4410a2ea34d50872b200384142acd5 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Fri, 21 Aug 2026 23:46:30 +0800 Subject: [PATCH 07/28] [doc, tests] test: document generic omni rollout hooks Signed-off-by: dongbo910220 <1275604947@qq.com> Co-authored-by: Michael-Zzq <52185141+Michael-Zzq@users.noreply.github.com> --- docs/api/pipelines.rst | 37 ++++++++++++++---- .../contributing/integrating_an_omni_model.md | 38 +++++++++++++++---- .../qwen3_tts/grpo_trainer_qwen3_tts.md | 1 + docs/index.md | 1 + examples/grpo_trainer/qwen3_tts/README.md | 2 + .../test_qwen3_tts_rollout_on_cpu.py | 12 ++++++ 6 files changed, 75 insertions(+), 16 deletions(-) create mode 120000 docs/examples/qwen3_tts/grpo_trainer_qwen3_tts.md diff --git a/docs/api/pipelines.rst b/docs/api/pipelines.rst index ecd54785a..4a5584e41 100644 --- a/docs/api/pipelines.rst +++ b/docs/api/pipelines.rst @@ -4,25 +4,32 @@ Pipelines Interface Last updated: |today| (API docstrings are auto-generated). A *pipeline* in VeRL-Omni packages everything needed to plug a particular -diffusion model architecture into the training loop: +model architecture into the training loop. Two adapter families are available: -- a **training-side adapter** subclassing +- autoregressive omni models use a training-side + :class:`~verl_omni.pipelines.model_base.OmniModelBase` and an optional + rollout-side :class:`~verl_omni.pipelines.model_base.OmniRolloutPipelineBase`; +- diffusion models use a training-side adapter subclassing :class:`~verl_omni.pipelines.model_base.DiffusionModelBase` that handles scheduler setup, model-input construction, and the per-step forward / reverse-sampling logic used by RL algorithms (e.g. FlowGRPO); -- an optional **rollout-side adapter** registered via +- their optional rollout-side adapter is registered via :class:`~verl_omni.pipelines.model_base.VllmOmniPipelineBase` that hooks into vLLM-Omni's diffusion serving stack to expose log-probabilities. -Adapters are auto-selected by matching the pair -``(DiffusionModelConfig.architecture, DiffusionModelConfig.algorithm)`` against the -registered ``(architecture, algorithm)`` key. The architecture is read from the -model's ``model_index.json``; the algorithm string is taken from the model config's -``actor_rollout_ref.model.algorithm`` value. +Autoregressive training adapters are selected by ``(architecture, model_stage)``; +their rollout adapters are selected by the vLLM-Omni ``pipeline_name``. Diffusion +adapters are selected by matching +``(DiffusionModelConfig.architecture, DiffusionModelConfig.algorithm)`` against a +registered ``(architecture, algorithm)`` key. Diffusion architecture is read from +``model_index.json`` and the algorithm from +``actor_rollout_ref.model.algorithm``. .. autosummary:: :nosignatures: + verl_omni.pipelines.model_base.OmniModelBase + verl_omni.pipelines.model_base.OmniRolloutPipelineBase verl_omni.pipelines.model_base.DiffusionModelBase verl_omni.pipelines.model_base.VllmOmniPipelineBase verl_omni.pipelines.qwen_image_flow_grpo.QwenImage @@ -33,6 +40,20 @@ model's ``model_index.json``; the algorithm string is taken from the model confi Model Base ~~~~~~~~~~~~~~~~~ +.. autoclass:: verl_omni.pipelines.model_base.OmniModelBase + :members: register, get_class, get_class_by_name, + get_strip_modules, configure_processor, configure_tokenizer, + configure_model, prepare_model_inputs + +.. autoclass:: verl_omni.pipelines.model_base.OmniRolloutPipelineBase + :members: register, get_class, + build_stage_configs, rollout_flags, weight_sync_stage_ids, + supports_cache_engine_sleep, get_pipeline_id, + ensure_pipeline_registered, get_engine_hf_overrides, + get_worker_extension_cls, initialize_rollout_workers, + get_stage_engine_extras, prepare_engine_prompt, + get_output_modalities, combine_engine_outputs + .. autoclass:: verl_omni.pipelines.model_base.DiffusionModelBase :members: register, get_class, build_scheduler, set_timesteps, diff --git a/docs/contributing/integrating_an_omni_model.md b/docs/contributing/integrating_an_omni_model.md index 1dd472aa8..f72738e15 100644 --- a/docs/contributing/integrating_an_omni_model.md +++ b/docs/contributing/integrating_an_omni_model.md @@ -13,10 +13,10 @@ under [`verl_omni/pipelines/`](https://github.com/verl-project/verl-omni/tree/ma Decide which **training stage** you want to train and how the model decomposes: - **Stage-split**: Multi-component omni models (thinker → talker → code2wav) - train only the text-understanding head during RL post-training. Other - components are stripped before FSDP wrapping to save memory. This is the - Qwen3-Omni pattern — the thinker is the autoregressive language model; talker - and codec are inference-only. + train one selected autoregressive stage during RL post-training. Other + components are stripped before FSDP wrapping to save memory. Qwen3-Omni + trains the thinker; Qwen3-TTS trains the talker's codec-0 policy while its + decoder remains rollout-only. - **Encoder-frozen**: Vision/audio encoders are typically frozen during RL training (`freeze_vision_tower=True`). The training adapter's `get_strip_modules` excludes them from the trainable set if they are separate @@ -70,6 +70,8 @@ adapt each implementation to your model's architecture: sequence alone cannot reconstruct the exact sampled trajectory. Missing required fields or inconsistent shapes should raise an actionable error; the adapter must not silently reconstruct a different trajectory. + Qwen3-TTS uses this hook to consume text tokens and all 16 codec codebooks + from its model-owned replay payload while optimizing codec-0 log-probabilities. Reference: [`verl_omni/pipelines/qwen3_omni/thinker_training_adapter.py`](../../verl_omni/pipelines/qwen3_omni/thinker_training_adapter.py) @@ -93,10 +95,21 @@ and implement: - **`get_pipeline_id(pipeline_mode)`**: Return the vLLM-Omni pipeline `model_type` string, used when auto-generating the deploy config YAML. -Optional overrides: `ensure_pipeline_registered` (register non-standard -pipeline variants with vLLM-Omni), `get_engine_hf_overrides` (HF config -overrides like `enable_audio_output: false`), `get_stage_engine_extras` -(per-stage overrides like `model_arch`). +Optional overrides fall into four groups: + +- Pipeline setup: `ensure_pipeline_registered`, `get_engine_hf_overrides`, + `get_stage_engine_extras`, `get_worker_extension_cls`, and + `initialize_rollout_workers`. +- Resource behavior: `weight_sync_stage_ids` and + `supports_cache_engine_sleep`. +- Request construction: `prepare_engine_prompt`. +- Multi-stage output retention: `get_output_modalities` and + `combine_engine_outputs`. + +Their defaults preserve the existing single-output AR behavior. Override only +the hooks required by the model. For example, Qwen3-TTS synchronizes actor +weights only to its talker stage and retains both codec and waveform outputs; +the decoder stage never receives actor weights. When training an omni model's autoregressive Talker stage, also override `postprocess_agent_loop_output`. Put the sampled policy sequence in @@ -207,6 +220,9 @@ KV is cheap relative to starving decode. Reference: [`examples/gspo_trainer/qwen3_omni/run_qwen3_omni_thinker_gspo_lora_v1.sh`](../../examples/gspo_trainer/qwen3_omni/run_qwen3_omni_thinker_gspo_lora_v1.sh) +For a talker-stage full-parameter example, see +[`examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh`](../../examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh). + ## 6. Common pitfalls These pitfalls are drawn from the Qwen3-Omni adapter. Some are @@ -234,3 +250,9 @@ model-specific — verify each against your own model's architecture. in `configure_tokenizer` and assign it to `tokenizer.chat_template`. verl's dataset loader calls `tokenizer.apply_chat_template()` and will fail without a template. + +- **Actor/rollout probability consistency**: Autoregressive codec policies may + combine several codebook embeddings before predicting the selected token. + Match actor, reference, rollout, and weight-sync dtypes, then verify selected + token log-probabilities before training. A nominal FP32 rollout fed BF16 + actor weights is not an FP32 consistency check. diff --git a/docs/examples/qwen3_tts/grpo_trainer_qwen3_tts.md b/docs/examples/qwen3_tts/grpo_trainer_qwen3_tts.md new file mode 120000 index 000000000..3bffcd1a3 --- /dev/null +++ b/docs/examples/qwen3_tts/grpo_trainer_qwen3_tts.md @@ -0,0 +1 @@ +../../../examples/grpo_trainer/qwen3_tts/README.md \ No newline at end of file diff --git a/docs/index.md b/docs/index.md index fca0839a9..a40236585 100644 --- a/docs/index.md +++ b/docs/index.md @@ -80,6 +80,7 @@ examples/mixgrpo_trainer.md examples/diffusionopd_trainer.md examples/flowgrpo_trainer_sd35_drm.md examples/bagel/flowgrpo_trainer_bagel.md +examples/qwen3_tts/grpo_trainer_qwen3_tts.md examples/qwen_image_edit/flowgrpo_trainer_qwen_image_edit.md examples/ltx2/flowgrpo_trainer_ltx2.md examples/minimax_h3/diffusionnft_trainer_minimax_h3.md diff --git a/examples/grpo_trainer/qwen3_tts/README.md b/examples/grpo_trainer/qwen3_tts/README.md index 2d94e3ef9..005606c8a 100644 --- a/examples/grpo_trainer/qwen3_tts/README.md +++ b/examples/grpo_trainer/qwen3_tts/README.md @@ -1,5 +1,7 @@ # Qwen3-TTS GRPO with an audio reward +Last updated: 08/21/2026. + This example full-parameter tunes the codec-0 policy of `Qwen/Qwen3-TTS-12Hz-0.6B-Base`. It uses verl's stock GRPO advantage, vanilla PPO policy loss, and optional direct reference-model KL. The other diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index 57b8e303a..ab7974e2c 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -23,6 +23,7 @@ from vllm import SamplingParams +from verl_omni.pipelines.model_base import OmniRolloutPipelineBase from verl_omni.pipelines.qwen3_tts import omni_rollout_adapter from verl_omni.pipelines.qwen3_tts.omni_rollout_adapter import Qwen3TTSRolloutAdapter from verl_omni.pipelines.qwen3_tts.talker_training_adapter import Qwen3TTSTalkerAdapter @@ -42,6 +43,17 @@ def __call__(self, text, **kwargs): return {"input_ids": list(range(len(text)))} +def test_optional_rollout_hooks_preserve_existing_ar_defaults(): + first, final = object(), object() + + assert OmniRolloutPipelineBase.weight_sync_stage_ids() is None + assert OmniRolloutPipelineBase.supports_cache_engine_sleep() + assert OmniRolloutPipelineBase.get_worker_extension_cls() is None + assert OmniRolloutPipelineBase.prepare_engine_prompt([], None, {}) is None + assert OmniRolloutPipelineBase.get_output_modalities() is None + assert OmniRolloutPipelineBase.combine_engine_outputs([first, final], {}) == (final, {}) + + def test_rollout_pipeline_registers_upstream_talker(monkeypatch): registered_pipelines = [] monkeypatch.setattr( From 932e80d1dfac7d65590649aaba1bbe6ce16daf80 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Sat, 22 Aug 2026 23:41:59 +0800 Subject: [PATCH 08/28] [rollout, tests] refactor: use native bucketed weight loading Signed-off-by: dongbo910220 <1275604947@qq.com> Co-authored-by: Michael-Zzq <52185141+Michael-Zzq@users.noreply.github.com> --- .../test_qwen3_tts_rollout_on_cpu.py | 38 ------------------- .../workers/rollout/vllm_rollout/utils.py | 28 ++------------ 2 files changed, 4 insertions(+), 62 deletions(-) diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index ab7974e2c..e0f1eb3c9 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -31,7 +31,6 @@ Qwen3TTSColocateWorkerExtension, _align_prompt_embedding_dtype, ) -from verl_omni.workers.rollout.vllm_rollout.utils import _receive_model_weight_buckets from verl_omni.workers.rollout.vllm_rollout.vllm_omni_async_server import vLLMOmniHttpServer @@ -199,43 +198,6 @@ def test_rollout_adapter_combines_policy_codes_and_waveform(): assert fields["tts_text"] == "first text" -def test_bucketed_weight_sync_rebuilds_derived_codec_table_once(monkeypatch): - class Model: - def __init__(self): - self.rebuilds = 0 - self.loads = 0 - self._stacked_codec_embed = object() - - def _build_stacked_codec_embed(self): - self.rebuilds += 1 - self._stacked_codec_embed = object() - - def load_weights(self, weights): - self.loads += 1 - self._build_stacked_codec_embed() - - class Receiver: - @staticmethod - def receive_weights(on_bucket_received): - on_bucket_received({"first": torch.tensor(1)}) - on_bucket_received({"second": torch.tensor(2)}) - - model = Model() - empty_cache_calls = [] - monkeypatch.setattr( - "verl_omni.workers.rollout.vllm_rollout.utils.get_torch_device", - lambda: SimpleNamespace(empty_cache=lambda: empty_cache_calls.append(True)), - ) - - _receive_model_weight_buckets(Receiver(), model) - - assert model.loads == 2 - assert model.rebuilds == 1 - assert empty_cache_calls == [True] - model._build_stacked_codec_embed() - assert model.rebuilds == 2 - - def test_server_prepares_stage_specific_sampling_params(): class Adapter: @staticmethod diff --git a/verl_omni/workers/rollout/vllm_rollout/utils.py b/verl_omni/workers/rollout/vllm_rollout/utils.py index c389fc9a0..13d2c2fa7 100644 --- a/verl_omni/workers/rollout/vllm_rollout/utils.py +++ b/verl_omni/workers/rollout/vllm_rollout/utils.py @@ -16,7 +16,7 @@ import time import torch -from verl.utils.device import get_torch_device, get_visible_devices_keyword +from verl.utils.device import get_visible_devices_keyword from verl.workers.rollout.vllm_rollout.utils import VLLM_LORA_INT_ID, VLLM_LORA_NAME, VLLM_LORA_PATH, set_death_signal from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension @@ -30,28 +30,6 @@ def _split_visible_devices(value: str) -> list[str]: """Split a visible-devices env value into stripped, non-empty entries.""" return [entry.strip() for entry in value.split(",") if entry.strip()] - - -def _receive_model_weight_buckets(receiver, model) -> None: - """Stream model weights and rebuild optional derived tensors once.""" - rebuild_derived_weights = getattr(model, "_build_stacked_codec_embed", None) - if callable(rebuild_derived_weights): - model._build_stacked_codec_embed = lambda: None - try: - receiver.receive_weights( - on_bucket_received=lambda weights, *args, **kwargs: model.load_weights(weights) - ) - finally: - if callable(rebuild_derived_weights): - model._build_stacked_codec_embed = rebuild_derived_weights - if callable(rebuild_derived_weights): - old_derived_weights = getattr(model, "_stacked_codec_embed", None) - model._stacked_codec_embed = None - del old_derived_weights - get_torch_device().empty_cache() - rebuild_derived_weights() - - class vLLMOmniColocateWorkerExtension(CustomPipelineWorkerExtension): """ The class for vLLM-Omni's worker to inherit from, in the colocate setting. @@ -211,7 +189,9 @@ def update_weights_from_ipc( from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader patch_vllm_moe_model_weight_loader(model) - _receive_model_weight_buckets(receiver, model) + receiver.receive_weights( + on_bucket_received=lambda weights, *args, **kwargs: model.load_weights(weights) + ) from vllm.model_executor.model_loader.utils import process_weights_after_loading process_weights_after_loading(model, model_config, self.device) From 0598b1a903a35a4bc01a8a6d2d45936fca685602 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Sun, 23 Aug 2026 11:09:56 +0800 Subject: [PATCH 09/28] [trainer] refactor: restore BF16 Qwen3-TTS GRPO path Signed-off-by: dongbo910220 <1275604947@qq.com> Co-authored-by: Michael-Zzq <52185141+Michael-Zzq@users.noreply.github.com> --- docs/api/pipelines.rst | 1 - .../contributing/integrating_an_omni_model.md | 10 +-- examples/grpo_trainer/qwen3_tts/README.md | 13 ++-- .../qwen3_tts/run_qwen3_tts_grpo.sh | 10 +-- .../test_qwen3_tts_rollout_on_cpu.py | 56 ----------------- tests/workers/test_omni_fsdp_engine_on_cpu.py | 14 ++--- verl_omni/pipelines/model_base.py | 10 --- .../qwen3_tts/omni_rollout_adapter.py | 13 ---- .../pipelines/qwen3_tts/worker_extension.py | 35 ----------- verl_omni/workers/engine/fsdp/omni_impl.py | 62 ++++++++----------- verl_omni/workers/engine_workers.py | 19 +----- 11 files changed, 48 insertions(+), 195 deletions(-) delete mode 100644 verl_omni/pipelines/qwen3_tts/worker_extension.py diff --git a/docs/api/pipelines.rst b/docs/api/pipelines.rst index 4a5584e41..cc62f26d1 100644 --- a/docs/api/pipelines.rst +++ b/docs/api/pipelines.rst @@ -50,7 +50,6 @@ Model Base build_stage_configs, rollout_flags, weight_sync_stage_ids, supports_cache_engine_sleep, get_pipeline_id, ensure_pipeline_registered, get_engine_hf_overrides, - get_worker_extension_cls, initialize_rollout_workers, get_stage_engine_extras, prepare_engine_prompt, get_output_modalities, combine_engine_outputs diff --git a/docs/contributing/integrating_an_omni_model.md b/docs/contributing/integrating_an_omni_model.md index f72738e15..d15adc90e 100644 --- a/docs/contributing/integrating_an_omni_model.md +++ b/docs/contributing/integrating_an_omni_model.md @@ -97,9 +97,8 @@ and implement: Optional overrides fall into four groups: -- Pipeline setup: `ensure_pipeline_registered`, `get_engine_hf_overrides`, - `get_stage_engine_extras`, `get_worker_extension_cls`, and - `initialize_rollout_workers`. +- Pipeline setup: `ensure_pipeline_registered`, `get_engine_hf_overrides`, and + `get_stage_engine_extras`. - Resource behavior: `weight_sync_stage_ids` and `supports_cache_engine_sleep`. - Request construction: `prepare_engine_prompt`. @@ -254,5 +253,6 @@ model-specific — verify each against your own model's architecture. - **Actor/rollout probability consistency**: Autoregressive codec policies may combine several codebook embeddings before predicting the selected token. Match actor, reference, rollout, and weight-sync dtypes, then verify selected - token log-probabilities before training. A nominal FP32 rollout fed BF16 - actor weights is not an FP32 consistency check. + token log-probabilities before training. The Qwen3-TTS recipe uses BF16 for + all four paths. Treat `diff_mean` and Pearson as consistency diagnostics, not + evidence of speech quality or bitwise agreement with an FP32 execution. diff --git a/examples/grpo_trainer/qwen3_tts/README.md b/examples/grpo_trainer/qwen3_tts/README.md index 005606c8a..ac81d13ac 100644 --- a/examples/grpo_trainer/qwen3_tts/README.md +++ b/examples/grpo_trainer/qwen3_tts/README.md @@ -1,6 +1,6 @@ # Qwen3-TTS GRPO with an audio reward -Last updated: 08/21/2026. +Last updated: 08/23/2026. This example full-parameter tunes the codec-0 policy of `Qwen/Qwen3-TTS-12Hz-0.6B-Base`. It uses verl's stock GRPO advantage, @@ -97,10 +97,13 @@ bash examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh The example defaults are `B=4`, `G=8`, `lr=2e-7`, direct `low_var_kl` with coefficient `0.12`, two GPUs, and 500 updates. These are recipe values, not algorithm requirements. `norm_adv_by_std_in_grpo` remains at the upstream -default. Actor, reference, rollout, and weight synchronization all use FP32: -the codec policy sums 16 codebook embeddings autoregressively, so a nominal -FP32 rollout with BF16 FSDP forward or BF16 synchronization does not reproduce -the same selected-token probabilities. +default. Actor, reference, rollout, and floating synchronized weights all use +BF16; synchronized integer buffers keep their integer dtype. Check selected-token +`diff_mean` and Pearson after synchronization as execution-consistency diagnostics, +not as evidence of speech quality or FP32-equivalent numerics. Prefer +`diff_mean < 0.005`; values from `0.005` to `0.01` require high Pearson and tail +inspection, while sustained values at or above `0.01` should stop the run for +investigation. For a two-update implementation smoke test: diff --git a/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh b/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh index 6facbd7fc..d58fe2960 100755 --- a/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh +++ b/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh @@ -70,9 +70,7 @@ export VLLM_USE_FLASHINFER_SAMPLER=0 "actor_rollout_ref.actor.data_loader_seed=${SEED}" \ "actor_rollout_ref.actor.fsdp_config.fsdp_size=${NUM_GPUS}" \ "actor_rollout_ref.actor.fsdp_config.seed=${SEED}" \ - actor_rollout_ref.actor.fsdp_config.model_dtype=float32 \ - actor_rollout_ref.actor.fsdp_config.dtype=float32 \ - '+actor_rollout_ref.actor.fsdp_config.mixed_precision={param_dtype:fp32,reduce_dtype:fp32,buffer_dtype:fp32}' \ + actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \ actor_rollout_ref.actor.fsdp_config.param_offload=false \ actor_rollout_ref.actor.fsdp_config.optimizer_offload=false \ actor_rollout_ref.actor.fsdp_config.use_orig_params=true \ @@ -84,7 +82,7 @@ export VLLM_USE_FLASHINFER_SAMPLER=0 actor_rollout_ref.rollout.temperature=1.0 \ actor_rollout_ref.rollout.top_p=1.0 \ actor_rollout_ref.rollout.top_k=-1 \ - actor_rollout_ref.rollout.dtype=float32 \ + actor_rollout_ref.rollout.dtype=bfloat16 \ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ actor_rollout_ref.rollout.gpu_memory_utilization=0.20 \ actor_rollout_ref.rollout.max_num_seqs=8 \ @@ -114,9 +112,7 @@ export VLLM_USE_FLASHINFER_SAMPLER=0 actor_rollout_ref.ref.strategy=fsdp \ "actor_rollout_ref.ref.fsdp_config.fsdp_size=${NUM_GPUS}" \ "actor_rollout_ref.ref.fsdp_config.seed=${SEED}" \ - actor_rollout_ref.ref.fsdp_config.model_dtype=float32 \ - actor_rollout_ref.ref.fsdp_config.dtype=float32 \ - '+actor_rollout_ref.ref.fsdp_config.mixed_precision={param_dtype:fp32,reduce_dtype:fp32,buffer_dtype:fp32}' \ + actor_rollout_ref.ref.fsdp_config.model_dtype=bfloat16 \ actor_rollout_ref.ref.fsdp_config.param_offload=false \ actor_rollout_ref.ref.fsdp_config.use_orig_params=true \ actor_rollout_ref.ref.fsdp_config.wrap_policy.min_num_params=0 \ diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index e0f1eb3c9..407321e06 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -27,10 +27,6 @@ from verl_omni.pipelines.qwen3_tts import omni_rollout_adapter from verl_omni.pipelines.qwen3_tts.omni_rollout_adapter import Qwen3TTSRolloutAdapter from verl_omni.pipelines.qwen3_tts.talker_training_adapter import Qwen3TTSTalkerAdapter -from verl_omni.pipelines.qwen3_tts.worker_extension import ( - Qwen3TTSColocateWorkerExtension, - _align_prompt_embedding_dtype, -) from verl_omni.workers.rollout.vllm_rollout.vllm_omni_async_server import vLLMOmniHttpServer @@ -47,7 +43,6 @@ def test_optional_rollout_hooks_preserve_existing_ar_defaults(): assert OmniRolloutPipelineBase.weight_sync_stage_ids() is None assert OmniRolloutPipelineBase.supports_cache_engine_sleep() - assert OmniRolloutPipelineBase.get_worker_extension_cls() is None assert OmniRolloutPipelineBase.prepare_engine_prompt([], None, {}) is None assert OmniRolloutPipelineBase.get_output_modalities() is None assert OmniRolloutPipelineBase.combine_engine_outputs([first, final], {}) == (final, {}) @@ -67,46 +62,6 @@ def test_rollout_pipeline_registers_upstream_talker(monkeypatch): assert registered_pipelines[0].model_arch == omni_rollout_adapter.QWEN3_TTS_PIPELINE.model_arch -def test_rollout_model_aligns_talker_and_prompt_builder_embedding_dtype(): - model = SimpleNamespace( - _embedding_dtype=torch.bfloat16, - _prompt_builder=SimpleNamespace(_embedding_dtype=torch.bfloat16), - ) - - _align_prompt_embedding_dtype(model, torch.float32) - - assert model._embedding_dtype == torch.float32 - assert model._prompt_builder._embedding_dtype == torch.float32 - - -def test_worker_extension_aligns_loaded_model_dtype(): - model = SimpleNamespace( - _embedding_dtype=torch.bfloat16, - _prompt_builder=SimpleNamespace(_embedding_dtype=torch.bfloat16), - ) - worker = SimpleNamespace( - _get_standard_weight_model_and_config=lambda: (model, SimpleNamespace(dtype=torch.float32)) - ) - - Qwen3TTSColocateWorkerExtension.align_qwen3_tts_prompt_embedding_dtype(worker) - - assert model._embedding_dtype == torch.float32 - assert model._prompt_builder._embedding_dtype == torch.float32 - - -@pytest.mark.asyncio -async def test_rollout_adapter_initializes_stage_zero_workers(): - calls = [] - - class Engine: - async def collective_rpc(self, **kwargs): - calls.append(kwargs) - - await Qwen3TTSRolloutAdapter.initialize_rollout_workers(Engine(), "full") - - assert calls == [{"method": "align_qwen3_tts_prompt_embedding_dtype", "stage_ids": [0]}] - - def test_rollout_adapter_builds_unique_prompt_and_scopes_weight_sync(tmp_path): speaker = tmp_path / "speaker.json" speaker.write_text("[0.0, 1.0]") @@ -124,17 +79,6 @@ def test_rollout_adapter_builds_unique_prompt_and_scopes_weight_sync(tmp_path): assert Qwen3TTSRolloutAdapter.weight_sync_stage_ids("full") == [0] assert not Qwen3TTSRolloutAdapter.supports_cache_engine_sleep("full") assert Qwen3TTSRolloutAdapter.get_output_modalities("full") == ["latent", "audio"] - assert Qwen3TTSRolloutAdapter.get_worker_extension_cls("full") == ( - "verl_omni.pipelines.qwen3_tts.worker_extension.Qwen3TTSColocateWorkerExtension" - ) - - -def test_server_selects_qwen3_tts_worker_extension(): - server = object.__new__(vLLMOmniHttpServer) - server._omni_rollout_adapter = Qwen3TTSRolloutAdapter - server._omni_pipeline_mode = "full" - - assert server._get_worker_extension_cls() == Qwen3TTSRolloutAdapter.get_worker_extension_cls("full") def test_rollout_adapter_requires_speaker_embedding(): diff --git a/tests/workers/test_omni_fsdp_engine_on_cpu.py b/tests/workers/test_omni_fsdp_engine_on_cpu.py index b9da529ed..7ccac1f1d 100644 --- a/tests/workers/test_omni_fsdp_engine_on_cpu.py +++ b/tests/workers/test_omni_fsdp_engine_on_cpu.py @@ -328,23 +328,21 @@ def prepare_model_inputs(cls, model_inputs, replay_batch, model_config): assert output_args == {"base": True} -def test_weight_sync_dtype_honors_float32_rollout(): +def test_weight_sync_casts_floating_dtensor_to_bfloat16(): omni_impl = _get_omni_impl_module() - tensor = torch.tensor([1.25], dtype=torch.bfloat16) + tensor = torch.tensor([1.25], dtype=torch.float32) - dtype = omni_impl.OmniFSDPEngine._resolve_weight_sync_dtype("float32") - synced = omni_impl.OmniFSDPEngine._cast_weight_for_sync(tensor, dtype) + synced = omni_impl.OmniFSDPEngine._cast_dtensor_weight_for_sync(tensor) - assert dtype is torch.float32 - assert synced.dtype is torch.float32 + assert synced.dtype is torch.bfloat16 assert synced.item() == pytest.approx(1.25) -def test_weight_sync_dtype_keeps_integer_buffers(): +def test_weight_sync_keeps_integer_dtensor_buffers(): omni_impl = _get_omni_impl_module() tensor = torch.tensor([1, 2], dtype=torch.int64) - synced = omni_impl.OmniFSDPEngine._cast_weight_for_sync(tensor, torch.float32) + synced = omni_impl.OmniFSDPEngine._cast_dtensor_weight_for_sync(tensor) assert synced is tensor assert synced.dtype is torch.int64 diff --git a/verl_omni/pipelines/model_base.py b/verl_omni/pipelines/model_base.py index d3b995f96..c8e1451c8 100644 --- a/verl_omni/pipelines/model_base.py +++ b/verl_omni/pipelines/model_base.py @@ -810,16 +810,6 @@ def get_engine_hf_overrides(cls, pipeline_mode: str = "thinker_only") -> dict: """ return {} - @classmethod - def get_worker_extension_cls(cls, pipeline_mode: str = "thinker_only") -> str | None: - """Return a model-specific vLLM worker extension, if required.""" - return None - - @classmethod - async def initialize_rollout_workers(cls, engine, pipeline_mode: str = "thinker_only") -> None: - """Run model-specific setup after all rollout workers are ready.""" - return - @classmethod def get_stage_engine_extras(cls, stage_id: int, pipeline_mode: str = "thinker_only") -> dict: """Return per-stage ``engine_extras`` to inject into the deploy config. diff --git a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py index 4c1bc2d3a..e0a1a5364 100644 --- a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py @@ -120,19 +120,6 @@ def supports_cache_engine_sleep(cls, pipeline_mode="full"): cls._check_mode(pipeline_mode) return False - @classmethod - def get_worker_extension_cls(cls, pipeline_mode="full"): - cls._check_mode(pipeline_mode) - return "verl_omni.pipelines.qwen3_tts.worker_extension.Qwen3TTSColocateWorkerExtension" - - @classmethod - async def initialize_rollout_workers(cls, engine, pipeline_mode="full"): - cls._check_mode(pipeline_mode) - await engine.collective_rpc( - method="align_qwen3_tts_prompt_embedding_dtype", - stage_ids=[0], - ) - @classmethod def get_stage_engine_extras(cls, stage_id, pipeline_mode="full"): cls._check_mode(pipeline_mode) diff --git a/verl_omni/pipelines/qwen3_tts/worker_extension.py b/verl_omni/pipelines/qwen3_tts/worker_extension.py deleted file mode 100644 index 380eedb1a..000000000 --- a/verl_omni/pipelines/qwen3_tts/worker_extension.py +++ /dev/null @@ -1,35 +0,0 @@ -# Copyright 2026 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Qwen3-TTS vLLM worker initialization.""" - -import torch - -from verl_omni.workers.rollout.vllm_rollout.utils import vLLMOmniColocateWorkerExtension - - -def _align_prompt_embedding_dtype(model, dtype: torch.dtype) -> None: - """Keep request embeddings compatible with the vLLM model input buffer.""" - model._embedding_dtype = dtype - model._prompt_builder._embedding_dtype = dtype - - -class Qwen3TTSColocateWorkerExtension(vLLMOmniColocateWorkerExtension): - """Apply Qwen3-TTS setup after the upstream model is loaded.""" - - def align_qwen3_tts_prompt_embedding_dtype(self) -> None: - standard = self._get_standard_weight_model_and_config() - if standard is None: - raise RuntimeError("Qwen3-TTS rollout worker has no loaded AR model") - model, model_config = standard - _align_prompt_embedding_dtype(model, model_config.dtype) diff --git a/verl_omni/workers/engine/fsdp/omni_impl.py b/verl_omni/workers/engine/fsdp/omni_impl.py index e2875b51c..d580c5de2 100644 --- a/verl_omni/workers/engine/fsdp/omni_impl.py +++ b/verl_omni/workers/engine/fsdp/omni_impl.py @@ -44,28 +44,11 @@ class OmniFSDPEngine(FSDPEngineWithLMHead): """FSDP engine for omni models""" @staticmethod - def _resolve_weight_sync_dtype(weight_sync_dtype): - if weight_sync_dtype is None or isinstance(weight_sync_dtype, torch.dtype): - return weight_sync_dtype - - from verl.utils.torch_dtypes import PrecisionType - - return PrecisionType.to_dtype(weight_sync_dtype) - - @staticmethod - def _cast_weight_for_sync(tensor: torch.Tensor, dtype: torch.dtype | None) -> torch.Tensor: - if dtype is not None and tensor.is_floating_point() and tensor.dtype != dtype: - return tensor.to(dtype=dtype, non_blocking=True) + def _cast_dtensor_weight_for_sync(tensor: torch.Tensor) -> torch.Tensor: + if tensor.is_floating_point() and tensor.dtype != torch.bfloat16: + return tensor.to(dtype=torch.bfloat16, non_blocking=True) return tensor - def _materialize_weight_for_sync(self, param, device, dtype: torch.dtype | None) -> torch.Tensor: - if isinstance(param, DTensor): - tensor = param.to(device, non_blocking=True).full_tensor() - dtype = torch.bfloat16 if dtype is None else dtype - else: - tensor = param - return self._cast_weight_for_sync(tensor, dtype) - def _run_with_manual_ref_offload(self, call): adapter_cls = getattr(self, "model_adapter_cls", None) manual_ref_offload = ( @@ -102,15 +85,8 @@ def prepare_model_inputs(self, micro_batch): ) return model_inputs, output_args - def get_per_tensor_param( - self, - layered_summon=False, - base_sync_done=False, - weight_sync_dtype=None, - **kwargs, - ): + def get_per_tensor_param(self, layered_summon=False, base_sync_done=False, **kwargs): log_gpu_memory_usage("Before load_fsdp_model_to_gpu", logger=logger) - sync_dtype = self._resolve_weight_sync_dtype(weight_sync_dtype) # FSDP2 CPUOffloadPolicy owns CPU<->GPU placement; calling model.to(device) here # leaves the module half-moved and crashes state_dict() below (verl#5995). The @@ -136,7 +112,7 @@ def get_per_tensor_param( if not base_sync_done: params = {replace_lora_wrapper(k, peft_config): v for k, v in params.items()} else: # merge lora - return self._merged_lora_per_tensor_param(sync_dtype), None + return self._merged_lora_per_tensor_param(), None else: params = self.module.state_dict() @@ -147,10 +123,19 @@ def get_per_tensor_param( offload_fsdp_model_to_cpu(self.module) log_gpu_memory_usage("After offload_fsdp_model_to_cpu", logger=logger) - device = get_device_id() # used when fsdp2 set cpu_offload_policy - per_tensor_param = ( - (name, self._materialize_weight_for_sync(param, device, sync_dtype)) for name, param in params.items() - ) + if peft_config is not None and base_sync_done: + per_tensor_param = params.items() + else: + device = get_device_id() # used when fsdp2 set cpu_offload_policy + per_tensor_param = ( + ( + name, + self._cast_dtensor_weight_for_sync(param.to(device, non_blocking=True).full_tensor()) + if isinstance(param, DTensor) + else param, + ) + for name, param in params.items() + ) if self._qat_enabled: from verl.utils.qat.quantizer import QATQuantizer @@ -178,17 +163,20 @@ def get_per_tensor_param( return per_tensor_param, peft_config_dict - def _merged_lora_per_tensor_param(self, weight_sync_dtype=None): + def _merged_lora_per_tensor_param(self): """Stream materialized merged weights before restoring the actor.""" device = get_device_id() - sync_dtype = self._resolve_weight_sync_dtype(weight_sync_dtype) try: with merged_lora_context(self.module, backup_adapters=True): params = normalize_peft_param_name(self.module.state_dict()) params = convert_weight_keys(params, getattr(self.module, "_fsdp_wrapped_module", self.module)) for name, param in params.items(): - materialized = self._materialize_weight_for_sync(param, device, sync_dtype) - yield name, materialized.detach().clone() + yield ( + name, + self._cast_dtensor_weight_for_sync(param.to(device, non_blocking=True).full_tensor()) + if isinstance(param, DTensor) + else param.detach().clone(), + ) finally: log_gpu_memory_usage("Before offload_fsdp_model_to_cpu", logger=logger) if self._is_offload_param: diff --git a/verl_omni/workers/engine_workers.py b/verl_omni/workers/engine_workers.py index 93a2978d8..ce1979343 100644 --- a/verl_omni/workers/engine_workers.py +++ b/verl_omni/workers/engine_workers.py @@ -792,7 +792,6 @@ def init_model(self): # 3. build rollout engine if "rollout" in self.role: rollout_config: RolloutConfig = omega_conf_to_dataclass(self.config.rollout) - self._rollout_weight_sync_dtype = rollout_config.dtype # TODO: move rollout_device_mesh into ServerAdapter # 3.1 build rollout device mesh (sglang need only) @@ -967,17 +966,6 @@ def _offload_actor_and_empty_cache(self, timings: Optional[dict] = None): if timings is not None: timings["offload_actor_to_cpu"] = time.perf_counter() - start - def _get_rollout_weight_sync_dtype(self): - """Use the dtype of the rollout instance that owns the target weights.""" - rollout_dtype = getattr(self, "_rollout_weight_sync_dtype", None) - if rollout_dtype is not None: - return rollout_dtype - rollout_config = getattr(getattr(self, "rollout", None), "config", None) - rollout_dtype = getattr(rollout_config, "dtype", None) - if rollout_dtype is not None: - return rollout_dtype - return self.config.rollout.get("dtype", None) - def _gather_lora_weights(self, timings: Optional[dict] = None): """Gather LoRA adapter params into a CPU dict, without offloading the actor. @@ -992,7 +980,6 @@ def _gather_lora_weights(self, timings: Optional[dict] = None): layered_summon=self.layered_summon, base_sync_done=True, adapter_name=self.config.rollout.rollout_adapter, - weight_sync_dtype=self._get_rollout_weight_sync_dtype(), ) lora_weights = {name: tensor for name, tensor in per_tensor_param} if timings is not None: @@ -1052,14 +1039,12 @@ async def update_weights(self, global_steps: int = None, mode: str = "auto"): per_tensor_param, _ = self.actor.engine.get_per_tensor_param( base_sync_done=True, adapter_name=self.config.rollout.rollout_adapter, - weight_sync_dtype=self._get_rollout_weight_sync_dtype(), ) await self.checkpoint_engine.send_weights(per_tensor_param) return per_tensor_param, _ = self.actor.engine.get_per_tensor_param( - adapter_name=self.config.rollout.rollout_adapter, - weight_sync_dtype=self._get_rollout_weight_sync_dtype(), + adapter_name=self.config.rollout.rollout_adapter ) await self.checkpoint_engine.send_weights(per_tensor_param) return @@ -1156,7 +1141,6 @@ async def update_weights(self, global_steps: int = None, mode: str = "auto"): layered_summon=self.layered_summon, base_sync_done=True, adapter_name=self.config.rollout.rollout_adapter, - weight_sync_dtype=self._get_rollout_weight_sync_dtype(), ) do_lora_base_sync = False @@ -1170,7 +1154,6 @@ async def update_weights(self, global_steps: int = None, mode: str = "auto"): layered_summon=self.layered_summon, base_sync_done=False, adapter_name=self.config.rollout.rollout_adapter, - weight_sync_dtype=self._get_rollout_weight_sync_dtype(), ) await self.rollout.update_weights( per_tensor_param_base, peft_config=peft_config, base_sync_done=False, global_steps=global_steps From bd152363bfb5c3b1b27aa28dc57e6f6e27bd23f7 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Tue, 25 Aug 2026 20:43:49 +0800 Subject: [PATCH 10/28] [rollout, model] refactor: generalize omni single-turn agent Signed-off-by: dongbo910220 <1275604947@qq.com> --- README.md | 12 +- docs/api/pipelines.rst | 3 +- .../qwen3_tts/run_qwen3_tts_grpo.sh | 2 +- .../test_qwen3_tts_rollout_on_cpu.py | 146 ++++++++++++++++++ verl_omni/pipelines/qwen3_tts/__init__.py | 3 +- verl_omni/pipelines/qwen3_tts/agent_loop.py | 61 -------- .../qwen3_tts/omni_rollout_adapter.py | 60 ++++++- 7 files changed, 217 insertions(+), 70 deletions(-) delete mode 100644 verl_omni/pipelines/qwen3_tts/agent_loop.py diff --git a/README.md b/README.md index 19c614da5..44681b0c7 100644 --- a/README.md +++ b/README.md @@ -191,16 +191,20 @@ Visit our [documentation](https://verl-omni.readthedocs.io/en/latest/index.html) ✅ - Qwen3-TTS - Audio-modality - Text → Audio + Qwen3-TTS + Audio-modality + Text → Audio DPO WIP - + GSPO WIP + + GRPO + ✅ + diff --git a/docs/api/pipelines.rst b/docs/api/pipelines.rst index cc62f26d1..4410cf2f0 100644 --- a/docs/api/pipelines.rst +++ b/docs/api/pipelines.rst @@ -51,7 +51,8 @@ Model Base supports_cache_engine_sleep, get_pipeline_id, ensure_pipeline_registered, get_engine_hf_overrides, get_stage_engine_extras, prepare_engine_prompt, - get_output_modalities, combine_engine_outputs + get_output_modalities, prepare_agent_sampling_params, + postprocess_agent_loop_output, combine_engine_outputs .. autoclass:: verl_omni.pipelines.model_base.DiffusionModelBase :members: register, get_class, diff --git a/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh b/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh index d58fe2960..193e727b0 100755 --- a/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh +++ b/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh @@ -96,7 +96,7 @@ export VLLM_USE_FLASHINFER_SAMPLER=0 actor_rollout_ref.rollout.enforce_eager=true \ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ actor_rollout_ref.rollout.agent.num_workers=8 \ - actor_rollout_ref.rollout.agent.default_agent_loop=qwen3_tts_single_turn \ + actor_rollout_ref.rollout.agent.default_agent_loop=omni_single_turn_agent \ "+actor_rollout_ref.rollout.engine_kwargs.vllm_omni.seed=${SEED}" \ +actor_rollout_ref.rollout.engine_kwargs.vllm_omni.output_mode=ar \ +actor_rollout_ref.rollout.engine_kwargs.vllm_omni.pipeline_name=qwen3_tts_rl \ diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index 407321e06..acc8f0671 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -13,6 +13,8 @@ # limitations under the License. """CPU contracts for Qwen3-TTS's multi-stage rollout integration.""" +import subprocess +import sys from types import SimpleNamespace import pytest @@ -23,6 +25,7 @@ from vllm import SamplingParams +from verl_omni.agent_loop.single_turn_agent_loop import OmniSingleTurnAgentLoop from verl_omni.pipelines.model_base import OmniRolloutPipelineBase from verl_omni.pipelines.qwen3_tts import omni_rollout_adapter from verl_omni.pipelines.qwen3_tts.omni_rollout_adapter import Qwen3TTSRolloutAdapter @@ -38,6 +41,22 @@ def __call__(self, text, **kwargs): return {"input_ids": list(range(len(text)))} +def test_external_module_import_registers_omni_agent_loop(): + code = """ +import vllm_omni.platforms as platforms +from vllm_omni.platforms.interface import UnspecifiedOmniPlatform + +platforms._current_omni_platform = UnspecifiedOmniPlatform() + +import verl_omni +from verl.experimental.agent_loop.agent_loop import _agent_loop_registry + +target = _agent_loop_registry[\"omni_single_turn_agent\"][\"_target_\"] +assert target == \"verl_omni.agent_loop.single_turn_agent_loop.OmniSingleTurnAgentLoop\" +""" + subprocess.run([sys.executable, "-c", code], check=True) + + def test_optional_rollout_hooks_preserve_existing_ar_defaults(): first, final = object(), object() @@ -45,9 +64,94 @@ def test_optional_rollout_hooks_preserve_existing_ar_defaults(): assert OmniRolloutPipelineBase.supports_cache_engine_sleep() assert OmniRolloutPipelineBase.prepare_engine_prompt([], None, {}) is None assert OmniRolloutPipelineBase.get_output_modalities() is None + sampling_params = {"temperature": 0.8} + assert ( + OmniRolloutPipelineBase.prepare_agent_sampling_params( + sampling_params, + rollout_config=None, + trainer_config=None, + agent_inputs={}, + ) + == sampling_params + ) + assert ( + OmniRolloutPipelineBase.postprocess_agent_loop_output( + final, + tokenizer=None, + response_length=8, + ) + is final + ) assert OmniRolloutPipelineBase.combine_engine_outputs([first, final], {}) == (final, {}) +def test_omni_single_turn_agent_resolves_registered_pipeline_adapter(): + rollout_config = SimpleNamespace(engine_kwargs={"vllm_omni": {"pipeline_name": "qwen3_tts_rl"}}) + + assert OmniSingleTurnAgentLoop._resolve_rollout_adapter(rollout_config) is Qwen3TTSRolloutAdapter + + missing_config = SimpleNamespace(engine_kwargs={"vllm_omni": {"pipeline_name": "missing"}}) + with pytest.raises(ValueError, match="requires a registered"): + OmniSingleTurnAgentLoop._resolve_rollout_adapter(missing_config) + + +@pytest.mark.asyncio +async def test_omni_single_turn_agent_delegates_policy_mapping_to_adapter(): + class Adapter: + @staticmethod + def prepare_agent_sampling_params(sampling_params, **kwargs): + assert kwargs["agent_inputs"]["session_id"] == 2 + return {**sampling_params, "seed": 123} + + @staticmethod + def postprocess_agent_loop_output(output, **kwargs): + assert kwargs["response_length"] == 4 + output.prompt_ids = [0] + return output + + class Server: + async def generate(self, **kwargs): + self.kwargs = kwargs + return SimpleNamespace( + token_ids=[101, 102], + log_probs=[-0.1, -0.2], + routed_experts=None, + num_preempted=0, + extra_fields={"audio": torch.ones(8)}, + ) + + loop = object.__new__(OmniSingleTurnAgentLoop) + loop.rollout_adapter = Adapter + loop.rollout_config = SimpleNamespace(response_length=4, full_determinism=True) + loop.response_length = 4 + loop.config = SimpleNamespace(data={"seed": 42}) + loop.server_manager = Server() + loop.tokenizer = _Tokenizer() + loop.process_multi_modal_info = lambda messages: _async_value({}) + loop.ct_build_initial_tokens = lambda *args, **kwargs: _async_value([11, 12]) + loop._assert_mm_supported = lambda has_multi_modal: None + loop._get_mm_processor_kwargs = lambda audios: {} + + result = await OmniSingleTurnAgentLoop.run.__wrapped__( + loop, + {"temperature": 0.8}, + priority=7, + raw_prompt=[{"role": "user", "content": "hello"}], + session_id=2, + ) + + assert loop.server_manager.kwargs["request_id"] == "det-7" + assert loop.server_manager.kwargs["sampling_params"]["seed"] == 123 + assert result.prompt_ids == [0] + assert result.response_ids == [101, 102] + assert result.response_logprobs == [-0.1, -0.2] + assert result.extra_fields["audio"].shape == (8,) + + +async def _async_value(value): + return value + + def test_rollout_pipeline_registers_upstream_talker(monkeypatch): registered_pipelines = [] monkeypatch.setattr( @@ -142,6 +246,48 @@ def test_rollout_adapter_combines_policy_codes_and_waveform(): assert fields["tts_text"] == "first text" +def test_rollout_adapter_prepares_sampling_and_actor_policy_sequence(): + rollout_config = SimpleNamespace(n=8, val_kwargs=SimpleNamespace(n=1)) + trainer_config = SimpleNamespace(data={"seed": 42}) + agent_inputs = { + "extra_info": {"split": "train", "id": "sample-7"}, + "session_id": 3, + "global_steps": 12, + "uid": "uid-7", + } + + seeded = Qwen3TTSRolloutAdapter.prepare_agent_sampling_params( + {"temperature": 0.8}, + rollout_config=rollout_config, + trainer_config=trainer_config, + agent_inputs=agent_inputs, + ) + assert seeded["seed"] == seeded["extra_args"]["tts_local_seed"] + + codes = torch.arange(5 * 16, dtype=torch.long).reshape(5, 16) + output = SimpleNamespace( + prompt_ids=[9, 8], + response_ids=[7, 6, 5, 4, 3], + response_mask=[1] * 5, + response_logprobs=[-0.1, -0.2, -0.3, -0.4, -0.5], + extra_fields={"tts_audio_codes": codes, "tts_text": "first text"}, + ) + + result = Qwen3TTSRolloutAdapter.postprocess_agent_loop_output( + output, + tokenizer=_Tokenizer(), + response_length=3, + ) + + assert result is output + assert result.prompt_ids == [0] + assert result.response_ids == codes[:3, 0].tolist() + assert result.response_mask == [1, 1, 1] + assert result.response_logprobs == [-0.1, -0.2, -0.3] + torch.testing.assert_close(result.extra_fields["tts_audio_codes"], codes[:3]) + assert result.extra_fields["tts_text_ids"] + + def test_server_prepares_stage_specific_sampling_params(): class Adapter: @staticmethod diff --git a/verl_omni/pipelines/qwen3_tts/__init__.py b/verl_omni/pipelines/qwen3_tts/__init__.py index 47fca4762..5831f3727 100644 --- a/verl_omni/pipelines/qwen3_tts/__init__.py +++ b/verl_omni/pipelines/qwen3_tts/__init__.py @@ -12,8 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .agent_loop import Qwen3TTSSingleTurnAgentLoop from .omni_rollout_adapter import Qwen3TTSRolloutAdapter from .talker_training_adapter import Qwen3TTSTalkerAdapter -__all__ = ["Qwen3TTSSingleTurnAgentLoop", "Qwen3TTSTalkerAdapter", "Qwen3TTSRolloutAdapter"] +__all__ = ["Qwen3TTSTalkerAdapter", "Qwen3TTSRolloutAdapter"] diff --git a/verl_omni/pipelines/qwen3_tts/agent_loop.py b/verl_omni/pipelines/qwen3_tts/agent_loop.py deleted file mode 100644 index b8a8bca18..000000000 --- a/verl_omni/pipelines/qwen3_tts/agent_loop.py +++ /dev/null @@ -1,61 +0,0 @@ -# Copyright 2026 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Single-turn agent loop that makes codec-0 the policy sequence.""" - -import torch -from verl.experimental.agent_loop.agent_loop import register -from verl.experimental.agent_loop.single_turn_agent_loop import SingleTurnAgentLoop - -from verl_omni.pipelines.qwen3_tts.rollout_utils import is_evaluation_split, with_rollout_generation_seed -from verl_omni.pipelines.qwen3_tts.talker_forward import TEXT_PROMPT_TRAILER_TOKENS, build_assistant_text - - -@register("qwen3_tts_single_turn") -class Qwen3TTSSingleTurnAgentLoop(SingleTurnAgentLoop): - async def run(self, sampling_params, **kwargs): - extra_info = kwargs.get("extra_info") - evaluation = is_evaluation_split(extra_info) - candidate_count = self.rollout_config.val_kwargs.n if evaluation else self.rollout_config.n - sampling_params = with_rollout_generation_seed( - sampling_params, - extra_info, - session_id=kwargs.get("session_id"), - global_steps=kwargs.get("global_steps"), - uid=kwargs.get("uid"), - base_seed=int(self.config.data.get("seed", 0)), - require_session_id=int(candidate_count) > 1, - ) - output = await super().run(sampling_params, **kwargs) - extra = output.extra_fields - codes, text = extra.get("tts_audio_codes"), extra.get("tts_text") - if codes is None or text is None: - raise RuntimeError("Qwen3-TTS rollout did not return codec codes and text.") - codes = torch.as_tensor(codes, dtype=torch.long) - if codes.ndim != 2 or codes.shape[-1] != 16: - raise ValueError("Qwen3-TTS codec codes must have shape (frames, 16).") - codes = codes[: self.response_length] - policy_ids = codes[:, 0].tolist() - if not policy_ids: - raise RuntimeError("Qwen3-TTS rollout returned an empty codec trajectory.") - if output.response_logprobs is not None: - if len(output.response_logprobs) < len(policy_ids): - raise RuntimeError("Qwen3-TTS rollout logprobs are shorter than the policy trajectory.") - output.response_logprobs = output.response_logprobs[: len(policy_ids)] - text_ids = self.tokenizer(build_assistant_text(str(text)), return_tensors="pt", padding=False)["input_ids"] - extra["tts_text_ids"] = text_ids[:, :-TEXT_PROMPT_TRAILER_TOKENS].reshape(-1).tolist() - extra["tts_audio_codes"] = codes - output.prompt_ids = [0] - output.response_ids = policy_ids - output.response_mask = [1] * len(policy_ids) - return output diff --git a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py index e0a1a5364..1dddbde88 100644 --- a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py @@ -25,8 +25,14 @@ from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE from verl_omni.pipelines.model_base import OmniRolloutPipelineBase -from verl_omni.pipelines.qwen3_tts.rollout_utils import align_audio_codes, append_tensor_chunk +from verl_omni.pipelines.qwen3_tts.rollout_utils import ( + align_audio_codes, + append_tensor_chunk, + is_evaluation_split, + with_rollout_generation_seed, +) from verl_omni.pipelines.qwen3_tts.talker_forward import ( + TEXT_PROMPT_TRAILER_TOKENS, build_assistant_text, load_speaker_xvector, require_auto_language, @@ -125,6 +131,58 @@ def get_stage_engine_extras(cls, stage_id, pipeline_mode="full"): cls._check_mode(pipeline_mode) return {"max_model_len": 65536, "max_num_batched_tokens": 65536} if stage_id == 1 else {} + @classmethod + def prepare_agent_sampling_params( + cls, + sampling_params, + *, + rollout_config, + trainer_config, + agent_inputs, + ): + extra_info = agent_inputs.get("extra_info") + evaluation = is_evaluation_split(extra_info) + candidate_count = rollout_config.val_kwargs.n if evaluation else rollout_config.n + return with_rollout_generation_seed( + sampling_params, + extra_info, + session_id=agent_inputs.get("session_id"), + global_steps=agent_inputs.get("global_steps"), + uid=agent_inputs.get("uid"), + base_seed=int(trainer_config.data.get("seed", 0)), + require_session_id=int(candidate_count) > 1, + ) + + @classmethod + def postprocess_agent_loop_output(cls, output, *, tokenizer, response_length): + extra = output.extra_fields + codes, text = extra.get("tts_audio_codes"), extra.get("tts_text") + if codes is None or text is None: + raise RuntimeError("Qwen3-TTS rollout did not return codec codes and text.") + codes = torch.as_tensor(codes, dtype=torch.long) + if codes.ndim != 2 or codes.shape[-1] != 16: + raise ValueError("Qwen3-TTS codec codes must have shape (frames, 16).") + codes = codes[:response_length] + policy_ids = codes[:, 0].tolist() + if not policy_ids: + raise RuntimeError("Qwen3-TTS rollout returned an empty codec trajectory.") + if output.response_logprobs is not None: + if len(output.response_logprobs) < len(policy_ids): + raise RuntimeError("Qwen3-TTS rollout logprobs are shorter than the policy trajectory.") + output.response_logprobs = output.response_logprobs[: len(policy_ids)] + text_ids = tokenizer(build_assistant_text(str(text)), return_tensors="pt", padding=False)["input_ids"] + text_ids = torch.as_tensor(text_ids, dtype=torch.long) + if text_ids.ndim == 1: + text_ids = text_ids.unsqueeze(0) + if text_ids.ndim != 2 or text_ids.shape[1] <= TEXT_PROMPT_TRAILER_TOKENS: + raise ValueError("Qwen3-TTS assistant text tokenization returned an invalid sequence.") + extra["tts_text_ids"] = text_ids[:, :-TEXT_PROMPT_TRAILER_TOKENS].reshape(-1).tolist() + extra["tts_audio_codes"] = codes + output.prompt_ids = [0] + output.response_ids = policy_ids + output.response_mask = [1] * len(policy_ids) + return output + @classmethod def prepare_engine_prompt(cls, prompt_ids, model_config, multi_modal_data, mm_processor_kwargs=None): text = model_config.tokenizer.decode(prompt_ids, skip_special_tokens=True).strip() From c7277cc15bf26e2b59619d1a7b308c271df844d2 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Tue, 25 Aug 2026 20:51:25 +0800 Subject: [PATCH 11/28] [rollout, tests] refactor: reuse standard single-turn flow Signed-off-by: dongbo910220 <1275604947@qq.com> --- .../test_qwen3_tts_rollout_on_cpu.py | 48 +++++++------------ 1 file changed, 18 insertions(+), 30 deletions(-) diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index acc8f0671..1220e601e 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -23,6 +23,7 @@ pytest.importorskip("verl") pytest.importorskip("vllm_omni") +from verl.experimental.agent_loop.single_turn_agent_loop import SingleTurnAgentLoop from vllm import SamplingParams from verl_omni.agent_loop.single_turn_agent_loop import OmniSingleTurnAgentLoop @@ -96,43 +97,36 @@ def test_omni_single_turn_agent_resolves_registered_pipeline_adapter(): @pytest.mark.asyncio -async def test_omni_single_turn_agent_delegates_policy_mapping_to_adapter(): +async def test_omni_single_turn_agent_delegates_policy_mapping_to_adapter(monkeypatch): class Adapter: - @staticmethod - def prepare_agent_sampling_params(sampling_params, **kwargs): - assert kwargs["agent_inputs"]["session_id"] == 2 - return {**sampling_params, "seed": 123} - @staticmethod def postprocess_agent_loop_output(output, **kwargs): assert kwargs["response_length"] == 4 output.prompt_ids = [0] return output - class Server: - async def generate(self, **kwargs): - self.kwargs = kwargs - return SimpleNamespace( - token_ids=[101, 102], - log_probs=[-0.1, -0.2], - routed_experts=None, - num_preempted=0, - extra_fields={"audio": torch.ones(8)}, - ) + upstream_output = SimpleNamespace( + prompt_ids=[11, 12], + response_ids=[101, 102], + response_mask=[1, 1], + response_logprobs=[-0.1, -0.2], + extra_fields={"audio": torch.ones(8)}, + ) + + async def upstream_run(_self, sampling_params, **kwargs): + assert sampling_params == {"temperature": 0.8} + assert kwargs["priority"] == 7 + return upstream_output + + monkeypatch.setattr(SingleTurnAgentLoop, "run", upstream_run) loop = object.__new__(OmniSingleTurnAgentLoop) loop.rollout_adapter = Adapter - loop.rollout_config = SimpleNamespace(response_length=4, full_determinism=True) + loop.rollout_config = SimpleNamespace(response_length=4) loop.response_length = 4 - loop.config = SimpleNamespace(data={"seed": 42}) - loop.server_manager = Server() loop.tokenizer = _Tokenizer() - loop.process_multi_modal_info = lambda messages: _async_value({}) - loop.ct_build_initial_tokens = lambda *args, **kwargs: _async_value([11, 12]) - loop._assert_mm_supported = lambda has_multi_modal: None - loop._get_mm_processor_kwargs = lambda audios: {} - result = await OmniSingleTurnAgentLoop.run.__wrapped__( + result = await OmniSingleTurnAgentLoop.run( loop, {"temperature": 0.8}, priority=7, @@ -140,18 +134,12 @@ async def generate(self, **kwargs): session_id=2, ) - assert loop.server_manager.kwargs["request_id"] == "det-7" - assert loop.server_manager.kwargs["sampling_params"]["seed"] == 123 assert result.prompt_ids == [0] assert result.response_ids == [101, 102] assert result.response_logprobs == [-0.1, -0.2] assert result.extra_fields["audio"].shape == (8,) -async def _async_value(value): - return value - - def test_rollout_pipeline_registers_upstream_talker(monkeypatch): registered_pipelines = [] monkeypatch.setattr( From cdbe1edf477b9dc070db1e98cab5c39e7ac252d2 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Tue, 25 Aug 2026 22:42:55 +0800 Subject: [PATCH 12/28] [model, tests] fix: restore Qwen3-TTS config default Signed-off-by: dongbo910220 <1275604947@qq.com> --- .../pipelines/test_qwen3_tts_transformers_compat_on_cpu.py | 3 ++- verl_omni/pipelines/qwen3_tts/talker_training_adapter.py | 6 +++++- verl_omni/pipelines/qwen3_tts/transformers_compat.py | 7 +++++++ 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py b/tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py index 6e3da9759..a3b61c07f 100644 --- a/tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py +++ b/tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py @@ -44,6 +44,7 @@ def test_qwen_tts_tiny_model_constructs_and_forwards_without_source_patch(): with compat.qwen3_tts_import_context(): config_module = importlib.import_module("qwen_tts.core.models.configuration_qwen3_tts") model_module = importlib.import_module("qwen_tts.core.models.modeling_qwen3_tts") + compat.patch_qwen3_tts_config_defaults(config_module.Qwen3TTSConfig) predictor = { "vocab_size": 32, @@ -72,7 +73,6 @@ def test_qwen_tts_tiny_model_constructs_and_forwards_without_source_patch(): "text_vocab_size": 80, "spk_id": {}, "codec_language_id": {}, - "pad_token_id": None, "rope_scaling": { "rope_type": "default", "type": "default", @@ -86,6 +86,7 @@ def test_qwen_tts_tiny_model_constructs_and_forwards_without_source_patch(): tts_model_type="custom", tokenizer_type="12hz", ) + assert config.talker_config.pad_token_id is None model = model_module.Qwen3TTSForConditionalGeneration(config) output = model.talker( inputs_embeds=torch.randn(2, 5, 8), diff --git a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py index f3dcc5cd4..9f7b96be4 100644 --- a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py @@ -25,7 +25,10 @@ require_auto_language, tts_actor_logits, ) -from verl_omni.pipelines.qwen3_tts.transformers_compat import qwen3_tts_import_context +from verl_omni.pipelines.qwen3_tts.transformers_compat import ( + patch_qwen3_tts_config_defaults, + qwen3_tts_import_context, +) logger = logging.getLogger(__name__) _PASSTHROUGH_TEMPLATE = "{% for message in messages %}{{ message['content'] }}{% endfor %}" @@ -65,6 +68,7 @@ def register_qwen3_tts_automodel() -> None: from transformers import AutoConfig, AutoModelForMultimodalLM + patch_qwen3_tts_config_defaults(config_cls) try: AutoConfig.register(getattr(config_cls, "model_type", "qwen3_tts"), config_cls) except ValueError: diff --git a/verl_omni/pipelines/qwen3_tts/transformers_compat.py b/verl_omni/pipelines/qwen3_tts/transformers_compat.py index 3784386ba..94fcd68a0 100644 --- a/verl_omni/pipelines/qwen3_tts/transformers_compat.py +++ b/verl_omni/pipelines/qwen3_tts/transformers_compat.py @@ -38,6 +38,13 @@ def wrapper(*args, input_embeds=None, inputs_embeds=None, cache_position=None, * return wrapper +def patch_qwen3_tts_config_defaults(config_cls) -> None: + """Restore the config default that qwen-tts expects from Transformers 4.x.""" + talker_config_cls = getattr(config_cls, "sub_configs", {}).get("talker_config") + if talker_config_cls is not None and not hasattr(talker_config_cls, "pad_token_id"): + talker_config_cls.pad_token_id = None + + @contextmanager def qwen3_tts_import_context(): """Expose the TF5 APIs expected while qwen-tts binds its imports. From a772a8c2b41ba0b8cf8d8ca03d53dc068beb77d0 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Fri, 28 Aug 2026 12:34:26 +0800 Subject: [PATCH 13/28] [doc] chore: cite TTS GRPO references Signed-off-by: dongbo910220 <1275604947@qq.com> --- examples/grpo_trainer/qwen3_tts/README.md | 43 ++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/examples/grpo_trainer/qwen3_tts/README.md b/examples/grpo_trainer/qwen3_tts/README.md index ac81d13ac..1cf77c550 100644 --- a/examples/grpo_trainer/qwen3_tts/README.md +++ b/examples/grpo_trainer/qwen3_tts/README.md @@ -1,6 +1,6 @@ # Qwen3-TTS GRPO with an audio reward -Last updated: 08/23/2026. +Last updated: 08/28/2026. This example full-parameter tunes the codec-0 policy of `Qwen/Qwen3-TTS-12Hz-0.6B-Base`. It uses verl's stock GRPO advantage, @@ -17,6 +17,37 @@ client because its published Transformers environment conflicts with the Transformers 5.x vLLM stack. The Trainer and `AudioRewardManager` do not contain SpeechJudge-specific branches, candidate masks, ASR gates, or custom losses. +## Algorithm background + +The rollout and update flow follows the TTS application of GRPO described in +[Group Relative Policy Optimization for Text-to-Speech with Large Language +Models](https://arxiv.org/abs/2509.18798): sample a group of speech-token +trajectories for each text prompt, decode every trajectory to a waveform, +compute scalar audio rewards, derive group-relative advantages, and replay the +sampled policy tokens for the GRPO update with an optional reference-model KL +penalty. That paper uses a specific ASR-based CER-and-NLL reward. This example +keeps the same GRPO structure but intentionally places reward computation behind +the generic audio HTTP protocol, so it is an integration of the paper-supported +algorithm rather than an exact reproduction of its reward function or recipe. + +The policy boundary follows the released Qwen3-TTS architecture. The +[Qwen3-TTS Technical Report](https://arxiv.org/abs/2601.15621) describes the +12 Hz tokenizer as a 16-layer multi-codebook representation: the first layer +captures semantic content, the other 15 RVQ layers add acoustic detail, the +Talker backbone predicts codec-0, and its MTP module predicts the residual +codebooks. This example therefore optimizes codec-0 as the policy sequence while +retaining the complete 16-codebook rollout for actor replay and waveform +decoding. Here, "teacher-forced" means that the actor replays the tokens sampled +during rollout as fixed history; it does not mean that ground-truth speech +tokens are used. + +[SpeechAlign](https://arxiv.org/abs/2404.05600) is a related multi-codebook +speech-alignment precedent: its autoregressive model generates the first of +eight RVQ codebooks and a pretrained non-autoregressive model supplies the +remaining layers. It supports treating the first codebook as the optimized +sequence while preserving residual acoustic codebooks, but it uses preference +optimization rather than the GRPO objective implemented here. + ## Install Install the engine before the training stack: @@ -118,3 +149,13 @@ This smoke proves rollout, finite audio reward, optimizer update, weight sync, and checkpoint wiring only. It is not evidence that GRPO improves held-out speech quality; that requires the complete fixed-validation curve and paired human listening evaluation. + +## References + +- Chang Liu, Ya-Jun Hu, Ying-Ying Gao, Shi-Lei Zhang, and Zhen-Hua Ling. + [Group Relative Policy Optimization for Text-to-Speech with Large Language + Models](https://arxiv.org/abs/2509.18798), 2025. +- Hangrui Hu et al. [Qwen3-TTS Technical + Report](https://arxiv.org/abs/2601.15621), 2026. +- Dong Zhang et al. [SpeechAlign: Aligning Speech Generation to Human + Preferences](https://arxiv.org/abs/2404.05600), 2024. From 1d95ccf2852a6e5e9fcc5d879281d98ccc512b5c Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Fri, 28 Aug 2026 16:39:36 +0800 Subject: [PATCH 14/28] [model, rollout, tests, doc] fix: address Qwen3-TTS review feedback Signed-off-by: dongbo910220 <1275604947@qq.com> --- docs/api/pipelines.rst | 8 +- .../contributing/integrating_an_omni_model.md | 15 +++- docs/start/http_scorer.md | 2 +- examples/grpo_trainer/qwen3_tts/README.md | 63 ++++++-------- tests/gpu_smoke/run_gpu_smoke_omni_e2e.sh | 4 + ...st_qwen3_tts_transformers_compat_on_cpu.py | 7 +- .../create_dummy_qwen3_tts_grpo_data.py | 73 ++++++++++++++++ tests/special_e2e/qwen3_tts_smoke_scorer.py | 83 +++++++++++++++++++ tests/special_e2e/run_qwen3_tts_grpo_smoke.sh | 79 ++++++++++++++++++ tests/trainer/omni/test_main_omni_on_cpu.py | 14 ++-- .../test_qwen3_tts_rollout_on_cpu.py | 10 +-- tests/workers/test_omni_fsdp_engine_on_cpu.py | 58 +++++++++++-- verl_omni/pipelines/model_base.py | 29 +++++-- .../qwen3_tts/omni_rollout_adapter.py | 26 +++--- .../qwen3_tts/talker_training_adapter.py | 61 +++++--------- .../qwen3_tts/transformers_compat.py | 13 +-- verl_omni/workers/config/omni/model.py | 32 +++++-- verl_omni/workers/engine/fsdp/omni_impl.py | 39 +++++---- 18 files changed, 455 insertions(+), 161 deletions(-) create mode 100644 tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py create mode 100644 tests/special_e2e/qwen3_tts_smoke_scorer.py create mode 100755 tests/special_e2e/run_qwen3_tts_grpo_smoke.sh diff --git a/docs/api/pipelines.rst b/docs/api/pipelines.rst index 4410cf2f0..5a04401f8 100644 --- a/docs/api/pipelines.rst +++ b/docs/api/pipelines.rst @@ -42,17 +42,17 @@ Model Base .. autoclass:: verl_omni.pipelines.model_base.OmniModelBase :members: register, get_class, get_class_by_name, + load_hf_config, get_model_class, get_strip_modules, configure_processor, configure_tokenizer, configure_model, prepare_model_inputs .. autoclass:: verl_omni.pipelines.model_base.OmniRolloutPipelineBase :members: register, get_class, build_stage_configs, rollout_flags, weight_sync_stage_ids, - supports_cache_engine_sleep, get_pipeline_id, - ensure_pipeline_registered, get_engine_hf_overrides, + get_pipeline_id, ensure_pipeline_registered, get_engine_hf_overrides, get_stage_engine_extras, prepare_engine_prompt, - get_output_modalities, prepare_agent_sampling_params, - postprocess_agent_loop_output, combine_engine_outputs + prepare_agent_sampling_params, postprocess_agent_loop_output, + combine_engine_outputs .. autoclass:: verl_omni.pipelines.model_base.DiffusionModelBase :members: register, get_class, diff --git a/docs/contributing/integrating_an_omni_model.md b/docs/contributing/integrating_an_omni_model.md index d15adc90e..5783176e4 100644 --- a/docs/contributing/integrating_an_omni_model.md +++ b/docs/contributing/integrating_an_omni_model.md @@ -60,6 +60,13 @@ adapt each implementation to your model's architecture: `module._no_split_modules` to the correct decoder layer class for FSDP. This method runs before FSDP wrapping and LoRA injection. +- **`load_hf_config(...)` and `get_model_class()`** (optional): Override these + only when the architecture is not supported by `AutoConfig` or + `AutoModelForMultimodalLM`. Returning `None` from `get_model_class` keeps the + default auto-model path. The FSDP engine still owns `from_pretrained`; + Qwen3-TTS selects the config and model classes published by `qwen-tts` + directly instead of registering them globally with Transformers. + - **`prepare_model_inputs(model_inputs, micro_batch, model_config)`** (optional): Validate model-native trajectory or conditioning data retained by rollout and add it to the actor forward inputs. Per-sample rollout data starts @@ -99,11 +106,11 @@ Optional overrides fall into four groups: - Pipeline setup: `ensure_pipeline_registered`, `get_engine_hf_overrides`, and `get_stage_engine_extras`. -- Resource behavior: `weight_sync_stage_ids` and - `supports_cache_engine_sleep`. +- Resource behavior: `weight_sync_stage_ids`. - Request construction: `prepare_engine_prompt`. -- Multi-stage output retention: `get_output_modalities` and - `combine_engine_outputs`. +- Multi-stage output assembly: `combine_engine_outputs`. The rollout server + derives retained output modalities from stages marked `final_output` in the + pipeline topology. Their defaults preserve the existing single-output AR behavior. Override only the hooks required by the model. For example, Qwen3-TTS synchronizes actor diff --git a/docs/start/http_scorer.md b/docs/start/http_scorer.md index cae6d9da2..d46347932 100644 --- a/docs/start/http_scorer.md +++ b/docs/start/http_scorer.md @@ -1,7 +1,7 @@ (http_scorer)= # Using an External HTTP Scorer Service -Last updated: 08/09/2026 +Last updated: 08/28/2026 VeRL-Omni ships a generic HTTP reward client (`verl_omni.utils.reward_score.http_scorer_client`) that sends generated images to an external scorer service over HTTP and returns the score. This is useful when your reward model is too large to co-locate with training, needs a different runtime (e.g., a separate GPU pool), or is shared across multiple experiments. diff --git a/examples/grpo_trainer/qwen3_tts/README.md b/examples/grpo_trainer/qwen3_tts/README.md index 1cf77c550..afad6e556 100644 --- a/examples/grpo_trainer/qwen3_tts/README.md +++ b/examples/grpo_trainer/qwen3_tts/README.md @@ -12,41 +12,22 @@ The launcher follows the V1 omni-model integration guide: it calls `verl_omni.trainer.main_omni` and expresses the recipe as CLI overrides on the standard `omni_trainer` config, without a model-specific Trainer or config tree. -SpeechJudge-BTRM is one possible scorer. It runs behind the generic audio HTTP -client because its published Transformers environment conflicts with the -Transformers 5.x vLLM stack. The Trainer and `AudioRewardManager` do not contain -SpeechJudge-specific branches, candidate masks, ASR gates, or custom losses. +SpeechJudge-BTRM is one possible pointwise scorer. SpeechJudge's published vLLM +entry point targets the pairwise generative GRM, while BTRM uses a scalar reward +head with Transformers. This example therefore keeps reward inference behind the +generic audio HTTP protocol instead of adding a SpeechJudge-specific Trainer +path. The scorer may run in a separate environment from the Transformers 5.x +vLLM training stack. ## Algorithm background -The rollout and update flow follows the TTS application of GRPO described in -[Group Relative Policy Optimization for Text-to-Speech with Large Language -Models](https://arxiv.org/abs/2509.18798): sample a group of speech-token -trajectories for each text prompt, decode every trajectory to a waveform, -compute scalar audio rewards, derive group-relative advantages, and replay the -sampled policy tokens for the GRPO update with an optional reference-model KL -penalty. That paper uses a specific ASR-based CER-and-NLL reward. This example -keeps the same GRPO structure but intentionally places reward computation behind -the generic audio HTTP protocol, so it is an integration of the paper-supported -algorithm rather than an exact reproduction of its reward function or recipe. - -The policy boundary follows the released Qwen3-TTS architecture. The -[Qwen3-TTS Technical Report](https://arxiv.org/abs/2601.15621) describes the -12 Hz tokenizer as a 16-layer multi-codebook representation: the first layer -captures semantic content, the other 15 RVQ layers add acoustic detail, the -Talker backbone predicts codec-0, and its MTP module predicts the residual -codebooks. This example therefore optimizes codec-0 as the policy sequence while -retaining the complete 16-codebook rollout for actor replay and waveform -decoding. Here, "teacher-forced" means that the actor replays the tokens sampled -during rollout as fixed history; it does not mean that ground-truth speech -tokens are used. - -[SpeechAlign](https://arxiv.org/abs/2404.05600) is a related multi-codebook -speech-alignment precedent: its autoregressive model generates the first of -eight RVQ codebooks and a pretrained non-autoregressive model supplies the -remaining layers. It supports treating the first codebook as the optimized -sequence while preserving residual acoustic codebooks, but it uses preference -optimization rather than the GRPO objective implemented here. +This recipe applies the paper's TTS GRPO flow to Qwen3-TTS: grouped codec-token +rollouts are decoded, scored, converted to group-relative advantages, and +replayed with optional reference KL. It optimizes codec-0, the autoregressive +policy sequence described by the Qwen3-TTS architecture, while retaining all 16 +codebooks for replay and waveform decoding. The HTTP scorer is configurable, so +this is not an exact reproduction of the paper's CER-and-NLL reward. See the +references below for the algorithm and multi-codebook design details. ## Install @@ -110,8 +91,11 @@ For SpeechJudge-BTRM, deploy the official [`AmphionTeam/SpeechJudge`](https://github.com/AmphionTeam/SpeechJudge) code and [`RMSnow/SpeechJudge-BTRM`](https://huggingface.co/RMSnow/SpeechJudge-BTRM) checkpoint in a separate environment, then expose its pointwise score through -this protocol. Pin the SpeechJudge source revision and runtime versions in the -service deployment. SpeechJudge-BTRM is licensed CC-BY-NC-4.0. +this protocol. The official [`main_grm_vllm.py`](https://github.com/AmphionTeam/SpeechJudge/blob/master/infer/main_grm_vllm.py) +runs a different, pairwise generative GRM path; the BTRM entry point is +[`main_btrm.py`](https://github.com/AmphionTeam/SpeechJudge/blob/master/infer/main_btrm.py). +Pin the SpeechJudge source revision and runtime versions in the service +deployment. SpeechJudge-BTRM is licensed CC-BY-NC-4.0. ## Train @@ -139,17 +123,22 @@ investigation. For a two-update implementation smoke test: ```bash -TOTAL_TRAINING_STEPS=2 TEST_FREQ=-1 SAVE_FREQ=1 RESUME_MODE=disable \ +TOTAL_TRAINING_STEPS=2 TEST_FREQ=-1 SAVE_FREQ=-1 RESUME_MODE=disable \ OUTPUT_DIR=outputs/qwen3_tts_grpo_smoke \ bash examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh \ trainer.val_before_train=false trainer.log_val_generations=0 ``` -This smoke proves rollout, finite audio reward, optimizer update, weight sync, -and checkpoint wiring only. It is not evidence that GRPO improves held-out +This smoke proves rollout, finite audio reward, optimizer update, and +post-update weight sync only. It is not evidence that GRPO improves held-out speech quality; that requires the complete fixed-validation curve and paired human listening evaluation. +The CI-oriented wrapper at +[`tests/special_e2e/run_qwen3_tts_grpo_smoke.sh`](../../../tests/special_e2e/run_qwen3_tts_grpo_smoke.sh) +creates deterministic fixtures, starts a duration-only test scorer, and runs two +updates with the official 0.6B Base model. + ## References - Chang Liu, Ya-Jun Hu, Ying-Ying Gao, Shi-Lei Zhang, and Zhen-Hua Ling. diff --git a/tests/gpu_smoke/run_gpu_smoke_omni_e2e.sh b/tests/gpu_smoke/run_gpu_smoke_omni_e2e.sh index 979213742..dd8d5431d 100644 --- a/tests/gpu_smoke/run_gpu_smoke_omni_e2e.sh +++ b/tests/gpu_smoke/run_gpu_smoke_omni_e2e.sh @@ -21,4 +21,8 @@ run_test 1 "Qwen3-Omni multimodal offline MLLM DPO LoRA e2e" \ env CUDA_VISIBLE_DEVICES="${CUDA_DEVICE_LIST}" NUM_GPUS=2 \ bash tests/special_e2e/run_qwen3_omni_multimodal_offline_mllm_dpo_lora_smoke.sh "${omni_trainer_args[@]}" +run_test 2 "Qwen3-TTS Talker full-parameter GRPO e2e" \ + env CUDA_VISIBLE_DEVICES="${CUDA_DEVICE_LIST}" NUM_GPUS=2 \ + bash tests/special_e2e/run_qwen3_tts_grpo_smoke.sh "${omni_trainer_args[@]}" + gpu_smoke_summary diff --git a/tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py b/tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py index a3b61c07f..724e7b532 100644 --- a/tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py +++ b/tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py @@ -19,7 +19,6 @@ import pytest import torch -from packaging.version import Version ROOT = Path(__file__).parents[2] @@ -35,15 +34,17 @@ def _load_compat_module(): def test_qwen_tts_tiny_model_constructs_and_forwards_without_source_patch(): transformers = pytest.importorskip("transformers") - if Version(transformers.__version__).major < 5: - pytest.skip("Transformers 5.x compatibility test") + assert int(transformers.__version__.split(".", maxsplit=1)[0]) >= 5 if importlib.util.find_spec("qwen_tts") is None: pytest.skip("qwen-tts is an optional dependency") compat = _load_compat_module() + rope_utils = importlib.import_module("transformers.modeling_rope_utils") + original_rope_functions = rope_utils.ROPE_INIT_FUNCTIONS with compat.qwen3_tts_import_context(): config_module = importlib.import_module("qwen_tts.core.models.configuration_qwen3_tts") model_module = importlib.import_module("qwen_tts.core.models.modeling_qwen3_tts") + assert rope_utils.ROPE_INIT_FUNCTIONS is original_rope_functions compat.patch_qwen3_tts_config_defaults(config_module.Qwen3TTSConfig) predictor = { diff --git a/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py b/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py new file mode 100644 index 000000000..9e2f9836a --- /dev/null +++ b/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Create deterministic Qwen3-TTS GRPO smoke data and speaker conditioning.""" + +import argparse +import json +from pathlib import Path + +import pandas as pd + +TRAIN_TEXTS = ( + "Please read this sentence at a calm and steady pace.", + "A short speech sample checks the complete training path.", + "Clear pronunciation makes this audio easy to inspect.", + "The weather is pleasant and the morning train is on time.", + "Four simple prompts are enough for one smoke-test batch.", + "This second batch verifies another optimizer update.", + "Generated speech is decoded before the reward is computed.", + "The final checkpoint confirms that training completed.", +) +VALIDATION_TEXTS = ( + "This is fixed validation sample one.", + "This is fixed validation sample two.", + "This is fixed validation sample three.", + "This is fixed validation sample four.", +) + + +def _row(text: str, sample_id: str, split: str, generation_seed: int | None = None) -> dict: + extra_info = {"id": sample_id, "split": split} + if generation_seed is not None: + extra_info["generation_seed"] = generation_seed + return { + "data_source": "tts", + "prompt": [{"role": "user", "content": text}], + "reward_model": {"style": "model", "ground_truth": text}, + "extra_info": extra_info, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, required=True) + args = parser.parse_args() + args.output_dir.mkdir(parents=True, exist_ok=True) + + train_rows = [_row(text, f"train-{index}", "train") for index, text in enumerate(TRAIN_TEXTS)] + validation_rows = [ + _row(text, f"validation-{index}", "validation", 1_000 + index) for index, text in enumerate(VALIDATION_TEXTS) + ] + pd.DataFrame(train_rows).to_parquet(args.output_dir / "train.parquet", index=False) + pd.DataFrame(validation_rows).to_parquet(args.output_dir / "validation.parquet", index=False) + + # Qwen3-TTS Base expects a 1024-dimensional speaker x-vector. A unit-norm + # deterministic fixture is sufficient for execution testing. + speaker = [1.0 / 32.0] * 1024 + (args.output_dir / "speaker.json").write_text(json.dumps(speaker), encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/tests/special_e2e/qwen3_tts_smoke_scorer.py b/tests/special_e2e/qwen3_tts_smoke_scorer.py new file mode 100644 index 000000000..e20b2162b --- /dev/null +++ b/tests/special_e2e/qwen3_tts_smoke_scorer.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Deterministic duration reward service for the Qwen3-TTS execution smoke.""" + +import argparse +import base64 +import json +import math +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + + +class _Handler(BaseHTTPRequestHandler): + server_version = "Qwen3TTSSmokeScorer/1.0" + + def _write_json(self, status: int, payload: dict) -> None: + body = json.dumps(payload, allow_nan=False).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self) -> None: + if self.path != "/health": + self._write_json(404, {"error": "not found"}) + return + self._write_json(200, {"status": "ready", "reward_mode": "duration_smoke_only"}) + + def do_POST(self) -> None: + if self.path != "/score": + self._write_json(404, {"error": "not found"}) + return + try: + content_length = int(self.headers.get("Content-Length", "0")) + payload = json.loads(self.rfile.read(content_length)) + if payload.get("protocol_version") != "1": + raise ValueError("protocol_version must be '1'") + num_samples = int(payload["num_samples"]) + sample_rate = int(payload["sample_rate"]) + waveform = base64.b64decode(payload["waveform_f32_base64"], validate=True) + if num_samples <= 0 or sample_rate <= 0 or len(waveform) != num_samples * 4: + raise ValueError("invalid float32 waveform shape") + duration_s = num_samples / sample_rate + if not math.isfinite(duration_s): + raise ValueError("duration must be finite") + result = {"score": duration_s, "duration_s": duration_s, "reward_mode": "duration_smoke_only"} + with self.server.log_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(result, allow_nan=False) + "\n") + self._write_json(200, result) + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: + self._write_json(400, {"error": str(exc)}) + + def log_message(self, format: str, *args) -> None: + return + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, required=True) + parser.add_argument("--log", type=Path, required=True) + args = parser.parse_args() + args.log.parent.mkdir(parents=True, exist_ok=True) + server = ThreadingHTTPServer((args.host, args.port), _Handler) + server.log_path = args.log + server.serve_forever() + + +if __name__ == "__main__": + main() diff --git a/tests/special_e2e/run_qwen3_tts_grpo_smoke.sh b/tests/special_e2e/run_qwen3_tts_grpo_smoke.sh new file mode 100755 index 000000000..2a5a7dfb4 --- /dev/null +++ b/tests/special_e2e/run_qwen3_tts_grpo_smoke.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# Qwen3-TTS full-parameter GRPO e2e smoke: real 0.6B model, two updates. +# This validates execution only; the duration scorer is not a quality reward. + +set -xeuo pipefail + +export NCCL_IB_DISABLE=1 +export CPATH=/usr/include${CPATH:+:${CPATH}} +export RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0 + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +cd "${REPO_ROOT}" + +NUM_GPUS="${NUM_GPUS:-2}" +[[ "${NUM_GPUS}" == "2" ]] || { echo "Qwen3-TTS smoke requires exactly two GPUs" >&2; exit 2; } +PYTHON_BIN="${PYTHON_BIN:-python3}" +MODEL_REPO="${MODEL_REPO:-Qwen/Qwen3-TTS-12Hz-0.6B-Base}" +WORK_DIR="${WORK_DIR:-${TMPDIR:-/tmp}/qwen3_tts_grpo_smoke_${USER:-user}_$$}" +DATA_DIR="${WORK_DIR}/data" +OUTPUT_DIR="${OUTPUT_DIR:-${WORK_DIR}/output}" +SCORER_LOG="${WORK_DIR}/scorer_requests.jsonl" +mkdir -p "${WORK_DIR}" "${OUTPUT_DIR}" + +if ! "${PYTHON_BIN}" -c 'import importlib.util; modules=("qwen_tts", "onnxruntime", "soundfile", "librosa", "sox"); raise SystemExit(any(importlib.util.find_spec(name) is None for name in modules))'; then + uv pip install --python "${PYTHON_BIN}" -e ".[tts]" + uv pip install --python "${PYTHON_BIN}" --no-deps qwen-tts==0.1.1 +fi + +MODEL_PATH="${MODEL_PATH:-}" +if [[ -z "${MODEL_PATH}" ]]; then + MODEL_PATH="$("${PYTHON_BIN}" -c \ + "from huggingface_hub import snapshot_download; print(snapshot_download('${MODEL_REPO}'))")" +fi +[[ -f "${MODEL_PATH}/config.json" ]] || { echo "Invalid MODEL_PATH: ${MODEL_PATH}" >&2; exit 2; } + +"${PYTHON_BIN}" tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py --output-dir "${DATA_DIR}" + +SCORER_PORT="${SCORER_PORT:-$("${PYTHON_BIN}" -c \ + 'import socket; s=socket.socket(); s.bind(("127.0.0.1", 0)); print(s.getsockname()[1]); s.close()')}" +"${PYTHON_BIN}" tests/special_e2e/qwen3_tts_smoke_scorer.py \ + --port "${SCORER_PORT}" --log "${SCORER_LOG}" >"${WORK_DIR}/scorer.log" 2>&1 & +SCORER_PID=$! +cleanup() { + kill "${SCORER_PID}" 2>/dev/null || true + wait "${SCORER_PID}" 2>/dev/null || true +} +trap cleanup EXIT + +for _ in $(seq 1 60); do + if curl -fsS "http://127.0.0.1:${SCORER_PORT}/health" >"${WORK_DIR}/scorer_health.json"; then + break + fi + sleep 1 +done +curl -fsS "http://127.0.0.1:${SCORER_PORT}/health" + +MODEL_PATH="${MODEL_PATH}" \ +TRAIN_FILE="${DATA_DIR}/train.parquet" \ +VAL_FILE="${DATA_DIR}/validation.parquet" \ +SPK_EMBED_PATH="${DATA_DIR}/speaker.json" \ +SCORER_URL="http://127.0.0.1:${SCORER_PORT}/score" \ +OUTPUT_DIR="${OUTPUT_DIR}" \ +NUM_GPUS="${NUM_GPUS}" \ +TOTAL_TRAINING_STEPS=2 \ +TEST_FREQ=-1 \ +SAVE_FREQ=-1 \ +RESUME_MODE=disable \ +PYTHON_BIN="${PYTHON_BIN}" \ +bash examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh \ + actor_rollout_ref.rollout.n=2 \ + actor_rollout_ref.rollout.agent.num_workers=2 \ + actor_rollout_ref.rollout.max_num_seqs=4 \ + trainer.val_before_train=false \ + trainer.log_val_generations=0 \ + "$@" + +[[ -s "${SCORER_LOG}" ]] +grep -q "training/global_step.*2" "${OUTPUT_DIR}/train.log" +echo "Qwen3-TTS GRPO e2e smoke passed; artifacts: ${WORK_DIR}" diff --git a/tests/trainer/omni/test_main_omni_on_cpu.py b/tests/trainer/omni/test_main_omni_on_cpu.py index 7fb313597..59f220d71 100644 --- a/tests/trainer/omni/test_main_omni_on_cpu.py +++ b/tests/trainer/omni/test_main_omni_on_cpu.py @@ -98,6 +98,10 @@ def test_omni_model_config_loads_tokenizer_and_processor_via_adapter(self, monke from verl_omni.workers.config.omni.model import OmniModelConfig mock_adapter = MagicMock() + mock_adapter.load_hf_config.return_value = SimpleNamespace( + tie_word_embeddings=False, + architectures=["arch"], + ) mock_adapter.configure_tokenizer.return_value = "tokenizer" mock_adapter.configure_processor.return_value = "processor" monkeypatch.setattr( @@ -106,11 +110,6 @@ def test_omni_model_config_loads_tokenizer_and_processor_via_adapter(self, monke ) monkeypatch.setattr(model_config_module, "resolve_model_local_dir", lambda path, use_shm=False: str(tmp_path)) monkeypatch.setattr(model_config_module, "copy_to_local", lambda path, use_shm=False: f"local:{path}") - monkeypatch.setattr( - model_config_module.AutoConfig, - "from_pretrained", - lambda *_args, **_kwargs: SimpleNamespace(tie_word_embeddings=False, architectures=["arch"]), - ) model_config = OmniModelConfig( path=str(tmp_path), @@ -123,6 +122,11 @@ def test_omni_model_config_loads_tokenizer_and_processor_via_adapter(self, monke assert model_config.tokenizer == "tokenizer" assert model_config.processor == "processor" + mock_adapter.load_hf_config.assert_called_once_with( + "local:" + str(tmp_path), + trust_remote_code=False, + attn_implementation="flash_attention_2", + ) mock_adapter.configure_tokenizer.assert_called_once_with("local:tokenizer-path", model_config) mock_adapter.configure_processor.assert_called_once_with(str(tmp_path), model_config) diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index 1220e601e..e8f121742 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -31,7 +31,10 @@ from verl_omni.pipelines.qwen3_tts import omni_rollout_adapter from verl_omni.pipelines.qwen3_tts.omni_rollout_adapter import Qwen3TTSRolloutAdapter from verl_omni.pipelines.qwen3_tts.talker_training_adapter import Qwen3TTSTalkerAdapter -from verl_omni.workers.rollout.vllm_rollout.vllm_omni_async_server import vLLMOmniHttpServer +from verl_omni.workers.rollout.vllm_rollout.vllm_omni_async_server import ( + _retained_output_modalities, + vLLMOmniHttpServer, +) class _Tokenizer: @@ -62,9 +65,7 @@ def test_optional_rollout_hooks_preserve_existing_ar_defaults(): first, final = object(), object() assert OmniRolloutPipelineBase.weight_sync_stage_ids() is None - assert OmniRolloutPipelineBase.supports_cache_engine_sleep() assert OmniRolloutPipelineBase.prepare_engine_prompt([], None, {}) is None - assert OmniRolloutPipelineBase.get_output_modalities() is None sampling_params = {"temperature": 0.8} assert ( OmniRolloutPipelineBase.prepare_agent_sampling_params( @@ -169,8 +170,7 @@ def test_rollout_adapter_builds_unique_prompt_and_scopes_weight_sync(tmp_path): assert first["additional_information"]["text"] == ["first text"] assert first["cache_salt"] != second["cache_salt"] assert Qwen3TTSRolloutAdapter.weight_sync_stage_ids("full") == [0] - assert not Qwen3TTSRolloutAdapter.supports_cache_engine_sleep("full") - assert Qwen3TTSRolloutAdapter.get_output_modalities("full") == ["latent", "audio"] + assert _retained_output_modalities(Qwen3TTSRolloutAdapter.build_stage_configs("full")) == ["latent", "audio"] def test_rollout_adapter_requires_speaker_embedding(): diff --git a/tests/workers/test_omni_fsdp_engine_on_cpu.py b/tests/workers/test_omni_fsdp_engine_on_cpu.py index 7ccac1f1d..2f1448c62 100644 --- a/tests/workers/test_omni_fsdp_engine_on_cpu.py +++ b/tests/workers/test_omni_fsdp_engine_on_cpu.py @@ -375,20 +375,20 @@ def test_collect_lora_params_import_not_from_verl(): # --------------------------------------------------------------------------- -# Qwen3-TTS reference-policy manual offload +# Qwen3-TTS reference-policy CPU-offload exception # --------------------------------------------------------------------------- -def _make_manual_ref_engine(omni_impl, *, enabled=True): +def _make_reference_engine(omni_impl, *, disable_cpu_offload=True): engine = object.__new__(omni_impl.OmniFSDPEngine) engine.engine_config = types.SimpleNamespace(forward_only=True, param_offload=False) - engine.model_adapter_cls = types.SimpleNamespace(requires_manual_ref_offload=enabled) + engine.model_adapter_cls = types.SimpleNamespace(disable_reference_cpu_offload=disable_cpu_offload) return engine def test_tts_ref_build_disables_forced_fsdp_cpu_offload_temporarily(): omni_impl = _get_omni_impl_module() - engine = _make_manual_ref_engine(omni_impl) + engine = _make_reference_engine(omni_impl) observed = [] def fake_build(module): @@ -403,9 +403,9 @@ def fake_build(module): assert engine.engine_config.forward_only is True -def test_tts_ref_to_uses_manual_load_path_and_restores_on_error(): +def test_tts_ref_to_bypasses_forced_offload_and_restores_on_error(): omni_impl = _get_omni_impl_module() - engine = _make_manual_ref_engine(omni_impl) + engine = _make_reference_engine(omni_impl) observed = [] def fake_to(*args, **kwargs): @@ -424,7 +424,7 @@ def fake_to(*args, **kwargs): def test_non_tts_ref_keeps_standard_forward_only_path(): omni_impl = _get_omni_impl_module() - engine = _make_manual_ref_engine(omni_impl, enabled=False) + engine = _make_reference_engine(omni_impl, disable_cpu_offload=False) observed = [] def fake_build(module): @@ -469,6 +469,7 @@ def test_build_module_calls_adapter_configure_model(architecture): fake_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] fake_adapter_cls = MagicMock() + fake_adapter_cls.get_model_class.return_value = None fake_configured_module = MagicMock(spec=torch.nn.Module) fake_configured_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] fake_adapter_cls.configure_model.return_value = fake_configured_module @@ -482,8 +483,9 @@ def test_build_module_calls_adapter_configure_model(architecture): patch.object(model_base_mod.OmniModelBase, "get_class_by_name", return_value=fake_adapter_cls) as mock_get_cls, patch.object(omni_impl, "get_init_weight_context_manager", return_value=MagicMock()), patch.object(omni_impl.warnings, "catch_warnings", return_value=MagicMock()), - patch("verl.utils.torch_dtypes.PrecisionType"), + patch("verl.utils.torch_dtypes.PrecisionType") as precision_type, ): + precision_type.to_dtype.side_effect = lambda value: value engine = object.__new__(omni_impl.OmniFSDPEngine) engine.model_config = model_config engine.engine_config = MagicMock() @@ -506,9 +508,49 @@ def test_build_module_calls_adapter_configure_model(architecture): ) fake_adapter_cls.configure_model.assert_called_once_with(fake_module, model_config) + assert engine.model_adapter_cls is fake_adapter_cls assert result is fake_configured_module +def test_build_module_uses_adapter_model_class_when_auto_model_is_incompatible(): + omni_impl = _get_omni_impl_module() + model_config = _make_mock_model_config(architecture="CustomOmniForConditionalGeneration") + loaded_module = MagicMock(spec=torch.nn.Module) + loaded_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] + configured_module = MagicMock(spec=torch.nn.Module) + configured_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] + adapter_cls = MagicMock() + custom_model_cls = MagicMock() + custom_model_cls.from_pretrained.return_value = loaded_module + adapter_cls.get_model_class.return_value = custom_model_cls + adapter_cls.configure_model.return_value = configured_module + model_base_mod = sys.modules["verl_omni.pipelines.model_base"] + + with ( + patch.object(model_base_mod.OmniModelBase, "get_class_by_name", return_value=adapter_cls), + patch.object(omni_impl, "get_init_weight_context_manager", return_value=MagicMock()), + patch.object(omni_impl.warnings, "catch_warnings", return_value=MagicMock()), + patch("verl.utils.torch_dtypes.PrecisionType") as precision_type, + ): + precision_type.to_dtype.side_effect = lambda value: value + engine = object.__new__(omni_impl.OmniFSDPEngine) + engine.model_config = model_config + engine.engine_config = MagicMock(model_dtype=None, forward_only=False) + engine.device_mesh = None + + result = engine._build_module() + + adapter_cls.get_model_class.assert_called_once_with() + custom_model_cls.from_pretrained.assert_called_once_with( + pretrained_model_name_or_path=model_config.local_path, + torch_dtype=torch.float32, + config=model_config.hf_config, + trust_remote_code=model_config.trust_remote_code, + ) + adapter_cls.configure_model.assert_called_once_with(loaded_module, model_config) + assert result is configured_module + + @pytest.mark.parametrize("option", ["use_liger", "use_fused_kernels"]) def test_build_module_rejects_unsupported_optimizations_before_model_load(option): omni_impl = _get_omni_impl_module() diff --git a/verl_omni/pipelines/model_base.py b/verl_omni/pipelines/model_base.py index c8e1451c8..db560872c 100644 --- a/verl_omni/pipelines/model_base.py +++ b/verl_omni/pipelines/model_base.py @@ -491,6 +491,7 @@ class Qwen3OmniThinkerAdapter(OmniModelBase): """ _registry: dict[tuple[str, str], type["OmniModelBase"]] = {} + disable_reference_cpu_offload = False @classmethod def register(cls, architecture: str, stage: str = "thinker"): @@ -558,6 +559,28 @@ def get_class_by_name( f"Set ``external_lib`` to load your training adapter." ) from None + @classmethod + def load_hf_config( + cls, + model_path: str, + *, + trust_remote_code: bool, + attn_implementation: str, + ): + """Load the Hugging Face config used by the FSDP engine.""" + from transformers import AutoConfig + + return AutoConfig.from_pretrained( + model_path, + trust_remote_code=trust_remote_code, + attn_implementation=attn_implementation, + ) + + @classmethod + def get_model_class(cls): + """Return a model class override, or ``None`` for the default auto model.""" + return None + @classmethod @abstractmethod def get_strip_modules(cls, model_config) -> list[str]: @@ -627,7 +650,6 @@ def configure_model(cls, module, model_config): Default implementation strips the submodules returned by ``get_strip_modules``. Override to also: - - Register the model class with ``AutoModelForCausalLM``. - Redirect ``forward()`` and embedding accessors to the trainable sub-component. - Force ``tie_word_embeddings=False`` for FSDP compatibility. @@ -763,11 +785,6 @@ def weight_sync_stage_ids(cls, pipeline_mode="thinker_only") -> list[int] | None """Return stages that receive actor weights, or all stages by default.""" return None - @classmethod - def supports_cache_engine_sleep(cls, pipeline_mode="thinker_only") -> bool: - """Return whether the pipeline supports rollout cache sleep and wake.""" - return True - @classmethod def get_pipeline_id(cls, pipeline_mode: str = "thinker_only") -> str: """Return the vLLM-Omni pipeline model_type for *pipeline_mode*. diff --git a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py index 1dddbde88..ea336fe91 100644 --- a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py @@ -61,16 +61,17 @@ def _completion(output): return completions[0] if completions else None -def _materialize(value): +def _copy_plain_containers(value): + """Copy Ray/shared-memory mapping containers without copying tensor payloads.""" if isinstance(value, Mapping): - return {key: _materialize(item) for key, item in value.items()} + return {key: _copy_plain_containers(item) for key, item in value.items()} if isinstance(value, list): - return [_materialize(item) for item in value] + return [_copy_plain_containers(item) for item in value] return value def talker2code2wav_token_only(source_outputs, prompt=None, _requires_multimodal_data=False): - """Materialize shared-memory Mapping payloads before the upstream processor mutates them.""" + """Give the mutating upstream processor ordinary Python containers.""" from vllm_omni.model_executor.stage_input_processors.qwen3_tts import ( talker2code2wav_token_only as upstream_processor, ) @@ -83,7 +84,7 @@ def talker2code2wav_token_only(source_outputs, prompt=None, _requires_multimodal completion_copy = copy.copy(completion) multimodal = getattr(completion, "multimodal_output", None) if isinstance(multimodal, Mapping): - completion_copy.multimodal_output = _materialize(multimodal) + completion_copy.multimodal_output = _copy_plain_containers(multimodal) source_copy.outputs.append(completion_copy) converted.append(source_copy) return upstream_processor(converted, prompt, _requires_multimodal_data) @@ -111,21 +112,12 @@ def ensure_pipeline_registered(cls, pipeline_mode="full"): cls._check_mode(pipeline_mode) register_pipeline(QWEN3_TTS_RL_PIPELINE) - @classmethod - def get_output_modalities(cls, pipeline_mode="full"): - cls._check_mode(pipeline_mode) - return ["latent", "audio"] - @classmethod def weight_sync_stage_ids(cls, pipeline_mode="full"): + """Sync actor weights only to stage 0; stage 1 is the frozen decoder.""" cls._check_mode(pipeline_mode) return [0] - @classmethod - def supports_cache_engine_sleep(cls, pipeline_mode="full"): - cls._check_mode(pipeline_mode) - return False - @classmethod def get_stage_engine_extras(cls, stage_id, pipeline_mode="full"): cls._check_mode(pipeline_mode) @@ -140,6 +132,7 @@ def prepare_agent_sampling_params( trainer_config, agent_inputs, ): + """Seed codec-0 and residual-codebook sampling for each GRPO candidate.""" extra_info = agent_inputs.get("extra_info") evaluation = is_evaluation_split(extra_info) candidate_count = rollout_config.val_kwargs.n if evaluation else rollout_config.n @@ -155,6 +148,7 @@ def prepare_agent_sampling_params( @classmethod def postprocess_agent_loop_output(cls, output, *, tokenizer, response_length): + """Map the 16-codebook rollout to codec-0 policy tokens and replay fields.""" extra = output.extra_fields codes, text = extra.get("tts_audio_codes"), extra.get("tts_text") if codes is None or text is None: @@ -185,6 +179,7 @@ def postprocess_agent_loop_output(cls, output, *, tokenizer, response_length): @classmethod def prepare_engine_prompt(cls, prompt_ids, model_config, multi_modal_data, mm_processor_kwargs=None): + """Build the Base-task prompt and fixed-speaker conditioning for vLLM-Omni.""" text = model_config.tokenizer.decode(prompt_ids, skip_special_tokens=True).strip() if not text: raise ValueError("Qwen3-TTS received an empty text prompt.") @@ -212,6 +207,7 @@ def prepare_engine_prompt(cls, prompt_ids, model_config, multi_modal_data, mm_pr @classmethod def combine_engine_outputs(cls, outputs, prompt): + """Combine stage-0 policy tokens with codec and waveform outputs.""" policy_output = None policy_length = -1 audio_codes = waveform = None diff --git a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py index 9f7b96be4..db6c6cbee 100644 --- a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py @@ -42,43 +42,6 @@ def _prepare_config_for_checkpoint(config) -> None: speaker_config.__dict__.pop("_dtype", None) -def register_qwen3_tts_automodel() -> None: - config_cls = model_cls = None - with qwen3_tts_import_context(): - for module_name in ("transformers", "qwen_tts.core.models.configuration_qwen3_tts"): - try: - config_cls = getattr(__import__(module_name, fromlist=["Qwen3TTSConfig"]), "Qwen3TTSConfig", None) - except ImportError: - continue - if config_cls is not None: - break - for module_name in ("transformers", "qwen_tts.core.models.modeling_qwen3_tts"): - try: - model_cls = getattr( - __import__(module_name, fromlist=["Qwen3TTSForConditionalGeneration"]), - "Qwen3TTSForConditionalGeneration", - None, - ) - except ImportError: - continue - if model_cls is not None: - break - if config_cls is None or model_cls is None: - return - - from transformers import AutoConfig, AutoModelForMultimodalLM - - patch_qwen3_tts_config_defaults(config_cls) - try: - AutoConfig.register(getattr(config_cls, "model_type", "qwen3_tts"), config_cls) - except ValueError: - pass - try: - AutoModelForMultimodalLM.register(config_cls, model_cls) - except ValueError: - pass - - def _speaker_embedding(model, batch_size, device, dtype): cached = getattr(model, "_verl_tts_speaker_embedding", None) if cached is None: @@ -145,7 +108,26 @@ def _set_input_embeddings(self, value): @OmniModelBase.register("Qwen3TTSForConditionalGeneration", stage="talker") class Qwen3TTSTalkerAdapter(OmniModelBase): - requires_manual_ref_offload = True + disable_reference_cpu_offload = True + + @classmethod + def load_hf_config(cls, model_path, *, trust_remote_code, attn_implementation): + with qwen3_tts_import_context(): + from qwen_tts.core.models.configuration_qwen3_tts import Qwen3TTSConfig + + patch_qwen3_tts_config_defaults(Qwen3TTSConfig) + return Qwen3TTSConfig.from_pretrained( + model_path, + trust_remote_code=trust_remote_code, + attn_implementation=attn_implementation, + ) + + @classmethod + def get_model_class(cls): + with qwen3_tts_import_context(): + from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSForConditionalGeneration + + return Qwen3TTSForConditionalGeneration @classmethod def get_strip_modules(cls, model_config): @@ -227,6 +209,3 @@ def prepare_model_inputs(cls, model_inputs, micro_batch, model_config): } ) return model_inputs - - -register_qwen3_tts_automodel() diff --git a/verl_omni/pipelines/qwen3_tts/transformers_compat.py b/verl_omni/pipelines/qwen3_tts/transformers_compat.py index 94fcd68a0..1b9300bf6 100644 --- a/verl_omni/pipelines/qwen3_tts/transformers_compat.py +++ b/verl_omni/pipelines/qwen3_tts/transformers_compat.py @@ -17,7 +17,6 @@ from functools import wraps import torch -from packaging.version import Version def _default_rope_init(config, device=None, **kwargs): @@ -52,19 +51,14 @@ def qwen3_tts_import_context(): qwen-tts modules retain the mask wrappers they import. The global Transformers functions are restored immediately afterwards. """ - import transformers - - if Version(transformers.__version__).major < 5: - yield - return - import transformers.masking_utils as masking_utils + import transformers.modeling_rope_utils as rope_utils import transformers.utils.generic as generic_utils - from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS original_check = generic_utils.check_model_inputs original_causal_mask = masking_utils.create_causal_mask original_sliding_mask = masking_utils.create_sliding_window_causal_mask + original_rope_functions = rope_utils.ROPE_INIT_FUNCTIONS def compatible_check_model_inputs(func=None): return original_check if func is None else original_check(func) @@ -72,10 +66,11 @@ def compatible_check_model_inputs(func=None): generic_utils.check_model_inputs = compatible_check_model_inputs masking_utils.create_causal_mask = _compatible_mask(original_causal_mask) masking_utils.create_sliding_window_causal_mask = _compatible_mask(original_sliding_mask) - ROPE_INIT_FUNCTIONS.setdefault("default", _default_rope_init) + rope_utils.ROPE_INIT_FUNCTIONS = {**original_rope_functions, "default": _default_rope_init} try: yield finally: generic_utils.check_model_inputs = original_check masking_utils.create_causal_mask = original_causal_mask masking_utils.create_sliding_window_causal_mask = original_sliding_mask + rope_utils.ROPE_INIT_FUNCTIONS = original_rope_functions diff --git a/verl_omni/workers/config/omni/model.py b/verl_omni/workers/config/omni/model.py index 3f625e7ae..270368165 100644 --- a/verl_omni/workers/config/omni/model.py +++ b/verl_omni/workers/config/omni/model.py @@ -20,7 +20,6 @@ from typing import Any, Optional from omegaconf import MISSING -from transformers import AutoConfig from verl.base_config import BaseConfig from verl.utils.fs import copy_to_local from verl.utils.import_utils import import_external_libs @@ -168,20 +167,35 @@ def __post_init__(self): # Build hf_config so the FSDP engine can load and wrap the model. self.local_hf_config_path = copy_to_local(self.hf_config_path, use_shm=self.use_shm) attn_implementation = self.override_config.get("attn_implementation", "flash_attention_2") - self.hf_config = AutoConfig.from_pretrained( - self.local_hf_config_path, - trust_remote_code=self.trust_remote_code, - attn_implementation=attn_implementation, - ) + from verl_omni.pipelines.model_base import OmniModelBase + + try: + adapter_cls = OmniModelBase.get_class_by_name(self.architecture, self.model_stage, self.external_lib) + except NotImplementedError: + adapter_cls = None + + if adapter_cls is None: + from transformers import AutoConfig + + self.hf_config = AutoConfig.from_pretrained( + self.local_hf_config_path, + trust_remote_code=self.trust_remote_code, + attn_implementation=attn_implementation, + ) + else: + self.hf_config = adapter_cls.load_hf_config( + self.local_hf_config_path, + trust_remote_code=self.trust_remote_code, + attn_implementation=attn_implementation, + ) self.share_embeddings_and_output_weights = getattr(self.hf_config, "tie_word_embeddings", False) self.architectures = getattr(self.hf_config, "architectures", None) if self.load_tokenizer: - from verl_omni.pipelines.model_base import OmniModelBase - self.local_tokenizer_path = copy_to_local(self.tokenizer_path, use_shm=self.use_shm) - adapter_cls = OmniModelBase.get_class_by_name(self.architecture, self.model_stage, self.external_lib) + if adapter_cls is None: + adapter_cls = OmniModelBase.get_class_by_name(self.architecture, self.model_stage, self.external_lib) self.tokenizer = adapter_cls.configure_tokenizer(self.local_tokenizer_path, self) self.processor = adapter_cls.configure_processor(self.local_path, self) diff --git a/verl_omni/workers/engine/fsdp/omni_impl.py b/verl_omni/workers/engine/fsdp/omni_impl.py index d580c5de2..0b1faf2b7 100644 --- a/verl_omni/workers/engine/fsdp/omni_impl.py +++ b/verl_omni/workers/engine/fsdp/omni_impl.py @@ -49,14 +49,23 @@ def _cast_dtensor_weight_for_sync(tensor: torch.Tensor) -> torch.Tensor: return tensor.to(dtype=torch.bfloat16, non_blocking=True) return tensor - def _run_with_manual_ref_offload(self, call): + def _run_without_forced_reference_cpu_offload(self, call): + """Bypass verl's forced reference CPU offload for adapters that require it. + + Qwen3-TTS invokes leaf embedding tables outside their wrapped decoder + module. FSDP1's forced ``CPUOffload`` leaves those tables on CPU while + the replay tensors are on CUDA, so reference log-probability computation + fails with a device mismatch. Temporarily presenting the reference as a + regular engine during FSDP construction and movement keeps the full + module on CUDA; it does not enable gradients for the reference model. + """ adapter_cls = getattr(self, "model_adapter_cls", None) - manual_ref_offload = ( - getattr(adapter_cls, "requires_manual_ref_offload", False) + disable_cpu_offload = ( + getattr(adapter_cls, "disable_reference_cpu_offload", False) and getattr(self.engine_config, "forward_only", False) and not getattr(self.engine_config, "param_offload", True) ) - if not manual_ref_offload: + if not disable_cpu_offload: return call() self.engine_config.forward_only = False @@ -67,11 +76,13 @@ def _run_with_manual_ref_offload(self, call): def _build_fsdp_module(self, module): parent_build = super()._build_fsdp_module - return self._run_with_manual_ref_offload(lambda: parent_build(module)) + return self._run_without_forced_reference_cpu_offload(lambda: parent_build(module)) def to(self, device, model=True, optimizer=True, grad=True): parent_to = super().to - return self._run_with_manual_ref_offload(lambda: parent_to(device, model=model, optimizer=optimizer, grad=grad)) + return self._run_without_forced_reference_cpu_offload( + lambda: parent_to(device, model=model, optimizer=optimizer, grad=grad) + ) def prepare_model_inputs(self, micro_batch): """Prepare standard LM inputs, then add model-native replay fields.""" @@ -200,6 +211,12 @@ def _build_module(self): self.model_config: OmniModelConfig architecture = self.model_config.architecture + adapter_cls = OmniModelBase.get_class_by_name( + architecture, + self.model_config.model_stage, + self.model_config.get("external_lib"), + ) + self.model_adapter_cls = adapter_cls torch_dtype = self.engine_config.model_dtype @@ -221,19 +238,13 @@ def _build_module(self): with init_context(), warnings.catch_warnings(): warnings.simplefilter("ignore") - module = AutoModelForMultimodalLM.from_pretrained( + model_cls = adapter_cls.get_model_class() or AutoModelForMultimodalLM + module = model_cls.from_pretrained( pretrained_model_name_or_path=self.model_config.local_path, torch_dtype=torch_dtype, config=self.model_config.hf_config, trust_remote_code=self.model_config.trust_remote_code, ) - - adapter_cls = OmniModelBase.get_class_by_name( - architecture, - self.model_config.model_stage, - self.model_config.get("external_lib"), - ) - self.model_adapter_cls = adapter_cls module = adapter_cls.configure_model(module, self.model_config) module.to(torch_dtype) From d03a3177573a1c9a8156e6ddcdeb72f677da741e Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Sat, 29 Aug 2026 00:50:10 +0800 Subject: [PATCH 15/28] [model, tests] fix: remove unneeded reference offload exception Signed-off-by: dongbo910220 <1275604947@qq.com> --- tests/workers/test_omni_fsdp_engine_on_cpu.py | 63 ------------------- verl_omni/pipelines/model_base.py | 1 - .../qwen3_tts/talker_training_adapter.py | 2 - verl_omni/workers/engine/fsdp/omni_impl.py | 35 ----------- 4 files changed, 101 deletions(-) diff --git a/tests/workers/test_omni_fsdp_engine_on_cpu.py b/tests/workers/test_omni_fsdp_engine_on_cpu.py index 2f1448c62..496b882b0 100644 --- a/tests/workers/test_omni_fsdp_engine_on_cpu.py +++ b/tests/workers/test_omni_fsdp_engine_on_cpu.py @@ -374,69 +374,6 @@ def test_collect_lora_params_import_not_from_verl(): ) -# --------------------------------------------------------------------------- -# Qwen3-TTS reference-policy CPU-offload exception -# --------------------------------------------------------------------------- - - -def _make_reference_engine(omni_impl, *, disable_cpu_offload=True): - engine = object.__new__(omni_impl.OmniFSDPEngine) - engine.engine_config = types.SimpleNamespace(forward_only=True, param_offload=False) - engine.model_adapter_cls = types.SimpleNamespace(disable_reference_cpu_offload=disable_cpu_offload) - return engine - - -def test_tts_ref_build_disables_forced_fsdp_cpu_offload_temporarily(): - omni_impl = _get_omni_impl_module() - engine = _make_reference_engine(omni_impl) - observed = [] - - def fake_build(module): - observed.append(engine.engine_config.forward_only) - return module - - module = object() - with patch.object(omni_impl.FSDPEngineWithLMHead, "_build_fsdp_module", side_effect=fake_build): - assert engine._build_fsdp_module(module) is module - - assert observed == [False] - assert engine.engine_config.forward_only is True - - -def test_tts_ref_to_bypasses_forced_offload_and_restores_on_error(): - omni_impl = _get_omni_impl_module() - engine = _make_reference_engine(omni_impl) - observed = [] - - def fake_to(*args, **kwargs): - observed.append((engine.engine_config.forward_only, args, kwargs)) - raise RuntimeError("test failure") - - with ( - patch.object(omni_impl.FSDPEngineWithLMHead, "to", side_effect=fake_to), - pytest.raises(RuntimeError, match="test failure"), - ): - engine.to("cuda", model=True, optimizer=False, grad=False) - - assert observed == [(False, ("cuda",), {"model": True, "optimizer": False, "grad": False})] - assert engine.engine_config.forward_only is True - - -def test_non_tts_ref_keeps_standard_forward_only_path(): - omni_impl = _get_omni_impl_module() - engine = _make_reference_engine(omni_impl, disable_cpu_offload=False) - observed = [] - - def fake_build(module): - observed.append(engine.engine_config.forward_only) - return module - - with patch.object(omni_impl.FSDPEngineWithLMHead, "_build_fsdp_module", side_effect=fake_build): - engine._build_fsdp_module(object()) - - assert observed == [True] - - # --------------------------------------------------------------------------- # ``_build_module`` calls adapter ``configure_model`` # --------------------------------------------------------------------------- diff --git a/verl_omni/pipelines/model_base.py b/verl_omni/pipelines/model_base.py index db560872c..e3a2f1835 100644 --- a/verl_omni/pipelines/model_base.py +++ b/verl_omni/pipelines/model_base.py @@ -491,7 +491,6 @@ class Qwen3OmniThinkerAdapter(OmniModelBase): """ _registry: dict[tuple[str, str], type["OmniModelBase"]] = {} - disable_reference_cpu_offload = False @classmethod def register(cls, architecture: str, stage: str = "thinker"): diff --git a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py index db6c6cbee..1815b09da 100644 --- a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py @@ -108,8 +108,6 @@ def _set_input_embeddings(self, value): @OmniModelBase.register("Qwen3TTSForConditionalGeneration", stage="talker") class Qwen3TTSTalkerAdapter(OmniModelBase): - disable_reference_cpu_offload = True - @classmethod def load_hf_config(cls, model_path, *, trust_remote_code, attn_implementation): with qwen3_tts_import_context(): diff --git a/verl_omni/workers/engine/fsdp/omni_impl.py b/verl_omni/workers/engine/fsdp/omni_impl.py index 0b1faf2b7..6a96f8af9 100644 --- a/verl_omni/workers/engine/fsdp/omni_impl.py +++ b/verl_omni/workers/engine/fsdp/omni_impl.py @@ -49,41 +49,6 @@ def _cast_dtensor_weight_for_sync(tensor: torch.Tensor) -> torch.Tensor: return tensor.to(dtype=torch.bfloat16, non_blocking=True) return tensor - def _run_without_forced_reference_cpu_offload(self, call): - """Bypass verl's forced reference CPU offload for adapters that require it. - - Qwen3-TTS invokes leaf embedding tables outside their wrapped decoder - module. FSDP1's forced ``CPUOffload`` leaves those tables on CPU while - the replay tensors are on CUDA, so reference log-probability computation - fails with a device mismatch. Temporarily presenting the reference as a - regular engine during FSDP construction and movement keeps the full - module on CUDA; it does not enable gradients for the reference model. - """ - adapter_cls = getattr(self, "model_adapter_cls", None) - disable_cpu_offload = ( - getattr(adapter_cls, "disable_reference_cpu_offload", False) - and getattr(self.engine_config, "forward_only", False) - and not getattr(self.engine_config, "param_offload", True) - ) - if not disable_cpu_offload: - return call() - - self.engine_config.forward_only = False - try: - return call() - finally: - self.engine_config.forward_only = True - - def _build_fsdp_module(self, module): - parent_build = super()._build_fsdp_module - return self._run_without_forced_reference_cpu_offload(lambda: parent_build(module)) - - def to(self, device, model=True, optimizer=True, grad=True): - parent_to = super().to - return self._run_without_forced_reference_cpu_offload( - lambda: parent_to(device, model=model, optimizer=optimizer, grad=grad) - ) - def prepare_model_inputs(self, micro_batch): """Prepare standard LM inputs, then add model-native replay fields.""" model_inputs, output_args = super().prepare_model_inputs(micro_batch) From 887a88e017d317de32ad0622f4d3620ec5f5157b Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Sat, 29 Aug 2026 12:56:25 +0800 Subject: [PATCH 16/28] [rollout, tests, doc] refactor: align Qwen3-TTS with AR strategy Signed-off-by: dongbo910220 <1275604947@qq.com> --- .../contributing/integrating_an_omni_model.md | 6 +- .../test_qwen3_tts_rollout_on_cpu.py | 70 ++++++--- .../qwen3_tts/omni_rollout_adapter.py | 5 +- .../vllm_rollout/vllm_omni_ar_strategy.py | 142 +++++++++++++++--- .../vllm_rollout/vllm_omni_async_server.py | 21 ++- .../vllm_rollout/vllm_omni_strategy_base.py | 12 +- 6 files changed, 199 insertions(+), 57 deletions(-) diff --git a/docs/contributing/integrating_an_omni_model.md b/docs/contributing/integrating_an_omni_model.md index 5783176e4..5ce4ba218 100644 --- a/docs/contributing/integrating_an_omni_model.md +++ b/docs/contributing/integrating_an_omni_model.md @@ -108,9 +108,9 @@ Optional overrides fall into four groups: `get_stage_engine_extras`. - Resource behavior: `weight_sync_stage_ids`. - Request construction: `prepare_engine_prompt`. -- Multi-stage output assembly: `combine_engine_outputs`. The rollout server - derives retained output modalities from stages marked `final_output` in the - pipeline topology. +- Multi-stage output assembly: `combine_engine_outputs`. The AR generation + strategy derives retained output modalities from stages marked + `final_output` in the pipeline topology. Their defaults preserve the existing single-output AR behavior. Override only the hooks required by the model. For example, Qwen3-TTS synchronizes actor diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index e8f121742..44a5758ca 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -31,10 +31,11 @@ from verl_omni.pipelines.qwen3_tts import omni_rollout_adapter from verl_omni.pipelines.qwen3_tts.omni_rollout_adapter import Qwen3TTSRolloutAdapter from verl_omni.pipelines.qwen3_tts.talker_training_adapter import Qwen3TTSTalkerAdapter -from verl_omni.workers.rollout.vllm_rollout.vllm_omni_async_server import ( +from verl_omni.workers.rollout.vllm_rollout.vllm_omni_ar_strategy import ( + ARStrategy, _retained_output_modalities, - vLLMOmniHttpServer, ) +from verl_omni.workers.rollout.vllm_rollout.vllm_omni_async_server import vLLMOmniHttpServer class _Tokenizer: @@ -173,6 +174,31 @@ def test_rollout_adapter_builds_unique_prompt_and_scopes_weight_sync(tmp_path): assert _retained_output_modalities(Qwen3TTSRolloutAdapter.build_stage_configs("full")) == ["latent", "audio"] +def test_ar_strategy_resolves_qwen3_tts_adapter_and_scopes_weight_sync(monkeypatch): + server = SimpleNamespace(_rollout_flags={}) + strategy = ARStrategy(server) + deploy_calls = [] + monkeypatch.setattr( + strategy, + "_write_deploy_config", + lambda engine_kwargs, pipeline_name, adapter_cls, pipeline_mode: deploy_calls.append( + (pipeline_name, adapter_cls, pipeline_mode) + ), + ) + engine_kwargs = { + "output_mode": "ar", + "pipeline_name": "qwen3_tts_rl", + "pipeline_mode": "full", + } + + strategy.preprocess_engine_kwargs(engine_kwargs) + + assert deploy_calls == [("qwen3_tts_rl", Qwen3TTSRolloutAdapter, "full")] + assert strategy._rollout_adapter is Qwen3TTSRolloutAdapter + assert strategy._weight_sync_stage_ids == [0] + assert engine_kwargs == {} + + def test_rollout_adapter_requires_speaker_embedding(): model_config = SimpleNamespace( tokenizer=_Tokenizer(), @@ -276,7 +302,7 @@ def test_rollout_adapter_prepares_sampling_and_actor_policy_sequence(): assert result.extra_fields["tts_text_ids"] -def test_server_prepares_stage_specific_sampling_params(): +def test_ar_strategy_prepares_stage_specific_sampling_params(): class Adapter: @staticmethod def prepare_engine_prompt(**kwargs): @@ -285,20 +311,21 @@ def prepare_engine_prompt(**kwargs): "additional_information": {"text": ["hello"]}, } - server = object.__new__(vLLMOmniHttpServer) - server._ar_mode = True - server._omni_rollout_adapter = Adapter - server._stage_sampling_constraints = {0: {}} - server.model_config = SimpleNamespace() - server.config = SimpleNamespace( - max_model_len=64, - prompt_length=16, - response_length=8, - repetition_penalty=1.0, + server = SimpleNamespace( + model_config=SimpleNamespace(), + config=SimpleNamespace( + max_model_len=64, + prompt_length=16, + response_length=8, + repetition_penalty=1.0, + ), + engine=SimpleNamespace(default_sampling_params_list=[SamplingParams(), SimpleNamespace(stage="decoder")]), ) - server.engine = SimpleNamespace(default_sampling_params_list=[SamplingParams(), SimpleNamespace(stage="decoder")]) + strategy = ARStrategy(server) + strategy._rollout_adapter = Adapter + strategy._stage_sampling_constraints = {0: {}} - prompt, params = server._preprocess_input( + prompt, params = strategy.preprocess_input( [5, 6], {"temperature": 0.8, "logprobs": True}, {}, @@ -315,7 +342,7 @@ def prepare_engine_prompt(**kwargs): @pytest.mark.asyncio -async def test_server_retains_requested_stage_outputs_and_targets_weight_sync(): +async def test_ar_strategy_retains_requested_stage_outputs_and_targets_weight_sync(): policy = SimpleNamespace(request_output=SimpleNamespace(outputs=[])) class Engine: @@ -338,13 +365,14 @@ def combine_engine_outputs(outputs, prompt): return policy, {"audio_sample_rate": 24_000} server = object.__new__(vLLMOmniHttpServer) - server._ar_mode = True server.engine = Engine() - server._rollout_output_modalities = ["latent", "audio"] - server._omni_rollout_adapter = Adapter - server._weight_sync_stage_ids = [0] + strategy = ARStrategy(server) + strategy._rollout_output_modalities = ["latent", "audio"] + strategy._rollout_adapter = Adapter + strategy._weight_sync_stage_ids = [0] + server._generate_strategy = strategy - result = await server._run_generation({"prompt_token_ids": [1]}, SamplingParams(), "request-0", None, 0) + result = await strategy.run_generation({"prompt_token_ids": [1]}, SamplingParams(), "request-0", None, 0) rpc_result = await server.collective_rpc("update_weights_from_ipc", kwargs={"base_sync_done": True}) assert result is policy diff --git a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py index ea336fe91..264a79309 100644 --- a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py @@ -62,7 +62,10 @@ def _completion(output): def _copy_plain_containers(value): - """Copy Ray/shared-memory mapping containers without copying tensor payloads.""" + """Convert Ray/shared-memory Mapping/list shells to built-ins for upstream strict dict checks. + + Tensor payloads are preserved rather than copied. + """ if isinstance(value, Mapping): return {key: _copy_plain_containers(item) for key, item in value.items()} if isinstance(value, list): diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py index c82630ccd..9b738ffe7 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import copy import json import logging import os @@ -45,6 +46,16 @@ def _drop_none_mapping_values(value: Any) -> Any: return value +def _retained_output_modalities(stages: list[Any]) -> list[str] | None: + """Request every modality when an AR pipeline exposes multiple final outputs.""" + final_output_types = [ + stage.final_output_type + for stage in stages + if getattr(stage, "final_output", False) and getattr(stage, "final_output_type", None) + ] + return list(dict.fromkeys(final_output_types)) if len(final_output_types) > 1 else None + + class ARStrategy(OmniStrategyBase): """Concrete AR/thinker strategy. @@ -56,6 +67,14 @@ class ARStrategy(OmniStrategyBase): rollout_config_cls = RolloutConfig model_config_cls = OmniModelConfig + def __init__(self, server: Any) -> None: + super().__init__(server) + self._rollout_adapter: type[OmniRolloutPipelineBase] | None = None + self._pipeline_mode = "thinker_only" + self._rollout_output_modalities: list[str] | None = None + self._weight_sync_stage_ids: list[int] | None = None + self._stage_sampling_constraints: dict[int, dict[str, Any]] = {} + def validate_configs(self) -> None: if self.server.config.max_model_len is None: self.server.config.max_model_len = self.server.config.prompt_length + self.server.config.response_length @@ -74,13 +93,15 @@ def preprocess_engine_kwargs(self, engine_kwargs: dict[str, Any]) -> None: engine_kwargs.pop("custom_pipeline", None) # TODO (mike): drop this later. It should be inferred from the model config. pipeline_name = engine_kwargs.pop("pipeline_name", None) - pipeline_mode = engine_kwargs.pop("pipeline_mode", "thinker_only") + self._pipeline_mode = engine_kwargs.pop("pipeline_mode", "thinker_only") adapter_cls = OmniRolloutPipelineBase.get_class(pipeline_name) if adapter_cls is not None: - self._write_deploy_config(engine_kwargs, pipeline_name, adapter_cls, pipeline_mode) - self.server._rollout_flags = adapter_cls.rollout_flags(pipeline_mode=pipeline_mode) - adapter_overrides = adapter_cls.get_engine_hf_overrides(pipeline_mode=pipeline_mode) + self._rollout_adapter = adapter_cls + self._write_deploy_config(engine_kwargs, pipeline_name, adapter_cls, self._pipeline_mode) + self.server._rollout_flags = adapter_cls.rollout_flags(pipeline_mode=self._pipeline_mode) + self._weight_sync_stage_ids = adapter_cls.weight_sync_stage_ids(pipeline_mode=self._pipeline_mode) + adapter_overrides = adapter_cls.get_engine_hf_overrides(pipeline_mode=self._pipeline_mode) if adapter_overrides: hf_overrides = engine_kwargs.get("hf_overrides", {}) if isinstance(hf_overrides, str): @@ -111,12 +132,29 @@ def _write_deploy_config( adapter_cls.ensure_pipeline_registered(pipeline_mode) stages = adapter_cls.build_stage_configs(pipeline_mode=pipeline_mode) pipeline_id = adapter_cls.get_pipeline_id(pipeline_mode) + self._rollout_output_modalities = _retained_output_modalities(stages) + self._stage_sampling_constraints = { + stage.stage_id: dict(getattr(stage, "sampling_constraints", {}) or {}) for stage in stages + } + stage_extras = { + stage.stage_id: dict(adapter_cls.get_stage_engine_extras(stage.stage_id, pipeline_mode=pipeline_mode)) + for stage in stages + } + capacity_fields = ("max_model_len", "max_num_batched_tokens") + if any(field in extras for extras in stage_extras.values() for field in capacity_fields): + for extras in stage_extras.values(): + for field in capacity_fields: + extras.setdefault(field, getattr(self.server.config, field)) + for field in capacity_fields: + engine_kwargs[field] = None device_control_env = get_visible_devices_keyword() visible_devices = os.environ.get(device_control_env, "") tp_size = self.server.config.tensor_model_parallel_size deploy_dict: dict[str, object] = {"pipeline": pipeline_id} + if "async_chunk" in engine_kwargs: + deploy_dict["async_chunk"] = bool(engine_kwargs["async_chunk"]) if visible_devices: device_count = len([device for device in visible_devices.split(",") if device.strip()]) @@ -128,7 +166,7 @@ def _write_deploy_config( "devices": devices, "tensor_parallel_size": tp_size, "text_encoder_tp_size": getattr(self.server.config, "text_encoder_tp_size", 1), - "engine_extras": adapter_cls.get_stage_engine_extras(stage_id, pipeline_mode=pipeline_mode), + "engine_extras": stage_extras[stage_id], } for stage_id in stage_ids ] @@ -151,6 +189,10 @@ def _write_deploy_config( engine_kwargs["deploy_config"] = deploy_path def prepare_engine_args(self, engine_args: dict[str, Any], args: Namespace) -> None: + if self._rollout_adapter is not None: + # The generated per-stage deploy config owns model_stage for + # heterogeneous pipelines such as Qwen3-TTS. + engine_args["model_stage"] = None for timeout_key in ("stage_init_timeout", "init_timeout"): timeout_value = getattr(args, timeout_key, None) if timeout_value is not None: @@ -159,6 +201,11 @@ def prepare_engine_args(self, engine_args: dict[str, Any], args: Namespace) -> N if isinstance(engine_args.get("compilation_config"), dict): engine_args["compilation_config"] = _drop_none_mapping_values(engine_args["compilation_config"]) + def collective_rpc_stage_ids(self, method: Any) -> list[int] | None: + if method in {"set_pending_lora_peft_config", "update_weights_from_ipc"}: + return self._weight_sync_stage_ids + return None + def preprocess_input( self, prompt_ids: list[int], @@ -170,15 +217,27 @@ def preprocess_input( mm_processor_kwargs: Optional[dict[str, Any]] = None, extra_prompt_ids: Optional[dict[str, list[int]]] = None, negative_extra_prompt_ids: Optional[dict[str, list[int]]] = None, - ) -> tuple[dict[str, Any], SamplingParams]: + ) -> tuple[dict[str, Any], SamplingParams | list[Any]]: if multi_modal_data: processor = getattr(self.server.model_config, "processor", None) if processor is not None and hasattr(processor, "dedup_pad_tokens"): prompt_ids = processor.dedup_pad_tokens(prompt_ids) - max_possible_tokens = self.server.config.max_model_len - len(prompt_ids) + + prompt = None + adapter_prepared_prompt = False + if self._rollout_adapter is not None: + prompt = self._rollout_adapter.prepare_engine_prompt( + prompt_ids=prompt_ids, + model_config=self.server.model_config, + multi_modal_data=multi_modal_data, + mm_processor_kwargs=mm_processor_kwargs, + ) + adapter_prepared_prompt = prompt is not None + effective_prompt_ids = prompt.get("prompt_token_ids", prompt_ids) if prompt is not None else prompt_ids + max_possible_tokens = self.server.config.max_model_len - len(effective_prompt_ids) if max_possible_tokens <= 0: raise ValueError( - f"Prompt length ({len(prompt_ids)}) meets or exceeds the model's maximum context length " + f"Prompt length ({len(effective_prompt_ids)}) meets or exceeds the model's maximum context length " f"({self.server.config.max_model_len}), leaving no space for generation." ) @@ -189,7 +248,7 @@ def preprocess_input( else: max_tokens = min( self.server.config.response_length, - self.server.config.prompt_length + self.server.config.response_length - len(prompt_ids), + self.server.config.prompt_length + self.server.config.response_length - len(effective_prompt_ids), ) max_tokens = max(0, min(max_tokens, max_possible_tokens)) @@ -201,13 +260,26 @@ def preprocess_input( else: sampling_params["logprobs"] = None sampling_params.setdefault("repetition_penalty", getattr(self.server.config, "repetition_penalty", 1.0)) - params = SamplingParams(max_tokens=max_tokens, **sampling_params) + policy_params = SamplingParams(max_tokens=max_tokens, **sampling_params) + engine = getattr(self.server, "engine", None) + default_params = list(getattr(engine, "default_sampling_params_list", []) or []) + if len(default_params) > 1: + params = copy.deepcopy(default_params) + constrained = self._stage_sampling_constraints.get(0, {}) + for field in {"max_tokens", *sampling_params} - constrained.keys(): + setattr(params[0], field, getattr(policy_params, field)) + else: + params = policy_params - prompt = {"prompt_token_ids": prompt_ids} + if prompt is None: + prompt = {"prompt_token_ids": prompt_ids} + additional_information = prompt.get("additional_information") + if isinstance(additional_information, dict): + additional_information.setdefault("max_new_tokens", [max_tokens]) if multi_modal_data: - prompt["multi_modal_data"] = multi_modal_data - if mm_processor_kwargs: - prompt["mm_processor_kwargs"] = mm_processor_kwargs + prompt.setdefault("multi_modal_data", multi_modal_data) + if mm_processor_kwargs and not adapter_prepared_prompt: + prompt.setdefault("mm_processor_kwargs", mm_processor_kwargs) return prompt, params async def run_generation( @@ -218,17 +290,35 @@ async def run_generation( lora_request: Optional[LoRARequest], priority: int, ) -> Any: - return await self._collect_last_output( - self.server.engine.generate( - prompt=prompt, - sampling_params_list=params, - request_id=request_id, - lora_request=lora_request, - priority=priority, - ) + generate_kwargs = dict( + prompt=prompt, + sampling_params_list=params, + request_id=request_id, + lora_request=lora_request, + priority=priority, ) - - def process_output(self, final_res: Any, params: SamplingParams, sampling_params: dict[str, Any]) -> TokenOutput: + if self._rollout_output_modalities is not None: + generate_kwargs["output_modalities"] = self._rollout_output_modalities + generator = self.server.engine.generate(**generate_kwargs) + if self._rollout_output_modalities is None: + return await self._collect_last_output(generator) + + outputs = [] + async for output in generator: + outputs.append(output) + if self._rollout_adapter is None: + return outputs[-1] if outputs else None + final_res, rollout_fields = self._rollout_adapter.combine_engine_outputs(outputs, prompt) + if final_res is not None and rollout_fields: + final_res._verl_omni_rollout_fields = rollout_fields + return final_res + + def process_output( + self, + final_res: Any, + params: SamplingParams | list[Any], + sampling_params: dict[str, Any], + ) -> TokenOutput: if final_res is None: raise RuntimeError("AR mode: vLLM-Omni engine yielded no output for the prompt.") @@ -237,9 +327,11 @@ def process_output(self, final_res: Any, params: SamplingParams, sampling_params raise RuntimeError("AR mode expects outputs with token IDs, but got None or empty.") extra_fields = {"global_steps": self.server.global_steps} + extra_fields.update(getattr(final_res, "_verl_omni_rollout_fields", {})) token_ids = req_output.outputs[0].token_ids log_probs = None - if params.logprobs is not None: + policy_params = params[0] if isinstance(params, list) else params + if policy_params.logprobs is not None: log_probs = [ logprobs[token_ids[index]].logprob for index, logprobs in enumerate(req_output.outputs[0].logprobs) ] diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py index 5bb28a7c9..f324036d7 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_async_server.py @@ -178,11 +178,6 @@ async def run_server(self, args: argparse.Namespace): engine_client = AsyncOmni(**engine_args) app = build_app(args) await omni_init_app_state(engine_client, app.state, args) - if self._omni_rollout_adapter is not None: - await self._omni_rollout_adapter.initialize_rollout_workers( - engine_client, - self._omni_pipeline_mode, - ) # Deploy config YAML is consumed by AsyncOmni above; clean up the temp dir. if getattr(self, "_temp_deploy_ctx", None) is not None: @@ -197,6 +192,22 @@ async def run_headless(self, args: argparse.Namespace): # TODO (mike): support multi node raise NotImplementedError("vLLM-Omni headless mode is not implemented yet.") + async def collective_rpc( + self, + method: Any, + timeout: float | None = None, + args: tuple = (), + kwargs: dict[str, Any] | None = None, + ): + """Dispatch a shared RPC to the stages selected by the active strategy.""" + return await self.engine.collective_rpc( + method=method, + timeout=timeout, + args=args, + kwargs=kwargs, + stage_ids=self._generate_strategy.collective_rpc_stage_ids(method), + ) + # ----------------------------------------------------------------------- # wake_up hook: Omni does not restore KV cache on wake-up # ----------------------------------------------------------------------- diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_strategy_base.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_strategy_base.py index 57bf23922..a5c396c9c 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_strategy_base.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_strategy_base.py @@ -56,8 +56,8 @@ class OmniStrategyBase(ABC): * optionally override the concrete hooks (:meth:`init_config`, :meth:`init_model_config`, :meth:`validate_configs`, :meth:`post_init`, :meth:`apply_quantization`, :meth:`override_generation_config`, - :meth:`preprocess_engine_kwargs`) when the mode needs behavior beyond the - shared defaults. + :meth:`preprocess_engine_kwargs`, :meth:`collective_rpc_stage_ids`) when + the mode needs behavior beyond the shared defaults. The two concrete subclasses are :class:`~verl_omni.workers.rollout.vllm_rollout.vllm_omni_ar_strategy.ARStrategy` @@ -155,6 +155,14 @@ def preprocess_engine_kwargs(self, engine_kwargs: dict[str, Any]) -> None: """ engine_kwargs.pop("output_mode", None) + def collective_rpc_stage_ids(self, method: Any) -> list[int] | None: + """Return pipeline stages targeted by a shared collective RPC. + + ``None`` preserves the engine default of broadcasting to every stage. + AR adapters can narrow actor weight synchronization to trainable stages. + """ + return None + @abstractmethod def prepare_engine_args(self, engine_args: dict[str, Any], args: Namespace) -> None: """Mutate ``engine_args`` in place with mode-specific engine arguments. From 115d688b306ec5da99374d68c71e4c4e4f2e5b67 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Mon, 31 Aug 2026 18:50:53 +0800 Subject: [PATCH 17/28] [model, rollout, tests, doc] refactor: address follow-up review Signed-off-by: dongbo910220 <1275604947@qq.com> --- docs/api/pipelines.rst | 4 +- .../contributing/integrating_an_omni_model.md | 10 +- examples/grpo_trainer/qwen3_tts/README.md | 27 ++-- pyproject.toml | 4 +- tests/pipelines/test_qwen3_tts_on_cpu.py | 73 +-------- ...pu.py => test_qwen3_tts_package_on_cpu.py} | 44 ++--- .../test_audio_reward_manager_on_cpu.py | 18 +-- .../create_dummy_qwen3_tts_grpo_data.py | 8 +- tests/special_e2e/qwen3_tts_dummy_reward.py | 23 +++ tests/special_e2e/qwen3_tts_smoke_scorer.py | 83 ---------- tests/special_e2e/run_qwen3_tts_grpo_smoke.sh | 32 +--- tests/trainer/omni/test_main_omni_on_cpu.py | 15 +- .../test_qwen3_tts_rollout_on_cpu.py | 46 ++---- tests/workers/test_omni_fsdp_engine_on_cpu.py | 26 +-- .../agent_loop/single_turn_agent_loop.py | 5 - verl_omni/pipelines/model_base.py | 29 +--- .../qwen3_tts/omni_rollout_adapter.py | 152 +++++------------- .../pipelines/qwen3_tts/rollout_utils.py | 91 +---------- .../pipelines/qwen3_tts/talker_forward.py | 38 +++-- .../qwen3_tts/talker_training_adapter.py | 102 ++++-------- .../qwen3_tts/transformers_compat.py | 76 --------- verl_omni/reward_loop/reward_manager/audio.py | 40 ++--- .../reward_score/audio_http_scorer_client.py | 8 +- verl_omni/workers/config/omni/model.py | 29 +--- verl_omni/workers/engine/fsdp/omni_impl.py | 10 +- .../vllm_rollout/vllm_omni_ar_strategy.py | 56 +++---- 26 files changed, 282 insertions(+), 767 deletions(-) rename tests/pipelines/{test_qwen3_tts_transformers_compat_on_cpu.py => test_qwen3_tts_package_on_cpu.py} (63%) create mode 100644 tests/special_e2e/qwen3_tts_dummy_reward.py delete mode 100644 tests/special_e2e/qwen3_tts_smoke_scorer.py delete mode 100644 verl_omni/pipelines/qwen3_tts/transformers_compat.py diff --git a/docs/api/pipelines.rst b/docs/api/pipelines.rst index 5a04401f8..4db785024 100644 --- a/docs/api/pipelines.rst +++ b/docs/api/pipelines.rst @@ -42,7 +42,7 @@ Model Base .. autoclass:: verl_omni.pipelines.model_base.OmniModelBase :members: register, get_class, get_class_by_name, - load_hf_config, get_model_class, + register_auto_classes, get_strip_modules, configure_processor, configure_tokenizer, configure_model, prepare_model_inputs @@ -51,7 +51,7 @@ Model Base build_stage_configs, rollout_flags, weight_sync_stage_ids, get_pipeline_id, ensure_pipeline_registered, get_engine_hf_overrides, get_stage_engine_extras, prepare_engine_prompt, - prepare_agent_sampling_params, postprocess_agent_loop_output, + postprocess_agent_loop_output, combine_engine_outputs .. autoclass:: verl_omni.pipelines.model_base.DiffusionModelBase diff --git a/docs/contributing/integrating_an_omni_model.md b/docs/contributing/integrating_an_omni_model.md index 5ce4ba218..7a9d2ab34 100644 --- a/docs/contributing/integrating_an_omni_model.md +++ b/docs/contributing/integrating_an_omni_model.md @@ -60,12 +60,10 @@ adapt each implementation to your model's architecture: `module._no_split_modules` to the correct decoder layer class for FSDP. This method runs before FSDP wrapping and LoRA injection. -- **`load_hf_config(...)` and `get_model_class()`** (optional): Override these - only when the architecture is not supported by `AutoConfig` or - `AutoModelForMultimodalLM`. Returning `None` from `get_model_class` keeps the - default auto-model path. The FSDP engine still owns `from_pretrained`; - Qwen3-TTS selects the config and model classes published by `qwen-tts` - directly instead of registering them globally with Transformers. +- **`register_auto_classes()`** (optional): Register classes supplied by an + optional model package with the appropriate Transformers Auto APIs. Qwen3-TTS + registers the official `qwen-tts` config and model with `AutoConfig` and + `AutoModelForTextToWaveform`; the FSDP engine still owns `from_pretrained`. - **`prepare_model_inputs(model_inputs, micro_batch, model_config)`** (optional): Validate model-native trajectory or conditioning data retained by diff --git a/examples/grpo_trainer/qwen3_tts/README.md b/examples/grpo_trainer/qwen3_tts/README.md index afad6e556..1032cc352 100644 --- a/examples/grpo_trainer/qwen3_tts/README.md +++ b/examples/grpo_trainer/qwen3_tts/README.md @@ -1,6 +1,6 @@ # Qwen3-TTS GRPO with an audio reward -Last updated: 08/28/2026. +Last updated: 08/31/2026. This example full-parameter tunes the codec-0 policy of `Qwen/Qwen3-TTS-12Hz-0.6B-Base`. It uses verl's stock GRPO advantage, @@ -37,12 +37,17 @@ Install the engine before the training stack: uv pip install -e ".[gpu]" --torch-backend=auto uv pip install "vllm-omni @ git+https://github.com/vllm-project/vllm-omni.git@$(cat .github/vllm_omni_pin.txt)" uv pip install -e ".[tts,train,dev]" -uv pip install qwen-tts==0.1.1 --no-deps +uv pip install --no-deps \ + "qwen-tts @ git+https://github.com/QwenLM/Qwen3-TTS.git@00969daa8064e23adc9e5f52cdf20cf247f94159" ``` -The Qwen3-TTS adapter contains a bounded import compatibility layer for -`qwen-tts==0.1.1` on the repository's Transformers 5.x stack. It does not edit -site-packages. The system `sox` executable is also required by qwen-tts. +The pinned Qwen3-TTS revision is the upstream Transformers 5 support change +from Qwen3-TTS PR #360. The released `qwen-tts==0.1.1` source targets +Transformers 4.57 and cannot be imported unchanged on this repository's +Transformers 5 stack. The adapter registers the upstream config and model with +`AutoConfig` and `AutoModelForTextToWaveform`; it does not carry a local +Transformers compatibility layer. The system `sox` executable is also required +by qwen-tts. ## Data @@ -58,9 +63,9 @@ Training and validation parquet rows use the normal verl format: ``` Use disjoint prompts. The default recipe evaluates the same complete 100-row -validation parquet at step 0 and every 20 updates. Give validation rows fixed -`extra_info.generation_seed` values to keep candidate sampling paired across -checkpoints. +validation parquet at step 0 and every 20 updates. It uses the rollout engine's +global seed; model-specific per-request seed derivation is intentionally outside +this integration. The concatenated replay layout also requires one fixed speaker embedding JSON. Generate it once with the official Qwen3-TTS Base model's @@ -136,8 +141,8 @@ human listening evaluation. The CI-oriented wrapper at [`tests/special_e2e/run_qwen3_tts_grpo_smoke.sh`](../../../tests/special_e2e/run_qwen3_tts_grpo_smoke.sh) -creates deterministic fixtures, starts a duration-only test scorer, and runs two -updates with the official 0.6B Base model. +creates deterministic fixtures, uses an in-process CPU duration reward, and runs +two updates with the official 0.6B Base model. ## References @@ -146,5 +151,7 @@ updates with the official 0.6B Base model. Models](https://arxiv.org/abs/2509.18798), 2025. - Hangrui Hu et al. [Qwen3-TTS Technical Report](https://arxiv.org/abs/2601.15621), 2026. +- QwenLM. [Qwen3-TTS PR #360: Support Transformers + 5](https://github.com/QwenLM/Qwen3-TTS/pull/360), 2026. - Dong Zhang et al. [SpeechAlign: Aligning Speech Generation to Human Preferences](https://arxiv.org/abs/2404.05600), 2024. diff --git a/pyproject.toml b/pyproject.toml index dc094060d..50e610fc8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,8 +58,8 @@ vllm-omni = [ audio = [ "qwen-omni-utils>=0.0.9", ] -# Install qwen-tts separately with --no-deps because qwen-tts 0.1.1 pins -# Transformers 4.57.3 while vLLM 0.27 requires Transformers 5.x. +# Install qwen-tts separately from the pinned upstream Transformers 5 support +# revision. The released 0.1.1 source and metadata target Transformers 4.57.3. tts = [ "einops>=0.8.0", "librosa>=0.10.2", diff --git a/tests/pipelines/test_qwen3_tts_on_cpu.py b/tests/pipelines/test_qwen3_tts_on_cpu.py index 597b326a0..18a759e43 100644 --- a/tests/pipelines/test_qwen3_tts_on_cpu.py +++ b/tests/pipelines/test_qwen3_tts_on_cpu.py @@ -18,7 +18,6 @@ from pathlib import Path from types import SimpleNamespace -import numpy as np import pytest import torch @@ -71,6 +70,8 @@ def test_only_validated_auto_language_layout_is_accepted(): assert forward.require_auto_language("auto") == "Auto" with pytest.raises(ValueError, match="supports only tts_language=Auto"): forward.require_auto_language("Chinese") + with pytest.raises(ValueError, match="supports only tts_language=Auto"): + forward.require_auto_language(None) def test_actor_logits_align_to_effective_codec0_response(monkeypatch): @@ -153,68 +154,8 @@ def test_codec_alignment_preserves_final_row_and_rejects_heuristic_match(): rollout.align_audio_codes(malformed, token_ids) -def test_rollout_chunk_accumulator_handles_cumulative_and_delta_outputs(): - first = torch.zeros(12, 16, dtype=torch.long) - cumulative = torch.cat((first, torch.ones(1, 16, dtype=torch.long))) - assert torch.equal(rollout.append_tensor_chunk(first, cumulative), cumulative) - - delta = torch.full((1, 16), 2, dtype=torch.long) - assert torch.equal(rollout.append_tensor_chunk(cumulative, delta), torch.cat((cumulative, delta))) - with pytest.raises(RuntimeError, match="changed shape"): - rollout.append_tensor_chunk(cumulative, torch.zeros(1, 15, dtype=torch.long)) - - -def test_validation_seed_covers_both_codec_samplers_and_candidates_without_mutation(): - original = {"temperature": 0.8, "extra_args": {"existing": "kept"}} - seeded = rollout.with_rollout_generation_seed( - original, - {"split": "validation", "generation_seed": np.int64(42017)}, - session_id=3, - global_steps=100, - require_session_id=True, - ) - - assert seeded == { - "temperature": 0.8, - "seed": 42020, - "extra_args": {"existing": "kept", "tts_local_seed": 42020}, - } - assert seeded == rollout.with_rollout_generation_seed( - original, - {"split": "validation", "generation_seed": np.int64(42017)}, - session_id=3, - global_steps=0, - require_session_id=True, - ) - assert original == {"temperature": 0.8, "extra_args": {"existing": "kept"}} - - gate_first = rollout.with_rollout_generation_seed( - {}, {"split": "gate", "id": "gate-7"}, session_id=0, global_steps=0, base_seed=42 - ) - gate_second = rollout.with_rollout_generation_seed( - {}, {"split": "gate", "id": "gate-7"}, session_id=1, global_steps=100, base_seed=42 - ) - assert gate_first["seed"] != gate_second["seed"] - assert gate_first == rollout.with_rollout_generation_seed( - {}, {"split": "gate", "id": "gate-7"}, session_id=0, global_steps=100, base_seed=42 - ) - - -def test_training_seeds_are_reproducible_and_group_diverse(): - kwargs = { - "extra_info": {"split": "train", "id": "sample-7"}, - "global_steps": 12, - "uid": "uid-7", - "base_seed": 42, - "require_session_id": True, - } - first = rollout.with_rollout_generation_seed({}, session_id=0, **kwargs) - repeated = rollout.with_rollout_generation_seed({}, session_id=0, **kwargs) - second = rollout.with_rollout_generation_seed({}, session_id=1, **kwargs) - next_step = rollout.with_rollout_generation_seed({}, session_id=0, **{**kwargs, "global_steps": 13}) - - assert first == repeated - assert len({first["seed"], second["seed"], next_step["seed"]}) == 3 - assert first["seed"] == first["extra_args"]["tts_local_seed"] - with pytest.raises(RuntimeError, match="session_id"): - rollout.with_rollout_generation_seed({}, session_id=None, **kwargs) +def test_talker_batch_and_logit_mask_reject_contract_mismatches(): + with pytest.raises(ValueError, match="matching non-empty"): + forward.build_talker_batch([], [], TOKENS, sub_codebook_vocab=2048) + with pytest.raises(ValueError, match="codebook_vocab"): + forward.mask_codec0_logits(torch.zeros((1, 2, 100)), 101, TOKENS.codec_eos) diff --git a/tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py b/tests/pipelines/test_qwen3_tts_package_on_cpu.py similarity index 63% rename from tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py rename to tests/pipelines/test_qwen3_tts_package_on_cpu.py index 724e7b532..c1cbd5be3 100644 --- a/tests/pipelines/test_qwen3_tts_transformers_compat_on_cpu.py +++ b/tests/pipelines/test_qwen3_tts_package_on_cpu.py @@ -11,41 +11,32 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Optional real-package check for qwen-tts on the repository TF5 stack.""" +"""Check the pinned upstream qwen-tts TF5 source on the repository stack.""" -import importlib import importlib.util -from pathlib import Path import pytest import torch -ROOT = Path(__file__).parents[2] - -def _load_compat_module(): - path = ROOT / "verl_omni/pipelines/qwen3_tts/transformers_compat.py" - spec = importlib.util.spec_from_file_location("qwen3_tts_transformers_compat_under_test", path) - module = importlib.util.module_from_spec(spec) - assert spec.loader is not None - spec.loader.exec_module(module) - return module - - -def test_qwen_tts_tiny_model_constructs_and_forwards_without_source_patch(): +def test_qwen_tts_registers_and_runs_without_a_transformers_compatibility_layer(): transformers = pytest.importorskip("transformers") - assert int(transformers.__version__.split(".", maxsplit=1)[0]) >= 5 if importlib.util.find_spec("qwen_tts") is None: pytest.skip("qwen-tts is an optional dependency") - compat = _load_compat_module() - rope_utils = importlib.import_module("transformers.modeling_rope_utils") - original_rope_functions = rope_utils.ROPE_INIT_FUNCTIONS - with compat.qwen3_tts_import_context(): - config_module = importlib.import_module("qwen_tts.core.models.configuration_qwen3_tts") - model_module = importlib.import_module("qwen_tts.core.models.modeling_qwen3_tts") - assert rope_utils.ROPE_INIT_FUNCTIONS is original_rope_functions - compat.patch_qwen3_tts_config_defaults(config_module.Qwen3TTSConfig) + from qwen_tts.core.models.configuration_qwen3_tts import Qwen3TTSConfig + from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSForConditionalGeneration + from transformers import AutoConfig, AutoModelForTextToWaveform + + assert int(transformers.__version__.split(".", maxsplit=1)[0]) >= 5 + AutoConfig.register("qwen3_tts", Qwen3TTSConfig, exist_ok=True) + AutoModelForTextToWaveform.register( + Qwen3TTSConfig, + Qwen3TTSForConditionalGeneration, + exist_ok=True, + ) + assert AutoConfig.for_model("qwen3_tts").__class__ is Qwen3TTSConfig + assert AutoModelForTextToWaveform._model_mapping[Qwen3TTSConfig] is Qwen3TTSForConditionalGeneration predictor = { "vocab_size": 32, @@ -81,14 +72,13 @@ def test_qwen_tts_tiny_model_constructs_and_forwards_without_source_patch(): "interleaved": True, }, } - config = config_module.Qwen3TTSConfig( + config = Qwen3TTSConfig( talker_config=talker, speaker_encoder_config={}, tts_model_type="custom", tokenizer_type="12hz", ) - assert config.talker_config.pad_token_id is None - model = model_module.Qwen3TTSForConditionalGeneration(config) + model = Qwen3TTSForConditionalGeneration(config) output = model.talker( inputs_embeds=torch.randn(2, 5, 8), attention_mask=torch.ones(2, 5, dtype=torch.long), diff --git a/tests/reward_loop/test_audio_reward_manager_on_cpu.py b/tests/reward_loop/test_audio_reward_manager_on_cpu.py index 547b1289d..bf33fbc32 100644 --- a/tests/reward_loop/test_audio_reward_manager_on_cpu.py +++ b/tests/reward_loop/test_audio_reward_manager_on_cpu.py @@ -150,19 +150,13 @@ def test_missing_sample_rate_fails_closed(): manager.loop.run_until_complete(manager.run_single(data)) -def test_list_of_chunks_is_concatenated(): - def compute_score(solution_audio, **kwargs): - waveform, sample_rate = solution_audio - np.testing.assert_array_equal(waveform, np.array([0.1, 0.2, 0.3], dtype=np.float32)) - assert sample_rate == 16_000 - return 0.5 - - manager = _manager(compute_score) - result = manager.loop.run_until_complete( - manager.run_single(_data([torch.tensor([0.1, 0.2]), torch.tensor([0.3])], 16_000)) - ) +def test_chunked_waveform_is_rejected_instead_of_guessed(): + manager = _manager(lambda **kwargs: 0.5) - assert result == {"reward_score": 0.5, "reward_extra_info": {"acc": 0.5}} + with pytest.raises(ValueError, match="could not convert"): + manager.loop.run_until_complete( + manager.run_single(_data([torch.tensor([0.1, 0.2]), torch.tensor([0.3])], 16_000)) + ) @pytest.mark.asyncio diff --git a/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py b/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py index 9e2f9836a..ac67dcd05 100644 --- a/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py +++ b/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py @@ -38,10 +38,8 @@ ) -def _row(text: str, sample_id: str, split: str, generation_seed: int | None = None) -> dict: +def _row(text: str, sample_id: str, split: str) -> dict: extra_info = {"id": sample_id, "split": split} - if generation_seed is not None: - extra_info["generation_seed"] = generation_seed return { "data_source": "tts", "prompt": [{"role": "user", "content": text}], @@ -57,9 +55,7 @@ def main() -> None: args.output_dir.mkdir(parents=True, exist_ok=True) train_rows = [_row(text, f"train-{index}", "train") for index, text in enumerate(TRAIN_TEXTS)] - validation_rows = [ - _row(text, f"validation-{index}", "validation", 1_000 + index) for index, text in enumerate(VALIDATION_TEXTS) - ] + validation_rows = [_row(text, f"validation-{index}", "validation") for index, text in enumerate(VALIDATION_TEXTS)] pd.DataFrame(train_rows).to_parquet(args.output_dir / "train.parquet", index=False) pd.DataFrame(validation_rows).to_parquet(args.output_dir / "validation.parquet", index=False) diff --git a/tests/special_e2e/qwen3_tts_dummy_reward.py b/tests/special_e2e/qwen3_tts_dummy_reward.py new file mode 100644 index 000000000..3cda6b744 --- /dev/null +++ b/tests/special_e2e/qwen3_tts_dummy_reward.py @@ -0,0 +1,23 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU-only reward used by the Qwen3-TTS execution smoke test.""" + +import numpy as np + + +def compute_score(solution_audio, **kwargs): + del kwargs + waveform, sample_rate = solution_audio + duration_s = np.asarray(waveform).size / sample_rate + return {"score": float(duration_s), "duration_s": float(duration_s)} diff --git a/tests/special_e2e/qwen3_tts_smoke_scorer.py b/tests/special_e2e/qwen3_tts_smoke_scorer.py deleted file mode 100644 index e20b2162b..000000000 --- a/tests/special_e2e/qwen3_tts_smoke_scorer.py +++ /dev/null @@ -1,83 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2026 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Deterministic duration reward service for the Qwen3-TTS execution smoke.""" - -import argparse -import base64 -import json -import math -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path - - -class _Handler(BaseHTTPRequestHandler): - server_version = "Qwen3TTSSmokeScorer/1.0" - - def _write_json(self, status: int, payload: dict) -> None: - body = json.dumps(payload, allow_nan=False).encode("utf-8") - self.send_response(status) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - def do_GET(self) -> None: - if self.path != "/health": - self._write_json(404, {"error": "not found"}) - return - self._write_json(200, {"status": "ready", "reward_mode": "duration_smoke_only"}) - - def do_POST(self) -> None: - if self.path != "/score": - self._write_json(404, {"error": "not found"}) - return - try: - content_length = int(self.headers.get("Content-Length", "0")) - payload = json.loads(self.rfile.read(content_length)) - if payload.get("protocol_version") != "1": - raise ValueError("protocol_version must be '1'") - num_samples = int(payload["num_samples"]) - sample_rate = int(payload["sample_rate"]) - waveform = base64.b64decode(payload["waveform_f32_base64"], validate=True) - if num_samples <= 0 or sample_rate <= 0 or len(waveform) != num_samples * 4: - raise ValueError("invalid float32 waveform shape") - duration_s = num_samples / sample_rate - if not math.isfinite(duration_s): - raise ValueError("duration must be finite") - result = {"score": duration_s, "duration_s": duration_s, "reward_mode": "duration_smoke_only"} - with self.server.log_path.open("a", encoding="utf-8") as handle: - handle.write(json.dumps(result, allow_nan=False) + "\n") - self._write_json(200, result) - except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: - self._write_json(400, {"error": str(exc)}) - - def log_message(self, format: str, *args) -> None: - return - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--host", default="127.0.0.1") - parser.add_argument("--port", type=int, required=True) - parser.add_argument("--log", type=Path, required=True) - args = parser.parse_args() - args.log.parent.mkdir(parents=True, exist_ok=True) - server = ThreadingHTTPServer((args.host, args.port), _Handler) - server.log_path = args.log - server.serve_forever() - - -if __name__ == "__main__": - main() diff --git a/tests/special_e2e/run_qwen3_tts_grpo_smoke.sh b/tests/special_e2e/run_qwen3_tts_grpo_smoke.sh index 2a5a7dfb4..67e4e37e6 100755 --- a/tests/special_e2e/run_qwen3_tts_grpo_smoke.sh +++ b/tests/special_e2e/run_qwen3_tts_grpo_smoke.sh @@ -1,6 +1,6 @@ #!/usr/bin/env bash # Qwen3-TTS full-parameter GRPO e2e smoke: real 0.6B model, two updates. -# This validates execution only; the duration scorer is not a quality reward. +# This validates execution only; the in-process CPU reward is not a quality reward. set -xeuo pipefail @@ -18,12 +18,12 @@ MODEL_REPO="${MODEL_REPO:-Qwen/Qwen3-TTS-12Hz-0.6B-Base}" WORK_DIR="${WORK_DIR:-${TMPDIR:-/tmp}/qwen3_tts_grpo_smoke_${USER:-user}_$$}" DATA_DIR="${WORK_DIR}/data" OUTPUT_DIR="${OUTPUT_DIR:-${WORK_DIR}/output}" -SCORER_LOG="${WORK_DIR}/scorer_requests.jsonl" mkdir -p "${WORK_DIR}" "${OUTPUT_DIR}" -if ! "${PYTHON_BIN}" -c 'import importlib.util; modules=("qwen_tts", "onnxruntime", "soundfile", "librosa", "sox"); raise SystemExit(any(importlib.util.find_spec(name) is None for name in modules))'; then +if ! "${PYTHON_BIN}" -c 'from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSForConditionalGeneration; import onnxruntime, soundfile, librosa, sox'; then uv pip install --python "${PYTHON_BIN}" -e ".[tts]" - uv pip install --python "${PYTHON_BIN}" --no-deps qwen-tts==0.1.1 + uv pip install --python "${PYTHON_BIN}" --force-reinstall --no-deps \ + "qwen-tts @ git+https://github.com/QwenLM/Qwen3-TTS.git@00969daa8064e23adc9e5f52cdf20cf247f94159" fi MODEL_PATH="${MODEL_PATH:-}" @@ -35,30 +35,11 @@ fi "${PYTHON_BIN}" tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py --output-dir "${DATA_DIR}" -SCORER_PORT="${SCORER_PORT:-$("${PYTHON_BIN}" -c \ - 'import socket; s=socket.socket(); s.bind(("127.0.0.1", 0)); print(s.getsockname()[1]); s.close()')}" -"${PYTHON_BIN}" tests/special_e2e/qwen3_tts_smoke_scorer.py \ - --port "${SCORER_PORT}" --log "${SCORER_LOG}" >"${WORK_DIR}/scorer.log" 2>&1 & -SCORER_PID=$! -cleanup() { - kill "${SCORER_PID}" 2>/dev/null || true - wait "${SCORER_PID}" 2>/dev/null || true -} -trap cleanup EXIT - -for _ in $(seq 1 60); do - if curl -fsS "http://127.0.0.1:${SCORER_PORT}/health" >"${WORK_DIR}/scorer_health.json"; then - break - fi - sleep 1 -done -curl -fsS "http://127.0.0.1:${SCORER_PORT}/health" - MODEL_PATH="${MODEL_PATH}" \ TRAIN_FILE="${DATA_DIR}/train.parquet" \ VAL_FILE="${DATA_DIR}/validation.parquet" \ SPK_EMBED_PATH="${DATA_DIR}/speaker.json" \ -SCORER_URL="http://127.0.0.1:${SCORER_PORT}/score" \ +SCORER_URL="http://unused.invalid/score" \ OUTPUT_DIR="${OUTPUT_DIR}" \ NUM_GPUS="${NUM_GPUS}" \ TOTAL_TRAINING_STEPS=2 \ @@ -70,10 +51,11 @@ bash examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh \ actor_rollout_ref.rollout.n=2 \ actor_rollout_ref.rollout.agent.num_workers=2 \ actor_rollout_ref.rollout.max_num_seqs=4 \ + "reward.custom_reward_function.path=${REPO_ROOT}/tests/special_e2e/qwen3_tts_dummy_reward.py" \ + reward.custom_reward_function.name=compute_score \ trainer.val_before_train=false \ trainer.log_val_generations=0 \ "$@" -[[ -s "${SCORER_LOG}" ]] grep -q "training/global_step.*2" "${OUTPUT_DIR}/train.log" echo "Qwen3-TTS GRPO e2e smoke passed; artifacts: ${WORK_DIR}" diff --git a/tests/trainer/omni/test_main_omni_on_cpu.py b/tests/trainer/omni/test_main_omni_on_cpu.py index 59f220d71..1f1675ed4 100644 --- a/tests/trainer/omni/test_main_omni_on_cpu.py +++ b/tests/trainer/omni/test_main_omni_on_cpu.py @@ -98,10 +98,6 @@ def test_omni_model_config_loads_tokenizer_and_processor_via_adapter(self, monke from verl_omni.workers.config.omni.model import OmniModelConfig mock_adapter = MagicMock() - mock_adapter.load_hf_config.return_value = SimpleNamespace( - tie_word_embeddings=False, - architectures=["arch"], - ) mock_adapter.configure_tokenizer.return_value = "tokenizer" mock_adapter.configure_processor.return_value = "processor" monkeypatch.setattr( @@ -110,6 +106,11 @@ def test_omni_model_config_loads_tokenizer_and_processor_via_adapter(self, monke ) monkeypatch.setattr(model_config_module, "resolve_model_local_dir", lambda path, use_shm=False: str(tmp_path)) monkeypatch.setattr(model_config_module, "copy_to_local", lambda path, use_shm=False: f"local:{path}") + monkeypatch.setattr( + model_config_module.AutoConfig, + "from_pretrained", + lambda *_args, **_kwargs: SimpleNamespace(tie_word_embeddings=False, architectures=["arch"]), + ) model_config = OmniModelConfig( path=str(tmp_path), @@ -122,11 +123,7 @@ def test_omni_model_config_loads_tokenizer_and_processor_via_adapter(self, monke assert model_config.tokenizer == "tokenizer" assert model_config.processor == "processor" - mock_adapter.load_hf_config.assert_called_once_with( - "local:" + str(tmp_path), - trust_remote_code=False, - attn_implementation="flash_attention_2", - ) + mock_adapter.register_auto_classes.assert_called_once_with() mock_adapter.configure_tokenizer.assert_called_once_with("local:tokenizer-path", model_config) mock_adapter.configure_processor.assert_called_once_with(str(tmp_path), model_config) diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index 44a5758ca..e0bf6d962 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -31,10 +31,7 @@ from verl_omni.pipelines.qwen3_tts import omni_rollout_adapter from verl_omni.pipelines.qwen3_tts.omni_rollout_adapter import Qwen3TTSRolloutAdapter from verl_omni.pipelines.qwen3_tts.talker_training_adapter import Qwen3TTSTalkerAdapter -from verl_omni.workers.rollout.vllm_rollout.vllm_omni_ar_strategy import ( - ARStrategy, - _retained_output_modalities, -) +from verl_omni.workers.rollout.vllm_rollout.vllm_omni_ar_strategy import ARStrategy from verl_omni.workers.rollout.vllm_rollout.vllm_omni_async_server import vLLMOmniHttpServer @@ -67,16 +64,6 @@ def test_optional_rollout_hooks_preserve_existing_ar_defaults(): assert OmniRolloutPipelineBase.weight_sync_stage_ids() is None assert OmniRolloutPipelineBase.prepare_engine_prompt([], None, {}) is None - sampling_params = {"temperature": 0.8} - assert ( - OmniRolloutPipelineBase.prepare_agent_sampling_params( - sampling_params, - rollout_config=None, - trainer_config=None, - agent_inputs={}, - ) - == sampling_params - ) assert ( OmniRolloutPipelineBase.postprocess_agent_loop_output( final, @@ -86,6 +73,8 @@ def test_optional_rollout_hooks_preserve_existing_ar_defaults(): is final ) assert OmniRolloutPipelineBase.combine_engine_outputs([first, final], {}) == (final, {}) + with pytest.raises(RuntimeError, match="no outputs"): + OmniRolloutPipelineBase.combine_engine_outputs([], {}) def test_omni_single_turn_agent_resolves_registered_pipeline_adapter(): @@ -171,7 +160,9 @@ def test_rollout_adapter_builds_unique_prompt_and_scopes_weight_sync(tmp_path): assert first["additional_information"]["text"] == ["first text"] assert first["cache_salt"] != second["cache_salt"] assert Qwen3TTSRolloutAdapter.weight_sync_stage_ids("full") == [0] - assert _retained_output_modalities(Qwen3TTSRolloutAdapter.build_stage_configs("full")) == ["latent", "audio"] + assert [ + stage.final_output_type for stage in Qwen3TTSRolloutAdapter.build_stage_configs("full") if stage.final_output + ] == ["latent", "audio"] def test_ar_strategy_resolves_qwen3_tts_adapter_and_scopes_weight_sync(monkeypatch): @@ -214,7 +205,7 @@ def test_talker_adapter_pads_exact_rollout_fields_for_actor_forward(): model_inputs = {"input_ids": torch.zeros(2, 6, dtype=torch.long)} micro_batch = { "extra_fields": [ - {"tts_text_ids": [1, 2], "tts_audio_codes": torch.ones(3, 16, dtype=torch.long)}, + {"tts_text_ids": [1, 2, 6], "tts_audio_codes": torch.ones(3, 16, dtype=torch.long)}, {"tts_text_ids": [3, 4, 5], "tts_audio_codes": torch.full((2, 16), 2, dtype=torch.long)}, ] } @@ -223,7 +214,7 @@ def test_talker_adapter_pads_exact_rollout_fields_for_actor_forward(): assert prepared["tts_text_ids"].shape == (2, 3) assert prepared["tts_audio_codes"].shape == (2, 3, 16) - assert prepared["text_len"].tolist() == [2, 3] + assert prepared["text_len"].tolist() == [3, 3] assert prepared["response_len"].tolist() == [3, 2] assert not prepared["tts_audio_codes"][1, 2].any() @@ -260,24 +251,7 @@ def test_rollout_adapter_combines_policy_codes_and_waveform(): assert fields["tts_text"] == "first text" -def test_rollout_adapter_prepares_sampling_and_actor_policy_sequence(): - rollout_config = SimpleNamespace(n=8, val_kwargs=SimpleNamespace(n=1)) - trainer_config = SimpleNamespace(data={"seed": 42}) - agent_inputs = { - "extra_info": {"split": "train", "id": "sample-7"}, - "session_id": 3, - "global_steps": 12, - "uid": "uid-7", - } - - seeded = Qwen3TTSRolloutAdapter.prepare_agent_sampling_params( - {"temperature": 0.8}, - rollout_config=rollout_config, - trainer_config=trainer_config, - agent_inputs=agent_inputs, - ) - assert seeded["seed"] == seeded["extra_args"]["tts_local_seed"] - +def test_rollout_adapter_prepares_actor_policy_sequence(): codes = torch.arange(5 * 16, dtype=torch.long).reshape(5, 16) output = SimpleNamespace( prompt_ids=[9, 8], @@ -343,7 +317,7 @@ def prepare_engine_prompt(**kwargs): @pytest.mark.asyncio async def test_ar_strategy_retains_requested_stage_outputs_and_targets_weight_sync(): - policy = SimpleNamespace(request_output=SimpleNamespace(outputs=[])) + policy = SimpleNamespace(outputs=[]) class Engine: def __init__(self): diff --git a/tests/workers/test_omni_fsdp_engine_on_cpu.py b/tests/workers/test_omni_fsdp_engine_on_cpu.py index 496b882b0..e99683491 100644 --- a/tests/workers/test_omni_fsdp_engine_on_cpu.py +++ b/tests/workers/test_omni_fsdp_engine_on_cpu.py @@ -58,6 +58,7 @@ def _make_mock_model_config(**overrides): cfg.model_stage = "thinker" cfg.local_path = "/fake/model/path" cfg.trust_remote_code = False + cfg.external_lib = None cfg.use_liger = False cfg.use_fused_kernels = False cfg.enable_gradient_checkpointing = False @@ -379,10 +380,11 @@ def test_collect_lora_params_import_not_from_verl(): # --------------------------------------------------------------------------- -def test_build_module_uses_auto_model_for_multimodal_lm(): - """``_build_module`` uses ``AutoModelForMultimodalLM``, not ``AutoModelForCausalLM``.""" +def test_build_module_uses_stage_specific_auto_model_classes(): + """Thinkers use multimodal auto models and talkers use text-to-waveform auto models.""" omni_impl = _get_omni_impl_module() assert omni_impl.AutoModelForMultimodalLM is not None + assert omni_impl.AutoModelForTextToWaveform is not None tree = _parse_omni_impl_ast() import_names = set() @@ -393,6 +395,9 @@ def test_build_module_uses_auto_model_for_multimodal_lm(): assert "AutoModelForMultimodalLM" in import_names, ( f"AutoModelForMultimodalLM not imported from transformers; imports: {import_names}" ) + assert "AutoModelForTextToWaveform" in import_names, ( + f"AutoModelForTextToWaveform not imported from transformers; imports: {import_names}" + ) assert "AutoModelForCausalLM" not in import_names, "AutoModelForCausalLM should NOT be imported from transformers" @@ -406,7 +411,6 @@ def test_build_module_calls_adapter_configure_model(architecture): fake_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] fake_adapter_cls = MagicMock() - fake_adapter_cls.get_model_class.return_value = None fake_configured_module = MagicMock(spec=torch.nn.Module) fake_configured_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] fake_adapter_cls.configure_model.return_value = fake_configured_module @@ -441,7 +445,7 @@ def test_build_module_calls_adapter_configure_model(architecture): mock_get_cls.assert_called_once_with( architecture, model_config.model_stage, - model_config.get("external_lib"), + model_config.external_lib, ) fake_adapter_cls.configure_model.assert_called_once_with(fake_module, model_config) @@ -449,22 +453,23 @@ def test_build_module_calls_adapter_configure_model(architecture): assert result is fake_configured_module -def test_build_module_uses_adapter_model_class_when_auto_model_is_incompatible(): +def test_build_module_uses_text_to_waveform_auto_model_for_talker(): omni_impl = _get_omni_impl_module() - model_config = _make_mock_model_config(architecture="CustomOmniForConditionalGeneration") + model_config = _make_mock_model_config( + architecture="Qwen3TTSForConditionalGeneration", + model_stage="talker", + ) loaded_module = MagicMock(spec=torch.nn.Module) loaded_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] configured_module = MagicMock(spec=torch.nn.Module) configured_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] adapter_cls = MagicMock() - custom_model_cls = MagicMock() - custom_model_cls.from_pretrained.return_value = loaded_module - adapter_cls.get_model_class.return_value = custom_model_cls adapter_cls.configure_model.return_value = configured_module model_base_mod = sys.modules["verl_omni.pipelines.model_base"] with ( patch.object(model_base_mod.OmniModelBase, "get_class_by_name", return_value=adapter_cls), + patch.object(omni_impl.AutoModelForTextToWaveform, "from_pretrained", return_value=loaded_module) as load, patch.object(omni_impl, "get_init_weight_context_manager", return_value=MagicMock()), patch.object(omni_impl.warnings, "catch_warnings", return_value=MagicMock()), patch("verl.utils.torch_dtypes.PrecisionType") as precision_type, @@ -477,8 +482,7 @@ def test_build_module_uses_adapter_model_class_when_auto_model_is_incompatible() result = engine._build_module() - adapter_cls.get_model_class.assert_called_once_with() - custom_model_cls.from_pretrained.assert_called_once_with( + load.assert_called_once_with( pretrained_model_name_or_path=model_config.local_path, torch_dtype=torch.float32, config=model_config.hf_config, diff --git a/verl_omni/agent_loop/single_turn_agent_loop.py b/verl_omni/agent_loop/single_turn_agent_loop.py index 6034716bf..59b44480c 100644 --- a/verl_omni/agent_loop/single_turn_agent_loop.py +++ b/verl_omni/agent_loop/single_turn_agent_loop.py @@ -11,8 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import logging -import os from typing import Any from uuid import uuid4 @@ -24,9 +22,6 @@ from verl_omni.agent_loop.diffusion_agent_loop import DiffusionAgentLoopOutput from verl_omni.pipelines.model_base import OmniRolloutPipelineBase -logger = logging.getLogger(__file__) -logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) - @register("omni_single_turn_agent") class OmniSingleTurnAgentLoop(SingleTurnAgentLoop): diff --git a/verl_omni/pipelines/model_base.py b/verl_omni/pipelines/model_base.py index e3a2f1835..54aa439b1 100644 --- a/verl_omni/pipelines/model_base.py +++ b/verl_omni/pipelines/model_base.py @@ -559,26 +559,9 @@ def get_class_by_name( ) from None @classmethod - def load_hf_config( - cls, - model_path: str, - *, - trust_remote_code: bool, - attn_implementation: str, - ): - """Load the Hugging Face config used by the FSDP engine.""" - from transformers import AutoConfig - - return AutoConfig.from_pretrained( - model_path, - trust_remote_code=trust_remote_code, - attn_implementation=attn_implementation, - ) - - @classmethod - def get_model_class(cls): - """Return a model class override, or ``None`` for the default auto model.""" - return None + def register_auto_classes(cls) -> None: + """Register optional model-package classes with Transformers auto APIs.""" + return @classmethod @abstractmethod @@ -802,7 +785,7 @@ def get_pipeline_id(cls, pipeline_mode: str = "thinker_only") -> str: for model_type, cls_ref in cls._registry.items(): if cls_ref is cls: return model_type - return "" + raise RuntimeError(f"{cls.__name__} is not registered as an omni rollout pipeline.") @classmethod def ensure_pipeline_registered(cls, pipeline_mode: str = "thinker_only") -> None: @@ -858,4 +841,6 @@ def get_output_modalities(cls, pipeline_mode: str = "thinker_only") -> list[str] @classmethod def combine_engine_outputs(cls, outputs: list, prompt: dict) -> tuple[Any, dict[str, Any]]: """Select the policy output and collect architecture-specific fields.""" - return (outputs[-1] if outputs else None), {} + if not outputs: + raise RuntimeError("The omni rollout engine returned no outputs.") + return outputs[-1], {} diff --git a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py index 264a79309..54b5b84c9 100644 --- a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py @@ -13,9 +13,7 @@ # limitations under the License. """Qwen3-TTS two-stage rollout adapter.""" -import copy import hashlib -from collections.abc import Mapping from dataclasses import replace from functools import lru_cache @@ -25,12 +23,7 @@ from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE from verl_omni.pipelines.model_base import OmniRolloutPipelineBase -from verl_omni.pipelines.qwen3_tts.rollout_utils import ( - align_audio_codes, - append_tensor_chunk, - is_evaluation_split, - with_rollout_generation_seed, -) +from verl_omni.pipelines.qwen3_tts.rollout_utils import align_audio_codes from verl_omni.pipelines.qwen3_tts.talker_forward import ( TEXT_PROMPT_TRAILER_TOKENS, build_assistant_text, @@ -39,13 +32,12 @@ ) _PIPELINE_ID = "qwen3_tts_rl" -_SYNC_PROCESSOR = "verl_omni.pipelines.qwen3_tts.omni_rollout_adapter.talker2code2wav_token_only" QWEN3_TTS_RL_PIPELINE = PipelineConfig( model_type=_PIPELINE_ID, model_arch=QWEN3_TTS_PIPELINE.model_arch, stages=( replace(QWEN3_TTS_PIPELINE.stages[0], final_output=True, final_output_type="latent"), - replace(QWEN3_TTS_PIPELINE.stages[1], sync_process_input_func=_SYNC_PROCESSOR), + QWEN3_TTS_PIPELINE.stages[1], ), ) @@ -55,44 +47,6 @@ def _load_speaker_vector(path: str) -> list[float]: return load_speaker_xvector(path).reshape(-1).tolist() -def _completion(output): - request_output = getattr(output, "request_output", None) or output - completions = getattr(request_output, "outputs", None) - return completions[0] if completions else None - - -def _copy_plain_containers(value): - """Convert Ray/shared-memory Mapping/list shells to built-ins for upstream strict dict checks. - - Tensor payloads are preserved rather than copied. - """ - if isinstance(value, Mapping): - return {key: _copy_plain_containers(item) for key, item in value.items()} - if isinstance(value, list): - return [_copy_plain_containers(item) for item in value] - return value - - -def talker2code2wav_token_only(source_outputs, prompt=None, _requires_multimodal_data=False): - """Give the mutating upstream processor ordinary Python containers.""" - from vllm_omni.model_executor.stage_input_processors.qwen3_tts import ( - talker2code2wav_token_only as upstream_processor, - ) - - converted = [] - for source_output in source_outputs: - source_copy = copy.copy(source_output) - source_copy.outputs = [] - for completion in getattr(source_output, "outputs", []): - completion_copy = copy.copy(completion) - multimodal = getattr(completion, "multimodal_output", None) - if isinstance(multimodal, Mapping): - completion_copy.multimodal_output = _copy_plain_containers(multimodal) - source_copy.outputs.append(completion_copy) - converted.append(source_copy) - return upstream_processor(converted, prompt, _requires_multimodal_data) - - @OmniRolloutPipelineBase.register(_PIPELINE_ID) class Qwen3TTSRolloutAdapter(OmniRolloutPipelineBase): @classmethod @@ -124,38 +78,21 @@ def weight_sync_stage_ids(cls, pipeline_mode="full"): @classmethod def get_stage_engine_extras(cls, stage_id, pipeline_mode="full"): cls._check_mode(pipeline_mode) - return {"max_model_len": 65536, "max_num_batched_tokens": 65536} if stage_id == 1 else {} - - @classmethod - def prepare_agent_sampling_params( - cls, - sampling_params, - *, - rollout_config, - trainer_config, - agent_inputs, - ): - """Seed codec-0 and residual-codebook sampling for each GRPO candidate.""" - extra_info = agent_inputs.get("extra_info") - evaluation = is_evaluation_split(extra_info) - candidate_count = rollout_config.val_kwargs.n if evaluation else rollout_config.n - return with_rollout_generation_seed( - sampling_params, - extra_info, - session_id=agent_inputs.get("session_id"), - global_steps=agent_inputs.get("global_steps"), - uid=agent_inputs.get("uid"), - base_seed=int(trainer_config.data.get("seed", 0)), - require_session_id=int(candidate_count) > 1, - ) + if stage_id == 0: + return {} + if stage_id == 1: + return {"max_model_len": 65536, "max_num_batched_tokens": 65536} + raise ValueError(f"Qwen3-TTS has no rollout stage {stage_id}.") @classmethod def postprocess_agent_loop_output(cls, output, *, tokenizer, response_length): """Map the 16-codebook rollout to codec-0 policy tokens and replay fields.""" extra = output.extra_fields - codes, text = extra.get("tts_audio_codes"), extra.get("tts_text") - if codes is None or text is None: + if "tts_audio_codes" not in extra or "tts_text" not in extra: raise RuntimeError("Qwen3-TTS rollout did not return codec codes and text.") + codes, text = extra["tts_audio_codes"], extra["tts_text"] + if not isinstance(text, str): + raise TypeError(f"Qwen3-TTS rollout text must be a string, got {type(text).__name__}.") codes = torch.as_tensor(codes, dtype=torch.long) if codes.ndim != 2 or codes.shape[-1] != 16: raise ValueError("Qwen3-TTS codec codes must have shape (frames, 16).") @@ -167,7 +104,7 @@ def postprocess_agent_loop_output(cls, output, *, tokenizer, response_length): if len(output.response_logprobs) < len(policy_ids): raise RuntimeError("Qwen3-TTS rollout logprobs are shorter than the policy trajectory.") output.response_logprobs = output.response_logprobs[: len(policy_ids)] - text_ids = tokenizer(build_assistant_text(str(text)), return_tensors="pt", padding=False)["input_ids"] + text_ids = tokenizer(build_assistant_text(text), return_tensors="pt", padding=False)["input_ids"] text_ids = torch.as_tensor(text_ids, dtype=torch.long) if text_ids.ndim == 1: text_ids = text_ids.unsqueeze(0) @@ -211,42 +148,41 @@ def prepare_engine_prompt(cls, prompt_ids, model_config, multi_modal_data, mm_pr @classmethod def combine_engine_outputs(cls, outputs, prompt): """Combine stage-0 policy tokens with codec and waveform outputs.""" - policy_output = None - policy_length = -1 - audio_codes = waveform = None - sample_rate = None - diagnostics = [] - for output in outputs: - completion = _completion(output) - if getattr(output, "stage_id", None) == 0 and completion is not None: - length = len(getattr(completion, "token_ids", None) or []) - if length >= policy_length: - policy_output, policy_length = output, length - multimodal = getattr(output, "multimodal_output", None) - diagnostics.append((getattr(output, "stage_id", None), type(multimodal).__name__)) - if not isinstance(multimodal, Mapping): - continue - codes = multimodal.get("codes") - if isinstance(codes, Mapping): - audio_codes = append_tensor_chunk(audio_codes, codes.get("audio")) - if getattr(output, "stage_id", None) == 1: - waveform = append_tensor_chunk( - waveform, multimodal.get("audio", multimodal.get("model_outputs")), flatten=True - ) - sample_rate = multimodal.get("sr", multimodal.get("audio_sample_rate", sample_rate)) - if policy_output is None: + policy_outputs = [output for output in outputs if output.stage_id == 0] + decoder_outputs = [output for output in outputs if output.stage_id == 1] + if not policy_outputs: raise RuntimeError("Qwen3-TTS rollout produced no stage-0 policy output.") - token_ids = list(getattr(_completion(policy_output), "token_ids", None) or []) - if audio_codes is None: - raise RuntimeError(f"Qwen3-TTS rollout produced no codec trajectory: {diagnostics}") + if not decoder_outputs: + raise RuntimeError("Qwen3-TTS rollout produced no stage-1 decoder output.") + + policy_output = policy_outputs[-1] + decoder_output = decoder_outputs[-1] + if len(policy_output.outputs) != 1: + raise RuntimeError( + f"Qwen3-TTS stage 0 must return exactly one completion, got {len(policy_output.outputs)}." + ) + token_ids = list(policy_output.outputs[0].token_ids) + if not token_ids: + raise RuntimeError("Qwen3-TTS rollout returned an empty stage-0 policy trajectory.") + + try: + audio_codes = torch.as_tensor(policy_output.multimodal_output["codes"]["audio"]).detach().cpu() + waveform = torch.as_tensor(decoder_output.multimodal_output["audio"]).detach().cpu().float().reshape(-1) + sample_rate = torch.as_tensor(decoder_output.multimodal_output["sr"]) + except (KeyError, TypeError, ValueError, RuntimeError) as error: + raise RuntimeError("Qwen3-TTS rollout output does not match the pinned two-stage contract.") from error + if waveform.numel() == 0: + raise RuntimeError("Qwen3-TTS stage 1 returned an empty waveform.") + if sample_rate.numel() != 1: + raise RuntimeError("Qwen3-TTS stage 1 must return one scalar sample rate.") + sample_rate_value = float(sample_rate.item()) + if sample_rate_value <= 0 or not sample_rate_value.is_integer(): + raise RuntimeError(f"Qwen3-TTS stage 1 returned an invalid sample rate: {sample_rate_value!r}.") + fields = { "tts_audio_codes": align_audio_codes(audio_codes, token_ids), "tts_text": prompt["additional_information"]["text"][0], + "audio": waveform, + "audio_sample_rate": int(sample_rate_value), } - if waveform is not None: - fields["audio"] = waveform.float().reshape(-1) - if sample_rate is not None: - if isinstance(sample_rate, list | tuple): - sample_rate = sample_rate[-1] - fields["audio_sample_rate"] = int(sample_rate.item() if hasattr(sample_rate, "item") else sample_rate) return policy_output, fields diff --git a/verl_omni/pipelines/qwen3_tts/rollout_utils.py b/verl_omni/pipelines/qwen3_tts/rollout_utils.py index 78add32dc..bedf3651f 100644 --- a/verl_omni/pipelines/qwen3_tts/rollout_utils.py +++ b/verl_omni/pipelines/qwen3_tts/rollout_utils.py @@ -11,99 +11,10 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Helpers for collecting Qwen3-TTS rollout outputs.""" - -import hashlib -from collections.abc import Mapping -from typing import Any +"""Helpers for validating Qwen3-TTS rollout outputs.""" import torch -_EVALUATION_SPLITS = {"gate", "official_test", "test", "val", "validation"} -_MAX_SAMPLING_SEED = 2**31 - 1 - - -def _scalar(value): - if hasattr(value, "item"): - try: - value = value.item() - except ValueError: - return None - return value - - -def _extra_info_mapping(extra_info) -> dict: - extra_info = _scalar(extra_info) - return dict(extra_info.items()) if isinstance(extra_info, Mapping) else {} - - -def is_evaluation_split(extra_info) -> bool: - info = _extra_info_mapping(extra_info) - return str(_scalar(info.get("split")) or "").lower() in _EVALUATION_SPLITS - - -def validation_generation_seed(extra_info) -> int | None: - info = _extra_info_mapping(extra_info) - if not is_evaluation_split(info): - return None - seed = _scalar(info.get("generation_seed")) - return None if seed is None else int(seed) - - -def rollout_generation_seed( - extra_info, - *, - session_id=None, - global_steps=None, - uid=None, - base_seed=0, - require_session_id=False, -) -> int: - """Derive a stable, group-diverse seed for both Qwen3-TTS samplers.""" - if require_session_id and session_id is None: - raise RuntimeError("Qwen3-TTS group sampling requires a per-candidate session_id when rollout.n > 1.") - candidate = int(_scalar(session_id) or 0) - explicit_seed = validation_generation_seed(extra_info) - if explicit_seed is not None: - return (explicit_seed + candidate) % _MAX_SAMPLING_SEED - - info = _extra_info_mapping(extra_info) - sample_id = _scalar(info.get("id", info.get("index", uid))) - step = "evaluation" if is_evaluation_split(info) else _scalar(global_steps) - payload = "\0".join(map(str, (int(base_seed), step, sample_id, candidate))).encode() - return int.from_bytes(hashlib.blake2b(payload, digest_size=8).digest(), "big") % _MAX_SAMPLING_SEED - - -def with_rollout_generation_seed(sampling_params, extra_info, **seed_kwargs): - """Seed codec-0 and residual-codebook sampling without mutating input.""" - seed = rollout_generation_seed(extra_info, **seed_kwargs) - seeded = dict(sampling_params) - residual_args = dict(seeded.get("extra_args") or {}) - residual_args["tts_local_seed"] = seed - seeded.update(seed=seed, extra_args=residual_args) - return seeded - - -def append_tensor_chunk(accumulated: torch.Tensor | None, value: Any, *, flatten=False): - if isinstance(value, list | tuple): - value = value[-1] if value else None - if value is None: - return accumulated - chunk = torch.as_tensor(value).detach().cpu() - if flatten: - chunk = chunk.reshape(-1) - if not chunk.numel(): - return accumulated - if accumulated is None: - return chunk - if chunk.shape[1:] != accumulated.shape[1:]: - raise RuntimeError(f"Rollout chunks changed shape from {tuple(accumulated.shape)} to {tuple(chunk.shape)}.") - if chunk.shape[0] >= accumulated.shape[0] and torch.equal(chunk[: accumulated.shape[0]], accumulated): - return chunk - if accumulated.shape[0] >= chunk.shape[0] and torch.equal(accumulated[: chunk.shape[0]], chunk): - return accumulated - return torch.cat((accumulated, chunk), dim=0) - def align_audio_codes(audio_codes: torch.Tensor, token_ids: list[int]) -> torch.Tensor: """Recover residual codebooks using codec-0 policy tokens as an exact invariant. diff --git a/verl_omni/pipelines/qwen3_tts/talker_forward.py b/verl_omni/pipelines/qwen3_tts/talker_forward.py index 924dec9a0..a1acc763a 100644 --- a/verl_omni/pipelines/qwen3_tts/talker_forward.py +++ b/verl_omni/pipelines/qwen3_tts/talker_forward.py @@ -78,9 +78,17 @@ def build_talker_batch( audio_codes, tokens, *, + sub_codebook_vocab: int, device=None, - sub_codebook_vocab=None, ) -> TalkerBatch: + if not text_ids or len(text_ids) != len(audio_codes): + raise ValueError("Qwen3-TTS teacher forcing requires matching non-empty text and codec batches.") + if sub_codebook_vocab <= 0: + raise ValueError(f"sub_codebook_vocab must be positive, got {sub_codebook_vocab}.") + if any(ids.reshape(-1).numel() < 3 for ids in text_ids): + raise ValueError("Qwen3-TTS teacher-forcing text sequences must contain at least three prefix tokens.") + if any(codes.ndim != 2 or codes.shape[-1] != NUM_CODEBOOKS for codes in audio_codes): + raise ValueError(f"Qwen3-TTS codec codes must have shape (frames, {NUM_CODEBOOKS}).") text_lens = [int(ids.reshape(-1).shape[0]) for ids in text_ids] codec_lens = [int(codes.shape[0]) for codes in audio_codes] speaker_slot = 6 @@ -97,9 +105,8 @@ def build_talker_batch( for index, (sample_text, sample_codes) in enumerate(zip(text_ids, audio_codes, strict=True)): ids = sample_text.reshape(-1).to(device=device, dtype=torch.long) codes = sample_codes.to(device=device, dtype=torch.long) - if sub_codebook_vocab is not None: - codes = codes.clone() - codes[:, 1:].clamp_(0, sub_codebook_vocab - 1) + if torch.any((codes[:, 1:] < 0) | (codes[:, 1:] >= sub_codebook_vocab)): + raise ValueError("Qwen3-TTS residual codec IDs are outside the code-predictor vocabulary.") text_len, codec_len = text_lens[index], codec_lens[index] input_ids[index, :3, 0] = ids[:3] @@ -141,7 +148,7 @@ def build_talker_batch( def require_auto_language(language) -> str: """Limit RL training to the prompt layout validated by the actor forward.""" - normalized = str(language or "Auto").strip() + normalized = str(language).strip() if normalized.lower() != "auto": raise ValueError( "Qwen3-TTS RL currently supports only tts_language=Auto; " @@ -152,9 +159,7 @@ def require_auto_language(language) -> str: def codec0_input_embeddings(talker, batch: TalkerBatch, speaker_embedding: torch.Tensor) -> torch.Tensor: ids = batch.input_ids - text_embeddings = talker.model.text_embedding(ids[:, :, 0]) - if getattr(talker, "text_projection", None) is not None: - text_embeddings = talker.text_projection(text_embeddings) + text_embeddings = talker.text_projection(talker.model.text_embedding(ids[:, :, 0])) codec_embeddings = talker.model.codec_embedding(ids[:, :, 1]) * batch.codec_embedding_mask codec_embeddings = codec_embeddings.clone() sample_indices = torch.arange(codec_embeddings.shape[0], device=codec_embeddings.device) @@ -180,10 +185,13 @@ def codec0_logits(talker, batch: TalkerBatch, speaker_embedding: torch.Tensor) - def mask_codec0_logits(logits: torch.Tensor, codebook_vocab: int, codec_eos_token_id: int) -> torch.Tensor: """Match the codec-token vocabulary exposed by the rollout model.""" + if not 1 < codebook_vocab <= logits.shape[-1]: + raise ValueError(f"codebook_vocab must be in [2, {logits.shape[-1]}], got {codebook_vocab}.") + if not 0 <= codec_eos_token_id < logits.shape[-1]: + raise ValueError(f"codec_eos_token_id must be in [0, {logits.shape[-1]}), got {codec_eos_token_id}.") valid = torch.zeros(logits.shape[-1], dtype=torch.bool, device=logits.device) - valid[1 : min(codebook_vocab, logits.shape[-1])] = True - if 0 <= codec_eos_token_id < logits.shape[-1]: - valid[codec_eos_token_id] = True + valid[1:codebook_vocab] = True + valid[codec_eos_token_id] = True return logits.masked_fill(~valid, -1e4) @@ -230,13 +238,13 @@ def tts_actor_logits( sub_vocab, int(model.config.talker_config.codec_eos_token_id), ) - output_vocab = max(logits.shape[-1], int(input_ids.max()) + 1) + if int(input_ids.min()) < 0 or int(input_ids.max()) >= logits.shape[-1]: + raise ValueError("Qwen3-TTS actor token IDs are outside the codec-0 logit vocabulary.") + output_vocab = logits.shape[-1] aligned = logits.new_zeros((batch_size, output_len, output_vocab)) for index, response_start in enumerate(response_starts): codec_len = batch.codec_lens[index] target = slice(response_start - 1, response_start - 1 + codec_len) - if output_vocab > logits.shape[-1]: - aligned[index, target, logits.shape[-1] :] = -1e4 source = batch.logit_start[index] - aligned[index, target, : logits.shape[-1]] = logits[index, source : source + codec_len] + aligned[index, target] = logits[index, source : source + codec_len] return aligned diff --git a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py index 1815b09da..8c4ed240c 100644 --- a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py @@ -15,8 +15,10 @@ import logging import types +from collections.abc import Mapping from typing import Any +import numpy as np import torch from verl_omni.pipelines.model_base import OmniModelBase @@ -25,43 +27,10 @@ require_auto_language, tts_actor_logits, ) -from verl_omni.pipelines.qwen3_tts.transformers_compat import ( - patch_qwen3_tts_config_defaults, - qwen3_tts_import_context, -) - -logger = logging.getLogger(__name__) -_PASSTHROUGH_TEMPLATE = "{% for message in messages %}{{ message['content'] }}{% endfor %}" -_TRAINABLE_PREFIXES = ("talker.model.", "talker.codec_head.") - - -def _prepare_config_for_checkpoint(config) -> None: - speaker_config = getattr(config, "speaker_encoder_config", None) - if speaker_config is not None: - speaker_config.__dict__.pop("dtype", None) - speaker_config.__dict__.pop("_dtype", None) def _speaker_embedding(model, batch_size, device, dtype): - cached = getattr(model, "_verl_tts_speaker_embedding", None) - if cached is None: - path = getattr(model.config, "tts_spk_embed_path", None) - if not path: - return None - cached = load_speaker_xvector(path) - model._verl_tts_speaker_embedding = cached - return cached.to(device=device, dtype=dtype).expand(batch_size, -1) - - -def _reinitialize_rope_buffers(model): - for submodule in model.modules(): - rope_init = getattr(submodule, "rope_init_fn", None) - inv_freq = getattr(submodule, "inv_freq", None) - if rope_init is None or not torch.is_tensor(inv_freq): - continue - new_inv_freq, scaling = rope_init(submodule.config, device=inv_freq.device) - submodule.inv_freq.data.copy_(new_inv_freq.to(device=inv_freq.device, dtype=inv_freq.dtype)) - submodule.attention_scaling = scaling + return model._verl_tts_speaker_embedding.to(device=device, dtype=dtype).expand(batch_size, -1) def _qwen3_tts_forward( @@ -76,13 +45,9 @@ def _qwen3_tts_forward( ): from transformers.modeling_outputs import CausalLMOutputWithPast - if any(value is None for value in (tts_text_ids, tts_audio_codes, response_len, text_len)): + required_inputs = (input_ids, attention_mask, tts_text_ids, tts_audio_codes, response_len, text_len) + if any(value is None for value in required_inputs): raise RuntimeError("Qwen3-TTS forward is missing exact rollout codec fields.") - if attention_mask is None: - attention_mask = torch.ones_like(input_ids) - if not getattr(self, "_verl_tts_rope_initialized", False): - _reinitialize_rope_buffers(self) - self._verl_tts_rope_initialized = True speaker = _speaker_embedding(self, input_ids.shape[0], input_ids.device, next(self.talker.parameters()).dtype) return CausalLMOutputWithPast( logits=tts_actor_logits( @@ -109,24 +74,18 @@ def _set_input_embeddings(self, value): @OmniModelBase.register("Qwen3TTSForConditionalGeneration", stage="talker") class Qwen3TTSTalkerAdapter(OmniModelBase): @classmethod - def load_hf_config(cls, model_path, *, trust_remote_code, attn_implementation): - with qwen3_tts_import_context(): - from qwen_tts.core.models.configuration_qwen3_tts import Qwen3TTSConfig - - patch_qwen3_tts_config_defaults(Qwen3TTSConfig) - return Qwen3TTSConfig.from_pretrained( - model_path, - trust_remote_code=trust_remote_code, - attn_implementation=attn_implementation, + def register_auto_classes(cls) -> None: + from qwen_tts.core.models.configuration_qwen3_tts import Qwen3TTSConfig + from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSForConditionalGeneration + from transformers import AutoConfig, AutoModelForTextToWaveform + + AutoConfig.register("qwen3_tts", Qwen3TTSConfig, exist_ok=True) + AutoModelForTextToWaveform.register( + Qwen3TTSConfig, + Qwen3TTSForConditionalGeneration, + exist_ok=True, ) - @classmethod - def get_model_class(cls): - with qwen3_tts_import_context(): - from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSForConditionalGeneration - - return Qwen3TTSForConditionalGeneration - @classmethod def get_strip_modules(cls, model_config): return ["speaker_encoder", "speech_tokenizer", "code2wav"] @@ -134,20 +93,20 @@ def get_strip_modules(cls, model_config): @classmethod def configure_model(cls, module, model_config): module = super().configure_model(module, model_config) - _prepare_config_for_checkpoint(module.config) module.config.tts_spk_embed_path = model_config.override_config.get("tts_spk_embed_path") module.config.tts_language = require_auto_language(model_config.override_config.get("tts_language", "Auto")) if not module.config.tts_spk_embed_path: raise ValueError("Qwen3-TTS GRPO requires tts_spk_embed_path for the validated non-streaming replay.") + module._verl_tts_speaker_embedding = load_speaker_xvector(module.config.tts_spk_embed_path) module.forward = types.MethodType(_qwen3_tts_forward, module) module.get_input_embeddings = types.MethodType(_get_input_embeddings, module) module.set_input_embeddings = types.MethodType(_set_input_embeddings, module) module._no_split_modules = ["Qwen3TTSTalkerDecoderLayer", "Qwen3TTSDecoderLayer"] trainable = 0 for name, parameter in module.named_parameters(): - parameter.requires_grad_(name.startswith(_TRAINABLE_PREFIXES)) + parameter.requires_grad_(name.startswith(("talker.model.", "talker.codec_head."))) trainable += int(parameter.requires_grad) - logger.info("Qwen3-TTS talker adapter enabled %d trainable parameter tensors", trainable) + logging.getLogger(__name__).info("Qwen3-TTS talker adapter enabled %d trainable parameter tensors", trainable) return module @classmethod @@ -159,34 +118,39 @@ def configure_tokenizer(cls, model_path: str, model_config): from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=model_config.trust_remote_code) - tokenizer.chat_template = _PASSTHROUGH_TEMPLATE + tokenizer.chat_template = "{% for message in messages %}{{ message['content'] }}{% endfor %}" if tokenizer.pad_token_id is None: - tokenizer.pad_token_id = tokenizer.eos_token_id or 0 - talker_config = getattr(model_config.hf_config, "talker_config", None) - if talker_config is not None: - talker_config.tie_word_embeddings = False + if tokenizer.eos_token_id is None: + raise ValueError("Qwen3-TTS tokenizer must define either pad_token_id or eos_token_id.") + tokenizer.pad_token_id = tokenizer.eos_token_id + model_config.hf_config.talker_config.tie_word_embeddings = False return tokenizer @classmethod def prepare_model_inputs(cls, model_inputs, micro_batch, model_config): del model_config - fields = micro_batch.get("extra_fields") - if fields is None: + if "extra_fields" not in micro_batch: raise RuntimeError( "Qwen3-TTS actor inputs require AgentLoopOutput.extra_fields; use the V1 agent-loop trainer path." ) - if hasattr(fields, "tolist"): + fields = micro_batch["extra_fields"] + if isinstance(fields, np.ndarray): fields = fields.tolist() - if isinstance(fields, dict): + if isinstance(fields, Mapping): fields = [fields] - fields = [getattr(item, "data", item) for item in fields] + if not isinstance(fields, list) or any(not isinstance(item, Mapping) for item in fields): + raise TypeError("Qwen3-TTS actor extra_fields must be a list of mappings.") if len(fields) != model_inputs["input_ids"].shape[0]: raise RuntimeError("Qwen3-TTS actor extra_fields do not match its batch size.") texts = [torch.as_tensor(item["tts_text_ids"], dtype=torch.long).reshape(-1) for item in fields] codes = [torch.as_tensor(item["tts_audio_codes"], dtype=torch.long) for item in fields] + if any(item.numel() < 3 for item in texts): + raise ValueError("Qwen3-TTS actor text fields must contain at least three prefix tokens.") if any(item.ndim != 2 or item.shape[-1] != 16 for item in codes): raise ValueError("Qwen3-TTS codec codes must have shape (frames, 16).") + if any(item.shape[0] == 0 for item in codes): + raise ValueError("Qwen3-TTS actor codec fields must contain at least one frame.") device = model_inputs["input_ids"].device text_buffer = torch.zeros((len(fields), max(item.numel() for item in texts)), dtype=torch.long, device=device) code_buffer = torch.zeros( diff --git a/verl_omni/pipelines/qwen3_tts/transformers_compat.py b/verl_omni/pipelines/qwen3_tts/transformers_compat.py deleted file mode 100644 index 1b9300bf6..000000000 --- a/verl_omni/pipelines/qwen3_tts/transformers_compat.py +++ /dev/null @@ -1,76 +0,0 @@ -# Copyright 2026 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Import-time compatibility for qwen-tts 0.1.1 on Transformers 5.x.""" - -from contextlib import contextmanager -from functools import wraps - -import torch - - -def _default_rope_init(config, device=None, **kwargs): - del kwargs - base = getattr(config, "rope_theta", 10_000.0) or 10_000.0 - dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads - positions = torch.arange(0, dim, 2, dtype=torch.int64, device=device).float() / dim - return 1.0 / (base**positions), 1.0 - - -def _compatible_mask(original): - @wraps(original) - def wrapper(*args, input_embeds=None, inputs_embeds=None, cache_position=None, **kwargs): - del cache_position - embeddings = inputs_embeds if inputs_embeds is not None else input_embeds - return original(*args, inputs_embeds=embeddings, **kwargs) - - return wrapper - - -def patch_qwen3_tts_config_defaults(config_cls) -> None: - """Restore the config default that qwen-tts expects from Transformers 4.x.""" - talker_config_cls = getattr(config_cls, "sub_configs", {}).get("talker_config") - if talker_config_cls is not None and not hasattr(talker_config_cls, "pad_token_id"): - talker_config_cls.pad_token_id = None - - -@contextmanager -def qwen3_tts_import_context(): - """Expose the TF5 APIs expected while qwen-tts binds its imports. - - qwen-tts modules retain the mask wrappers they import. The global - Transformers functions are restored immediately afterwards. - """ - import transformers.masking_utils as masking_utils - import transformers.modeling_rope_utils as rope_utils - import transformers.utils.generic as generic_utils - - original_check = generic_utils.check_model_inputs - original_causal_mask = masking_utils.create_causal_mask - original_sliding_mask = masking_utils.create_sliding_window_causal_mask - original_rope_functions = rope_utils.ROPE_INIT_FUNCTIONS - - def compatible_check_model_inputs(func=None): - return original_check if func is None else original_check(func) - - generic_utils.check_model_inputs = compatible_check_model_inputs - masking_utils.create_causal_mask = _compatible_mask(original_causal_mask) - masking_utils.create_sliding_window_causal_mask = _compatible_mask(original_sliding_mask) - rope_utils.ROPE_INIT_FUNCTIONS = {**original_rope_functions, "default": _default_rope_init} - try: - yield - finally: - generic_utils.check_model_inputs = original_check - masking_utils.create_causal_mask = original_causal_mask - masking_utils.create_sliding_window_causal_mask = original_sliding_mask - rope_utils.ROPE_INIT_FUNCTIONS = original_rope_functions diff --git a/verl_omni/reward_loop/reward_manager/audio.py b/verl_omni/reward_loop/reward_manager/audio.py index ceb49d2e8..6b1f80a27 100644 --- a/verl_omni/reward_loop/reward_manager/audio.py +++ b/verl_omni/reward_loop/reward_manager/audio.py @@ -15,6 +15,7 @@ import inspect import math +from collections.abc import Mapping import numpy as np import torch @@ -37,13 +38,16 @@ def __init__(self, config, tokenizer, compute_score, reward_router_address=None, def _mapping(value): if isinstance(value, np.ndarray) and value.shape == (): value = value.item() - return dict(value.items()) if hasattr(value, "items") else {} + if value is None: + return {} + if not isinstance(value, Mapping): + raise TypeError(f"Audio reward metadata must be a mapping, got {type(value).__name__}.") + return dict(value) @classmethod - def _extract_audio(cls, data_item, extra_info): - batch = data_item.non_tensor_batch - audio = extra_info.get("audio", batch.get("audio")) - sample_rate = extra_info.get("audio_sample_rate", batch.get("audio_sample_rate")) + def _extract_audio(cls, extra_info): + audio = extra_info.get("audio") + sample_rate = extra_info.get("audio_sample_rate") if audio is None: raise KeyError("Audio reward requires extra_info['audio'] from the rollout.") if sample_rate is None: @@ -52,15 +56,7 @@ def _extract_audio(cls, data_item, extra_info): try: waveform = torch.as_tensor(audio).detach().float().cpu() except (TypeError, ValueError, RuntimeError) as exc: - if not isinstance(audio, list | tuple): - raise ValueError("Audio reward could not convert the waveform to numeric samples.") from exc - try: - chunks = [torch.as_tensor(chunk).detach().float().cpu().reshape(-1) for chunk in audio] - except (TypeError, ValueError, RuntimeError) as chunk_exc: - raise ValueError( - "Audio reward could not convert all waveform chunks to numeric samples." - ) from chunk_exc - waveform = torch.cat(chunks) if chunks else torch.empty(0) + raise ValueError("Audio reward could not convert the waveform to numeric samples.") from exc while waveform.ndim > 1 and waveform.shape[0] == 1: waveform = waveform[0] if waveform.ndim == 2: @@ -72,15 +68,11 @@ def _extract_audio(cls, data_item, extra_info): if not torch.isfinite(waveform).all(): raise ValueError("Audio reward received a waveform containing NaN or infinity.") - if isinstance(sample_rate, list | tuple): - if len(sample_rate) != 1: - raise ValueError("Audio reward requires exactly one sample rate per waveform.") - sample_rate = sample_rate[0] - if hasattr(sample_rate, "item"): - try: - sample_rate = sample_rate.item() - except (RuntimeError, ValueError) as exc: - raise ValueError("Audio reward requires one scalar sample rate per waveform.") from exc + if isinstance(sample_rate, np.ndarray | torch.Tensor): + sample_rate_count = sample_rate.size if isinstance(sample_rate, np.ndarray) else sample_rate.numel() + if sample_rate_count != 1: + raise ValueError("Audio reward requires one scalar sample rate per waveform.") + sample_rate = sample_rate.item() if isinstance(sample_rate, bool) or not isinstance(sample_rate, int | float): raise TypeError(f"Audio sample rate must be numeric, got {type(sample_rate).__name__}.") if not math.isfinite(float(sample_rate)) or float(sample_rate) <= 0 or float(sample_rate) != int(sample_rate): @@ -97,7 +89,7 @@ async def run_single(self, data: DataProto) -> dict: extra_info["num_turns"] = batch.get("__num_turns__", extra_info.get("num_turns")) extra_info["global_steps"] = batch.get("global_steps", extra_info.get("global_steps", 0)) ground_truth = batch["reward_model"]["ground_truth"] - audio = self._extract_audio(item, extra_info) + audio = self._extract_audio(extra_info) kwargs = { "data_source": batch["data_source"], "solution_audio": audio, diff --git a/verl_omni/utils/reward_score/audio_http_scorer_client.py b/verl_omni/utils/reward_score/audio_http_scorer_client.py index 29ff3d3d9..0ec5b064e 100644 --- a/verl_omni/utils/reward_score/audio_http_scorer_client.py +++ b/verl_omni/utils/reward_score/audio_http_scorer_client.py @@ -21,8 +21,6 @@ import aiohttp import numpy as np -PROTOCOL_VERSION = "1" - class _RetryableHTTPError(RuntimeError): pass @@ -60,12 +58,14 @@ def _serialize_request(solution_audio, ground_truth: str, extra_info: dict | Non raise TypeError("Audio HTTP scorer sample rate must be numeric.") if not math.isfinite(float(sample_rate)) or float(sample_rate) <= 0 or int(sample_rate) != sample_rate: raise ValueError(f"Audio HTTP scorer sample rate must be a positive integer, got {sample_rate!r}.") + if not isinstance(ground_truth, str): + raise TypeError(f"Audio HTTP scorer prompt must be a string, got {type(ground_truth).__name__}.") return { - "protocol_version": PROTOCOL_VERSION, + "protocol_version": "1", "waveform_f32_base64": base64.b64encode(waveform.tobytes()).decode("ascii"), "num_samples": int(waveform.size), "sample_rate": int(sample_rate), - "prompt": str(ground_truth or ""), + "prompt": ground_truth, "metadata": _scalar_metadata(extra_info), } diff --git a/verl_omni/workers/config/omni/model.py b/verl_omni/workers/config/omni/model.py index 270368165..85ebd4fa1 100644 --- a/verl_omni/workers/config/omni/model.py +++ b/verl_omni/workers/config/omni/model.py @@ -20,6 +20,7 @@ from typing import Any, Optional from omegaconf import MISSING +from transformers import AutoConfig from verl.base_config import BaseConfig from verl.utils.fs import copy_to_local from verl.utils.import_utils import import_external_libs @@ -169,33 +170,19 @@ def __post_init__(self): attn_implementation = self.override_config.get("attn_implementation", "flash_attention_2") from verl_omni.pipelines.model_base import OmniModelBase - try: - adapter_cls = OmniModelBase.get_class_by_name(self.architecture, self.model_stage, self.external_lib) - except NotImplementedError: - adapter_cls = None - - if adapter_cls is None: - from transformers import AutoConfig - - self.hf_config = AutoConfig.from_pretrained( - self.local_hf_config_path, - trust_remote_code=self.trust_remote_code, - attn_implementation=attn_implementation, - ) - else: - self.hf_config = adapter_cls.load_hf_config( - self.local_hf_config_path, - trust_remote_code=self.trust_remote_code, - attn_implementation=attn_implementation, - ) + adapter_cls = OmniModelBase.get_class_by_name(self.architecture, self.model_stage, self.external_lib) + adapter_cls.register_auto_classes() + self.hf_config = AutoConfig.from_pretrained( + self.local_hf_config_path, + trust_remote_code=self.trust_remote_code, + attn_implementation=attn_implementation, + ) self.share_embeddings_and_output_weights = getattr(self.hf_config, "tie_word_embeddings", False) self.architectures = getattr(self.hf_config, "architectures", None) if self.load_tokenizer: self.local_tokenizer_path = copy_to_local(self.tokenizer_path, use_shm=self.use_shm) - if adapter_cls is None: - adapter_cls = OmniModelBase.get_class_by_name(self.architecture, self.model_stage, self.external_lib) self.tokenizer = adapter_cls.configure_tokenizer(self.local_tokenizer_path, self) self.processor = adapter_cls.configure_processor(self.local_path, self) diff --git a/verl_omni/workers/engine/fsdp/omni_impl.py b/verl_omni/workers/engine/fsdp/omni_impl.py index 6a96f8af9..b11c0f1eb 100644 --- a/verl_omni/workers/engine/fsdp/omni_impl.py +++ b/verl_omni/workers/engine/fsdp/omni_impl.py @@ -18,7 +18,7 @@ import torch from torch.distributed.tensor import DTensor -from transformers import AutoModelForMultimodalLM +from transformers import AutoModelForMultimodalLM, AutoModelForTextToWaveform from verl.utils.debug import log_gpu_memory_usage from verl.utils.device import get_device_id from verl.utils.fsdp_utils import ( @@ -179,7 +179,7 @@ def _build_module(self): adapter_cls = OmniModelBase.get_class_by_name( architecture, self.model_config.model_stage, - self.model_config.get("external_lib"), + self.model_config.external_lib, ) self.model_adapter_cls = adapter_cls @@ -203,8 +203,10 @@ def _build_module(self): with init_context(), warnings.catch_warnings(): warnings.simplefilter("ignore") - model_cls = adapter_cls.get_model_class() or AutoModelForMultimodalLM - module = model_cls.from_pretrained( + auto_model_cls = ( + AutoModelForTextToWaveform if self.model_config.model_stage == "talker" else AutoModelForMultimodalLM + ) + module = auto_model_cls.from_pretrained( pretrained_model_name_or_path=self.model_config.local_path, torch_dtype=torch_dtype, config=self.model_config.hf_config, diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py index 9b738ffe7..f3c8ac03f 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py @@ -35,8 +35,6 @@ logger = logging.getLogger(__file__) logger.setLevel(logging.INFO) -_WORKER_EXTENSION = "verl_omni.workers.rollout.vllm_rollout.utils.vLLMOmniColocateWorkerExtension" - def _drop_none_mapping_values(value: Any) -> Any: if isinstance(value, dict): @@ -46,16 +44,6 @@ def _drop_none_mapping_values(value: Any) -> Any: return value -def _retained_output_modalities(stages: list[Any]) -> list[str] | None: - """Request every modality when an AR pipeline exposes multiple final outputs.""" - final_output_types = [ - stage.final_output_type - for stage in stages - if getattr(stage, "final_output", False) and getattr(stage, "final_output_type", None) - ] - return list(dict.fromkeys(final_output_types)) if len(final_output_types) > 1 else None - - class ARStrategy(OmniStrategyBase): """Concrete AR/thinker strategy. @@ -86,7 +74,7 @@ def override_generation_config(self) -> dict[str, Any]: return vLLMHttpServer._get_override_generation_config(self.server) def worker_extension_cls(self, device_type: str) -> str: - return _WORKER_EXTENSION + return "verl_omni.workers.rollout.vllm_rollout.utils.vLLMOmniColocateWorkerExtension" def preprocess_engine_kwargs(self, engine_kwargs: dict[str, Any]) -> None: super().preprocess_engine_kwargs(engine_kwargs) @@ -109,8 +97,8 @@ def preprocess_engine_kwargs(self, engine_kwargs: dict[str, Any]) -> None: hf_overrides.update(adapter_overrides) engine_kwargs["hf_overrides"] = hf_overrides - stage_init_timeout = engine_kwargs.get("stage_init_timeout") or engine_kwargs.get("stage-init-timeout") - init_timeout = engine_kwargs.get("init_timeout") or engine_kwargs.get("init-timeout") + stage_init_timeout = engine_kwargs.get("stage_init_timeout") + init_timeout = engine_kwargs.get("init_timeout") if stage_init_timeout is not None and init_timeout is None: engine_kwargs["init_timeout"] = max(int(stage_init_timeout), 600) @@ -132,10 +120,11 @@ def _write_deploy_config( adapter_cls.ensure_pipeline_registered(pipeline_mode) stages = adapter_cls.build_stage_configs(pipeline_mode=pipeline_mode) pipeline_id = adapter_cls.get_pipeline_id(pipeline_mode) - self._rollout_output_modalities = _retained_output_modalities(stages) - self._stage_sampling_constraints = { - stage.stage_id: dict(getattr(stage, "sampling_constraints", {}) or {}) for stage in stages - } + final_output_types = [stage.final_output_type for stage in stages if stage.final_output] + self._rollout_output_modalities = ( + list(dict.fromkeys(final_output_types)) if len(final_output_types) > 1 else None + ) + self._stage_sampling_constraints = {stage.stage_id: dict(stage.sampling_constraints) for stage in stages} stage_extras = { stage.stage_id: dict(adapter_cls.get_stage_engine_extras(stage.stage_id, pipeline_mode=pipeline_mode)) for stage in stages @@ -261,11 +250,11 @@ def preprocess_input( sampling_params["logprobs"] = None sampling_params.setdefault("repetition_penalty", getattr(self.server.config, "repetition_penalty", 1.0)) policy_params = SamplingParams(max_tokens=max_tokens, **sampling_params) - engine = getattr(self.server, "engine", None) - default_params = list(getattr(engine, "default_sampling_params_list", []) or []) - if len(default_params) > 1: - params = copy.deepcopy(default_params) - constrained = self._stage_sampling_constraints.get(0, {}) + if self._rollout_adapter is not None: + params = copy.deepcopy(self.server.engine.default_sampling_params_list) + if len(params) <= 1: + raise RuntimeError("An omni rollout adapter requires per-stage sampling parameters.") + constrained = self._stage_sampling_constraints[0] for field in {"max_tokens", *sampling_params} - constrained.keys(): setattr(params[0], field, getattr(policy_params, field)) else: @@ -307,10 +296,9 @@ async def run_generation( async for output in generator: outputs.append(output) if self._rollout_adapter is None: - return outputs[-1] if outputs else None + raise RuntimeError("Retaining multiple stage outputs requires a registered rollout adapter.") final_res, rollout_fields = self._rollout_adapter.combine_engine_outputs(outputs, prompt) - if final_res is not None and rollout_fields: - final_res._verl_omni_rollout_fields = rollout_fields + final_res._verl_omni_rollout_fields = rollout_fields return final_res def process_output( @@ -322,23 +310,23 @@ def process_output( if final_res is None: raise RuntimeError("AR mode: vLLM-Omni engine yielded no output for the prompt.") - req_output = getattr(final_res, "request_output", None) or final_res - if not req_output.outputs: + if not final_res.outputs: raise RuntimeError("AR mode expects outputs with token IDs, but got None or empty.") extra_fields = {"global_steps": self.server.global_steps} - extra_fields.update(getattr(final_res, "_verl_omni_rollout_fields", {})) - token_ids = req_output.outputs[0].token_ids + if self._rollout_adapter is not None: + extra_fields.update(final_res._verl_omni_rollout_fields) + token_ids = final_res.outputs[0].token_ids log_probs = None policy_params = params[0] if isinstance(params, list) else params if policy_params.logprobs is not None: log_probs = [ - logprobs[token_ids[index]].logprob for index, logprobs in enumerate(req_output.outputs[0].logprobs) + logprobs[token_ids[index]].logprob for index, logprobs in enumerate(final_res.outputs[0].logprobs) ] - finish_reason = req_output.outputs[0].finish_reason + finish_reason = final_res.outputs[0].finish_reason stop_reason = self._map_stop_reason(finish_reason) - num_preempted = self._extract_num_preempted(req_output) + num_preempted = self._extract_num_preempted(final_res) return TokenOutput( token_ids=token_ids, From 94dd2d197ca65276989092c9395b6dd7009b79e9 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Tue, 1 Sep 2026 13:31:09 +0800 Subject: [PATCH 18/28] [model, rollout, tests] refactor: use generic Talker replay interface Adopt the shared Talker hooks from PR #504 and namespace the Qwen3-TTS replay payload. Preserve the V1 TransferQueue extra_fields shell explicitly and cover it with the real AgentLoopOutput conversion path. Signed-off-by: dongbo910220 <1275604947@qq.com> --- .../test_qwen3_tts_rollout_on_cpu.py | 91 +++++++------------ verl_omni/pipelines/model_base.py | 5 - .../qwen3_tts/omni_rollout_adapter.py | 12 ++- .../pipelines/qwen3_tts/rollout_utils.py | 2 + .../qwen3_tts/talker_training_adapter.py | 38 ++++---- .../workers/rollout/vllm_rollout/utils.py | 2 + 6 files changed, 70 insertions(+), 80 deletions(-) diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index e0bf6d962..a2c8c2126 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -19,17 +19,20 @@ import pytest import torch +from tensordict import TensorDict pytest.importorskip("verl") pytest.importorskip("vllm_omni") -from verl.experimental.agent_loop.single_turn_agent_loop import SingleTurnAgentLoop +from verl.experimental.agent_loop.agent_loop import AgentLoopMetrics, AgentLoopOutput +from verl.utils.tensordict_utils import list_of_dict_to_tensordict from vllm import SamplingParams from verl_omni.agent_loop.single_turn_agent_loop import OmniSingleTurnAgentLoop from verl_omni.pipelines.model_base import OmniRolloutPipelineBase from verl_omni.pipelines.qwen3_tts import omni_rollout_adapter from verl_omni.pipelines.qwen3_tts.omni_rollout_adapter import Qwen3TTSRolloutAdapter +from verl_omni.pipelines.qwen3_tts.rollout_utils import QWEN3_TTS_REPLAY_KEY from verl_omni.pipelines.qwen3_tts.talker_training_adapter import Qwen3TTSTalkerAdapter from verl_omni.workers.rollout.vllm_rollout.vllm_omni_ar_strategy import ARStrategy from verl_omni.workers.rollout.vllm_rollout.vllm_omni_async_server import vLLMOmniHttpServer @@ -87,50 +90,6 @@ def test_omni_single_turn_agent_resolves_registered_pipeline_adapter(): OmniSingleTurnAgentLoop._resolve_rollout_adapter(missing_config) -@pytest.mark.asyncio -async def test_omni_single_turn_agent_delegates_policy_mapping_to_adapter(monkeypatch): - class Adapter: - @staticmethod - def postprocess_agent_loop_output(output, **kwargs): - assert kwargs["response_length"] == 4 - output.prompt_ids = [0] - return output - - upstream_output = SimpleNamespace( - prompt_ids=[11, 12], - response_ids=[101, 102], - response_mask=[1, 1], - response_logprobs=[-0.1, -0.2], - extra_fields={"audio": torch.ones(8)}, - ) - - async def upstream_run(_self, sampling_params, **kwargs): - assert sampling_params == {"temperature": 0.8} - assert kwargs["priority"] == 7 - return upstream_output - - monkeypatch.setattr(SingleTurnAgentLoop, "run", upstream_run) - - loop = object.__new__(OmniSingleTurnAgentLoop) - loop.rollout_adapter = Adapter - loop.rollout_config = SimpleNamespace(response_length=4) - loop.response_length = 4 - loop.tokenizer = _Tokenizer() - - result = await OmniSingleTurnAgentLoop.run( - loop, - {"temperature": 0.8}, - priority=7, - raw_prompt=[{"role": "user", "content": "hello"}], - session_id=2, - ) - - assert result.prompt_ids == [0] - assert result.response_ids == [101, 102] - assert result.response_logprobs == [-0.1, -0.2] - assert result.extra_fields["audio"].shape == (8,) - - def test_rollout_pipeline_registers_upstream_talker(monkeypatch): registered_pipelines = [] monkeypatch.setattr( @@ -203,12 +162,22 @@ def test_rollout_adapter_requires_speaker_embedding(): def test_talker_adapter_pads_exact_rollout_fields_for_actor_forward(): model_inputs = {"input_ids": torch.zeros(2, 6, dtype=torch.long)} - micro_batch = { - "extra_fields": [ - {"tts_text_ids": [1, 2, 6], "tts_audio_codes": torch.ones(3, 16, dtype=torch.long)}, - {"tts_text_ids": [3, 4, 5], "tts_audio_codes": torch.full((2, 16), 2, dtype=torch.long)}, + payloads = [ + {"text_ids": [1, 2, 6], "audio_codes": torch.ones(3, 16, dtype=torch.long)}, + {"text_ids": [3, 4, 5], "audio_codes": torch.full((2, 16), 2, dtype=torch.long)}, + ] + micro_batch = list_of_dict_to_tensordict( + [ + AgentLoopOutput( + prompt_ids=[1], + response_ids=[2], + response_mask=[1], + metrics=AgentLoopMetrics(), + extra_fields={QWEN3_TTS_REPLAY_KEY: item}, + ).as_dict() + for item in payloads ] - } + ) prepared = Qwen3TTSTalkerAdapter.prepare_model_inputs(model_inputs, micro_batch, None) @@ -219,11 +188,12 @@ def test_talker_adapter_pads_exact_rollout_fields_for_actor_forward(): assert not prepared["tts_audio_codes"][1, 2].any() -def test_talker_adapter_requires_v1_agent_loop_extra_fields(): +def test_talker_adapter_requires_namespaced_replay_payload(): model_inputs = {"input_ids": torch.zeros(1, 4, dtype=torch.long)} + micro_batch = TensorDict({}, batch_size=[1]) - with pytest.raises(RuntimeError, match="V1 agent-loop trainer"): - Qwen3TTSTalkerAdapter.prepare_model_inputs(model_inputs, {}, None) + with pytest.raises(RuntimeError, match=QWEN3_TTS_REPLAY_KEY): + Qwen3TTSTalkerAdapter.prepare_model_inputs(model_inputs, micro_batch, None) def test_rollout_adapter_combines_policy_codes_and_waveform(): @@ -258,7 +228,12 @@ def test_rollout_adapter_prepares_actor_policy_sequence(): response_ids=[7, 6, 5, 4, 3], response_mask=[1] * 5, response_logprobs=[-0.1, -0.2, -0.3, -0.4, -0.5], - extra_fields={"tts_audio_codes": codes, "tts_text": "first text"}, + extra_fields={ + "tts_audio_codes": codes, + "tts_text": "first text", + "audio": torch.ones(2400), + "audio_sample_rate": 24_000, + }, ) result = Qwen3TTSRolloutAdapter.postprocess_agent_loop_output( @@ -272,8 +247,12 @@ def test_rollout_adapter_prepares_actor_policy_sequence(): assert result.response_ids == codes[:3, 0].tolist() assert result.response_mask == [1, 1, 1] assert result.response_logprobs == [-0.1, -0.2, -0.3] - torch.testing.assert_close(result.extra_fields["tts_audio_codes"], codes[:3]) - assert result.extra_fields["tts_text_ids"] + assert "tts_audio_codes" not in result.extra_fields + assert "tts_text" not in result.extra_fields + replay = result.extra_fields[QWEN3_TTS_REPLAY_KEY] + torch.testing.assert_close(replay["audio_codes"], codes[:3]) + assert replay["text_ids"] + assert result.extra_fields["audio_sample_rate"] == 24_000 def test_ar_strategy_prepares_stage_specific_sampling_params(): diff --git a/verl_omni/pipelines/model_base.py b/verl_omni/pipelines/model_base.py index 54aa439b1..0db96f99d 100644 --- a/verl_omni/pipelines/model_base.py +++ b/verl_omni/pipelines/model_base.py @@ -833,11 +833,6 @@ def prepare_engine_prompt( """Build an architecture-specific rollout prompt when required.""" return None - @classmethod - def get_output_modalities(cls, pipeline_mode: str = "thinker_only") -> list[str] | None: - """Return intermediate modalities that must be retained by the engine.""" - return None - @classmethod def combine_engine_outputs(cls, outputs: list, prompt: dict) -> tuple[Any, dict[str, Any]]: """Select the policy output and collect architecture-specific fields.""" diff --git a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py index 54b5b84c9..e0f54fe27 100644 --- a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py @@ -23,7 +23,7 @@ from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE from verl_omni.pipelines.model_base import OmniRolloutPipelineBase -from verl_omni.pipelines.qwen3_tts.rollout_utils import align_audio_codes +from verl_omni.pipelines.qwen3_tts.rollout_utils import QWEN3_TTS_REPLAY_KEY, align_audio_codes from verl_omni.pipelines.qwen3_tts.talker_forward import ( TEXT_PROMPT_TRAILER_TOKENS, build_assistant_text, @@ -90,6 +90,8 @@ def postprocess_agent_loop_output(cls, output, *, tokenizer, response_length): extra = output.extra_fields if "tts_audio_codes" not in extra or "tts_text" not in extra: raise RuntimeError("Qwen3-TTS rollout did not return codec codes and text.") + if QWEN3_TTS_REPLAY_KEY in extra: + raise RuntimeError(f"Qwen3-TTS rollout unexpectedly returned reserved field {QWEN3_TTS_REPLAY_KEY!r}.") codes, text = extra["tts_audio_codes"], extra["tts_text"] if not isinstance(text, str): raise TypeError(f"Qwen3-TTS rollout text must be a string, got {type(text).__name__}.") @@ -110,8 +112,12 @@ def postprocess_agent_loop_output(cls, output, *, tokenizer, response_length): text_ids = text_ids.unsqueeze(0) if text_ids.ndim != 2 or text_ids.shape[1] <= TEXT_PROMPT_TRAILER_TOKENS: raise ValueError("Qwen3-TTS assistant text tokenization returned an invalid sequence.") - extra["tts_text_ids"] = text_ids[:, :-TEXT_PROMPT_TRAILER_TOKENS].reshape(-1).tolist() - extra["tts_audio_codes"] = codes + extra[QWEN3_TTS_REPLAY_KEY] = { + "text_ids": text_ids[:, :-TEXT_PROMPT_TRAILER_TOKENS].reshape(-1).tolist(), + "audio_codes": codes, + } + del extra["tts_audio_codes"] + del extra["tts_text"] output.prompt_ids = [0] output.response_ids = policy_ids output.response_mask = [1] * len(policy_ids) diff --git a/verl_omni/pipelines/qwen3_tts/rollout_utils.py b/verl_omni/pipelines/qwen3_tts/rollout_utils.py index bedf3651f..2f9a85283 100644 --- a/verl_omni/pipelines/qwen3_tts/rollout_utils.py +++ b/verl_omni/pipelines/qwen3_tts/rollout_utils.py @@ -15,6 +15,8 @@ import torch +QWEN3_TTS_REPLAY_KEY = "qwen3_tts_talker_replay" + def align_audio_codes(audio_codes: torch.Tensor, token_ids: list[int]) -> torch.Tensor: """Recover residual codebooks using codec-0 policy tokens as an exact invariant. diff --git a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py index 8c4ed240c..8ea5b90c4 100644 --- a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py @@ -18,10 +18,11 @@ from collections.abc import Mapping from typing import Any -import numpy as np import torch +from verl.utils import tensordict_utils as tu from verl_omni.pipelines.model_base import OmniModelBase +from verl_omni.pipelines.qwen3_tts.rollout_utils import QWEN3_TTS_REPLAY_KEY from verl_omni.pipelines.qwen3_tts.talker_forward import ( load_speaker_xvector, require_auto_language, @@ -129,22 +130,27 @@ def configure_tokenizer(cls, model_path: str, model_config): @classmethod def prepare_model_inputs(cls, model_inputs, micro_batch, model_config): del model_config - if "extra_fields" not in micro_batch: - raise RuntimeError( - "Qwen3-TTS actor inputs require AgentLoopOutput.extra_fields; use the V1 agent-loop trainer path." - ) - fields = micro_batch["extra_fields"] - if isinstance(fields, np.ndarray): - fields = fields.tolist() - if isinstance(fields, Mapping): - fields = [fields] - if not isinstance(fields, list) or any(not isinstance(item, Mapping) for item in fields): + # The online V1 trainer retains AgentLoopOutput.extra_fields as one mapping per sample. + sample_extra_fields = tu.get(micro_batch, "extra_fields") + if sample_extra_fields is None: + raise RuntimeError(f"Qwen3-TTS actor inputs require the {QWEN3_TTS_REPLAY_KEY!r} replay payload.") + if not isinstance(sample_extra_fields, list) or any( + not isinstance(item, Mapping) for item in sample_extra_fields + ): raise TypeError("Qwen3-TTS actor extra_fields must be a list of mappings.") - if len(fields) != model_inputs["input_ids"].shape[0]: - raise RuntimeError("Qwen3-TTS actor extra_fields do not match its batch size.") - - texts = [torch.as_tensor(item["tts_text_ids"], dtype=torch.long).reshape(-1) for item in fields] - codes = [torch.as_tensor(item["tts_audio_codes"], dtype=torch.long) for item in fields] + if len(sample_extra_fields) != model_inputs["input_ids"].shape[0]: + raise RuntimeError("Qwen3-TTS actor extra_fields do not match the actor batch size.") + if any(QWEN3_TTS_REPLAY_KEY not in item for item in sample_extra_fields): + raise RuntimeError(f"Qwen3-TTS actor inputs require the {QWEN3_TTS_REPLAY_KEY!r} replay payload.") + + fields = [item[QWEN3_TTS_REPLAY_KEY] for item in sample_extra_fields] + if any(not isinstance(item, Mapping) for item in fields): + raise TypeError(f"Qwen3-TTS {QWEN3_TTS_REPLAY_KEY!r} must be a list of mappings.") + if any("text_ids" not in item or "audio_codes" not in item for item in fields): + raise RuntimeError("Qwen3-TTS replay payloads must contain text_ids and audio_codes.") + + texts = [torch.as_tensor(item["text_ids"], dtype=torch.long).reshape(-1) for item in fields] + codes = [torch.as_tensor(item["audio_codes"], dtype=torch.long) for item in fields] if any(item.numel() < 3 for item in texts): raise ValueError("Qwen3-TTS actor text fields must contain at least three prefix tokens.") if any(item.ndim != 2 or item.shape[-1] != 16 for item in codes): diff --git a/verl_omni/workers/rollout/vllm_rollout/utils.py b/verl_omni/workers/rollout/vllm_rollout/utils.py index 13d2c2fa7..b03205af0 100644 --- a/verl_omni/workers/rollout/vllm_rollout/utils.py +++ b/verl_omni/workers/rollout/vllm_rollout/utils.py @@ -30,6 +30,8 @@ def _split_visible_devices(value: str) -> list[str]: """Split a visible-devices env value into stripped, non-empty entries.""" return [entry.strip() for entry in value.split(",") if entry.strip()] + + class vLLMOmniColocateWorkerExtension(CustomPipelineWorkerExtension): """ The class for vLLM-Omni's worker to inherit from, in the colocate setting. From 37af0369f1fb2ca528d83aa5f94b3ebc3ec10792 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Tue, 1 Sep 2026 14:29:26 +0800 Subject: [PATCH 19/28] [model, rollout, tests] refactor: enforce fail-closed Talker contracts Remove one-use module assignments, require explicit Qwen3-TTS language and adapter prompt fields, and reject ambiguous codec alignment. Add focused CPU coverage for each fail-closed contract. Signed-off-by: dongbo910220 <1275604947@qq.com> --- tests/pipelines/test_qwen3_tts_on_cpu.py | 12 ++++-- .../create_dummy_qwen3_tts_grpo_data.py | 37 +++++++++-------- .../test_qwen3_tts_rollout_on_cpu.py | 40 +++++++++++++++++++ tests/workers/test_omni_fsdp_engine_on_cpu.py | 13 +----- verl_omni/pipelines/model_base.py | 4 +- .../qwen3_tts/omni_rollout_adapter.py | 2 +- .../pipelines/qwen3_tts/rollout_utils.py | 8 +++- .../qwen3_tts/talker_training_adapter.py | 2 +- verl_omni/workers/config/omni/model.py | 3 -- .../vllm_rollout/vllm_omni_ar_strategy.py | 13 +++++- 10 files changed, 91 insertions(+), 43 deletions(-) diff --git a/tests/pipelines/test_qwen3_tts_on_cpu.py b/tests/pipelines/test_qwen3_tts_on_cpu.py index 18a759e43..fc00b3488 100644 --- a/tests/pipelines/test_qwen3_tts_on_cpu.py +++ b/tests/pipelines/test_qwen3_tts_on_cpu.py @@ -21,11 +21,9 @@ import pytest import torch -ROOT = Path(__file__).parents[2] - def _load(name: str, relative_path: str): - spec = importlib.util.spec_from_file_location(name, ROOT / relative_path) + spec = importlib.util.spec_from_file_location(name, Path(__file__).parents[2] / relative_path) module = importlib.util.module_from_spec(spec) assert spec.loader is not None sys.modules[name] = module @@ -154,6 +152,14 @@ def test_codec_alignment_preserves_final_row_and_rejects_heuristic_match(): rollout.align_audio_codes(malformed, token_ids) +def test_codec_alignment_rejects_ambiguous_exact_matches(): + raw = torch.zeros(2, 16, dtype=torch.long) + raw[:, 0] = 101 + + with pytest.raises(RuntimeError, match="Ambiguous Qwen3-TTS codec alignment"): + rollout.align_audio_codes(raw, [101, 2150]) + + def test_talker_batch_and_logit_mask_reject_contract_mismatches(): with pytest.raises(ValueError, match="matching non-empty"): forward.build_talker_batch([], [], TOKENS, sub_codebook_vocab=2048) diff --git a/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py b/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py index ac67dcd05..8ae71226f 100644 --- a/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py +++ b/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py @@ -20,23 +20,6 @@ import pandas as pd -TRAIN_TEXTS = ( - "Please read this sentence at a calm and steady pace.", - "A short speech sample checks the complete training path.", - "Clear pronunciation makes this audio easy to inspect.", - "The weather is pleasant and the morning train is on time.", - "Four simple prompts are enough for one smoke-test batch.", - "This second batch verifies another optimizer update.", - "Generated speech is decoded before the reward is computed.", - "The final checkpoint confirms that training completed.", -) -VALIDATION_TEXTS = ( - "This is fixed validation sample one.", - "This is fixed validation sample two.", - "This is fixed validation sample three.", - "This is fixed validation sample four.", -) - def _row(text: str, sample_id: str, split: str) -> dict: extra_info = {"id": sample_id, "split": split} @@ -54,8 +37,24 @@ def main() -> None: args = parser.parse_args() args.output_dir.mkdir(parents=True, exist_ok=True) - train_rows = [_row(text, f"train-{index}", "train") for index, text in enumerate(TRAIN_TEXTS)] - validation_rows = [_row(text, f"validation-{index}", "validation") for index, text in enumerate(VALIDATION_TEXTS)] + train_texts = ( + "Please read this sentence at a calm and steady pace.", + "A short speech sample checks the complete training path.", + "Clear pronunciation makes this audio easy to inspect.", + "The weather is pleasant and the morning train is on time.", + "Four simple prompts are enough for one smoke-test batch.", + "This second batch verifies another optimizer update.", + "Generated speech is decoded before the reward is computed.", + "The final checkpoint confirms that training completed.", + ) + validation_texts = ( + "This is fixed validation sample one.", + "This is fixed validation sample two.", + "This is fixed validation sample three.", + "This is fixed validation sample four.", + ) + train_rows = [_row(text, f"train-{index}", "train") for index, text in enumerate(train_texts)] + validation_rows = [_row(text, f"validation-{index}", "validation") for index, text in enumerate(validation_texts)] pd.DataFrame(train_rows).to_parquet(args.output_dir / "train.parquet", index=False) pd.DataFrame(validation_rows).to_parquet(args.output_dir / "validation.parquet", index=False) diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index a2c8c2126..48153bc3c 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -160,6 +160,21 @@ def test_rollout_adapter_requires_speaker_embedding(): Qwen3TTSRolloutAdapter.prepare_engine_prompt([1], model_config, {}) +def test_qwen3_tts_adapters_require_explicit_language(tmp_path): + speaker = tmp_path / "speaker.json" + speaker.write_text("[0.0, 1.0]") + model_config = SimpleNamespace( + tokenizer=_Tokenizer(), + override_config={"tts_spk_embed_path": str(speaker)}, + hf_config=SimpleNamespace(talker_config=SimpleNamespace(codec_eos_token_id=2150)), + ) + + with pytest.raises(ValueError, match="supports only tts_language=Auto"): + Qwen3TTSRolloutAdapter.prepare_engine_prompt([1], model_config, {}) + with pytest.raises(ValueError, match="supports only tts_language=Auto"): + Qwen3TTSTalkerAdapter.configure_model(SimpleNamespace(config=SimpleNamespace()), model_config) + + def test_talker_adapter_pads_exact_rollout_fields_for_actor_forward(): model_inputs = {"input_ids": torch.zeros(2, 6, dtype=torch.long)} payloads = [ @@ -294,6 +309,31 @@ def prepare_engine_prompt(**kwargs): assert params[1].stage == "decoder" +@pytest.mark.parametrize( + ("adapter_prompt", "message"), + [ + ({"additional_information": {"text": ["hello"]}}, "must contain prompt_token_ids"), + ({"prompt_token_ids": "1,2"}, "list of integers"), + ([1, 2], "must return a dict or None"), + ], +) +def test_ar_strategy_rejects_invalid_adapter_prompt(adapter_prompt, message): + class Adapter: + @staticmethod + def prepare_engine_prompt(**kwargs): + return adapter_prompt + + server = SimpleNamespace( + model_config=SimpleNamespace(), + config=SimpleNamespace(max_model_len=64, prompt_length=16, response_length=8), + ) + strategy = ARStrategy(server) + strategy._rollout_adapter = Adapter + + with pytest.raises((RuntimeError, TypeError), match=message): + strategy.preprocess_input([5, 6], {}, {}, None, None) + + @pytest.mark.asyncio async def test_ar_strategy_retains_requested_stage_outputs_and_targets_weight_sync(): policy = SimpleNamespace(outputs=[]) diff --git a/tests/workers/test_omni_fsdp_engine_on_cpu.py b/tests/workers/test_omni_fsdp_engine_on_cpu.py index e99683491..a7d6bcde9 100644 --- a/tests/workers/test_omni_fsdp_engine_on_cpu.py +++ b/tests/workers/test_omni_fsdp_engine_on_cpu.py @@ -85,26 +85,18 @@ def _get(key, default=None): # --------------------------------------------------------------------------- -_omni_impl_cache = None - - def _get_omni_impl_module(): """Load ``omni_impl.py`` with pre-mocked ``sys.modules`` to bypass CUDA. - Returns a cached module after the first call. + Reuses the module from ``sys.modules`` after the first call. """ - global _omni_impl_cache - if _omni_impl_cache is not None: - return _omni_impl_cache - # If the module was already imported during pytest collection (e.g. by # ``test_omni_fsdp_merge_on_cpu.py``), reuse it to avoid re-running the # ``@EngineRegistry.register`` decorator and triggering a duplicate-key # assertion. _OMNI_IMPL_FQN = "verl_omni.workers.engine.fsdp.omni_impl" if _OMNI_IMPL_FQN in sys.modules: - _omni_impl_cache = sys.modules[_OMNI_IMPL_FQN] - return _omni_impl_cache + return sys.modules[_OMNI_IMPL_FQN] root_mod = sys.modules.setdefault("verl_omni", types.ModuleType("verl_omni")) root_mod.__path__ = [_VERL_OMNI_DIR] @@ -151,7 +143,6 @@ def _get_omni_impl_module(): sys.modules["verl_omni.workers.engine.fsdp.omni_impl"] = omni_impl spec.loader.exec_module(omni_impl) - _omni_impl_cache = omni_impl return omni_impl diff --git a/verl_omni/pipelines/model_base.py b/verl_omni/pipelines/model_base.py index 0db96f99d..dd6893958 100644 --- a/verl_omni/pipelines/model_base.py +++ b/verl_omni/pipelines/model_base.py @@ -22,8 +22,6 @@ from verl_omni.workers.config import DiffusionModelConfig -logger = logging.getLogger(__name__) - class DiffusionModelBase(ABC): """Abstract base class for diffusion model training helpers. @@ -64,7 +62,7 @@ def get_class(cls, model_config: DiffusionModelConfig) -> type["DiffusionModelBa algorithm = model_config.algorithm if architecture in {"QwenImagePipeline", "QwenImageEditPlusPipeline"}: - logger.info( + logging.getLogger(__name__).info( "Applying monkey-patch for QwenImageTransformer2DModel Ulysses SP " "This workaround will be removed once we upgrade to a diffusers release that " "includes the upstream fix." diff --git a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py index e0f54fe27..4ac274f75 100644 --- a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py @@ -132,7 +132,7 @@ def prepare_engine_prompt(cls, prompt_ids, model_config, multi_modal_data, mm_pr speaker_path = model_config.override_config.get("tts_spk_embed_path") if not speaker_path: raise ValueError("Qwen3-TTS GRPO requires tts_spk_embed_path for the validated non-streaming replay.") - language = require_auto_language(model_config.override_config.get("tts_language", "Auto")) + language = require_auto_language(model_config.override_config.get("tts_language")) additional_information = { "task_type": ["Base"], "text": [text], diff --git a/verl_omni/pipelines/qwen3_tts/rollout_utils.py b/verl_omni/pipelines/qwen3_tts/rollout_utils.py index 2f9a85283..24ce2949a 100644 --- a/verl_omni/pipelines/qwen3_tts/rollout_utils.py +++ b/verl_omni/pipelines/qwen3_tts/rollout_utils.py @@ -50,7 +50,13 @@ def align_audio_codes(audio_codes: torch.Tensor, token_ids: list[int]) -> torch. f"response_length={len(token_ids)}, raw_codec_rows={raw_codes.shape[0]}." ) - copy_length, start = max(candidates) + copy_length = max(length for length, _ in candidates) + best_starts = [start for length, start in candidates if length == copy_length] + if len(best_starts) != 1: + raise RuntimeError( + "Ambiguous Qwen3-TTS codec alignment: multiple residual-codebook spans match the sampled policy." + ) + start = best_starts[0] codes = raw_codes.new_zeros((len(token_ids), 16)) if copy_length: codes[:copy_length] = raw_codes[start : start + copy_length] diff --git a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py index 8ea5b90c4..073747b1a 100644 --- a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py @@ -95,7 +95,7 @@ def get_strip_modules(cls, model_config): def configure_model(cls, module, model_config): module = super().configure_model(module, model_config) module.config.tts_spk_embed_path = model_config.override_config.get("tts_spk_embed_path") - module.config.tts_language = require_auto_language(model_config.override_config.get("tts_language", "Auto")) + module.config.tts_language = require_auto_language(model_config.override_config.get("tts_language")) if not module.config.tts_spk_embed_path: raise ValueError("Qwen3-TTS GRPO requires tts_spk_embed_path for the validated non-streaming replay.") module._verl_tts_speaker_embedding = load_speaker_xvector(module.config.tts_spk_embed_path) diff --git a/verl_omni/workers/config/omni/model.py b/verl_omni/workers/config/omni/model.py index 85ebd4fa1..daa2bf497 100644 --- a/verl_omni/workers/config/omni/model.py +++ b/verl_omni/workers/config/omni/model.py @@ -14,7 +14,6 @@ """Configuration dataclass for omni (thinker/talker) model training.""" import json -import logging import os from dataclasses import dataclass, field from typing import Any, Optional @@ -30,8 +29,6 @@ __all__ = ["OmniModelConfig"] -logger = logging.getLogger(__name__) - @dataclass class OmniModelConfig(BaseConfig): diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py index f3c8ac03f..88e6b7271 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py @@ -222,7 +222,18 @@ def preprocess_input( mm_processor_kwargs=mm_processor_kwargs, ) adapter_prepared_prompt = prompt is not None - effective_prompt_ids = prompt.get("prompt_token_ids", prompt_ids) if prompt is not None else prompt_ids + if prompt is not None: + if not isinstance(prompt, dict): + raise TypeError(f"An omni rollout adapter must return a dict or None, got {type(prompt).__name__}.") + if "prompt_token_ids" not in prompt: + raise RuntimeError("An adapter-prepared omni prompt must contain prompt_token_ids.") + effective_prompt_ids = prompt["prompt_token_ids"] + if not isinstance(effective_prompt_ids, list) or any( + isinstance(token_id, bool) or not isinstance(token_id, int) for token_id in effective_prompt_ids + ): + raise TypeError("An adapter-prepared omni prompt must contain prompt_token_ids as a list of integers.") + else: + effective_prompt_ids = prompt_ids max_possible_tokens = self.server.config.max_model_len - len(effective_prompt_ids) if max_possible_tokens <= 0: raise ValueError( From c5fafa7fc7627a91083248277397141f94b21cef Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Tue, 1 Sep 2026 21:48:27 +0800 Subject: [PATCH 20/28] [rollout, tests] fix: preserve AR strategy compatibility Restore pre-existing shared logger and worker-extension definitions, preserve underscore and hyphenated timeout arguments, and cover the timeout behavior with focused CPU tests. Signed-off-by: dongbo910220 <1275604947@qq.com> --- .../test_vllm_omni_strategy_on_cpu.py | 22 ++++++++++++++++++- .../agent_loop/single_turn_agent_loop.py | 5 +++++ .../vllm_rollout/vllm_omni_ar_strategy.py | 8 ++++--- 3 files changed, 31 insertions(+), 4 deletions(-) diff --git a/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py index a1771f3a5..131c379d8 100644 --- a/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py @@ -160,7 +160,7 @@ def test_ar_strategy_preserves_output_conversion(): finish_reason="length", num_preempted=2, ) - final_res = SimpleNamespace(request_output=SimpleNamespace(outputs=[completion])) + final_res = SimpleNamespace(outputs=[completion]) output = strategy.process_output( final_res, @@ -195,6 +195,26 @@ def test_ar_strategy_preserves_engine_kwarg_normalization(monkeypatch): } +@pytest.mark.parametrize( + ("timeout_kwargs", "expected"), + [ + ({"stage-init-timeout": 45}, {"stage-init-timeout": 45, "init-timeout": 600}), + ( + {"stage_init_timeout": 45, "init-timeout": 90}, + {"stage-init-timeout": 45, "init-timeout": 90}, + ), + ], +) +def test_ar_strategy_preserves_hyphenated_timeout_kwargs(monkeypatch, timeout_kwargs, expected): + monkeypatch.setattr(ar_strategy_module.OmniRolloutPipelineBase, "get_class", lambda pipeline_name: None) + strategy = ARStrategy(SimpleNamespace()) + engine_kwargs = {"pipeline_name": "missing", **timeout_kwargs} + + strategy.preprocess_engine_kwargs(engine_kwargs) + + assert engine_kwargs == expected + + def test_ar_strategy_preserves_engine_argument_normalization(): server = SimpleNamespace(config=SimpleNamespace(logprobs_mode="raw_logprobs")) strategy = ARStrategy(server) diff --git a/verl_omni/agent_loop/single_turn_agent_loop.py b/verl_omni/agent_loop/single_turn_agent_loop.py index 59b44480c..6034716bf 100644 --- a/verl_omni/agent_loop/single_turn_agent_loop.py +++ b/verl_omni/agent_loop/single_turn_agent_loop.py @@ -11,6 +11,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import logging +import os from typing import Any from uuid import uuid4 @@ -22,6 +24,9 @@ from verl_omni.agent_loop.diffusion_agent_loop import DiffusionAgentLoopOutput from verl_omni.pipelines.model_base import OmniRolloutPipelineBase +logger = logging.getLogger(__file__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + @register("omni_single_turn_agent") class OmniSingleTurnAgentLoop(SingleTurnAgentLoop): diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py index 88e6b7271..4b8b42c1a 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py @@ -35,6 +35,8 @@ logger = logging.getLogger(__file__) logger.setLevel(logging.INFO) +_WORKER_EXTENSION = "verl_omni.workers.rollout.vllm_rollout.utils.vLLMOmniColocateWorkerExtension" + def _drop_none_mapping_values(value: Any) -> Any: if isinstance(value, dict): @@ -74,7 +76,7 @@ def override_generation_config(self) -> dict[str, Any]: return vLLMHttpServer._get_override_generation_config(self.server) def worker_extension_cls(self, device_type: str) -> str: - return "verl_omni.workers.rollout.vllm_rollout.utils.vLLMOmniColocateWorkerExtension" + return _WORKER_EXTENSION def preprocess_engine_kwargs(self, engine_kwargs: dict[str, Any]) -> None: super().preprocess_engine_kwargs(engine_kwargs) @@ -97,8 +99,8 @@ def preprocess_engine_kwargs(self, engine_kwargs: dict[str, Any]) -> None: hf_overrides.update(adapter_overrides) engine_kwargs["hf_overrides"] = hf_overrides - stage_init_timeout = engine_kwargs.get("stage_init_timeout") - init_timeout = engine_kwargs.get("init_timeout") + stage_init_timeout = engine_kwargs.get("stage_init_timeout") or engine_kwargs.get("stage-init-timeout") + init_timeout = engine_kwargs.get("init_timeout") or engine_kwargs.get("init-timeout") if stage_init_timeout is not None and init_timeout is None: engine_kwargs["init_timeout"] = max(int(stage_init_timeout), 600) From 9d2fe6f8c24e7c6f23be0d75b17f47a8d36d1559 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Tue, 1 Sep 2026 22:22:20 +0800 Subject: [PATCH 21/28] [model, rollout, tests] fix: restore upstream-owned behavior Restore pre-existing loggers, test caching, configuration access, pipeline defaults, and AR output normalization that were changed by an over-broad review cleanup. Signed-off-by: dongbo910220 <1275604947@qq.com> --- .../rollout_vllm/test_vllm_omni_strategy_on_cpu.py | 2 +- tests/workers/test_omni_fsdp_engine_on_cpu.py | 14 +++++++++++--- verl_omni/pipelines/model_base.py | 6 ++++-- verl_omni/workers/config/omni/model.py | 3 +++ verl_omni/workers/engine/fsdp/omni_impl.py | 2 +- .../rollout/vllm_rollout/vllm_omni_ar_strategy.py | 11 ++++++----- 6 files changed, 26 insertions(+), 12 deletions(-) diff --git a/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py index 131c379d8..87bb4cd2b 100644 --- a/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py @@ -160,7 +160,7 @@ def test_ar_strategy_preserves_output_conversion(): finish_reason="length", num_preempted=2, ) - final_res = SimpleNamespace(outputs=[completion]) + final_res = SimpleNamespace(request_output=SimpleNamespace(outputs=[completion])) output = strategy.process_output( final_res, diff --git a/tests/workers/test_omni_fsdp_engine_on_cpu.py b/tests/workers/test_omni_fsdp_engine_on_cpu.py index a7d6bcde9..d724cd66d 100644 --- a/tests/workers/test_omni_fsdp_engine_on_cpu.py +++ b/tests/workers/test_omni_fsdp_engine_on_cpu.py @@ -84,19 +84,26 @@ def _get(key, default=None): # Isolated module loader # --------------------------------------------------------------------------- +_omni_impl_cache = None + def _get_omni_impl_module(): """Load ``omni_impl.py`` with pre-mocked ``sys.modules`` to bypass CUDA. - Reuses the module from ``sys.modules`` after the first call. + Returns a cached module after the first call. """ + global _omni_impl_cache + if _omni_impl_cache is not None: + return _omni_impl_cache + # If the module was already imported during pytest collection (e.g. by # ``test_omni_fsdp_merge_on_cpu.py``), reuse it to avoid re-running the # ``@EngineRegistry.register`` decorator and triggering a duplicate-key # assertion. _OMNI_IMPL_FQN = "verl_omni.workers.engine.fsdp.omni_impl" if _OMNI_IMPL_FQN in sys.modules: - return sys.modules[_OMNI_IMPL_FQN] + _omni_impl_cache = sys.modules[_OMNI_IMPL_FQN] + return _omni_impl_cache root_mod = sys.modules.setdefault("verl_omni", types.ModuleType("verl_omni")) root_mod.__path__ = [_VERL_OMNI_DIR] @@ -143,6 +150,7 @@ def _get_omni_impl_module(): sys.modules["verl_omni.workers.engine.fsdp.omni_impl"] = omni_impl spec.loader.exec_module(omni_impl) + _omni_impl_cache = omni_impl return omni_impl @@ -436,7 +444,7 @@ def test_build_module_calls_adapter_configure_model(architecture): mock_get_cls.assert_called_once_with( architecture, model_config.model_stage, - model_config.external_lib, + model_config.get("external_lib"), ) fake_adapter_cls.configure_model.assert_called_once_with(fake_module, model_config) diff --git a/verl_omni/pipelines/model_base.py b/verl_omni/pipelines/model_base.py index dd6893958..8c77bf0e4 100644 --- a/verl_omni/pipelines/model_base.py +++ b/verl_omni/pipelines/model_base.py @@ -22,6 +22,8 @@ from verl_omni.workers.config import DiffusionModelConfig +logger = logging.getLogger(__name__) + class DiffusionModelBase(ABC): """Abstract base class for diffusion model training helpers. @@ -62,7 +64,7 @@ def get_class(cls, model_config: DiffusionModelConfig) -> type["DiffusionModelBa algorithm = model_config.algorithm if architecture in {"QwenImagePipeline", "QwenImageEditPlusPipeline"}: - logging.getLogger(__name__).info( + logger.info( "Applying monkey-patch for QwenImageTransformer2DModel Ulysses SP " "This workaround will be removed once we upgrade to a diffusers release that " "includes the upstream fix." @@ -783,7 +785,7 @@ def get_pipeline_id(cls, pipeline_mode: str = "thinker_only") -> str: for model_type, cls_ref in cls._registry.items(): if cls_ref is cls: return model_type - raise RuntimeError(f"{cls.__name__} is not registered as an omni rollout pipeline.") + return "" @classmethod def ensure_pipeline_registered(cls, pipeline_mode: str = "thinker_only") -> None: diff --git a/verl_omni/workers/config/omni/model.py b/verl_omni/workers/config/omni/model.py index daa2bf497..85ebd4fa1 100644 --- a/verl_omni/workers/config/omni/model.py +++ b/verl_omni/workers/config/omni/model.py @@ -14,6 +14,7 @@ """Configuration dataclass for omni (thinker/talker) model training.""" import json +import logging import os from dataclasses import dataclass, field from typing import Any, Optional @@ -29,6 +30,8 @@ __all__ = ["OmniModelConfig"] +logger = logging.getLogger(__name__) + @dataclass class OmniModelConfig(BaseConfig): diff --git a/verl_omni/workers/engine/fsdp/omni_impl.py b/verl_omni/workers/engine/fsdp/omni_impl.py index b11c0f1eb..f2181e2e5 100644 --- a/verl_omni/workers/engine/fsdp/omni_impl.py +++ b/verl_omni/workers/engine/fsdp/omni_impl.py @@ -179,7 +179,7 @@ def _build_module(self): adapter_cls = OmniModelBase.get_class_by_name( architecture, self.model_config.model_stage, - self.model_config.external_lib, + self.model_config.get("external_lib"), ) self.model_adapter_cls = adapter_cls diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py index 4b8b42c1a..58e60a22b 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py @@ -323,23 +323,24 @@ def process_output( if final_res is None: raise RuntimeError("AR mode: vLLM-Omni engine yielded no output for the prompt.") - if not final_res.outputs: + req_output = getattr(final_res, "request_output", None) or final_res + if not req_output.outputs: raise RuntimeError("AR mode expects outputs with token IDs, but got None or empty.") extra_fields = {"global_steps": self.server.global_steps} if self._rollout_adapter is not None: extra_fields.update(final_res._verl_omni_rollout_fields) - token_ids = final_res.outputs[0].token_ids + token_ids = req_output.outputs[0].token_ids log_probs = None policy_params = params[0] if isinstance(params, list) else params if policy_params.logprobs is not None: log_probs = [ - logprobs[token_ids[index]].logprob for index, logprobs in enumerate(final_res.outputs[0].logprobs) + logprobs[token_ids[index]].logprob for index, logprobs in enumerate(req_output.outputs[0].logprobs) ] - finish_reason = final_res.outputs[0].finish_reason + finish_reason = req_output.outputs[0].finish_reason stop_reason = self._map_stop_reason(finish_reason) - num_preempted = self._extract_num_preempted(final_res) + num_preempted = self._extract_num_preempted(req_output) return TokenOutput( token_ids=token_ids, From 2ea5103e520aa8cf1f2ff4f7aeeb7285e423f5d6 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Wed, 2 Sep 2026 00:23:30 +0800 Subject: [PATCH 22/28] [rollout, tests] fix: preserve single-output AR pipelines Restrict multi-output sampling and replay behavior to pipelines that actually retain multiple stage outputs. Cover the existing Qwen3-Omni thinker-only contract and keep the Qwen3-TTS fixture explicit. Signed-off-by: dongbo910220 <1275604947@qq.com> --- .../test_qwen3_tts_rollout_on_cpu.py | 1 + .../test_vllm_omni_strategy_on_cpu.py | 45 +++++++++++++++++++ .../vllm_rollout/vllm_omni_ar_strategy.py | 10 ++--- 3 files changed, 51 insertions(+), 5 deletions(-) diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index 48153bc3c..21f1425da 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -291,6 +291,7 @@ def prepare_engine_prompt(**kwargs): ) strategy = ARStrategy(server) strategy._rollout_adapter = Adapter + strategy._rollout_output_modalities = ["latent", "audio"] strategy._stage_sampling_constraints = {0: {}} prompt, params = strategy.preprocess_input( diff --git a/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py index 87bb4cd2b..17743d824 100644 --- a/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py @@ -19,6 +19,7 @@ import pytest import torch +from verl_omni.pipelines.qwen3_omni.omni_rollout_adapter import Qwen3OmniRolloutAdapter from verl_omni.pipelines.rollout_media import DiffusionIOSpec, MediaSpec from verl_omni.workers.rollout.vllm_rollout import vllm_omni_ar_strategy as ar_strategy_module from verl_omni.workers.rollout.vllm_rollout import vllm_omni_async_server as server_module @@ -238,6 +239,50 @@ def test_ar_strategy_preserves_engine_argument_normalization(): } +def test_ar_strategy_preserves_qwen3_omni_thinker_only_contract(): + server = SimpleNamespace( + config=SimpleNamespace( + max_model_len=8, + prompt_length=4, + response_length=4, + repetition_penalty=1.0, + logprobs_mode="processed_logprobs", + ), + model_config=SimpleNamespace(processor=None), + global_steps=12, + ) + strategy = ARStrategy(server) + strategy._rollout_adapter = Qwen3OmniRolloutAdapter + strategy._rollout_output_modalities = None + + engine_args = {"model_stage": "thinker"} + strategy.prepare_engine_args(engine_args, Namespace(stage_init_timeout=None, init_timeout=None)) + assert engine_args["model_stage"] == "thinker" + + prompt, params = strategy.preprocess_input( + prompt_ids=[1, 2], + sampling_params={"max_new_tokens": 2, "logprobs": True}, + multi_modal_data={}, + lora_request=None, + negative_prompt_ids=None, + ) + assert prompt == {"prompt_token_ids": [1, 2]} + assert isinstance(params, ar_strategy_module.SamplingParams) + + completion = SimpleNamespace( + token_ids=[7], + logprobs=[{7: SimpleNamespace(logprob=-0.25)}], + finish_reason="stop", + num_preempted=0, + ) + output = strategy.process_output( + SimpleNamespace(request_output=SimpleNamespace(outputs=[completion])), + params=params, + sampling_params={}, + ) + assert output.extra_fields == {"global_steps": 12} + + def test_diffusion_strategy_preserves_engine_argument_preparation(monkeypatch): imported = [] monkeypatch.setattr(diffusion_strategy_module, "import_external_libs", imported.append) diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py index 58e60a22b..660ba8531 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py @@ -180,9 +180,9 @@ def _write_deploy_config( engine_kwargs["deploy_config"] = deploy_path def prepare_engine_args(self, engine_args: dict[str, Any], args: Namespace) -> None: - if self._rollout_adapter is not None: + if self._rollout_output_modalities is not None: # The generated per-stage deploy config owns model_stage for - # heterogeneous pipelines such as Qwen3-TTS. + # multi-output pipelines such as Qwen3-TTS. engine_args["model_stage"] = None for timeout_key in ("stage_init_timeout", "init_timeout"): timeout_value = getattr(args, timeout_key, None) @@ -263,10 +263,10 @@ def preprocess_input( sampling_params["logprobs"] = None sampling_params.setdefault("repetition_penalty", getattr(self.server.config, "repetition_penalty", 1.0)) policy_params = SamplingParams(max_tokens=max_tokens, **sampling_params) - if self._rollout_adapter is not None: + if self._rollout_output_modalities is not None: params = copy.deepcopy(self.server.engine.default_sampling_params_list) if len(params) <= 1: - raise RuntimeError("An omni rollout adapter requires per-stage sampling parameters.") + raise RuntimeError("A multi-output omni rollout requires per-stage sampling parameters.") constrained = self._stage_sampling_constraints[0] for field in {"max_tokens", *sampling_params} - constrained.keys(): setattr(params[0], field, getattr(policy_params, field)) @@ -328,7 +328,7 @@ def process_output( raise RuntimeError("AR mode expects outputs with token IDs, but got None or empty.") extra_fields = {"global_steps": self.server.global_steps} - if self._rollout_adapter is not None: + if self._rollout_output_modalities is not None: extra_fields.update(final_res._verl_omni_rollout_fields) token_ids = req_output.outputs[0].token_ids log_probs = None From e5a55b3d61505fea41328004eabff6a0eb26bd34 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Thu, 3 Sep 2026 17:06:44 +0800 Subject: [PATCH 23/28] fix: harden Qwen3-TTS rollout and smoke contracts Co-authored-by: OpenAI Codex Signed-off-by: dongbo910220 <1275604947@qq.com> --- .github/actions/gpu-smoke-prepare/action.yml | 7 +- .github/qwen_tts_pin.txt | 1 + .github/workflows/gpu_smoke.yml | 2 + .../contributing/integrating_an_omni_model.md | 6 +- examples/grpo_trainer/qwen3_tts/README.md | 18 ++- .../qwen3_tts/run_qwen3_tts_grpo.sh | 4 +- pyproject.toml | 2 +- tests/gpu_smoke/select_gpu_smoke_groups.py | 7 + .../test_qwen3_tts_package_on_cpu.py | 102 +++++++++++++- .../test_audio_reward_manager_on_cpu.py | 41 +++++- .../build_qwen3_tts_tiny_random.py | 124 ++++++++++++++++++ .../create_dummy_qwen3_tts_grpo_data.py | 20 ++- tests/special_e2e/run_qwen3_tts_grpo_smoke.sh | 43 ++++-- .../test_gpu_smoke_selector_on_cpu.py | 44 +++++++ .../test_audio_http_scorer_client_on_cpu.py | 17 ++- .../test_qwen3_tts_rollout_on_cpu.py | 29 +++- .../test_vllm_omni_strategy_on_cpu.py | 43 ++++++ tests/workers/test_omni_fsdp_engine_on_cpu.py | 59 +++++++-- verl_omni/pipelines/model_base.py | 8 +- .../qwen3_tts/omni_rollout_adapter.py | 2 + .../qwen3_tts/talker_training_adapter.py | 5 + verl_omni/reward_loop/reward_manager/audio.py | 20 ++- .../reward_score/audio_http_scorer_client.py | 38 +++--- verl_omni/workers/engine/fsdp/omni_impl.py | 13 +- .../vllm_rollout/vllm_omni_ar_strategy.py | 24 +++- 25 files changed, 594 insertions(+), 85 deletions(-) create mode 100644 .github/qwen_tts_pin.txt create mode 100644 tests/special_e2e/build_qwen3_tts_tiny_random.py create mode 100644 tests/special_sanity/test_gpu_smoke_selector_on_cpu.py diff --git a/.github/actions/gpu-smoke-prepare/action.yml b/.github/actions/gpu-smoke-prepare/action.yml index 128beb0b8..d4403f342 100644 --- a/.github/actions/gpu-smoke-prepare/action.yml +++ b/.github/actions/gpu-smoke-prepare/action.yml @@ -20,7 +20,7 @@ runs: export UV_CACHE_DIR="${UV_CACHE_DIR:-${HOME}/.cache/uv}" git config --global http.postBuffer 524288000 || true # Base image PyPI mirror may lag pypi.org (e.g. kernels, fa3-fwd). - uv pip install --system --break-system-packages ".[gpu,dev,audio]" + uv pip install --system --break-system-packages ".[gpu,dev,audio,tts]" uv pip install --system --break-system-packages "vllm-omni @ git+https://github.com/vllm-project/vllm-omni.git@$(cat .github/vllm_omni_pin.txt)" uv pip install --system --break-system-packages --no-deps --reinstall "verl @ git+https://github.com/verl-project/verl.git@$(cat .github/verl_pin.txt)" uv pip install --system --break-system-packages TransferQueue==0.1.8 @@ -31,6 +31,11 @@ runs: uv pip install --system --break-system-packages veomni==0.1.11 --no-deps uv pip install --system --break-system-packages torchcodec librosa soundfile av audioread uv pip install --system --break-system-packages "transformers[mistral-common]==5.14.1" + # The pinned Qwen3-TTS TF5 source declares Transformers >=5.15.1, while + # this repository deliberately caps Transformers at 5.14.1. Install its + # exact tested source without letting that metadata replace the CI stack. + uv pip install --system --break-system-packages --no-deps \ + "qwen-tts @ git+https://github.com/QwenLM/Qwen3-TTS.git@$(cat .github/qwen_tts_pin.txt)" # GPU-smoke-only override; the project GPU extra remains on kernels==0.16.0. uv pip install --system --break-system-packages "kernels==0.14.1" # NCCL checkpoint engine (diffusion v1 separate_async weight sync). diff --git a/.github/qwen_tts_pin.txt b/.github/qwen_tts_pin.txt new file mode 100644 index 000000000..c0af4ea17 --- /dev/null +++ b/.github/qwen_tts_pin.txt @@ -0,0 +1 @@ +00969daa8064e23adc9e5f52cdf20cf247f94159 diff --git a/.github/workflows/gpu_smoke.yml b/.github/workflows/gpu_smoke.yml index 00e89aeaf..975755710 100644 --- a/.github/workflows/gpu_smoke.yml +++ b/.github/workflows/gpu_smoke.yml @@ -14,6 +14,7 @@ on: - "tests/special_e2e/**" - "pyproject.toml" - .github/workflows/gpu_smoke.yml + - .github/qwen_tts_pin.txt - .github/vllm_omni_pin.txt - .github/actions/gpu-smoke-prepare/** - .github/actions/gpu-smoke-upload-logs/** @@ -31,6 +32,7 @@ on: - "tests/special_e2e/**" - "pyproject.toml" - .github/workflows/gpu_smoke.yml + - .github/qwen_tts_pin.txt - .github/vllm_omni_pin.txt - .github/actions/gpu-smoke-prepare/** - .github/actions/gpu-smoke-upload-logs/** diff --git a/docs/contributing/integrating_an_omni_model.md b/docs/contributing/integrating_an_omni_model.md index 7a9d2ab34..804e039fc 100644 --- a/docs/contributing/integrating_an_omni_model.md +++ b/docs/contributing/integrating_an_omni_model.md @@ -61,8 +61,10 @@ adapt each implementation to your model's architecture: This method runs before FSDP wrapping and LoRA injection. - **`register_auto_classes()`** (optional): Register classes supplied by an - optional model package with the appropriate Transformers Auto APIs. Qwen3-TTS - registers the official `qwen-tts` config and model with `AutoConfig` and + optional model package with the appropriate Transformers Auto APIs. The model + config resolves one `(architecture, stage)` adapter before calling this hook; + the base implementation is a no-op. Qwen3-TTS registers the official + `qwen-tts` config and model with `AutoConfig` and `AutoModelForTextToWaveform`; the FSDP engine still owns `from_pretrained`. - **`prepare_model_inputs(model_inputs, micro_batch, model_config)`** diff --git a/examples/grpo_trainer/qwen3_tts/README.md b/examples/grpo_trainer/qwen3_tts/README.md index 1032cc352..b620da7a9 100644 --- a/examples/grpo_trainer/qwen3_tts/README.md +++ b/examples/grpo_trainer/qwen3_tts/README.md @@ -38,16 +38,20 @@ uv pip install -e ".[gpu]" --torch-backend=auto uv pip install "vllm-omni @ git+https://github.com/vllm-project/vllm-omni.git@$(cat .github/vllm_omni_pin.txt)" uv pip install -e ".[tts,train,dev]" uv pip install --no-deps \ - "qwen-tts @ git+https://github.com/QwenLM/Qwen3-TTS.git@00969daa8064e23adc9e5f52cdf20cf247f94159" + "qwen-tts @ git+https://github.com/QwenLM/Qwen3-TTS.git@$(cat .github/qwen_tts_pin.txt)" ``` The pinned Qwen3-TTS revision is the upstream Transformers 5 support change -from Qwen3-TTS PR #360. The released `qwen-tts==0.1.1` source targets -Transformers 4.57 and cannot be imported unchanged on this repository's -Transformers 5 stack. The adapter registers the upstream config and model with -`AutoConfig` and `AutoModelForTextToWaveform`; it does not carry a local -Transformers compatibility layer. The system `sox` executable is also required -by qwen-tts. +from Qwen3-TTS PR #360. Its package metadata requires Transformers 5.15.1 or +newer, while this repository intentionally caps Transformers at 5.14.1. The +`--no-deps` flag preserves that repository-wide cap; the `tts` extra explicitly +owns the runtime dependencies, including `torchaudio==2.11.0` to match vLLM's +Torch pin, and CI tests the exact Qwen3-TTS revision from +`.github/qwen_tts_pin.txt` on this stack. The released `qwen-tts==0.1.1` source +targets Transformers 4.57 and cannot be imported unchanged here. The adapter +registers the upstream config and model with `AutoConfig` and +`AutoModelForTextToWaveform`; it does not carry a local Transformers +compatibility layer. The system `sox` executable is also required by qwen-tts. ## Data diff --git a/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh b/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh index 193e727b0..cf79effa8 100755 --- a/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh +++ b/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh @@ -123,9 +123,9 @@ export VLLM_USE_FLASHINFER_SAMPLER=0 reward.custom_reward_function.path=pkg://verl_omni.utils.reward_score.audio_http_scorer_client \ reward.custom_reward_function.name=compute_score \ "+reward.custom_reward_function.reward_kwargs.server_url=${SCORER_URL}" \ - +reward.custom_reward_function.reward_kwargs.timeout_s=120.0 \ + +reward.custom_reward_function.reward_kwargs.timeout=120.0 \ +reward.custom_reward_function.reward_kwargs.max_retries=2 \ - +reward.custom_reward_function.reward_kwargs.retry_backoff_s=0.5 \ + +reward.custom_reward_function.reward_kwargs.retry_backoff=0.5 \ reward.reward_manager.source=importlib \ reward.reward_manager.name=AudioRewardManager \ reward.reward_manager.module.path=pkg://verl_omni.reward_loop.reward_manager \ diff --git a/pyproject.toml b/pyproject.toml index 50e610fc8..b82b9b74a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,7 +66,7 @@ tts = [ "onnxruntime>=1.20.0", "soundfile>=0.12.1", "sox>=1.5.0", - "torchaudio", + "torchaudio==2.11.0", ] # CUDA rollout backend (vllm) + actor FA3 (kernels) + liger-kernel. Install in step 1 on GPU. gpu = [ diff --git a/tests/gpu_smoke/select_gpu_smoke_groups.py b/tests/gpu_smoke/select_gpu_smoke_groups.py index 5f7a38cbe..d06d59452 100644 --- a/tests/gpu_smoke/select_gpu_smoke_groups.py +++ b/tests/gpu_smoke/select_gpu_smoke_groups.py @@ -66,11 +66,18 @@ class SmokeGroup: "verl_omni/workers/**", ), "ci-e2e-omni": ( + ".github/qwen_tts_pin.txt", + "examples/grpo_trainer/qwen3_tts/**", "tests/gpu_smoke/run_gpu_smoke_omni_e2e.sh", + "tests/pipelines/test_qwen3_tts*", "tests/special_e2e/*omni*", + "tests/special_e2e/*qwen3_tts*", + "tests/utils/reward_score/test_audio_http_scorer_client_on_cpu.py", "verl_omni/models/transformers/qwen3_omni_thinker.py", + "verl_omni/pipelines/qwen3_tts/**", "verl_omni/trainer/config/omni/**", "verl_omni/trainer/omni/**", + "verl_omni/utils/reward_score/audio_http_scorer_client.py", ), "ci-e2e-diffusion": ( "tests/gpu_smoke/run_gpu_smoke_diffusion_e2e.sh", diff --git a/tests/pipelines/test_qwen3_tts_package_on_cpu.py b/tests/pipelines/test_qwen3_tts_package_on_cpu.py index c1cbd5be3..a6c6fa069 100644 --- a/tests/pipelines/test_qwen3_tts_package_on_cpu.py +++ b/tests/pipelines/test_qwen3_tts_package_on_cpu.py @@ -14,12 +14,61 @@ """Check the pinned upstream qwen-tts TF5 source on the repository stack.""" import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace import pytest import torch -def test_qwen_tts_registers_and_runs_without_a_transformers_compatibility_layer(): +def _load_smoke_helper(filename: str): + path = Path(__file__).parents[1] / f"special_e2e/{filename}.py" + spec = importlib.util.spec_from_file_location(f"{filename}_under_test", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_smoke_speaker_dimension_matches_model_config(tmp_path): + builder = _load_smoke_helper("create_dummy_qwen3_tts_grpo_data") + model_config_path = tmp_path / "config.json" + model_config_path.write_text( + json.dumps( + { + "talker_config": {"hidden_size": 128}, + "speaker_encoder_config": {"enc_dim": 128}, + } + ), + encoding="utf-8", + ) + + assert builder._speaker_dimension(model_config_path) == 128 + + model_config_path.write_text( + json.dumps( + { + "talker_config": {"hidden_size": 128}, + "speaker_encoder_config": {"enc_dim": 1024}, + } + ), + encoding="utf-8", + ) + with pytest.raises(ValueError, match="speaker encoder output must match"): + builder._speaker_dimension(model_config_path) + + +def test_smoke_tiny_model_mrope_section_matches_head_dimension(): + builder = _load_smoke_helper("build_qwen3_tts_tiny_random") + + section = builder._scaled_mrope_section([24, 20, 20], head_dim=64) + + assert section == [12, 10, 10] + assert sum(section) == 64 // 2 + + +def test_qwen_tts_registers_and_runs_without_a_transformers_compatibility_layer(tmp_path, monkeypatch): transformers = pytest.importorskip("transformers") if importlib.util.find_spec("qwen_tts") is None: pytest.skip("qwen-tts is an optional dependency") @@ -47,7 +96,7 @@ def test_qwen_tts_registers_and_runs_without_a_transformers_compatibility_layer( "num_key_value_heads": 1, "head_dim": 4, "max_position_embeddings": 64, - "num_code_groups": 4, + "num_code_groups": 16, "layer_types": ["full_attention"], "pad_token_id": None, } @@ -60,9 +109,15 @@ def test_qwen_tts_registers_and_runs_without_a_transformers_compatibility_layer( "num_attention_heads": 2, "num_key_value_heads": 1, "max_position_embeddings": 64, - "num_code_groups": 4, + "num_code_groups": 16, "text_hidden_size": 8, "text_vocab_size": 80, + "codec_eos_token_id": 50, + "codec_nothink_id": 51, + "codec_think_bos_id": 52, + "codec_think_eos_id": 53, + "codec_pad_id": 54, + "codec_bos_id": 55, "spk_id": {}, "codec_language_id": {}, "rope_scaling": { @@ -77,6 +132,9 @@ def test_qwen_tts_registers_and_runs_without_a_transformers_compatibility_layer( speaker_encoder_config={}, tts_model_type="custom", tokenizer_type="12hz", + tts_pad_token_id=60, + tts_bos_token_id=61, + tts_eos_token_id=62, ) model = Qwen3TTSForConditionalGeneration(config) output = model.talker( @@ -87,3 +145,41 @@ def test_qwen_tts_registers_and_runs_without_a_transformers_compatibility_layer( ) assert output.logits.shape == (2, 5, 64) + + import vllm_omni.platforms as platforms + from vllm_omni.platforms.interface import UnspecifiedOmniPlatform + + monkeypatch.setattr(platforms, "_current_omni_platform", UnspecifiedOmniPlatform()) + from verl_omni.pipelines.qwen3_tts.talker_forward import ( + TalkerTokens, + build_talker_batch, + codec0_input_embeddings, + ) + + codes = torch.randint(1, 31, (3, 16), dtype=torch.long) + batch = build_talker_batch( + [torch.tensor([1, 2, 3])], + [codes], + TalkerTokens.from_config(model.config), + sub_codebook_vocab=32, + ) + embeddings = codec0_input_embeddings(model.talker, batch, torch.zeros(1, 8)) + assert embeddings.shape == (*batch.input_ids.shape[:2], 8) + assert torch.isfinite(embeddings).all() + + from verl_omni.pipelines.qwen3_tts.talker_training_adapter import Qwen3TTSTalkerAdapter + + speaker_path = tmp_path / "speaker.json" + speaker_path.write_text(json.dumps([0.0] * 8), encoding="utf-8") + configured = Qwen3TTSTalkerAdapter.configure_model( + model, + SimpleNamespace( + use_remove_padding=False, + override_config={"tts_spk_embed_path": str(speaker_path), "tts_language": "Auto"}, + ), + ) + trainable_names = {name for name, parameter in configured.named_parameters() if parameter.requires_grad} + assert trainable_names + assert all(name.startswith(("talker.model.", "talker.codec_head.")) for name in trainable_names) + assert any(not parameter.requires_grad for parameter in configured.parameters()) + assert configured.get_input_embeddings() is configured.talker.model.codec_embedding diff --git a/tests/reward_loop/test_audio_reward_manager_on_cpu.py b/tests/reward_loop/test_audio_reward_manager_on_cpu.py index bf33fbc32..24bca57a2 100644 --- a/tests/reward_loop/test_audio_reward_manager_on_cpu.py +++ b/tests/reward_loop/test_audio_reward_manager_on_cpu.py @@ -44,18 +44,21 @@ def _manager(compute_score): return AudioRewardManager(_config(), MagicMock(), compute_score=compute_score) -def _data(audio=None, sample_rate=24_000): +def _data(audio=None, sample_rate=24_000, *, layout="streamed"): fields = {} if audio is not None: fields = {"audio": audio, "audio_sample_rate": sample_rate} + non_tensors = { + "data_source": ["tts_reward"], + "reward_model": [{"ground_truth": "ni3 hao3"}], + "extra_info": [{"id": "sample-0"}], + "tool_extra_fields": [fields if layout == "streamed" else {}], + } + if layout == "finalized" and audio is not None: + non_tensors.update({"audio": [audio], "audio_sample_rate": [sample_rate]}) return DataProto.from_dict( tensors={"responses": torch.zeros(1, 4, dtype=torch.long)}, - non_tensors={ - "data_source": ["tts_reward"], - "reward_model": [{"ground_truth": "ni3 hao3"}], - "extra_info": [{"id": "sample-0"}], - "tool_extra_fields": [fields], - }, + non_tensors=non_tensors, ) @@ -113,6 +116,7 @@ def compute_score(data_source, solution_audio, ground_truth, extra_info): assert sample_rate == 24_000 assert ground_truth == "ni3 hao3" assert extra_info["id"] == "sample-0" + assert "global_steps" not in extra_info return {"score": 0.75, "pinyin_error_rate": 0.1} manager = _manager(compute_score) @@ -124,6 +128,22 @@ def compute_score(data_source, solution_audio, ground_truth, extra_info): } +def test_run_single_reads_finalized_top_level_audio_layout(): + def compute_score(solution_audio, extra_info, **kwargs): + waveform, sample_rate = solution_audio + np.testing.assert_array_equal(waveform, np.ones(8, dtype=np.float32)) + assert sample_rate == 16_000 + assert extra_info["id"] == "sample-0" + return 0.5 + + manager = _manager(compute_score) + result = manager.loop.run_until_complete( + manager.run_single(_data(np.ones(8, dtype=np.float32), 16_000, layout="finalized")) + ) + + assert result["reward_score"] == 0.5 + + @pytest.mark.parametrize( ("data", "message"), [ @@ -159,6 +179,13 @@ def test_chunked_waveform_is_rejected_instead_of_guessed(): ) +def test_two_dimensional_waveform_is_rejected_instead_of_downmixed_on_the_wrong_axis(): + manager = _manager(lambda **kwargs: 0.5) + + with pytest.raises(ValueError, match="one mono waveform"): + manager.loop.run_until_complete(manager.run_single(_data(np.zeros((128, 2)), 16_000))) + + @pytest.mark.asyncio async def test_async_score_function_is_supported(): async def compute_score(solution_audio, **kwargs): diff --git a/tests/special_e2e/build_qwen3_tts_tiny_random.py b/tests/special_e2e/build_qwen3_tts_tiny_random.py new file mode 100644 index 000000000..74c3eddd7 --- /dev/null +++ b/tests/special_e2e/build_qwen3_tts_tiny_random.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Build a tiny random 16-codebook Qwen3-TTS checkpoint for GPU smoke tests.""" + +import argparse +import json +import shutil +from pathlib import Path + +import torch +from safetensors.torch import save_file + + +def _copy_files(source: Path, destination: Path, names: tuple[str, ...]) -> None: + for name in names: + source_file = source / name + if source_file.exists(): + shutil.copy2(source_file, destination / name) + + +def _save_model(model, output_dir: Path) -> None: + state_dict = {name: tensor.detach().cpu().contiguous().clone() for name, tensor in model.state_dict().items()} + save_file(state_dict, output_dir / "model.safetensors", metadata={"format": "pt"}) + + +def _write_config(source: Path, destination: Path, updates) -> None: + config = json.loads(source.read_text(encoding="utf-8")) + updates(config) + destination.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8") + + +def _scaled_mrope_section(section: list[int], head_dim: int) -> list[int]: + target = head_dim // 2 + total = sum(section) + scaled = [value * target // total for value in section] + scaled[-1] += target - sum(scaled) + return scaled + + +def _set_model_codebooks(config: dict) -> None: + talker_config = config["talker_config"] + talker_config["num_code_groups"] = 16 + talker_config["code_predictor_config"]["num_code_groups"] = 16 + rope_scaling = talker_config["rope_scaling"] + rope_scaling["mrope_section"] = _scaled_mrope_section(rope_scaling["mrope_section"], talker_config["head_dim"]) + + +def _set_tokenizer_quantizers(config: dict) -> None: + config["dtype"] = "bfloat16" + config["encoder_valid_num_quantizers"] = 16 + config["encoder_config"]["num_quantizers"] = 16 + config["decoder_config"]["num_quantizers"] = 16 + + +def build(source_model_path: Path, output_dir: Path, seed: int = 42) -> Path: + from qwen_tts.core.models.configuration_qwen3_tts import Qwen3TTSConfig + from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSForConditionalGeneration + from qwen_tts.core.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config + from qwen_tts.core.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Model + + torch.manual_seed(seed) + output_dir.mkdir(parents=True, exist_ok=True) + + model_config = Qwen3TTSConfig.from_pretrained(source_model_path) + model_config.talker_config.num_code_groups = 16 + model_config.talker_config.code_predictor_config.num_code_groups = 16 + rope_scaling = model_config.talker_config.rope_scaling + rope_scaling["mrope_section"] = _scaled_mrope_section( + rope_scaling["mrope_section"], model_config.talker_config.head_dim + ) + model = Qwen3TTSForConditionalGeneration(model_config).to(torch.bfloat16) + _save_model(model, output_dir) + _write_config( + source_model_path / "config.json", + output_dir / "config.json", + _set_model_codebooks, + ) + _copy_files( + source_model_path, + output_dir, + ("generation_config.json", "merges.txt", "preprocessor_config.json", "tokenizer_config.json", "vocab.json"), + ) + + tokenizer_source = source_model_path / "speech_tokenizer" + tokenizer_output = output_dir / "speech_tokenizer" + tokenizer_output.mkdir(exist_ok=True) + tokenizer_config = Qwen3TTSTokenizerV2Config.from_pretrained(tokenizer_source) + tokenizer_config.encoder_valid_num_quantizers = 16 + tokenizer_config.encoder_config.num_quantizers = 16 + tokenizer_config.decoder_config.num_quantizers = 16 + tokenizer = Qwen3TTSTokenizerV2Model(tokenizer_config).to(torch.bfloat16) + _save_model(tokenizer, tokenizer_output) + _write_config( + tokenizer_source / "config.json", + tokenizer_output / "config.json", + _set_tokenizer_quantizers, + ) + _copy_files(tokenizer_source, tokenizer_output, ("configuration.json", "preprocessor_config.json")) + return output_dir + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--source-model-path", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + print(build(args.source_model_path, args.output_dir, args.seed)) + + +if __name__ == "__main__": + main() diff --git a/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py b/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py index 8ae71226f..91d1a8a07 100644 --- a/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py +++ b/tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py @@ -16,11 +16,25 @@ import argparse import json +import math from pathlib import Path import pandas as pd +def _speaker_dimension(model_config_path: Path) -> int: + model_config = json.loads(model_config_path.read_text(encoding="utf-8")) + talker_dimension = model_config["talker_config"]["hidden_size"] + speaker_dimension = model_config["speaker_encoder_config"]["enc_dim"] + if talker_dimension != speaker_dimension: + raise ValueError( + f"speaker encoder output must match the talker hidden size: {speaker_dimension} != {talker_dimension}" + ) + if not isinstance(speaker_dimension, int) or speaker_dimension <= 0: + raise ValueError(f"speaker dimension must be a positive integer, got {speaker_dimension!r}") + return speaker_dimension + + def _row(text: str, sample_id: str, split: str) -> dict: extra_info = {"id": sample_id, "split": split} return { @@ -34,6 +48,7 @@ def _row(text: str, sample_id: str, split: str) -> dict: def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--model-config", type=Path, required=True) args = parser.parse_args() args.output_dir.mkdir(parents=True, exist_ok=True) @@ -58,9 +73,8 @@ def main() -> None: pd.DataFrame(train_rows).to_parquet(args.output_dir / "train.parquet", index=False) pd.DataFrame(validation_rows).to_parquet(args.output_dir / "validation.parquet", index=False) - # Qwen3-TTS Base expects a 1024-dimensional speaker x-vector. A unit-norm - # deterministic fixture is sufficient for execution testing. - speaker = [1.0 / 32.0] * 1024 + speaker_dimension = _speaker_dimension(args.model_config) + speaker = [1.0 / math.sqrt(speaker_dimension)] * speaker_dimension (args.output_dir / "speaker.json").write_text(json.dumps(speaker), encoding="utf-8") diff --git a/tests/special_e2e/run_qwen3_tts_grpo_smoke.sh b/tests/special_e2e/run_qwen3_tts_grpo_smoke.sh index 67e4e37e6..80fa691b0 100755 --- a/tests/special_e2e/run_qwen3_tts_grpo_smoke.sh +++ b/tests/special_e2e/run_qwen3_tts_grpo_smoke.sh @@ -1,5 +1,5 @@ #!/usr/bin/env bash -# Qwen3-TTS full-parameter GRPO e2e smoke: real 0.6B model, two updates. +# Qwen3-TTS full-parameter GRPO e2e smoke: tiny model, two updates. # This validates execution only; the in-process CPU reward is not a quality reward. set -xeuo pipefail @@ -12,28 +12,39 @@ REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" cd "${REPO_ROOT}" NUM_GPUS="${NUM_GPUS:-2}" -[[ "${NUM_GPUS}" == "2" ]] || { echo "Qwen3-TTS smoke requires exactly two GPUs" >&2; exit 2; } +[[ "${NUM_GPUS}" =~ ^[0-9]+$ && "${NUM_GPUS}" -ge 2 ]] || { + echo "Qwen3-TTS smoke requires at least two GPUs" >&2 + exit 2 +} PYTHON_BIN="${PYTHON_BIN:-python3}" -MODEL_REPO="${MODEL_REPO:-Qwen/Qwen3-TTS-12Hz-0.6B-Base}" +MODEL_REPO="${MODEL_REPO:-optimum-intel-internal-testing/tiny-random-qwen3-tts}" +MODEL_REVISION="${MODEL_REVISION:-6374d605b31381cac6f9577f5e742af2f76ba79c}" WORK_DIR="${WORK_DIR:-${TMPDIR:-/tmp}/qwen3_tts_grpo_smoke_${USER:-user}_$$}" DATA_DIR="${WORK_DIR}/data" OUTPUT_DIR="${OUTPUT_DIR:-${WORK_DIR}/output}" mkdir -p "${WORK_DIR}" "${OUTPUT_DIR}" -if ! "${PYTHON_BIN}" -c 'from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSForConditionalGeneration; import onnxruntime, soundfile, librosa, sox'; then - uv pip install --python "${PYTHON_BIN}" -e ".[tts]" - uv pip install --python "${PYTHON_BIN}" --force-reinstall --no-deps \ - "qwen-tts @ git+https://github.com/QwenLM/Qwen3-TTS.git@00969daa8064e23adc9e5f52cdf20cf247f94159" -fi +"${PYTHON_BIN}" -c \ + 'from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSForConditionalGeneration; import onnxruntime, soundfile, librosa, sox' \ + || { echo "Qwen3-TTS smoke dependencies must be installed by gpu-smoke-prepare" >&2; exit 2; } MODEL_PATH="${MODEL_PATH:-}" if [[ -z "${MODEL_PATH}" ]]; then - MODEL_PATH="$("${PYTHON_BIN}" -c \ - "from huggingface_hub import snapshot_download; print(snapshot_download('${MODEL_REPO}'))")" + SOURCE_MODEL_PATH="$("${PYTHON_BIN}" -c \ + 'import sys; from huggingface_hub import snapshot_download; print(snapshot_download(sys.argv[1], revision=sys.argv[2]))' \ + "${MODEL_REPO}" "${MODEL_REVISION}")" + # The published tiny checkpoint has four codebooks; the production actor + # contract has 16, so rebuild its small architecture with random 16-codebook weights. + MODEL_PATH="${WORK_DIR}/tiny-random-qwen3-tts-16-codebooks" + "${PYTHON_BIN}" tests/special_e2e/build_qwen3_tts_tiny_random.py \ + --source-model-path "${SOURCE_MODEL_PATH}" \ + --output-dir "${MODEL_PATH}" fi [[ -f "${MODEL_PATH}/config.json" ]] || { echo "Invalid MODEL_PATH: ${MODEL_PATH}" >&2; exit 2; } -"${PYTHON_BIN}" tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py --output-dir "${DATA_DIR}" +"${PYTHON_BIN}" tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py \ + --output-dir "${DATA_DIR}" \ + --model-config "${MODEL_PATH}/config.json" MODEL_PATH="${MODEL_PATH}" \ TRAIN_FILE="${DATA_DIR}/train.parquet" \ @@ -57,5 +68,13 @@ bash examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh \ trainer.log_val_generations=0 \ "$@" -grep -q "training/global_step.*2" "${OUTPUT_DIR}/train.log" +"${PYTHON_BIN}" - "${OUTPUT_DIR}/train.log" <<'PY' +import re +import sys +from pathlib import Path + +steps = [int(step) for step in re.findall(r"training/global_step:(\d+)(?!\d)", Path(sys.argv[1]).read_text())] +if not steps or max(steps) != 2: + raise SystemExit(f"Expected training/global_step to reach exactly 2, observed {steps!r}") +PY echo "Qwen3-TTS GRPO e2e smoke passed; artifacts: ${WORK_DIR}" diff --git a/tests/special_sanity/test_gpu_smoke_selector_on_cpu.py b/tests/special_sanity/test_gpu_smoke_selector_on_cpu.py new file mode 100644 index 000000000..cce4a4a77 --- /dev/null +++ b/tests/special_sanity/test_gpu_smoke_selector_on_cpu.py @@ -0,0 +1,44 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +import sys +from pathlib import Path + + +def _load_selector(): + path = Path(__file__).parents[1] / "gpu_smoke/select_gpu_smoke_groups.py" + spec = importlib.util.spec_from_file_location("gpu_smoke_selector_under_test", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_qwen3_tts_smoke_files_select_only_omni_e2e_group(): + selector = _load_selector() + + selected = selector.select_group_names( + [ + ".github/qwen_tts_pin.txt", + "examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh", + "tests/special_e2e/build_qwen3_tts_tiny_random.py", + "tests/special_e2e/run_qwen3_tts_grpo_smoke.sh", + "tests/special_e2e/create_dummy_qwen3_tts_grpo_data.py", + "tests/special_e2e/qwen3_tts_dummy_reward.py", + ] + ) + + assert selected == ["ci-e2e-omni"] diff --git a/tests/utils/reward_score/test_audio_http_scorer_client_on_cpu.py b/tests/utils/reward_score/test_audio_http_scorer_client_on_cpu.py index 87c25c24a..43b606313 100644 --- a/tests/utils/reward_score/test_audio_http_scorer_client_on_cpu.py +++ b/tests/utils/reward_score/test_audio_http_scorer_client_on_cpu.py @@ -94,9 +94,10 @@ async def score(request): kwargs = { "solution_audio": (np.zeros(32, dtype=np.float32), 24_000), "ground_truth": "target text", + "data_source": "tts_reward", "server_url": server_url, "max_retries": 2, - "retry_backoff_s": 0, + "retry_backoff": 0, } try: result = await client.compute_score(**kwargs) @@ -141,9 +142,9 @@ def test_response_validation_fails_closed(payload, message): @pytest.mark.parametrize( ("kwargs", "message"), [ - ({"timeout_s": float("nan")}, "finite number"), + ({"timeout": float("nan")}, "finite number"), ({"max_retries": 1.5}, "integer"), - ({"retry_backoff_s": float("inf")}, "finite number"), + ({"retry_backoff": float("inf")}, "finite number"), ], ) def test_invalid_retry_configuration_is_rejected(kwargs, message): @@ -156,3 +157,13 @@ def test_invalid_retry_configuration_is_rejected(kwargs, message): with pytest.raises(ValueError, match=message): asyncio.run(call) + + +def test_unknown_reward_configuration_is_not_silently_swallowed(): + with pytest.raises(TypeError, match="unexpected keyword argument"): + client.compute_score( + solution_audio=(np.zeros(1, dtype=np.float32), 24_000), + ground_truth="text", + server_url="http://127.0.0.1:1/score", + timeout_s=1.0, + ) diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index 21f1425da..766fee727 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -65,6 +65,7 @@ def test_external_module_import_registers_omni_agent_loop(): def test_optional_rollout_hooks_preserve_existing_ar_defaults(): first, final = object(), object() + assert OmniRolloutPipelineBase.supports_async_chunk is True assert OmniRolloutPipelineBase.weight_sync_stage_ids() is None assert OmniRolloutPipelineBase.prepare_engine_prompt([], None, {}) is None assert ( @@ -75,7 +76,9 @@ def test_optional_rollout_hooks_preserve_existing_ar_defaults(): ) is final ) - assert OmniRolloutPipelineBase.combine_engine_outputs([first, final], {}) == (final, {}) + assert OmniRolloutPipelineBase.combine_engine_outputs([final], {}) == (final, {}) + with pytest.raises(NotImplementedError, match="multiple final outputs"): + OmniRolloutPipelineBase.combine_engine_outputs([first, final], {}) with pytest.raises(RuntimeError, match="no outputs"): OmniRolloutPipelineBase.combine_engine_outputs([], {}) @@ -139,6 +142,7 @@ def test_ar_strategy_resolves_qwen3_tts_adapter_and_scopes_weight_sync(monkeypat "output_mode": "ar", "pipeline_name": "qwen3_tts_rl", "pipeline_mode": "full", + "async_chunk": False, } strategy.preprocess_engine_kwargs(engine_kwargs) @@ -146,7 +150,21 @@ def test_ar_strategy_resolves_qwen3_tts_adapter_and_scopes_weight_sync(monkeypat assert deploy_calls == [("qwen3_tts_rl", Qwen3TTSRolloutAdapter, "full")] assert strategy._rollout_adapter is Qwen3TTSRolloutAdapter assert strategy._weight_sync_stage_ids == [0] - assert engine_kwargs == {} + assert engine_kwargs == {"async-chunk": False} + + +@pytest.mark.parametrize("async_chunk", [None, True]) +def test_ar_strategy_requires_non_chunked_qwen3_tts_rollout(async_chunk): + strategy = ARStrategy(SimpleNamespace(_rollout_flags={})) + engine_kwargs = { + "pipeline_name": "qwen3_tts_rl", + "pipeline_mode": "full", + } + if async_chunk is not None: + engine_kwargs["async_chunk"] = async_chunk + + with pytest.raises(ValueError, match="requires async_chunk=false"): + strategy.preprocess_engine_kwargs(engine_kwargs) def test_rollout_adapter_requires_speaker_embedding(): @@ -175,6 +193,13 @@ def test_qwen3_tts_adapters_require_explicit_language(tmp_path): Qwen3TTSTalkerAdapter.configure_model(SimpleNamespace(config=SimpleNamespace()), model_config) +def test_talker_adapter_rejects_remove_padding_before_model_configuration(): + model_config = SimpleNamespace(use_remove_padding=True) + + with pytest.raises(ValueError, match="use_remove_padding=false"): + Qwen3TTSTalkerAdapter.configure_model(SimpleNamespace(), model_config) + + def test_talker_adapter_pads_exact_rollout_fields_for_actor_forward(): model_inputs = {"input_ids": torch.zeros(2, 6, dtype=torch.long)} payloads = [ diff --git a/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py index 17743d824..43caf124c 100644 --- a/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_vllm_omni_strategy_on_cpu.py @@ -13,11 +13,13 @@ # limitations under the License. from argparse import Namespace +from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock import pytest import torch +import yaml from verl_omni.pipelines.qwen3_omni.omni_rollout_adapter import Qwen3OmniRolloutAdapter from verl_omni.pipelines.rollout_media import DiffusionIOSpec, MediaSpec @@ -283,6 +285,47 @@ def test_ar_strategy_preserves_qwen3_omni_thinker_only_contract(): assert output.extra_fields == {"global_steps": 12} +def test_ar_strategy_writes_qwen3_omni_thinker_only_deploy_config(monkeypatch): + monkeypatch.setattr(ar_strategy_module, "get_visible_devices_keyword", lambda: "CUDA_VISIBLE_DEVICES") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") + server = SimpleNamespace( + config=SimpleNamespace( + tensor_model_parallel_size=1, + text_encoder_tp_size=1, + max_model_len=8, + max_num_batched_tokens=8, + ), + _rollout_flags={}, + ) + strategy = ARStrategy(server) + engine_kwargs = { + "pipeline_name": "qwen3_omni_moe", + "pipeline_mode": "thinker_only", + } + + strategy.preprocess_engine_kwargs(engine_kwargs) + + deploy_path = engine_kwargs["deploy-config"] + deploy = yaml.safe_load(Path(deploy_path).read_text(encoding="utf-8")) + assert deploy["pipeline"] == Qwen3OmniRolloutAdapter.get_pipeline_id("thinker_only") + assert [stage["stage_id"] for stage in deploy["stages"]] == [0] + assert strategy._rollout_output_modalities is None + server._temp_deploy_ctx.cleanup() + + +def test_ar_strategy_rejects_qwen3_omni_full_multi_output_without_combiner(): + server = SimpleNamespace(_rollout_flags={}) + strategy = ARStrategy(server) + + with pytest.raises(ValueError, match="multiple final pipeline outputs"): + strategy.preprocess_engine_kwargs( + { + "pipeline_name": "qwen3_omni_moe", + "pipeline_mode": "full", + } + ) + + def test_diffusion_strategy_preserves_engine_argument_preparation(monkeypatch): imported = [] monkeypatch.setattr(diffusion_strategy_module, "import_external_libs", imported.append) diff --git a/tests/workers/test_omni_fsdp_engine_on_cpu.py b/tests/workers/test_omni_fsdp_engine_on_cpu.py index d724cd66d..ec3ab0121 100644 --- a/tests/workers/test_omni_fsdp_engine_on_cpu.py +++ b/tests/workers/test_omni_fsdp_engine_on_cpu.py @@ -379,11 +379,10 @@ def test_collect_lora_params_import_not_from_verl(): # --------------------------------------------------------------------------- -def test_build_module_uses_stage_specific_auto_model_classes(): - """Thinkers use multimodal auto models and talkers use text-to-waveform auto models.""" +def test_build_module_uses_adapter_selected_auto_model_class(): + """Adapters select non-default auto model classes without relying on stage names.""" omni_impl = _get_omni_impl_module() assert omni_impl.AutoModelForMultimodalLM is not None - assert omni_impl.AutoModelForTextToWaveform is not None tree = _parse_omni_impl_ast() import_names = set() @@ -394,22 +393,27 @@ def test_build_module_uses_stage_specific_auto_model_classes(): assert "AutoModelForMultimodalLM" in import_names, ( f"AutoModelForMultimodalLM not imported from transformers; imports: {import_names}" ) - assert "AutoModelForTextToWaveform" in import_names, ( - f"AutoModelForTextToWaveform not imported from transformers; imports: {import_names}" - ) + assert "AutoModelForTextToWaveform" not in import_names assert "AutoModelForCausalLM" not in import_names, "AutoModelForCausalLM should NOT be imported from transformers" -@pytest.mark.parametrize("architecture", ["Qwen3OmniMoeForConditionalGeneration"]) -def test_build_module_calls_adapter_configure_model(architecture): +@pytest.mark.parametrize( + ("architecture", "model_stage"), + [ + ("Qwen3OmniMoeForConditionalGeneration", "thinker"), + ("FutureOmniForConditionalGeneration", "talker"), + ], +) +def test_build_module_calls_adapter_configure_model(architecture, model_stage): """Mock ``from_pretrained``; verify ``adapter_cls.configure_model(module, cfg)``.""" omni_impl = _get_omni_impl_module() - model_config = _make_mock_model_config(architecture=architecture) + model_config = _make_mock_model_config(architecture=architecture, model_stage=model_stage) fake_module = MagicMock(spec=torch.nn.Module) fake_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] fake_adapter_cls = MagicMock() + fake_adapter_cls.auto_model_class = None fake_configured_module = MagicMock(spec=torch.nn.Module) fake_configured_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] fake_adapter_cls.configure_model.return_value = fake_configured_module @@ -431,6 +435,7 @@ def test_build_module_calls_adapter_configure_model(architecture): engine.engine_config = MagicMock() engine.engine_config.model_dtype = None engine.engine_config.forward_only = False + engine.engine_config.strategy = "fsdp2" engine.device_mesh = None result = engine._build_module() @@ -453,6 +458,8 @@ def test_build_module_calls_adapter_configure_model(architecture): def test_build_module_uses_text_to_waveform_auto_model_for_talker(): + from transformers import AutoModelForTextToWaveform + omni_impl = _get_omni_impl_module() model_config = _make_mock_model_config( architecture="Qwen3TTSForConditionalGeneration", @@ -463,12 +470,13 @@ def test_build_module_uses_text_to_waveform_auto_model_for_talker(): configured_module = MagicMock(spec=torch.nn.Module) configured_module.named_parameters.return_value = [("weight", torch.nn.Parameter(torch.randn(2, 2)))] adapter_cls = MagicMock() + adapter_cls.auto_model_class = AutoModelForTextToWaveform adapter_cls.configure_model.return_value = configured_module model_base_mod = sys.modules["verl_omni.pipelines.model_base"] with ( patch.object(model_base_mod.OmniModelBase, "get_class_by_name", return_value=adapter_cls), - patch.object(omni_impl.AutoModelForTextToWaveform, "from_pretrained", return_value=loaded_module) as load, + patch.object(AutoModelForTextToWaveform, "from_pretrained", return_value=loaded_module) as load, patch.object(omni_impl, "get_init_weight_context_manager", return_value=MagicMock()), patch.object(omni_impl.warnings, "catch_warnings", return_value=MagicMock()), patch("verl.utils.torch_dtypes.PrecisionType") as precision_type, @@ -477,6 +485,7 @@ def test_build_module_uses_text_to_waveform_auto_model_for_talker(): engine = object.__new__(omni_impl.OmniFSDPEngine) engine.model_config = model_config engine.engine_config = MagicMock(model_dtype=None, forward_only=False) + engine.engine_config.strategy = "fsdp2" engine.device_mesh = None result = engine._build_module() @@ -491,6 +500,36 @@ def test_build_module_uses_text_to_waveform_auto_model_for_talker(): assert result is configured_module +def test_build_module_rejects_mixed_frozen_parameters_without_fsdp1_orig_params(): + omni_impl = _get_omni_impl_module() + model_config = _make_mock_model_config() + loaded_module = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)) + configured_module = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)) + configured_module[0].requires_grad_(False) + adapter_cls = MagicMock() + adapter_cls.auto_model_class = None + adapter_cls.configure_model.return_value = configured_module + model_base_mod = sys.modules["verl_omni.pipelines.model_base"] + + with ( + patch.object(model_base_mod.OmniModelBase, "get_class_by_name", return_value=adapter_cls), + patch.object(omni_impl.AutoModelForMultimodalLM, "from_pretrained", return_value=loaded_module), + patch.object(omni_impl, "get_init_weight_context_manager", return_value=MagicMock()), + patch.object(omni_impl.warnings, "catch_warnings", return_value=MagicMock()), + patch("verl.utils.torch_dtypes.PrecisionType") as precision_type, + ): + precision_type.to_dtype.side_effect = lambda value: value + engine = object.__new__(omni_impl.OmniFSDPEngine) + engine.model_config = model_config + engine.engine_config = MagicMock(model_dtype=None, forward_only=False) + engine.engine_config.strategy = "fsdp" + engine.engine_config.use_orig_params = False + engine.device_mesh = None + + with pytest.raises(ValueError, match="use_orig_params=true"): + engine._build_module() + + @pytest.mark.parametrize("option", ["use_liger", "use_fused_kernels"]) def test_build_module_rejects_unsupported_optimizations_before_model_load(option): omni_impl = _get_omni_impl_module() diff --git a/verl_omni/pipelines/model_base.py b/verl_omni/pipelines/model_base.py index 8c77bf0e4..f3bb8265e 100644 --- a/verl_omni/pipelines/model_base.py +++ b/verl_omni/pipelines/model_base.py @@ -491,6 +491,7 @@ class Qwen3OmniThinkerAdapter(OmniModelBase): """ _registry: dict[tuple[str, str], type["OmniModelBase"]] = {} + auto_model_class: Any = None @classmethod def register(cls, architecture: str, stage: str = "thinker"): @@ -686,6 +687,7 @@ class Qwen3OmniRolloutAdapter(OmniRolloutPipelineBase): """ _registry: dict[str, type["OmniRolloutPipelineBase"]] = {} + supports_async_chunk = True @classmethod def register(cls, model_type: str): @@ -838,4 +840,8 @@ def combine_engine_outputs(cls, outputs: list, prompt: dict) -> tuple[Any, dict[ """Select the policy output and collect architecture-specific fields.""" if not outputs: raise RuntimeError("The omni rollout engine returned no outputs.") - return outputs[-1], {} + if len(outputs) != 1: + raise NotImplementedError( + "An omni rollout adapter with multiple final outputs must implement combine_engine_outputs()." + ) + return outputs[0], {} diff --git a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py index 4ac274f75..6222bece1 100644 --- a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py @@ -49,6 +49,8 @@ def _load_speaker_vector(path: str) -> list[float]: @OmniRolloutPipelineBase.register(_PIPELINE_ID) class Qwen3TTSRolloutAdapter(OmniRolloutPipelineBase): + supports_async_chunk = False + @classmethod def _check_mode(cls, pipeline_mode): if pipeline_mode != "full": diff --git a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py index 073747b1a..82c5b7954 100644 --- a/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/talker_training_adapter.py @@ -19,6 +19,7 @@ from typing import Any import torch +from transformers import AutoModelForTextToWaveform from verl.utils import tensordict_utils as tu from verl_omni.pipelines.model_base import OmniModelBase @@ -74,6 +75,8 @@ def _set_input_embeddings(self, value): @OmniModelBase.register("Qwen3TTSForConditionalGeneration", stage="talker") class Qwen3TTSTalkerAdapter(OmniModelBase): + auto_model_class = AutoModelForTextToWaveform + @classmethod def register_auto_classes(cls) -> None: from qwen_tts.core.models.configuration_qwen3_tts import Qwen3TTSConfig @@ -93,6 +96,8 @@ def get_strip_modules(cls, model_config): @classmethod def configure_model(cls, module, model_config): + if getattr(model_config, "use_remove_padding", False): + raise ValueError("Qwen3-TTS Talker training requires actor_rollout_ref.model.use_remove_padding=false.") module = super().configure_model(module, model_config) module.config.tts_spk_embed_path = model_config.override_config.get("tts_spk_embed_path") module.config.tts_language = require_auto_language(model_config.override_config.get("tts_language")) diff --git a/verl_omni/reward_loop/reward_manager/audio.py b/verl_omni/reward_loop/reward_manager/audio.py index 6b1f80a27..57ff9c64c 100644 --- a/verl_omni/reward_loop/reward_manager/audio.py +++ b/verl_omni/reward_loop/reward_manager/audio.py @@ -54,15 +54,18 @@ def _extract_audio(cls, extra_info): raise KeyError("Audio reward requires extra_info['audio_sample_rate'] from the rollout.") try: + if isinstance(audio, np.ndarray) and audio.dtype == object: + audio = np.asarray(audio, dtype=np.float32) waveform = torch.as_tensor(audio).detach().float().cpu() except (TypeError, ValueError, RuntimeError) as exc: raise ValueError("Audio reward could not convert the waveform to numeric samples.") from exc while waveform.ndim > 1 and waveform.shape[0] == 1: waveform = waveform[0] - if waveform.ndim == 2: - waveform = waveform.mean(dim=0) - elif waveform.ndim != 1: - raise ValueError(f"Expected audio shape (T,) or (C,T), got {tuple(waveform.shape)}.") + if waveform.ndim != 1: + raise ValueError( + f"Expected one mono waveform with shape (T,) or leading singleton dimensions, " + f"got {tuple(waveform.shape)}." + ) if waveform.numel() == 0: raise ValueError("Audio reward received an empty waveform.") if not torch.isfinite(waveform).all(): @@ -86,8 +89,13 @@ async def run_single(self, data: DataProto) -> dict: batch = item.non_tensor_batch extra_info = self._mapping(batch.get("extra_info", {})) extra_info.update(self._mapping(batch.get("tool_extra_fields"))) - extra_info["num_turns"] = batch.get("__num_turns__", extra_info.get("num_turns")) - extra_info["global_steps"] = batch.get("global_steps", extra_info.get("global_steps", 0)) + for key in ("audio", "audio_sample_rate"): + if key in batch and batch[key] is not None: + extra_info[key] = batch[key] + if "__num_turns__" in batch: + extra_info["num_turns"] = batch["__num_turns__"] + if "global_steps" in batch: + extra_info["global_steps"] = batch["global_steps"] ground_truth = batch["reward_model"]["ground_truth"] audio = self._extract_audio(extra_info) kwargs = { diff --git a/verl_omni/utils/reward_score/audio_http_scorer_client.py b/verl_omni/utils/reward_score/audio_http_scorer_client.py index 0ec5b064e..4f7a5a236 100644 --- a/verl_omni/utils/reward_score/audio_http_scorer_client.py +++ b/verl_omni/utils/reward_score/audio_http_scorer_client.py @@ -110,13 +110,13 @@ async def _session() -> aiohttp.ClientSession: return session -async def _request_score(server_url: str, payload: dict, timeout_s: float) -> dict: +async def _request_score(server_url: str, payload: dict, timeout: float) -> dict: session = await _session() try: async with session.post( server_url, json=payload, - timeout=aiohttp.ClientTimeout(total=timeout_s), + timeout=aiohttp.ClientTimeout(total=timeout), ) as response: if response.status != 200: detail = await response.text() @@ -129,7 +129,7 @@ async def _request_score(server_url: str, payload: dict, timeout_s: float) -> di except (aiohttp.ContentTypeError, ValueError) as exc: raise RuntimeError("Audio scorer returned malformed JSON.") from exc except asyncio.TimeoutError as exc: - raise _RetryableHTTPError(f"Audio scorer timed out after {timeout_s} seconds.") from exc + raise _RetryableHTTPError(f"Audio scorer timed out after {timeout} seconds.") from exc return _validate_response(result) @@ -137,39 +137,39 @@ async def compute_score( solution_audio, ground_truth: str, extra_info: dict | None = None, + data_source: str | None = None, *, server_url: str, - timeout_s: float = 120.0, + timeout: float = 120.0, max_retries: int = 2, - retry_backoff_s: float = 0.5, - **kwargs, + retry_backoff: float = 0.5, ) -> dict: """Send one waveform to an external scorer and return its finite score.""" - del kwargs - if isinstance(timeout_s, bool) or not isinstance(timeout_s, int | float) or not math.isfinite(float(timeout_s)): - raise ValueError("timeout_s must be a finite number.") - if timeout_s <= 0: - raise ValueError("timeout_s must be positive.") + del data_source + if isinstance(timeout, bool) or not isinstance(timeout, int | float) or not math.isfinite(float(timeout)): + raise ValueError("timeout must be a finite number.") + if timeout <= 0: + raise ValueError("timeout must be positive.") if isinstance(max_retries, bool) or not isinstance(max_retries, int): raise ValueError("max_retries must be an integer.") if max_retries < 0: raise ValueError("max_retries must be non-negative.") if ( - isinstance(retry_backoff_s, bool) - or not isinstance(retry_backoff_s, int | float) - or not math.isfinite(float(retry_backoff_s)) + isinstance(retry_backoff, bool) + or not isinstance(retry_backoff, int | float) + or not math.isfinite(float(retry_backoff)) ): - raise ValueError("retry_backoff_s must be a finite number.") - if retry_backoff_s < 0: - raise ValueError("retry_backoff_s must be non-negative.") + raise ValueError("retry_backoff must be a finite number.") + if retry_backoff < 0: + raise ValueError("retry_backoff must be non-negative.") payload = _serialize_request(solution_audio, ground_truth, extra_info) last_error = None for attempt in range(max_retries + 1): try: - return await _request_score(server_url, payload, timeout_s) + return await _request_score(server_url, payload, timeout) except (_RetryableHTTPError, aiohttp.ClientConnectionError, aiohttp.ClientPayloadError) as exc: last_error = exc if attempt < max_retries: - await asyncio.sleep(retry_backoff_s * (2**attempt)) + await asyncio.sleep(retry_backoff * (2**attempt)) raise RuntimeError(f"Audio scoring failed after {max_retries + 1} attempts: {last_error}") from last_error diff --git a/verl_omni/workers/engine/fsdp/omni_impl.py b/verl_omni/workers/engine/fsdp/omni_impl.py index f2181e2e5..6fbdf92c5 100644 --- a/verl_omni/workers/engine/fsdp/omni_impl.py +++ b/verl_omni/workers/engine/fsdp/omni_impl.py @@ -18,7 +18,7 @@ import torch from torch.distributed.tensor import DTensor -from transformers import AutoModelForMultimodalLM, AutoModelForTextToWaveform +from transformers import AutoModelForMultimodalLM from verl.utils.debug import log_gpu_memory_usage from verl.utils.device import get_device_id from verl.utils.fsdp_utils import ( @@ -203,9 +203,7 @@ def _build_module(self): with init_context(), warnings.catch_warnings(): warnings.simplefilter("ignore") - auto_model_cls = ( - AutoModelForTextToWaveform if self.model_config.model_stage == "talker" else AutoModelForMultimodalLM - ) + auto_model_cls = getattr(adapter_cls, "auto_model_class", None) or AutoModelForMultimodalLM module = auto_model_cls.from_pretrained( pretrained_model_name_or_path=self.model_config.local_path, torch_dtype=torch_dtype, @@ -214,6 +212,13 @@ def _build_module(self): ) module = adapter_cls.configure_model(module, self.model_config) + if self.engine_config.strategy == "fsdp" and not self.engine_config.use_orig_params: + trainability = {parameter.requires_grad for parameter in module.parameters()} + if len(trainability) > 1: + raise ValueError( + "FSDP1 requires use_orig_params=true when a model adapter freezes only part of the model." + ) + module.to(torch_dtype) if self.model_config.enable_gradient_checkpointing: diff --git a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py index 660ba8531..fc7c70e7f 100644 --- a/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py +++ b/verl_omni/workers/rollout/vllm_rollout/vllm_omni_ar_strategy.py @@ -87,6 +87,14 @@ def preprocess_engine_kwargs(self, engine_kwargs: dict[str, Any]) -> None: adapter_cls = OmniRolloutPipelineBase.get_class(pipeline_name) if adapter_cls is not None: + async_chunk = engine_kwargs.get("async_chunk", engine_kwargs.get("async-chunk", True)) + if not isinstance(async_chunk, bool): + raise TypeError(f"async_chunk must be a boolean, got {type(async_chunk).__name__}.") + if async_chunk and not adapter_cls.supports_async_chunk: + raise ValueError( + f"{adapter_cls.__name__} requires async_chunk=false because chunked stage outputs " + "cannot be replayed by its actor adapter." + ) self._rollout_adapter = adapter_cls self._write_deploy_config(engine_kwargs, pipeline_name, adapter_cls, self._pipeline_mode) self.server._rollout_flags = adapter_cls.rollout_flags(pipeline_mode=self._pipeline_mode) @@ -126,6 +134,17 @@ def _write_deploy_config( self._rollout_output_modalities = ( list(dict.fromkeys(final_output_types)) if len(final_output_types) > 1 else None ) + adapter_combiner = getattr(adapter_cls.combine_engine_outputs, "__func__", adapter_cls.combine_engine_outputs) + default_combiner = getattr( + OmniRolloutPipelineBase.combine_engine_outputs, + "__func__", + OmniRolloutPipelineBase.combine_engine_outputs, + ) + if self._rollout_output_modalities is not None and adapter_combiner is default_combiner: + raise ValueError( + f"{adapter_cls.__name__} exposes multiple final pipeline outputs but does not implement " + "combine_engine_outputs(); refusing to guess which output contains policy token IDs." + ) self._stage_sampling_constraints = {stage.stage_id: dict(stage.sampling_constraints) for stage in stages} stage_extras = { stage.stage_id: dict(adapter_cls.get_stage_engine_extras(stage.stage_id, pipeline_mode=pipeline_mode)) @@ -144,8 +163,9 @@ def _write_deploy_config( tp_size = self.server.config.tensor_model_parallel_size deploy_dict: dict[str, object] = {"pipeline": pipeline_id} - if "async_chunk" in engine_kwargs: - deploy_dict["async_chunk"] = bool(engine_kwargs["async_chunk"]) + async_chunk = engine_kwargs.get("async_chunk", engine_kwargs.get("async-chunk")) + if async_chunk is not None: + deploy_dict["async_chunk"] = async_chunk if visible_devices: device_count = len([device for device in visible_devices.split(",") if device.strip()]) From c17bc7f3c8813f9d4e56b16339661b663e86acba Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Fri, 4 Sep 2026 01:45:29 +0800 Subject: [PATCH 24/28] [trainer, doc] fix: use FP32 master weights for Qwen3-TTS GRPO Co-authored-by: OpenAI Codex Signed-off-by: dongbo910220 <1275604947@qq.com> --- examples/grpo_trainer/qwen3_tts/README.md | 20 +++++++++---------- .../qwen3_tts/run_qwen3_tts_grpo.sh | 13 ++++++++---- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/examples/grpo_trainer/qwen3_tts/README.md b/examples/grpo_trainer/qwen3_tts/README.md index b620da7a9..f4204eb32 100644 --- a/examples/grpo_trainer/qwen3_tts/README.md +++ b/examples/grpo_trainer/qwen3_tts/README.md @@ -1,6 +1,6 @@ # Qwen3-TTS GRPO with an audio reward -Last updated: 08/31/2026. +Last updated: 09/04/2026. This example full-parameter tunes the codec-0 policy of `Qwen/Qwen3-TTS-12Hz-0.6B-Base`. It uses verl's stock GRPO advantage, @@ -118,16 +118,14 @@ OUTPUT_DIR=/path/to/output \ bash examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh ``` -The example defaults are `B=4`, `G=8`, `lr=2e-7`, direct `low_var_kl` with -coefficient `0.12`, two GPUs, and 500 updates. These are recipe values, not -algorithm requirements. `norm_adv_by_std_in_grpo` remains at the upstream -default. Actor, reference, rollout, and floating synchronized weights all use -BF16; synchronized integer buffers keep their integer dtype. Check selected-token -`diff_mean` and Pearson after synchronization as execution-consistency diagnostics, -not as evidence of speech quality or FP32-equivalent numerics. Prefer -`diff_mean < 0.005`; values from `0.005` to `0.01` require high Pearson and tail -inspection, while sustained values at or above `0.01` should stop the run for -investigation. +The example defaults are `B=4`, `G=8`, `lr=1e-6` with 10 warmup steps and a +constant schedule, direct `low_var_kl` with coefficient `0.12`, two GPUs, and +500 updates. These are recipe values, not algorithm requirements. +`norm_adv_by_std_in_grpo` remains at the upstream default. The actor and +reference keep persistent parameters in FP32, while FSDP uses BF16 parameters +for forward and backward computation with FP32 gradient reduction and buffers. +The actor's AdamW state therefore remains FP32, and rollout inference remains +BF16. For a two-update implementation smoke test: diff --git a/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh b/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh index cf79effa8..1a42a6b1d 100755 --- a/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh +++ b/examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh @@ -63,14 +63,17 @@ export VLLM_USE_FLASHINFER_SAMPLER=0 actor_rollout_ref.actor.loss_agg_mode=seq-mean-token-mean \ actor_rollout_ref.actor.use_torch_compile=false \ actor_rollout_ref.actor.policy_loss.loss_mode=vanilla \ - actor_rollout_ref.actor.optim.lr=2.0e-7 \ - actor_rollout_ref.actor.optim.lr_warmup_steps=0 \ + actor_rollout_ref.actor.optim.lr=1.0e-6 \ + actor_rollout_ref.actor.optim.lr_warmup_steps=10 \ + actor_rollout_ref.actor.optim.lr_scheduler_type=constant \ actor_rollout_ref.actor.optim.weight_decay=0.0 \ actor_rollout_ref.actor.optim.clip_grad=1.0 \ "actor_rollout_ref.actor.data_loader_seed=${SEED}" \ "actor_rollout_ref.actor.fsdp_config.fsdp_size=${NUM_GPUS}" \ "actor_rollout_ref.actor.fsdp_config.seed=${SEED}" \ - actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \ + actor_rollout_ref.actor.fsdp_config.model_dtype=float32 \ + actor_rollout_ref.actor.fsdp_config.dtype=bfloat16 \ + '+actor_rollout_ref.actor.fsdp_config.mixed_precision={param_dtype:bf16,reduce_dtype:fp32,buffer_dtype:fp32}' \ actor_rollout_ref.actor.fsdp_config.param_offload=false \ actor_rollout_ref.actor.fsdp_config.optimizer_offload=false \ actor_rollout_ref.actor.fsdp_config.use_orig_params=true \ @@ -112,7 +115,9 @@ export VLLM_USE_FLASHINFER_SAMPLER=0 actor_rollout_ref.ref.strategy=fsdp \ "actor_rollout_ref.ref.fsdp_config.fsdp_size=${NUM_GPUS}" \ "actor_rollout_ref.ref.fsdp_config.seed=${SEED}" \ - actor_rollout_ref.ref.fsdp_config.model_dtype=bfloat16 \ + actor_rollout_ref.ref.fsdp_config.model_dtype=float32 \ + actor_rollout_ref.ref.fsdp_config.dtype=bfloat16 \ + '+actor_rollout_ref.ref.fsdp_config.mixed_precision={param_dtype:bf16,reduce_dtype:fp32,buffer_dtype:fp32}' \ actor_rollout_ref.ref.fsdp_config.param_offload=false \ actor_rollout_ref.ref.fsdp_config.use_orig_params=true \ actor_rollout_ref.ref.fsdp_config.wrap_policy.min_num_params=0 \ From 7968c76ee169ea021941497ad44447216802da99 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Wed, 9 Sep 2026 05:49:55 +0800 Subject: [PATCH 25/28] [rollout, ci, tests] fix: preserve Qwen3-TTS Code2Wav placeholders Co-authored-by: GitHub Copilot Signed-off-by: dongbo910220 <1275604947@qq.com> --- .github/actions/gpu-smoke-prepare/action.yml | 5 +- tests/gpu_smoke/select_gpu_smoke_groups.py | 2 - .../test_qwen3_tts_rollout_on_cpu.py | 26 +++++-- .../qwen3_tts/omni_rollout_adapter.py | 70 ++++++++++++++++++- 4 files changed, 94 insertions(+), 9 deletions(-) diff --git a/.github/actions/gpu-smoke-prepare/action.yml b/.github/actions/gpu-smoke-prepare/action.yml index 5a0a81521..0cc7d6c34 100644 --- a/.github/actions/gpu-smoke-prepare/action.yml +++ b/.github/actions/gpu-smoke-prepare/action.yml @@ -30,7 +30,10 @@ runs: # TODO: rm --no-deps when VeOmni supports the vLLM torch pin (torch 2.13 as of vLLM 0.28) uv pip install --system --break-system-packages veomni==0.1.11 --no-deps uv pip install --system --break-system-packages torchcodec librosa soundfile av audioread - uv pip install --system --break-system-packages "transformers[mistral-common]==5.14.1" + # vllm-omni currently resolves accelerate 1.12, while the pinned + # Transformers stack and this project require accelerate >=1.14. + uv pip install --system --break-system-packages \ + "transformers[mistral-common]==5.14.1" "accelerate>=1.14.0" # The pinned Qwen3-TTS TF5 source declares Transformers >=5.15.1, while # this repository deliberately caps Transformers at 5.14.1. Install its # exact tested source without letting that metadata replace the CI stack. diff --git a/tests/gpu_smoke/select_gpu_smoke_groups.py b/tests/gpu_smoke/select_gpu_smoke_groups.py index 72c980e96..764a79d8c 100644 --- a/tests/gpu_smoke/select_gpu_smoke_groups.py +++ b/tests/gpu_smoke/select_gpu_smoke_groups.py @@ -72,12 +72,10 @@ class SmokeGroup: "tests/pipelines/test_qwen3_tts*", "tests/special_e2e/*omni*", "tests/special_e2e/*qwen3_tts*", - "tests/utils/reward_score/test_audio_http_scorer_client_on_cpu.py", "verl_omni/models/transformers/qwen3_omni_thinker.py", "verl_omni/pipelines/qwen3_tts/**", "verl_omni/trainer/config/omni/**", "verl_omni/trainer/omni/**", - "verl_omni/utils/reward_score/audio_http_scorer_client.py", ), "ci-e2e-diffusion": ( "tests/gpu_smoke/run_gpu_smoke_diffusion_e2e.sh", diff --git a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py index 461dd0c6e..b412bffba 100644 --- a/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py +++ b/tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py @@ -31,7 +31,10 @@ from verl_omni.agent_loop.single_turn_agent_loop import OmniSingleTurnAgentLoop from verl_omni.pipelines.model_base import OmniRolloutPipelineBase from verl_omni.pipelines.qwen3_tts import omni_rollout_adapter -from verl_omni.pipelines.qwen3_tts.omni_rollout_adapter import Qwen3TTSRolloutAdapter +from verl_omni.pipelines.qwen3_tts.omni_rollout_adapter import ( + Qwen3TTSRolloutAdapter, + prepare_code2wav_input_for_policy_replay, +) from verl_omni.pipelines.qwen3_tts.rollout_utils import QWEN3_TTS_REPLAY_KEY from verl_omni.pipelines.qwen3_tts.talker_training_adapter import Qwen3TTSTalkerAdapter from verl_omni.workers.rollout.vllm_rollout.vllm_omni_ar_strategy import ARStrategy @@ -131,9 +134,24 @@ def test_rollout_adapter_builds_unique_prompt_and_scopes_weight_sync(tmp_path): assert first["additional_information"]["text"] == ["first text"] assert first["cache_salt"] != second["cache_salt"] assert Qwen3TTSRolloutAdapter.weight_sync_stage_ids("full") == [0] - assert [ - stage.final_output_type for stage in Qwen3TTSRolloutAdapter.build_stage_configs("full") if stage.final_output - ] == ["latent", "audio"] + stages = Qwen3TTSRolloutAdapter.build_stage_configs("full") + assert [stage.final_output_type for stage in stages if stage.final_output] == ["latent", "audio"] + assert stages[0].sampling_constraints["min_tokens"] == 2 + assert stages[1].sync_process_input_func.endswith(".prepare_code2wav_input_for_policy_replay") + + +def test_code2wav_placeholder_uses_retained_policy_token_count_without_mutation(): + completion = SimpleNamespace( + cumulative_token_ids=[101, 201, 202, 2150], + multimodal_output=None, + ) + source_output = SimpleNamespace(finished=True, outputs=[completion]) + + prepared = prepare_code2wav_input_for_policy_replay([source_output]) + + assert len(prepared) == 1 + assert len(prepared[0]["prompt_token_ids"]) == 3 * 16 + assert completion.multimodal_output is None def test_ar_strategy_resolves_qwen3_tts_adapter(monkeypatch): diff --git a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py index cb86419cc..266e0338e 100644 --- a/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py +++ b/verl_omni/pipelines/qwen3_tts/omni_rollout_adapter.py @@ -16,6 +16,7 @@ import hashlib from dataclasses import replace from functools import lru_cache +from types import SimpleNamespace import torch from vllm_omni.config.pipeline_registry import register_pipeline @@ -25,6 +26,7 @@ from verl_omni.pipelines.model_base import OmniRolloutPipelineBase from verl_omni.pipelines.qwen3_tts.rollout_utils import QWEN3_TTS_REPLAY_KEY, align_audio_codes from verl_omni.pipelines.qwen3_tts.talker_forward import ( + NUM_CODEBOOKS, TEXT_PROMPT_TRAILER_TOKENS, build_assistant_text, load_speaker_xvector, @@ -32,12 +34,76 @@ ) _PIPELINE_ID = "qwen3_tts_rl" + + +def prepare_code2wav_input_for_policy_replay( + source_outputs: list, + prompt=None, + _requires_multimodal_data: bool = False, +) -> list: + """Size Code2Wav placeholders from the retained Talker policy trajectory. + + The pinned sync engine keeps codec values on the worker connector while + the orchestrator sees only cumulative policy tokens. Code2Wav still needs + a placeholder sized to 16 codebooks for every generated codec frame. + """ + from vllm_omni.model_executor.stage_input_processors.qwen3_tts import talker2code2wav_token_only + + normalized_outputs = [] + for source_output in source_outputs: + completions = getattr(source_output, "outputs", None) + if not getattr(source_output, "finished", False): + normalized_outputs.append(source_output) + continue + if not completions: + raise RuntimeError("Qwen3-TTS Talker finished without a completion.") + + completion = completions[0] + token_ids = list(getattr(completion, "cumulative_token_ids", None) or []) + if len(token_ids) < 2: + raise RuntimeError("Qwen3-TTS Talker produced no codec frames for Code2Wav.") + frame_count = len(token_ids) - 1 + # This tensor supplies only the placeholder shape. The worker + # connector sends the actual 16-codebook values to Code2Wav. + multimodal_output = { + "codes": {"audio": torch.ones((frame_count, NUM_CODEBOOKS), dtype=torch.long)}, + } + + completion_proxy = SimpleNamespace( + cumulative_token_ids=token_ids, + multimodal_output=multimodal_output, + ) + normalized_outputs.append( + SimpleNamespace( + finished=source_output.finished, + outputs=[completion_proxy], + ) + ) + + return talker2code2wav_token_only( + normalized_outputs, + prompt=prompt, + _requires_multimodal_data=_requires_multimodal_data, + ) + + QWEN3_TTS_RL_PIPELINE = PipelineConfig( model_type=_PIPELINE_ID, model_arch=QWEN3_TTS_PIPELINE.model_arch, stages=( - replace(QWEN3_TTS_PIPELINE.stages[0], final_output=True, final_output_type="latent"), - QWEN3_TTS_PIPELINE.stages[1], + replace( + QWEN3_TTS_PIPELINE.stages[0], + final_output=True, + final_output_type="latent", + sampling_constraints={ + **QWEN3_TTS_PIPELINE.stages[0].sampling_constraints, + "min_tokens": 2, + }, + ), + replace( + QWEN3_TTS_PIPELINE.stages[1], + sync_process_input_func=f"{__name__}.prepare_code2wav_input_for_policy_replay", + ), ), ) From 902de65323db82590b7f9457219bc1c231c6c8a8 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Wed, 9 Sep 2026 13:58:27 +0800 Subject: [PATCH 26/28] [ci, cfg, doc] refactor: simplify Qwen3-TTS dependencies Install qwen-tts through the shared omni optional dependency group, remove the repository-specific source pin, and keep GPU smoke setup and documentation aligned with the package metadata. Co-authored-by: GitHub Copilot Signed-off-by: dongbo910220 <1275604947@qq.com> --- .github/actions/gpu-smoke-prepare/action.yml | 11 +--------- .github/qwen_tts_pin.txt | 1 - .github/workflows/gpu_smoke.yml | 2 -- examples/grpo_trainer/qwen3_tts/README.md | 20 ++++++------------- pyproject.toml | 5 ++--- tests/gpu_smoke/select_gpu_smoke_groups.py | 1 - .../test_qwen3_tts_package_on_cpu.py | 2 +- .../test_gpu_smoke_selector_on_cpu.py | 1 - 8 files changed, 10 insertions(+), 33 deletions(-) delete mode 100644 .github/qwen_tts_pin.txt diff --git a/.github/actions/gpu-smoke-prepare/action.yml b/.github/actions/gpu-smoke-prepare/action.yml index 0cc7d6c34..b2694e32b 100644 --- a/.github/actions/gpu-smoke-prepare/action.yml +++ b/.github/actions/gpu-smoke-prepare/action.yml @@ -20,7 +20,7 @@ runs: export UV_CACHE_DIR="${UV_CACHE_DIR:-${HOME}/.cache/uv}" git config --global http.postBuffer 524288000 || true # Base image PyPI mirror may lag pypi.org (e.g. kernels, fa3-fwd). - uv pip install --system --break-system-packages ".[gpu,dev,audio,tts]" + uv pip install --system --break-system-packages ".[gpu,dev,audio,omni]" uv pip install --system --break-system-packages "vllm-omni @ git+https://github.com/vllm-project/vllm-omni.git@$(cat .github/vllm_omni_pin.txt)" uv pip install --system --break-system-packages --no-deps --reinstall "verl @ git+https://github.com/verl-project/verl.git@$(cat .github/verl_pin.txt)" uv pip install --system --break-system-packages TransferQueue==0.1.8 @@ -30,15 +30,6 @@ runs: # TODO: rm --no-deps when VeOmni supports the vLLM torch pin (torch 2.13 as of vLLM 0.28) uv pip install --system --break-system-packages veomni==0.1.11 --no-deps uv pip install --system --break-system-packages torchcodec librosa soundfile av audioread - # vllm-omni currently resolves accelerate 1.12, while the pinned - # Transformers stack and this project require accelerate >=1.14. - uv pip install --system --break-system-packages \ - "transformers[mistral-common]==5.14.1" "accelerate>=1.14.0" - # The pinned Qwen3-TTS TF5 source declares Transformers >=5.15.1, while - # this repository deliberately caps Transformers at 5.14.1. Install its - # exact tested source without letting that metadata replace the CI stack. - uv pip install --system --break-system-packages --no-deps \ - "qwen-tts @ git+https://github.com/QwenLM/Qwen3-TTS.git@$(cat .github/qwen_tts_pin.txt)" # NCCL checkpoint engine (diffusion v1 separate_async weight sync). uv pip install --system --break-system-packages pyzmq uv pip install --system --break-system-packages cupy-cuda12x || uv pip install --system --break-system-packages cupy-cuda13x diff --git a/.github/qwen_tts_pin.txt b/.github/qwen_tts_pin.txt deleted file mode 100644 index c0af4ea17..000000000 --- a/.github/qwen_tts_pin.txt +++ /dev/null @@ -1 +0,0 @@ -00969daa8064e23adc9e5f52cdf20cf247f94159 diff --git a/.github/workflows/gpu_smoke.yml b/.github/workflows/gpu_smoke.yml index a386adeb3..06888ce8d 100644 --- a/.github/workflows/gpu_smoke.yml +++ b/.github/workflows/gpu_smoke.yml @@ -14,7 +14,6 @@ on: - "tests/special_e2e/**" - "pyproject.toml" - .github/workflows/gpu_smoke.yml - - .github/qwen_tts_pin.txt - .github/vllm_omni_pin.txt - .github/verl_pin.txt - .github/actions/gpu-smoke-prepare/** @@ -33,7 +32,6 @@ on: - "tests/special_e2e/**" - "pyproject.toml" - .github/workflows/gpu_smoke.yml - - .github/qwen_tts_pin.txt - .github/vllm_omni_pin.txt - .github/verl_pin.txt - .github/actions/gpu-smoke-prepare/** diff --git a/examples/grpo_trainer/qwen3_tts/README.md b/examples/grpo_trainer/qwen3_tts/README.md index b494a24c0..d8d372ba7 100644 --- a/examples/grpo_trainer/qwen3_tts/README.md +++ b/examples/grpo_trainer/qwen3_tts/README.md @@ -36,22 +36,14 @@ Install the engine before the training stack: ```bash uv pip install -e ".[gpu]" --torch-backend=auto uv pip install "vllm-omni @ git+https://github.com/vllm-project/vllm-omni.git@$(cat .github/vllm_omni_pin.txt)" -uv pip install -e ".[tts,train,dev]" -uv pip install --no-deps \ - "qwen-tts @ git+https://github.com/QwenLM/Qwen3-TTS.git@$(cat .github/qwen_tts_pin.txt)" +uv pip install -e ".[omni,train,dev]" ``` -The pinned Qwen3-TTS revision is the upstream Transformers 5 support change -from Qwen3-TTS PR #360. Its package metadata requires Transformers 5.15.1 or -newer, while this repository intentionally caps Transformers at 5.14.1. The -`--no-deps` flag preserves that repository-wide cap; the `tts` extra explicitly -owns the runtime dependencies, including `torchaudio==2.11.0` to match vLLM's -Torch pin, and CI tests the exact Qwen3-TTS revision from -`.github/qwen_tts_pin.txt` on this stack. The released `qwen-tts==0.1.1` source -targets Transformers 4.57 and cannot be imported unchanged here. The adapter -registers the upstream config and model with `AutoConfig` and -`AutoModelForTextToWaveform`; it does not carry a local Transformers -compatibility layer. The system `sox` executable is also required by qwen-tts. +The `omni` extra installs qwen-tts and its runtime dependencies, including +`torchaudio==2.11.0` to match vLLM's Torch pin. The adapter registers the +upstream config and model with `AutoConfig` and `AutoModelForTextToWaveform`; it +does not carry a local Transformers compatibility layer. The system `sox` +executable is also required by qwen-tts. ## Data diff --git a/pyproject.toml b/pyproject.toml index 51d8a4e8b..8007f89f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,9 +59,8 @@ audio = [ "qwen-omni-utils>=0.0.9", "audioread", ] -# Install qwen-tts separately from the pinned upstream Transformers 5 support -# revision. The released 0.1.1 source and metadata target Transformers 4.57.3. -tts = [ +omni = [ + "qwen-tts", "einops>=0.8.0", "librosa>=0.10.2", "onnxruntime>=1.20.0", diff --git a/tests/gpu_smoke/select_gpu_smoke_groups.py b/tests/gpu_smoke/select_gpu_smoke_groups.py index 0e5de914f..1b9bf2c9d 100644 --- a/tests/gpu_smoke/select_gpu_smoke_groups.py +++ b/tests/gpu_smoke/select_gpu_smoke_groups.py @@ -66,7 +66,6 @@ class SmokeGroup: "verl_omni/workers/**", ), "ci-e2e-omni": ( - ".github/qwen_tts_pin.txt", "examples/grpo_trainer/qwen3_tts/**", "tests/gpu_smoke/run_gpu_smoke_omni_e2e.sh", "tests/pipelines/test_qwen3_tts*", diff --git a/tests/pipelines/test_qwen3_tts_package_on_cpu.py b/tests/pipelines/test_qwen3_tts_package_on_cpu.py index a6c6fa069..49d45ae52 100644 --- a/tests/pipelines/test_qwen3_tts_package_on_cpu.py +++ b/tests/pipelines/test_qwen3_tts_package_on_cpu.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Check the pinned upstream qwen-tts TF5 source on the repository stack.""" +"""Check the optional qwen-tts package on the repository stack.""" import importlib.util import json diff --git a/tests/special_sanity/test_gpu_smoke_selector_on_cpu.py b/tests/special_sanity/test_gpu_smoke_selector_on_cpu.py index cce4a4a77..ceb42597a 100644 --- a/tests/special_sanity/test_gpu_smoke_selector_on_cpu.py +++ b/tests/special_sanity/test_gpu_smoke_selector_on_cpu.py @@ -32,7 +32,6 @@ def test_qwen3_tts_smoke_files_select_only_omni_e2e_group(): selected = selector.select_group_names( [ - ".github/qwen_tts_pin.txt", "examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh", "tests/special_e2e/build_qwen3_tts_tiny_random.py", "tests/special_e2e/run_qwen3_tts_grpo_smoke.sh", From 0b54a8d653ba29a5ecb3f4df1e1b8fcb18617891 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Wed, 9 Sep 2026 17:36:50 +0800 Subject: [PATCH 27/28] [ci, cfg, doc] fix: install compatible Qwen-TTS source Use the tested upstream Transformers 5 Qwen-TTS revision without allowing its newer dependency metadata to replace the repository stack. Exercise the real package contract in CPU CI and route pin changes through the Omni GPU smoke. Co-authored-by: GitHub Copilot Signed-off-by: dongbo910220 <1275604947@qq.com> --- .github/actions/gpu-smoke-prepare/action.yml | 8 ++++++++ .github/qwen_tts_pin.txt | 1 + .github/workflows/cpu_unit_tests.yml | 7 ++++++- .github/workflows/gpu_smoke.yml | 2 ++ examples/grpo_trainer/qwen3_tts/README.md | 18 +++++++++++++----- pyproject.toml | 3 ++- tests/gpu_smoke/select_gpu_smoke_groups.py | 1 + .../pipelines/test_qwen3_tts_package_on_cpu.py | 2 +- .../test_gpu_smoke_selector_on_cpu.py | 1 + 9 files changed, 35 insertions(+), 8 deletions(-) create mode 100644 .github/qwen_tts_pin.txt diff --git a/.github/actions/gpu-smoke-prepare/action.yml b/.github/actions/gpu-smoke-prepare/action.yml index b2694e32b..2a033ca7f 100644 --- a/.github/actions/gpu-smoke-prepare/action.yml +++ b/.github/actions/gpu-smoke-prepare/action.yml @@ -30,6 +30,14 @@ runs: # TODO: rm --no-deps when VeOmni supports the vLLM torch pin (torch 2.13 as of vLLM 0.28) uv pip install --system --break-system-packages veomni==0.1.11 --no-deps uv pip install --system --break-system-packages torchcodec librosa soundfile av audioread + # Re-apply repository constraints after vllm-omni installs its runtime + # dependencies, which otherwise downgrade Accelerate to 1.12. + uv pip install --system --break-system-packages \ + "transformers[mistral-common]==5.14.1" "accelerate>=1.14.0" + # The released qwen-tts package targets Transformers 4.57.3. Install the + # tested upstream Transformers 5 source without changing this repo's stack. + uv pip install --system --break-system-packages --no-deps --reinstall \ + "qwen-tts @ git+https://github.com/QwenLM/Qwen3-TTS.git@$(cat .github/qwen_tts_pin.txt)" # NCCL checkpoint engine (diffusion v1 separate_async weight sync). uv pip install --system --break-system-packages pyzmq uv pip install --system --break-system-packages cupy-cuda12x || uv pip install --system --break-system-packages cupy-cuda13x diff --git a/.github/qwen_tts_pin.txt b/.github/qwen_tts_pin.txt new file mode 100644 index 000000000..c0af4ea17 --- /dev/null +++ b/.github/qwen_tts_pin.txt @@ -0,0 +1 @@ +00969daa8064e23adc9e5f52cdf20cf247f94159 diff --git a/.github/workflows/cpu_unit_tests.yml b/.github/workflows/cpu_unit_tests.yml index 4dd2b7485..8ee56e0ae 100644 --- a/.github/workflows/cpu_unit_tests.yml +++ b/.github/workflows/cpu_unit_tests.yml @@ -10,6 +10,7 @@ on: - "tests/**/*_on_cpu.py" - "pyproject.toml" - .github/workflows/cpu_unit_tests.yml + - .github/qwen_tts_pin.txt - .github/vllm_omni_pin.txt - .github/verl_pin.txt pull_request: @@ -22,6 +23,7 @@ on: - "tests/**/*_on_cpu.py" - "pyproject.toml" - .github/workflows/cpu_unit_tests.yml + - .github/qwen_tts_pin.txt - .github/vllm_omni_pin.txt - .github/verl_pin.txt @@ -56,6 +58,7 @@ jobs: cache: pip cache-dependency-path: | pyproject.toml + .github/qwen_tts_pin.txt .github/vllm_omni_pin.txt .github/verl_pin.txt - name: Install dependencies @@ -64,7 +67,9 @@ jobs: pip install "vllm-omni @ git+https://github.com/vllm-project/vllm-omni.git@$(cat .github/vllm_omni_pin.txt)" pip install TransferQueue==0.1.8 pip install --no-deps "verl @ git+https://github.com/verl-project/verl.git@$(cat .github/verl_pin.txt)" - pip install ".[dev]" + pip install ".[omni,dev]" + pip install --no-deps \ + "qwen-tts @ git+https://github.com/QwenLM/Qwen3-TTS.git@$(cat .github/qwen_tts_pin.txt)" pip install --no-deps -e . - name: Verify the documented train install resolves run: | diff --git a/.github/workflows/gpu_smoke.yml b/.github/workflows/gpu_smoke.yml index 06888ce8d..a386adeb3 100644 --- a/.github/workflows/gpu_smoke.yml +++ b/.github/workflows/gpu_smoke.yml @@ -14,6 +14,7 @@ on: - "tests/special_e2e/**" - "pyproject.toml" - .github/workflows/gpu_smoke.yml + - .github/qwen_tts_pin.txt - .github/vllm_omni_pin.txt - .github/verl_pin.txt - .github/actions/gpu-smoke-prepare/** @@ -32,6 +33,7 @@ on: - "tests/special_e2e/**" - "pyproject.toml" - .github/workflows/gpu_smoke.yml + - .github/qwen_tts_pin.txt - .github/vllm_omni_pin.txt - .github/verl_pin.txt - .github/actions/gpu-smoke-prepare/** diff --git a/examples/grpo_trainer/qwen3_tts/README.md b/examples/grpo_trainer/qwen3_tts/README.md index d8d372ba7..57ef3b4be 100644 --- a/examples/grpo_trainer/qwen3_tts/README.md +++ b/examples/grpo_trainer/qwen3_tts/README.md @@ -37,13 +37,21 @@ Install the engine before the training stack: uv pip install -e ".[gpu]" --torch-backend=auto uv pip install "vllm-omni @ git+https://github.com/vllm-project/vllm-omni.git@$(cat .github/vllm_omni_pin.txt)" uv pip install -e ".[omni,train,dev]" +uv pip install --no-deps --reinstall \ + "qwen-tts @ git+https://github.com/QwenLM/Qwen3-TTS.git@$(cat .github/qwen_tts_pin.txt)" ``` -The `omni` extra installs qwen-tts and its runtime dependencies, including -`torchaudio==2.11.0` to match vLLM's Torch pin. The adapter registers the -upstream config and model with `AutoConfig` and `AutoModelForTextToWaveform`; it -does not carry a local Transformers compatibility layer. The system `sox` -executable is also required by qwen-tts. +The pinned Qwen3-TTS revision is the upstream Transformers 5 support change +from Qwen3-TTS PR #360. Its package metadata requires Transformers 5.15.1 or +newer, while this repository intentionally caps Transformers at 5.14.1. The +`--no-deps` flag preserves that repository-wide cap; the `omni` extra owns the +runtime dependencies, including `torchaudio==2.11.0` to match vLLM's Torch pin, +and CI tests the exact Qwen3-TTS revision from `.github/qwen_tts_pin.txt` on this +stack. The released `qwen-tts==0.1.1` source targets Transformers 4.57.3 and +cannot be imported unchanged here. The adapter registers the upstream config +and model with `AutoConfig` and `AutoModelForTextToWaveform`; it does not carry +a local Transformers compatibility layer. The system `sox` executable is also +required by qwen-tts. ## Data diff --git a/pyproject.toml b/pyproject.toml index 8007f89f3..08fbe060f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,8 +59,9 @@ audio = [ "qwen-omni-utils>=0.0.9", "audioread", ] +# Install qwen-tts separately from the pinned upstream Transformers 5 support +# revision. The released 0.1.1 source and metadata target Transformers 4.57.3. omni = [ - "qwen-tts", "einops>=0.8.0", "librosa>=0.10.2", "onnxruntime>=1.20.0", diff --git a/tests/gpu_smoke/select_gpu_smoke_groups.py b/tests/gpu_smoke/select_gpu_smoke_groups.py index 1b9bf2c9d..0e5de914f 100644 --- a/tests/gpu_smoke/select_gpu_smoke_groups.py +++ b/tests/gpu_smoke/select_gpu_smoke_groups.py @@ -66,6 +66,7 @@ class SmokeGroup: "verl_omni/workers/**", ), "ci-e2e-omni": ( + ".github/qwen_tts_pin.txt", "examples/grpo_trainer/qwen3_tts/**", "tests/gpu_smoke/run_gpu_smoke_omni_e2e.sh", "tests/pipelines/test_qwen3_tts*", diff --git a/tests/pipelines/test_qwen3_tts_package_on_cpu.py b/tests/pipelines/test_qwen3_tts_package_on_cpu.py index 49d45ae52..a6c6fa069 100644 --- a/tests/pipelines/test_qwen3_tts_package_on_cpu.py +++ b/tests/pipelines/test_qwen3_tts_package_on_cpu.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Check the optional qwen-tts package on the repository stack.""" +"""Check the pinned upstream qwen-tts TF5 source on the repository stack.""" import importlib.util import json diff --git a/tests/special_sanity/test_gpu_smoke_selector_on_cpu.py b/tests/special_sanity/test_gpu_smoke_selector_on_cpu.py index ceb42597a..cce4a4a77 100644 --- a/tests/special_sanity/test_gpu_smoke_selector_on_cpu.py +++ b/tests/special_sanity/test_gpu_smoke_selector_on_cpu.py @@ -32,6 +32,7 @@ def test_qwen3_tts_smoke_files_select_only_omni_e2e_group(): selected = selector.select_group_names( [ + ".github/qwen_tts_pin.txt", "examples/grpo_trainer/qwen3_tts/run_qwen3_tts_grpo.sh", "tests/special_e2e/build_qwen3_tts_tiny_random.py", "tests/special_e2e/run_qwen3_tts_grpo_smoke.sh", From 1bf92cb3f5c858bd9c4c46cda0f86d37ec2fe787 Mon Sep 17 00:00:00 2001 From: dongbo910220 <1275604947@qq.com> Date: Wed, 9 Sep 2026 18:23:16 +0800 Subject: [PATCH 28/28] [ci] refactor: use project dependency overrides in smoke setup Remove the redundant Transformers and Accelerate reinstall from the GPU smoke action. The existing uv override-dependencies already resolve the repository-supported versions. Co-authored-by: GitHub Copilot Signed-off-by: dongbo910220 <1275604947@qq.com> --- .github/actions/gpu-smoke-prepare/action.yml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/.github/actions/gpu-smoke-prepare/action.yml b/.github/actions/gpu-smoke-prepare/action.yml index 2a033ca7f..4e610f60a 100644 --- a/.github/actions/gpu-smoke-prepare/action.yml +++ b/.github/actions/gpu-smoke-prepare/action.yml @@ -30,10 +30,6 @@ runs: # TODO: rm --no-deps when VeOmni supports the vLLM torch pin (torch 2.13 as of vLLM 0.28) uv pip install --system --break-system-packages veomni==0.1.11 --no-deps uv pip install --system --break-system-packages torchcodec librosa soundfile av audioread - # Re-apply repository constraints after vllm-omni installs its runtime - # dependencies, which otherwise downgrade Accelerate to 1.12. - uv pip install --system --break-system-packages \ - "transformers[mistral-common]==5.14.1" "accelerate>=1.14.0" # The released qwen-tts package targets Transformers 4.57.3. Install the # tested upstream Transformers 5 source without changing this repo's stack. uv pip install --system --break-system-packages --no-deps --reinstall \