[3/N][model, rollout, reward] feat: add Qwen3-TTS GRPO support - #428
Conversation
|
@knlnguyen1802 can you take a look |
321e886 to
fa6ecbd
Compare
updated |
|
before review, may I ask how grpo is applied exactly to TTS output, any paper for reference, and what is 16-codebook teacher-forced inputs? A detail background can be provided so we dont provide wrong review comment |
|
And If I understand correctly the idea is partially borrowed from #282 ? correct me if I am wrong~ |
Thanks for asking.
|
The two PRs overlap at the high-level goal of Qwen3-TTS post-training, but their scopes and integration paths are different. The main methodological reference for this PR was Group Relative Policy Optimization for Text-to-Speech with Large Language Models, as cited in the preceding reply.
Therefore, this PR is intentionally narrower and focuses on integrating Qwen3-TTS into the existing generic omni GRPO path. |
Thanks for your explanation. |
910484b to
0bc9083
Compare
References updated. Thanks! |
| @@ -132,7 +168,7 @@ def _merged_lora_per_tensor_param(self): | |||
| for name, param in params.items(): | |||
| yield ( | |||
| name, | |||
| param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True) | |||
| self._cast_dtensor_weight_for_sync(param.to(device, non_blocking=True).full_tensor()) | |||
| if isinstance(param, DTensor) | |||
| else param.detach().clone(), | |||
| ) | |||
| @@ -192,6 +228,7 @@ def _build_module(self): | |||
| self.model_config.model_stage, | |||
| self.model_config.get("external_lib"), | |||
| ) | |||
| self.model_adapter_cls = adapter_cls | |||
There was a problem hiding this comment.
seems irrelevant change to your pr
There was a problem hiding this comment.
self.model_adapter_cls is used by the FSDP engine to reconstruct model-specific replay inputs in prepare_model_inputs. It is also now selected before from_pretrained, allowing Qwen3-TTS to provide its official model class without modifying the global AutoModel registry.
|
|
||
| @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 prepare_agent_sampling_params( | ||
| cls, | ||
| sampling_params: dict[str, Any], | ||
| *, | ||
| rollout_config, | ||
| trainer_config, | ||
| agent_inputs: dict[str, Any], | ||
| ) -> dict[str, Any]: | ||
| """Prepare per-request sampling parameters for an omni agent loop.""" | ||
| return dict(sampling_params) | ||
|
|
||
| @classmethod | ||
| def postprocess_agent_loop_output(cls, output, *, tokenizer, response_length: int): | ||
| """Map an engine trajectory to the policy sequence used by the actor.""" | ||
| return output | ||
|
|
||
| @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), {} |
There was a problem hiding this comment.
consider minimize the hook introduced here.
We add hook if it is really necessary.
There was a problem hiding this comment.
Agreed. I removed get_output_modalities; the server now derives retained modalities from the pipeline stage metadata. The remaining hooks cover Qwen3-TTS-specific prompt construction, dual-sampler candidate seeding, codec-0 policy replay mapping, and multi-stage output assembly, which cannot be inferred from the generic topology.
Co-authored-by: GitHub Copilot Signed-off-by: dongbo910220 <1275604947@qq.com>
There was a problem hiding this comment.
Reviewed at head 1b2d302. I validated the three CPU contract files locally in an env matching this PR's pins (vllm 0.28.0, vllm-omni 0.28.0rc2.dev68+gded893462 == .github/vllm_omni_pin.txt ded8934, transformers 5.14.1):
pytest tests/pipelines/test_qwen3_tts_on_cpu.py \
tests/pipelines/test_qwen3_tts_package_on_cpu.py \
tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.py
# 36 passed, 1 skipped (skip = optional qwen_tts package absent)
The implementation itself looks solid and fail-closed: the actor/rollout codec-0 trajectory is guarded frame-by-frame (tts_actor_logits), align_audio_codes treats codec-0 policy tokens as an exact invariant and rejects ambiguous spans, mask_codec0_logits uses -1e4 (not -inf) and validates the vocab/eos bounds, and the scope is deliberately narrowed to the validated tts_language=Auto non-streaming layout. It also touches no generic engine/strategy files — the shared omni hooks it uses already landed on main (e.g. #504) — so the blast radius is contained. No blocking code-correctness issues found.
Requesting changes for one accuracy item plus a couple of tracking follow-ups:
1. (must fix) PR description no longer matches the diff. The "What does this PR do?" section still lists "a generic AudioRewardManager and fail-closed JSON HTTP client", and the Test section lists test_qwen3_tts_transformers_compat_on_cpu.py, tests/reward_loop/test_audio_reward_manager_on_cpu.py, tests/utils/reward_score/test_audio_http_scorer_client_on_cpu.py, and tests/workers/test_omni_fsdp_engine_on_cpu.py. None of those are in this three-dot diff — verl_omni/reward_loop/reward_manager/audio.py, verl_omni/utils/reward_score/audio_http_scorer_client.py, and those tests already live on main from an earlier slice. Please trim the description/Test section to the actual [3/N] scope (Qwen3-TTS pipeline + recipe + smoke + the 3 CPU test files), and note the reward manager/HTTP client as already-merged, so the Test commands are reproducible against this PR.
2. (follow-up) Re-request review on the server-touching thread. The vllm_omni_async_server.py thread was left as "postpone review until #444 lands" — #444/#504 have since merged, so this is ready for another look.
3. (nit) Two current threads look addressed but are still open: the recipe e2e smoke (a 2-GPU/2-step Talker GRPO smoke was added) and talker_forward.py generality (now via the shared Talker hooks from #504). Worth resolving or confirming.
| ] | ||
| # 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 = [ |
There was a problem hiding this comment.
Changed the optional dependency group from tts to omni.
| # 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)" |
There was a problem hiding this comment.
-
revert the change
accelerate>=1.14.0 already in .toml -
add qwen-tts in .toml instead.
There was a problem hiding this comment.
Moved qwen-tts into the omni extra and removed the redundant CI reinstall steps. The repository's uv overrides retain Transformers 5.14.1 and Accelerate 1.14.0; dependency resolution and the Qwen3-TTS CPU contracts pass with the released package.
| - "tests/special_e2e/**" | ||
| - "pyproject.toml" | ||
| - .github/workflows/gpu_smoke.yml | ||
| - .github/qwen_tts_pin.txt |
There was a problem hiding this comment.
drop the pin. we will not track qwen-tts pin.
Preserve the Qwen3-TTS GRPO catalogue entry while incorporating the finalized Qwen3-Omni OPD labels and the latest upstream CI and governance changes. Co-authored-by: GitHub Copilot Signed-off-by: dongbo910220 <1275604947@qq.com>
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>
|
@NancyFyong Updated the PR description and Test section to match the current |
zhtmike
left a comment
zhtmike
left a comment
There was a problem hiding this comment.
Looks good.
Are you interested to contribute qwen3-omni talker grpo as well?
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>
|
@zhtmike The final-head GPU smoke exposed a compatibility issue: the released I restored the previously validated Qwen-TTS source revision on top of the current head and installed it with |
|
@zhtmike Thanks for the suggestion. I am interested in Qwen3-Omni Talker GRPO, but I currently do not have enough GPU capacity to validate the full 30B-A3B Talker path end to end, so I would rather not commit to it before I can test it properly. I would still be happy to help with smaller, well-scoped parts that fit my hardware. In the near term, I may focus on work I can validate thoroughly, including the HPSv3 reward micro-batching optimization in #449 and cleaning up my locally reproduced Qwen3-TTS Hindi LoRA/GRPO work into a contribution. |
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>
|
All CI green. Wonderful work! |
5 participants
What does this PR do?
Adds full-parameter stock-GRPO training support for the codec-0 policy of
Qwen/Qwen3-TTS-12Hz-0.6B-Basethrough the existing V1verl_omni.trainer.main_omnipath.This
[3/N]PR adds a model-specific Qwen3-TTS Talker training adapter, a two-stage rollout adapter, a full-parameter GRPO recipe, a two-GPU smoke test, focused CPU contract tests, and documentation.The recipe keeps the actor/reference master parameters and optimizer state in FP32, uses BF16 for FSDP forward/backward computation and rollout, keeps gradient reduction and buffers in FP32, and uses LR
1e-6with 10 warmup steps and optional direct reference-model KL.It does not add a new Trainer, algorithm, policy loss, dataset format, ASR gate, candidate mask, or SpeechJudge-specific code path. This model-specific slice reuses existing shared rollout and reward utilities rather than duplicating them.
Related RFC: #90.
Test
CUDA_VISIBLE_DEVICES='' pytest -q \ tests/pipelines/test_qwen3_tts_on_cpu.py \ tests/pipelines/test_qwen3_tts_package_on_cpu.py \ tests/workers/rollout/rollout_vllm/test_qwen3_tts_rollout_on_cpu.pyResult:
all passed.The unchanged runtime implementation completed a two-GPU, two-step Qwen3-TTS Talker GRPO smoke on two RTX 5090 GPUs. Final-head GitHub GPU CI is still required before merge.
We also completed 200 Qwen3-TTS GRPO updates and evaluated the same fixed 100-sample validation set at step 0 and every 20 steps.
Validation SpeechJudge mean increased from
15.5760to19.0116, with a paired mean gain of+3.4356(95% CI[+2.1739, +4.7301]) and a paired median gain of+2.1563.Design & Code Changes
qwen-ttsthrough theomnioptional dependency group without a repository-specific source pin and selects the Qwen3-TTS two-GPU smoke through the existing GPU-smoke workflow.AI assistance was used;
Co-authored-by: @Michael-Zzq