[RL] Qwen3-Omni thinker: audio-input GRPO

[yxs]

Summary

End-to-end audio-input GRPO for Qwen3-Omni-30B-A3B (thinker): prompts carry audio (AVQA multiple-choice QA), rollouts run on an externally launched sglang-omni server, and the trainer holds only the extracted thinker text backbone, injecting frozen-audio-tower embeddings at audio placeholder positions.

upstream miles is not taking modality-specific RL for now (radixark#1700).

Design

Rollout. Generation runs on an externally launched sglang-omni server. The adapter sends processor-expanded input_ids plus base64-serialized audio tensors and speaks the omni server's API.

Training. The Megatron-side model is just the thinker text backbone (Qwen3-MoE). A frozen copy of the omni audio encoder turns each audio clip into embedding vectors, which replace the <|AUDIO|> placeholder tokens in the input embeddings. Sequence-parallel aware: each TP rank scatters only its own sequence chunk. The encoder stays out of the optimizer, DDP, and checkpoints.

Weight sync. After each optimizer step, Megatron weights are converted to HF thinker.* names and broadcast to the server over NCCL.

Example & data.examples/omni_thinker/run_qwen3_omni_thinker_grpo.py is the full driver (download → extract → convert → train).

Validation

off-policy

--sync-mode skip: the server keeps its initial weights; TIS reweights samples to absorb the train/rollout gap. Debug tier (256-sample subset, 6 rounds): the pipeline runs end to end, rewards grade correctly, and the two implementations agree on logprobs, mis_rs_catastrophic = 0, tis_truncate ≤ 0.78%.

on-policy

--sync-mode distributed: thinker weights are broadcast to the server after every optimizer step.

• Full AVQA epoch: 5120 samples, 160 GRPO steps, 161/161 weight refits succeeded (~5–11 s each for the 30B thinker), ~69 s/step, zero NCCL errors.

• Reward 0.807 → 0.867 (20-step means):

• reward curve

• Train/infer consistency stays healthy the whole run: train_rollout_logprob_abs_diff 0.002–0.02.
• Audio ablation on the post-epoch server (greedy, weight_version = 161): the same 256 questions score 89.8% with audio vs 50.4% with audio stripped.

note: On-policy requires the server co-located with the trainer (same container) and launched with NCCL_CUMEM_ENABLE=0 (miles pins the trainer to 0; the sglang-omni server entry does not set it, and the transport mismatch fails the first broadcast). Split-container deployments fail in the NCCL data plane. The audio tower is frozen, only the thinker text backbone trains.