Add SGLang-Omni RL rollout integration - #2
Closed
Hayden727 wants to merge 24 commits into
Closed
Conversation
Foundational miles-side plumbing for Omni RL (Qwen3-Omni-Thinker AR + higgs TTS) rollout against the sglang-omni /generate backend: - miles/utils/processing_utils.py: encode_audio_for_rollout_engine (WAV base64 data URI), the audio analog of encode_image_for_rollout_engine. - miles_plugins/omni/rollout_contract.py: strict sampling-param whitelist/alias (sglang-omni RolloutSamplingParams forbids extra keys), /generate payload builder, response parser with loud length validation, and a sample-accumulation helper following the miles loss-mask convention (mask spans response tokens only). - miles_plugins/omni/omni_generate_fn.py: OmniGenerateFn, a class-based --custom-generate-function-path entrypoint mirroring single_turn generate. - tests/fast/test_omni_rollout_contract.py: positive + negative unit tests. Omni-specific code stays under miles_plugins/omni with narrow public imports so it can later be extracted; only the generic audio encoder lands in miles core.
…iring Round-1 review fixes for the miles-side omni rollout integration: - OmniGenerateFn now mirrors single_turn/compute_request_payload halt semantics: partial-rollout resume shrinks max_new_tokens and rollout_max_context_len is enforced (returns TRUNCATED with no budget, never over-generates). - Decoded response audio moves to sample.metadata['generated_audio'] (reward-facing) instead of multimodal_train_inputs, which the training path moves to GPU and concatenates (would crash on a dict). - Input audio is encoded from sample.multimodal_inputs and sent as audio_data; the generic compute_request_payload path gains the same audio_data support (encode_audios_for_rollout_engine), parallel to image_data. - encode_audio_for_rollout_engine rejects out-of-range integer PCM; parse_generate_response is strict (entries must be exactly [log_prob, token_id]). - tests: new test_omni_generate_fn.py loads the hook via load_generate_function, stubs post, and asserts the emitted request + sample (payload cleaning, budget truncation, audio encoding); test_omni_rollout_contract.py updated for metadata audio + strictness.
…fied audio keys Round-2 review fixes completing the AC-3 miles-side audio scaffold: - apply_response_to_sample now appends loss-mask entries whenever a mask already exists (partial-rollout off-policy masking pre-sets loss_mask=[0]*old_len), not only when update_loss_mask=True. Keeps len(loss_mask)==response_length so convert_samples_to_train_data does not crash on resumed rollouts. - compute_prompt_ids_from_sample routes through call_processor (so audio_kwargs are actually injected) instead of calling state.processor directly, and normalizes processor input_ids to a JSON-safe list[int]. - one canonical extract_audio_inputs(multimodal_inputs) accepting 'audios'/'audio' is shared by compute_request_payload and OmniGenerateFn. - tests: loadable-hook resume test (off-policy mask + new tokens stays aligned), apply append-to-existing-mask, generic compute_request_payload audio (both keys), and a strict xfail marker documenting deferred mm_data audio-INPUT token expansion.
Extracts the Qwen3-Omni thinker into a standalone HF checkpoint (strips the thinker. prefix, writes thinker_config as config.json). The thinker config carries vision_config, so miles FSDP get_model_cls() loads it via AutoModelForImageTextToText; a forward pass on text input produces LM logits. RAM-bounded shard-by-shard copy so the 30B model never needs to fit in memory. Enables GATE-A thinker-only RL training without modifying miles core model loading.
async compute_math_reward(args, sample, **kwargs) -> 1.0 if the response contains the gold answer (sample.label), else 0.0. Numeric-aware (12 == 12.0), deterministic and dependency-free. Loaded via --custom-rm-path.
Self-contained single-GPU harness demonstrating all four closed-loop components on the real Qwen3-Omni Thinker: rollout (sglang-omni /generate) -> reward (math) -> GRPO advantage -> LoRA policy-gradient update, over multiple steps. Plus a tiny math_smoke dataset. Behavior policy is the served base thinker (rollouts not yet weight-synced; that is the documented next integration step).
…nker Extends the GATE-A loop with the 4th component done properly: per-step LoRA-merged weight broadcast into the served sglang-omni thinker stage via /init_weights_update_group + /update_weights_from_distributed (stages=[thinker]), so rollouts become on-policy. The thinker load_weights accepts plain model.* names, so extracted-thinker names sync directly.
Add torch.cuda.set_device(0) before model load / weight-sync, matching sglang-omni's E2E refit trainer, so NCCL broadcast uses the correct device. The NCCL rendezvous with the served thinker was proven to succeed (WEIGHT_UPDATE_GROUP_READY); final broadcast validation is pending a stable free GPU (shared host churn).
Allow a fresh NCCL group name per run so a re-run is not wedged by a stale weight-update group left on the server by a prior interrupted attempt.
… work The on-policy weight-sync broadcast failed with 'Cuda invalid argument' because the trainer and server are separate single-GPU-masked processes; NCCL tried direct P2P between their physical GPUs and could not resolve the masked peer device. Setting NCCL_P2P_DISABLE=1 on both ends (SHM transport) fixes it. Verified end-to-end: 4 GRPO steps on the real Qwen3-Omni-30B Thinker, synced_params=160/step, non-diverging loss, stable reward — full on-policy GATE-A closed loop.
compute_tts_reward / TtsCompositeReward: transcribe generated audio with Whisper, score content via CER vs target text, combined with hard audio-validity guards (decode success, duration bounds, non-silence RMS floor). Failed decode -> deterministic low reward, never crashes (DEC-2 composite design). Plus a tiny TTS smoke dataset. Pure logic verified locally (CER, normalization, WAV decode, all guards).
gate_b_loop.py demonstrates the first 3 closed-loop components on the real Higgs-audio model via sglang-omni: rollout (pretok /generate -> codec tokens + logprobs + decodable audio, AC-2 now PASS) -> composite Whisper-ASR-CER reward + audio guards -> GRPO advantage. The 4th (LoRA update + NCCL weight-sync) mirrors GATE-A. Bundles the verified sglang-omni Higgs codec-token-logprob patch (examples/omni_gate_b/sglang_omni_patches/, 4 files: model_runner/payload_types/request_builders/vocoder_scheduler).
Differentiable teacher-forced forward (transformers Qwen3 backbone + fused codec embedding/head loaded from the Higgs ckpt) recomputing codebook-0 logprobs over a sampled codec sequence, for GATE-B RL. Verified vs the served bf16 model: mean|delta|~0.05, max~0.19 -- an fp32 trainer gives the SAME residual, so it is the server's bf16/sglang-kernel numeric floor, not a reconstruction error (exp(0.2)~1.22 sits at the GRPO clip boundary, first-ratio-after-sync only). Unblocks GRPO + tts_engine weight-sync.
…-sync) Closes the TTS RL loop end to end, mirroring gate_a_full.py: rollout (full codec tokens + codebook-0 logprobs + audio) -> composite ASR reward -> GRPO advantage -> LoRA update (HiggsTtsActor recomputes new logprobs) -> NCCL broadcast of LoRA-merged body.* weights into the served tts_engine stage (server fuses q/k/v on load). Verified live: 3-step run, WEIGHT_UPDATE_GROUP_READY, 72 params synced/step, non-diverging. (Reward saturates on the easy smoke set -> adv 0; proves the loop + weight-sync, not policy improvement, same caveat as GATE-A.) Also: higgs_actor uses get_input_embeddings() for LoRA-wrap safety.
TEMP/MAX_NEW envs to induce reward variance for a non-saturating run. Verified live at TEMP=1.4: mean_reward 0.908 (variance -> non-zero GRPO advantage), avg_loss -0.0004 (real gradient), 72 changed body.* params synced/step to tts_engine, non-diverging. Demonstrates the GATE-B loop actually trains, not just plumbing.
…en copy) Captures the d6fdbf4f sglang-omni fix (record true codebook-0 logprob instead of the text-vocab zeros placeholder). Supersedes the stale higgs_codec_logprob.patch, which encodes the broken placeholder approach (logprob == ln(1/151936)). Canonical source is the sglang-omni repo (Hayden727/sglang-omni hayden/higgs-rl-rollout).
It encodes the placeholder bug (codec logprob == ln(1/151936)); superseded by higgs_codec_logprob_fix.patch and the d6fdbf4f fix in the sglang-omni repo.
…ning runs Adds math_harder.jsonl (10 moderate problems, base acc ~50%) and tts_harder.jsonl (8 pronunciation-hard sentences, base CER ~0.11) so reward is non-saturated, plus a MAX_NEW env on gate_a_full.py. Verified live: GATE-A reward 0.50->~0.70 (24-step run stable, plateaus ~0.68); GATE-B CER 0.108->~0.062 -- genuine GRPO learning signal with on-policy weight-sync, vs the saturated smoke sets.
… names
Per code-style (no plan-marker terms like GATE-/Phase- in code), rename by actual role:
examples/omni_gate_a -> examples/thinker_text_rl
gate_a_lora_smoke.py -> lora_grpo_smoke.py
gate_a_full.py -> onpolicy_grpo_weight_sync.py
gate_a_trainer.sh -> fsdp_trainer_launch.sh
examples/omni_gate_b -> examples/higgs_tts_rl
gate_b_loop.py -> rollout_reward_advantage.py
gate_b_full.py -> onpolicy_grpo_weight_sync.py
gate_b_parity_probe.py -> logprob_parity_probe.py
Also scrub GATE-A/GATE-B from docstrings/prints/GROUP_NAME and fix internal path refs
(incl. miles_plugins/omni docstrings). Datasets and miles_plugins module names unchanged
(already function-named).
Applies the docstring/print/GROUP_NAME/path-reference cleanup that the rename commit missed (a bad git-add pathspec had left these unstaged), and tracks the renamed fsdp_trainer_launch.sh + math_harder_msgs.jsonl.
--rollout-external / --rollout-external-engine-addrs were defined but unused: RolloutManager always launched internal sglang engines (assert num_gpus>0). Add an external branch that skips internal server launch and points sglang_router_ip/port at the supplied external addr, so a custom generate function can drive an external engine (e.g. sglang-omni for an omni pipeline). Verified: RolloutManager init now passes in external mode and the FSDP actor loads the model; FSDP2 sharding of the 30B-A3B omni thinker is a separate experimental-backend issue.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3186a5a169
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
JingwenGu0829
left a comment
JingwenGu0829
left a comment
Collaborator
There was a problem hiding this comment.
Did a rough look and had the following findings. Will take a more detailed look this weekend.
Owner
|
close for now, the follow up pr will be #4 ... |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this PR adds
This PR adds the Miles-side integration used to validate SGLang-Omni RL rollouts for:
The SGLang-Omni server-side companion PR is:
Main changes
/generatepayload construction and strict response parsing.output_codebook_tokensandomni_rolloutmetadata for Higgs TTS rollouts.Validation
Fast tests in the
hongccc/sglang-omni:devcontainer onhyper01-omni:End-to-end smoke validation:
thinkerstage,synced_params=160./generateaudio/codebook rollout, Whisper-CER reward, LoRA update, per-step NCCL sync to the servedtts_enginestage,synced_params=72.