From 041ad2df3a8925ffcc1bb30bca887196f3fa0d63 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Fri, 19 Jun 2026 19:06:33 +0800 Subject: [PATCH 01/24] feat: add omni rollout contract and audio encode for sglang-omni RL 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. --- miles/utils/processing_utils.py | 27 ++++ miles_plugins/omni/__init__.py | 6 + miles_plugins/omni/omni_generate_fn.py | 59 +++++++ miles_plugins/omni/rollout_contract.py | 183 ++++++++++++++++++++++ tests/fast/test_omni_rollout_contract.py | 190 +++++++++++++++++++++++ 5 files changed, 465 insertions(+) create mode 100644 miles_plugins/omni/__init__.py create mode 100644 miles_plugins/omni/omni_generate_fn.py create mode 100644 miles_plugins/omni/rollout_contract.py create mode 100644 tests/fast/test_omni_rollout_contract.py diff --git a/miles/utils/processing_utils.py b/miles/utils/processing_utils.py index 855a06a8fb2..05b8dd0cdaa 100644 --- a/miles/utils/processing_utils.py +++ b/miles/utils/processing_utils.py @@ -3,6 +3,7 @@ import io import logging import os +import wave from pathlib import Path from huggingface_hub import hf_hub_download @@ -173,3 +174,29 @@ def encode_image_for_rollout_engine(image) -> str: image.save(buffer, format="PNG") image_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8") return f"data:image/png;base64,{image_base64}" + + +def encode_audio_for_rollout_engine(audio, sampling_rate: int) -> str: + """Encode a mono waveform as a base64 WAV data URI for the rollout engine. + + Mirrors ``encode_image_for_rollout_engine`` for the audio modality. Accepts a + 1-D array of float samples in [-1, 1] (converted to 16-bit PCM) or int16 samples. + """ + import numpy as np + + samples = np.asarray(audio) + if samples.ndim != 1: + raise ValueError(f"expected a 1-D mono waveform, got shape {samples.shape}") + if samples.dtype.kind == "f": + samples = (np.clip(samples, -1.0, 1.0) * 32767.0).astype(np.int16) + elif samples.dtype != np.int16: + samples = samples.astype(np.int16) + + buffer = io.BytesIO() + with wave.open(buffer, "wb") as wav_file: + wav_file.setnchannels(1) + wav_file.setsampwidth(2) + wav_file.setframerate(int(sampling_rate)) + wav_file.writeframes(samples.tobytes()) + audio_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8") + return f"data:audio/wav;base64,{audio_base64}" diff --git a/miles_plugins/omni/__init__.py b/miles_plugins/omni/__init__.py new file mode 100644 index 00000000000..1502e920eed --- /dev/null +++ b/miles_plugins/omni/__init__.py @@ -0,0 +1,6 @@ +"""Omni RL rollout integration for the sglang-omni inference backend. + +This package keeps omni-specific rollout glue out of generic miles core. It is loaded +through path-string hooks (``--custom-generate-function-path``) and imports only public +miles entrypoints so it can later be extracted into a standalone distribution. +""" diff --git a/miles_plugins/omni/omni_generate_fn.py b/miles_plugins/omni/omni_generate_fn.py new file mode 100644 index 00000000000..2cdad936135 --- /dev/null +++ b/miles_plugins/omni/omni_generate_fn.py @@ -0,0 +1,59 @@ +"""Per-sample generate function that drives rollout against the sglang-omni backend. + +Load it via ``--custom-generate-function-path miles_plugins.omni.omni_generate_fn.OmniGenerateFn``. +It mirrors the stock single-turn generate path but speaks the omni ``/generate`` contract: +a whitelisted sampling-param payload, request ``metadata`` for response matching, and a +response parser that captures generated tokens, behavior-policy log-probs, decoded audio +(for TTS rewards), and ``weight_version`` provenance. +""" + +from __future__ import annotations + +from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput +from miles.rollout.generate_utils.generate_endpoint_utils import compute_prompt_ids_from_sample +from miles.utils.http_utils import post +from miles.utils.types import Sample + +from .rollout_contract import apply_response_to_sample, build_generate_payload, parse_generate_response + + +class OmniGenerateFn: + """Class-based generate function for omni (Thinker AR / TTS) rollout.""" + + async def __call__(self, input: GenerateFnInput) -> GenerateFnOutput: + args = input.args + sample = input.sample + sampling_params = input.sampling_params + assert sample.status in {Sample.Status.PENDING, Sample.Status.ABORTED}, f"{sample.status=}" + + url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate" + + prompt_ids = compute_prompt_ids_from_sample(input.state, sample) + # Partial-rollout resume: continue from already-generated tokens. + input_ids = sample.tokens if len(sample.response) > 0 else prompt_ids + + payload = build_generate_payload( + input_ids, + sampling_params, + metadata=_request_metadata(sample), + output_modalities=sample.metadata.get("output_modalities"), + ) + + output = await post(url, payload) + + result = parse_generate_response(output) + apply_response_to_sample(sample, prompt_ids, result) + # Reuse the existing meta_info handling for status / weight_version / prefix-cache stats. + sample.update_from_meta_info(args, output["meta_info"]) + + return GenerateFnOutput(samples=sample) + + +def _request_metadata(sample: Sample) -> dict: + """Identifiers echoed back by the backend so responses can be matched to a rollout.""" + fields = { + "group_index": sample.group_index, + "index": sample.index, + "session_id": sample.session_id, + } + return {k: v for k, v in fields.items() if v is not None} diff --git a/miles_plugins/omni/rollout_contract.py b/miles_plugins/omni/rollout_contract.py new file mode 100644 index 00000000000..45418bb5448 --- /dev/null +++ b/miles_plugins/omni/rollout_contract.py @@ -0,0 +1,183 @@ +"""Typed request/response contract for the sglang-omni ``/generate`` rollout endpoint. + +The omni backend exposes a stricter rollout schema than the stock sglang ``/generate``: +its sampling params reject unknown keys (``extra="forbid"``), so miles' default sampling +params (which carry keys such as ``skip_special_tokens`` or ``sampling_seed``) must be +whitelisted and aliased before they are sent. The response carries the generated tokens +and their behavior-policy log-probs inside ``meta_info.output_token_logprobs`` (one +``[log_prob, token_id]`` pair per generated token), optional decoded ``audio`` for TTS +rewards, and ``weight_version`` provenance. + +All functions here are pure and side-effect free except :func:`apply_response_to_sample`, +which accumulates generated tokens onto a sample following the existing miles rollout +convention. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from miles.utils.types import Sample + +# Sampling-param keys accepted by the omni rollout endpoint. Anything else is dropped so +# the request is not rejected by the backend's strict (extra-forbidding) schema. +OMNI_SAMPLING_PARAM_KEYS: frozenset[str] = frozenset( + { + "temperature", + "top_p", + "top_k", + "min_p", + "repetition_penalty", + "stop", + "stop_token_ids", + "seed", + "max_new_tokens", + "max_tokens", + } +) + +# miles uses some legacy names that map onto the omni schema's canonical fields. +OMNI_SAMPLING_PARAM_ALIASES: dict[str, str] = {"sampling_seed": "seed"} + + +def clean_sampling_params(sampling_params: dict[str, Any]) -> dict[str, Any]: + """Project miles sampling params onto the keys the omni endpoint accepts. + + Unknown keys are dropped, known aliases are renamed, and ``None`` values are removed + so optional fields fall back to backend defaults instead of failing validation. + """ + cleaned: dict[str, Any] = {} + for key, value in (sampling_params or {}).items(): + target = OMNI_SAMPLING_PARAM_ALIASES.get(key, key) + if target in OMNI_SAMPLING_PARAM_KEYS and value is not None: + cleaned[target] = value + return cleaned + + +def build_generate_payload( + input_ids: list[int], + sampling_params: dict[str, Any], + *, + metadata: dict[str, Any] | None = None, + output_modalities: list[str] | None = None, + return_logprob: bool = True, + audio_data: list[str] | None = None, +) -> dict[str, Any]: + """Build an omni ``/generate`` request body from pre-tokenized inputs. + + The trainer always sends ``input_ids`` (it computes gradients on these exact tokens), + requests log-probs by default, and may echo ``metadata`` so responses can be matched + back to a rollout batch. + """ + payload: dict[str, Any] = { + "input_ids": list(input_ids), + "sampling_params": clean_sampling_params(sampling_params), + "return_logprob": return_logprob, + } + if metadata: + payload["metadata"] = metadata + if output_modalities is not None: + payload["output_modalities"] = output_modalities + if audio_data is not None: + payload["audio_data"] = audio_data + return payload + + +@dataclass +class OmniRolloutResult: + """Parsed view of an omni ``/generate`` response, ready to apply to a sample.""" + + response_tokens: list[int] + response_log_probs: list[float] + text: str = "" + finish_reason: dict[str, Any] = field(default_factory=dict) + weight_version: str | None = None + cached_tokens: int = 0 + prompt_tokens: int = 0 + completion_tokens: int = 0 + audio: dict[str, Any] | None = None + + +def parse_generate_response(response: dict[str, Any]) -> OmniRolloutResult: + """Parse an omni ``/generate`` response into :class:`OmniRolloutResult`. + + Raises ``ValueError`` (loudly, never silently truncating) when the response is + malformed or when the per-token log-prob count disagrees with ``completion_tokens``. + """ + if "meta_info" not in response: + raise ValueError("omni /generate response is missing 'meta_info'") + meta = response["meta_info"] + + token_logprobs = meta.get("output_token_logprobs") or [] + response_tokens: list[int] = [] + response_log_probs: list[float] = [] + for i, item in enumerate(token_logprobs): + if not isinstance(item, (list, tuple)) or len(item) < 2: + raise ValueError( + f"output_token_logprobs[{i}] is malformed: {item!r}; expected [log_prob, token_id]" + ) + response_log_probs.append(float(item[0])) + response_tokens.append(int(item[1])) + + completion_tokens = meta.get("completion_tokens") + if completion_tokens is not None and len(response_tokens) != completion_tokens: + raise ValueError( + f"output_token_logprobs length ({len(response_tokens)}) " + f"!= completion_tokens ({completion_tokens})" + ) + + if "finish_reason" not in meta: + raise ValueError("omni /generate meta_info is missing 'finish_reason'") + + return OmniRolloutResult( + response_tokens=response_tokens, + response_log_probs=response_log_probs, + text=response.get("text", "") or "", + finish_reason=meta["finish_reason"], + weight_version=meta.get("weight_version"), + cached_tokens=int(meta.get("cached_tokens") or 0), + prompt_tokens=int(meta.get("prompt_tokens") or 0), + completion_tokens=int(completion_tokens if completion_tokens is not None else len(response_tokens)), + audio=response.get("audio"), + ) + + +def apply_response_to_sample( + sample: Sample, + prompt_ids: list[int], + result: OmniRolloutResult, + *, + update_loss_mask: bool = False, +) -> Sample: + """Accumulate parsed generation onto ``sample`` (tokens, log-probs, loss mask, audio). + + Follows the miles convention where ``loss_mask`` and ``rollout_log_probs`` span only + the generated (completion) tokens (length == ``response_length``); the prompt is + excluded by lying outside the mask rather than by leading zeros. Standard meta_info + handling (status, weight-version, prefix-cache stats) stays with the caller via the + existing ``Sample.update_from_meta_info`` so this stays backend-agnostic and testable + without trainer ``args``. + """ + if not sample.tokens: + sample.tokens = list(prompt_ids) + + sample.tokens = sample.tokens + result.response_tokens + sample.response_length += len(result.response_tokens) + sample.response += result.text + + if sample.rollout_log_probs is None: + sample.rollout_log_probs = [] + sample.rollout_log_probs += result.response_log_probs + + if update_loss_mask: + if sample.loss_mask is None: + sample.loss_mask = [] + sample.loss_mask += [1] * len(result.response_tokens) + + if result.audio is not None: + if sample.multimodal_train_inputs is None: + sample.multimodal_train_inputs = {} + sample.multimodal_train_inputs["audio"] = result.audio + + return sample diff --git a/tests/fast/test_omni_rollout_contract.py b/tests/fast/test_omni_rollout_contract.py new file mode 100644 index 00000000000..0b110b26c82 --- /dev/null +++ b/tests/fast/test_omni_rollout_contract.py @@ -0,0 +1,190 @@ +"""Unit tests for the omni rollout contract and the audio rollout-encode helper. + +These exercise the miles-side glue for the sglang-omni ``/generate`` backend without any +GPU, network, or model: payload whitelisting, response parsing/alignment, sample +accumulation, and WAV encoding. +""" + +import base64 +import io +import wave + +import numpy as np +import pytest + +from miles.utils.processing_utils import encode_audio_for_rollout_engine +from miles.utils.types import Sample +from miles_plugins.omni.rollout_contract import ( + OMNI_SAMPLING_PARAM_KEYS, + apply_response_to_sample, + build_generate_payload, + clean_sampling_params, + parse_generate_response, +) + + +def _response(token_logprobs, *, completion_tokens=None, finish="stop", **meta): + meta_info = { + "finish_reason": {"type": finish}, + "output_token_logprobs": token_logprobs, + **meta, + } + if completion_tokens is not None: + meta_info["completion_tokens"] = completion_tokens + return {"text": meta.pop("text", ""), "meta_info": meta_info} + + +# --- sampling param whitelisting ------------------------------------------------------- + + +def test_clean_sampling_params_drops_unknown_keys_and_aliases_seed(): + raw = { + "temperature": 0.7, + "top_p": 0.95, + "max_new_tokens": 128, + "sampling_seed": 1234, # legacy alias -> seed + "skip_special_tokens": True, # not accepted by omni schema + "no_stop_trim": False, # not accepted by omni schema + "spaces_between_special_tokens": True, # not accepted by omni schema + "top_k": None, # None dropped + } + cleaned = clean_sampling_params(raw) + assert cleaned == {"temperature": 0.7, "top_p": 0.95, "max_new_tokens": 128, "seed": 1234} + assert set(cleaned).issubset(OMNI_SAMPLING_PARAM_KEYS) + assert "skip_special_tokens" not in cleaned + assert "sampling_seed" not in cleaned + + +def test_build_generate_payload_shape_and_metadata(): + payload = build_generate_payload( + [1, 2, 3], + {"temperature": 1.0, "skip_special_tokens": True}, + metadata={"index": 5}, + output_modalities=["audio"], + ) + assert payload["input_ids"] == [1, 2, 3] + assert payload["return_logprob"] is True + assert payload["sampling_params"] == {"temperature": 1.0} # forbidden key removed + assert payload["metadata"] == {"index": 5} + assert payload["output_modalities"] == ["audio"] + # empty metadata must not be emitted + assert "metadata" not in build_generate_payload([1], {}) + + +# --- response parsing ------------------------------------------------------------------ + + +def test_parse_generate_response_aligns_tokens_and_logprobs(): + resp = _response( + [[-0.1, 10], [-0.2, 11], [-0.3, 12]], + completion_tokens=3, + weight_version="42", + cached_tokens=7, + prompt_tokens=9, + ) + result = parse_generate_response(resp) + assert result.response_tokens == [10, 11, 12] + assert result.response_log_probs == [-0.1, -0.2, -0.3] + assert result.weight_version == "42" + assert result.cached_tokens == 7 and isinstance(result.cached_tokens, int) + assert result.completion_tokens == 3 + + +def test_parse_generate_response_captures_audio_and_text(): + resp = _response([[-0.5, 99]], completion_tokens=1, text="hi") + resp["text"] = "hi" + resp["audio"] = {"format": "wav", "sample_rate": 24000, "data": ""} + result = parse_generate_response(resp) + assert result.audio == {"format": "wav", "sample_rate": 24000, "data": ""} + assert result.text == "hi" + + +def test_parse_generate_response_empty_completion_is_not_an_error(): + result = parse_generate_response(_response([], completion_tokens=0)) + assert result.response_tokens == [] + assert result.response_log_probs == [] + + +def test_parse_generate_response_length_mismatch_raises(): + with pytest.raises(ValueError, match="completion_tokens"): + parse_generate_response(_response([[-0.1, 10]], completion_tokens=5)) + + +def test_parse_generate_response_malformed_item_raises(): + with pytest.raises(ValueError, match="malformed"): + parse_generate_response(_response([[-0.1]], completion_tokens=1)) + + +def test_parse_generate_response_missing_meta_info_raises(): + with pytest.raises(ValueError, match="meta_info"): + parse_generate_response({"text": ""}) + + +def test_parse_generate_response_missing_finish_reason_raises(): + with pytest.raises(ValueError, match="finish_reason"): + parse_generate_response({"meta_info": {"output_token_logprobs": []}}) + + +# --- sample accumulation --------------------------------------------------------------- + + +def test_apply_response_to_sample_aligns_and_validates(): + sample = Sample(prompt="p", tokens=[]) + prompt_ids = [1, 2, 3] + result = parse_generate_response( + _response([[-0.1, 10], [-0.2, 11]], completion_tokens=2, weight_version="3") + ) + apply_response_to_sample(sample, prompt_ids, result, update_loss_mask=True) + + assert sample.tokens == [1, 2, 3, 10, 11] + assert sample.response_length == 2 + assert sample.rollout_log_probs == [-0.1, -0.2] + # miles convention: loss_mask spans only the response tokens + assert sample.loss_mask == [1, 1] + assert len(sample.loss_mask) == sample.response_length + assert len(sample.rollout_log_probs) == sample.response_length + sample.validate() # must not raise + + +def test_apply_response_to_sample_captures_audio_into_train_inputs(): + sample = Sample(prompt="p", tokens=[]) + result = parse_generate_response(_response([[-0.1, 5]], completion_tokens=1)) + result.audio = {"format": "wav", "data": ""} + apply_response_to_sample(sample, [1, 2], result) + assert sample.multimodal_train_inputs["audio"] == {"format": "wav", "data": ""} + + +def test_apply_response_to_sample_multi_turn_accumulates(): + sample = Sample(prompt="p", tokens=[]) + first = parse_generate_response(_response([[-0.1, 10]], completion_tokens=1)) + apply_response_to_sample(sample, [1, 2], first, update_loss_mask=True) + # second turn: tokens already present, continue appending + second = parse_generate_response(_response([[-0.2, 20], [-0.3, 21]], completion_tokens=2)) + apply_response_to_sample(sample, [1, 2], second, update_loss_mask=True) + + assert sample.tokens == [1, 2, 10, 20, 21] + assert sample.response_length == 3 + assert sample.rollout_log_probs == [-0.1, -0.2, -0.3] + assert sample.loss_mask == [1, 1, 1] + + +# --- audio encode helper --------------------------------------------------------------- + + +def test_encode_audio_for_rollout_engine_roundtrips_wav(): + sampling_rate = 24000 + waveform = np.linspace(-1.0, 1.0, num=480, dtype=np.float32) + uri = encode_audio_for_rollout_engine(waveform, sampling_rate) + assert uri.startswith("data:audio/wav;base64,") + + raw = base64.b64decode(uri.split(",", 1)[1]) + with wave.open(io.BytesIO(raw), "rb") as wav_file: + assert wav_file.getnchannels() == 1 + assert wav_file.getsampwidth() == 2 + assert wav_file.getframerate() == sampling_rate + assert wav_file.getnframes() == 480 + + +def test_encode_audio_for_rollout_engine_rejects_multichannel(): + with pytest.raises(ValueError, match="mono"): + encode_audio_for_rollout_engine(np.zeros((2, 100), dtype=np.float32), 16000) From db08da25a1c6fda6b65f089f42ee8cc6ba8fa904 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Fri, 19 Jun 2026 19:31:54 +0800 Subject: [PATCH 02/24] fix(omni): harden OmniGenerateFn budget/audio handling + audio_data wiring 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. --- .../generate_utils/generate_endpoint_utils.py | 7 +- miles/utils/processing_utils.py | 27 +++- miles_plugins/omni/omni_generate_fn.py | 55 ++++++- miles_plugins/omni/rollout_contract.py | 19 ++- tests/fast/test_omni_generate_fn.py | 143 ++++++++++++++++++ tests/fast/test_omni_rollout_contract.py | 29 +++- 6 files changed, 261 insertions(+), 19 deletions(-) create mode 100644 tests/fast/test_omni_generate_fn.py diff --git a/miles/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py index d50098e686b..48526350197 100644 --- a/miles/rollout/generate_utils/generate_endpoint_utils.py +++ b/miles/rollout/generate_utils/generate_endpoint_utils.py @@ -9,7 +9,10 @@ import pybase64 from miles.utils.lora import LORA_ADAPTER_NAME, is_lora_enabled -from miles.utils.processing_utils import encode_image_for_rollout_engine +from miles.utils.processing_utils import ( + encode_audios_for_rollout_engine, + encode_image_for_rollout_engine, +) from miles.utils.types import Sample @@ -60,6 +63,8 @@ def compute_request_payload( payload["lora_path"] = LORA_ADAPTER_NAME if image_data := (multimodal_inputs or {}).get("images"): payload["image_data"] = [encode_image_for_rollout_engine(image) for image in image_data] + if audio_data := (multimodal_inputs or {}).get("audios"): + payload["audio_data"] = encode_audios_for_rollout_engine(audio_data) return payload, None diff --git a/miles/utils/processing_utils.py b/miles/utils/processing_utils.py index 05b8dd0cdaa..fd31e2f602c 100644 --- a/miles/utils/processing_utils.py +++ b/miles/utils/processing_utils.py @@ -189,8 +189,17 @@ def encode_audio_for_rollout_engine(audio, sampling_rate: int) -> str: raise ValueError(f"expected a 1-D mono waveform, got shape {samples.shape}") if samples.dtype.kind == "f": samples = (np.clip(samples, -1.0, 1.0) * 32767.0).astype(np.int16) - elif samples.dtype != np.int16: + elif samples.dtype == np.int16: + pass + elif samples.dtype.kind in ("i", "u"): + # Reject integer PCM that would silently wrap when narrowed to int16. + if samples.size and (int(samples.min()) < -32768 or int(samples.max()) > 32767): + raise ValueError( + "integer audio samples must already fit int16 PCM range [-32768, 32767]" + ) samples = samples.astype(np.int16) + else: + raise ValueError(f"unsupported audio dtype {samples.dtype!r}; expected float or integer PCM") buffer = io.BytesIO() with wave.open(buffer, "wb") as wav_file: @@ -200,3 +209,19 @@ def encode_audio_for_rollout_engine(audio, sampling_rate: int) -> str: wav_file.writeframes(samples.tobytes()) audio_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8") return f"data:audio/wav;base64,{audio_base64}" + + +def encode_audios_for_rollout_engine(audios) -> list[str]: + """Encode a list of waveform entries to base64 WAV data URIs for the rollout engine. + + Each entry is a ``(waveform, sampling_rate)`` pair or a ``{"array", "sampling_rate"}`` + dict, mirroring how ``images`` are carried in ``Sample.multimodal_inputs``. + """ + encoded: list[str] = [] + for item in audios: + if isinstance(item, dict): + waveform, sampling_rate = item["array"], item["sampling_rate"] + else: + waveform, sampling_rate = item + encoded.append(encode_audio_for_rollout_engine(waveform, sampling_rate)) + return encoded diff --git a/miles_plugins/omni/omni_generate_fn.py b/miles_plugins/omni/omni_generate_fn.py index 2cdad936135..0d02104e764 100644 --- a/miles_plugins/omni/omni_generate_fn.py +++ b/miles_plugins/omni/omni_generate_fn.py @@ -1,10 +1,11 @@ """Per-sample generate function that drives rollout against the sglang-omni backend. Load it via ``--custom-generate-function-path miles_plugins.omni.omni_generate_fn.OmniGenerateFn``. -It mirrors the stock single-turn generate path but speaks the omni ``/generate`` contract: -a whitelisted sampling-param payload, request ``metadata`` for response matching, and a -response parser that captures generated tokens, behavior-policy log-probs, decoded audio -(for TTS rewards), and ``weight_version`` provenance. +It mirrors the stock single-turn generate path (including partial-rollout budget and +context-length halting) but speaks the omni ``/generate`` contract: a whitelisted +sampling-param payload, encoded input audio, request ``metadata`` for response matching, +and a response parser that captures generated tokens, behavior-policy log-probs, decoded +audio (for TTS rewards), and ``weight_version`` provenance. """ from __future__ import annotations @@ -12,6 +13,7 @@ from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput from miles.rollout.generate_utils.generate_endpoint_utils import compute_prompt_ids_from_sample from miles.utils.http_utils import post +from miles.utils.processing_utils import encode_audios_for_rollout_engine from miles.utils.types import Sample from .rollout_contract import apply_response_to_sample, build_generate_payload, parse_generate_response @@ -23,20 +25,32 @@ class OmniGenerateFn: async def __call__(self, input: GenerateFnInput) -> GenerateFnOutput: args = input.args sample = input.sample - sampling_params = input.sampling_params + sampling_params = dict(input.sampling_params) # copied; max_new_tokens is adjusted below assert sample.status in {Sample.Status.PENDING, Sample.Status.ABORTED}, f"{sample.status=}" url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate" prompt_ids = compute_prompt_ids_from_sample(input.state, sample) - # Partial-rollout resume: continue from already-generated tokens. - input_ids = sample.tokens if len(sample.response) > 0 else prompt_ids + # Partial-rollout resume: continue from already-generated tokens and shrink the + # remaining budget by what was already produced (mirrors single_turn.generate). + if len(sample.response) > 0: + input_ids = sample.tokens + if sampling_params.get("max_new_tokens") is not None: + sampling_params["max_new_tokens"] -= len(sample.tokens) - len(prompt_ids) + else: + input_ids = prompt_ids + + halt_status = _clamp_max_new_tokens(args, sampling_params, len(input_ids)) + if halt_status is not None: + sample.status = halt_status + return GenerateFnOutput(samples=sample) payload = build_generate_payload( input_ids, sampling_params, metadata=_request_metadata(sample), output_modalities=sample.metadata.get("output_modalities"), + audio_data=_encode_input_audio(sample), ) output = await post(url, payload) @@ -49,6 +63,33 @@ async def __call__(self, input: GenerateFnInput) -> GenerateFnOutput: return GenerateFnOutput(samples=sample) +def _clamp_max_new_tokens(args, sampling_params: dict, prompt_len: int) -> Sample.Status | None: + """Cap ``max_new_tokens`` by the context budget; return a halt status if none remains. + + Mirrors ``compute_request_payload`` so the omni path enforces the same limits as the + stock generate path. + """ + max_new_tokens = sampling_params.get("max_new_tokens") + if max_new_tokens is None: + max_new_tokens = args.rollout_max_response_len + if context_len := getattr(args, "rollout_max_context_len", None): + max_new_tokens = min(max_new_tokens, context_len - prompt_len) + if max_new_tokens <= 0: + return Sample.Status.TRUNCATED + sampling_params["max_new_tokens"] = max_new_tokens + return None + + +def _encode_input_audio(sample: Sample) -> list[str] | None: + """Encode input-side audio from ``sample.multimodal_inputs`` for the request payload.""" + if not sample.multimodal_inputs: + return None + audios = sample.multimodal_inputs.get("audios") or sample.multimodal_inputs.get("audio") + if not audios: + return None + return encode_audios_for_rollout_engine(audios) + + def _request_metadata(sample: Sample) -> dict: """Identifiers echoed back by the backend so responses can be matched to a rollout.""" fields = { diff --git a/miles_plugins/omni/rollout_contract.py b/miles_plugins/omni/rollout_contract.py index 45418bb5448..2ab087caa47 100644 --- a/miles_plugins/omni/rollout_contract.py +++ b/miles_plugins/omni/rollout_contract.py @@ -113,7 +113,7 @@ def parse_generate_response(response: dict[str, Any]) -> OmniRolloutResult: response_tokens: list[int] = [] response_log_probs: list[float] = [] for i, item in enumerate(token_logprobs): - if not isinstance(item, (list, tuple)) or len(item) < 2: + if not isinstance(item, (list, tuple)) or len(item) != 2: raise ValueError( f"output_token_logprobs[{i}] is malformed: {item!r}; expected [log_prob, token_id]" ) @@ -154,10 +154,12 @@ def apply_response_to_sample( Follows the miles convention where ``loss_mask`` and ``rollout_log_probs`` span only the generated (completion) tokens (length == ``response_length``); the prompt is - excluded by lying outside the mask rather than by leading zeros. Standard meta_info - handling (status, weight-version, prefix-cache stats) stays with the caller via the - existing ``Sample.update_from_meta_info`` so this stays backend-agnostic and testable - without trainer ``args``. + excluded by lying outside the mask rather than by leading zeros. Decoded response + ``audio`` is stored in ``sample.metadata`` (reward-facing), never in + ``multimodal_train_inputs``. Standard meta_info handling (status, weight-version, + prefix-cache stats) stays with the caller via the existing + ``Sample.update_from_meta_info`` so this stays backend-agnostic and testable without + trainer ``args``. """ if not sample.tokens: sample.tokens = list(prompt_ids) @@ -176,8 +178,9 @@ def apply_response_to_sample( sample.loss_mask += [1] * len(result.response_tokens) if result.audio is not None: - if sample.multimodal_train_inputs is None: - sample.multimodal_train_inputs = {} - sample.multimodal_train_inputs["audio"] = result.audio + # Decoded response audio is reward-facing (e.g. ASR scoring), not a model-forward + # tensor. Keep it out of multimodal_train_inputs, which the training path moves to + # GPU and concatenates; store it in metadata instead. + sample.metadata["generated_audio"] = result.audio return sample diff --git a/tests/fast/test_omni_generate_fn.py b/tests/fast/test_omni_generate_fn.py new file mode 100644 index 00000000000..87ca6d574b9 --- /dev/null +++ b/tests/fast/test_omni_generate_fn.py @@ -0,0 +1,143 @@ +"""Integration tests for the loadable OmniGenerateFn hook. + +Loads the class through the same path-string loader rollout uses, stubs the HTTP +transport, and asserts the exact request emitted to the omni ``/generate`` endpoint plus +the resulting sample. Exercises the highest-risk path (the real ``__call__``), unlike the +pure-helper tests in test_omni_rollout_contract.py. +""" + +import asyncio +from types import SimpleNamespace + +import numpy as np +import pytest + +import miles_plugins.omni.omni_generate_fn as omni_mod +from miles.rollout.base_types import GenerateFnInput +from miles.rollout.inference_rollout.compatibility import load_generate_function +from miles.utils.types import Sample + +_HOOK_PATH = "miles_plugins.omni.omni_generate_fn.OmniGenerateFn" + + +class _FakeTokenizer: + def encode(self, text, add_special_tokens=False): + return [1, 2, 3] + + +def _fake_state(*, max_context_len=0): + args = SimpleNamespace( + sglang_router_ip="127.0.0.1", + sglang_router_port=8000, + rollout_max_response_len=128, + rollout_max_context_len=max_context_len, + sglang_speculative_algorithm=None, + ) + return SimpleNamespace(args=args, tokenizer=_FakeTokenizer(), processor=None) + + +def _canned_response(): + return { + "text": "hello", + "audio": {"format": "wav", "data": ""}, + "meta_info": { + "finish_reason": {"type": "stop"}, + "output_token_logprobs": [[-0.1, 10], [-0.2, 11]], + "completion_tokens": 2, + "weight_version": "7", + "cached_tokens": 0, + "prompt_tokens": 3, + }, + } + + +def test_omni_generate_fn_emits_payload_and_applies_response(monkeypatch): + captured = {} + + async def fake_post(url, payload, **kwargs): + captured["url"] = url + captured["payload"] = payload + return _canned_response() + + monkeypatch.setattr(omni_mod, "post", fake_post) + + fn = load_generate_function(_HOOK_PATH) + assert fn is not None + + sample = Sample(prompt="hi", index=5, group_index=2) + inp = GenerateFnInput( + state=_fake_state(), + sample=sample, + sampling_params={ + "temperature": 0.7, + "skip_special_tokens": True, # dropped by the omni schema + "sampling_seed": 9, # aliased -> seed + "max_new_tokens": 64, + }, + evaluation=False, + ) + + out = asyncio.run(fn(inp)) + result_sample = out.samples + + payload = captured["payload"] + assert captured["url"] == "http://127.0.0.1:8000/generate" + assert payload["input_ids"] == [1, 2, 3] + assert payload["return_logprob"] is True + assert payload["sampling_params"] == {"temperature": 0.7, "seed": 9, "max_new_tokens": 64} + assert payload["metadata"] == {"group_index": 2, "index": 5} + assert "audio_data" not in payload # no input audio on this sample + + assert result_sample.tokens == [1, 2, 3, 10, 11] + assert result_sample.response_length == 2 + assert result_sample.rollout_log_probs == [-0.1, -0.2] + assert result_sample.response == "hello" + # generated audio is reward-facing -> metadata, never multimodal_train_inputs + assert result_sample.metadata["generated_audio"] == {"format": "wav", "data": ""} + assert result_sample.multimodal_train_inputs is None + assert result_sample.weight_versions == ["7"] + assert result_sample.status == Sample.Status.COMPLETED + + +def test_omni_generate_fn_truncates_when_no_context_budget(monkeypatch): + async def fail_post(url, payload, **kwargs): + raise AssertionError("post must not be called when there is no token budget") + + monkeypatch.setattr(omni_mod, "post", fail_post) + + fn = load_generate_function(_HOOK_PATH) + sample = Sample(prompt="hi") + inp = GenerateFnInput( + state=_fake_state(max_context_len=3), # prompt is 3 tokens -> 0 budget left + sample=sample, + sampling_params={"max_new_tokens": 64}, + evaluation=False, + ) + + out = asyncio.run(fn(inp)) + assert out.samples.status == Sample.Status.TRUNCATED + + +def test_omni_generate_fn_encodes_input_audio(monkeypatch): + captured = {} + + async def fake_post(url, payload, **kwargs): + captured["payload"] = payload + return _canned_response() + + monkeypatch.setattr(omni_mod, "post", fake_post) + + fn = load_generate_function(_HOOK_PATH) + sample = Sample(prompt="hi") + sample.multimodal_inputs = {"audios": [(np.zeros(160, dtype=np.float32), 16000)]} + inp = GenerateFnInput( + state=_fake_state(), + sample=sample, + sampling_params={"max_new_tokens": 32}, + evaluation=False, + ) + + asyncio.run(fn(inp)) + audio_data = captured["payload"]["audio_data"] + assert len(audio_data) == 1 + assert audio_data[0].startswith("data:audio/wav;base64,") diff --git a/tests/fast/test_omni_rollout_contract.py b/tests/fast/test_omni_rollout_contract.py index 0b110b26c82..06cccae1a1c 100644 --- a/tests/fast/test_omni_rollout_contract.py +++ b/tests/fast/test_omni_rollout_contract.py @@ -115,6 +115,12 @@ def test_parse_generate_response_malformed_item_raises(): parse_generate_response(_response([[-0.1]], completion_tokens=1)) +def test_parse_generate_response_rejects_overlong_logprob_entry(): + # strict contract: each entry must be exactly [log_prob, token_id] + with pytest.raises(ValueError, match="malformed"): + parse_generate_response(_response([[-0.1, 10, "extra"]], completion_tokens=1)) + + def test_parse_generate_response_missing_meta_info_raises(): with pytest.raises(ValueError, match="meta_info"): parse_generate_response({"text": ""}) @@ -146,12 +152,14 @@ def test_apply_response_to_sample_aligns_and_validates(): sample.validate() # must not raise -def test_apply_response_to_sample_captures_audio_into_train_inputs(): +def test_apply_response_to_sample_stores_audio_in_metadata_not_train_inputs(): sample = Sample(prompt="p", tokens=[]) result = parse_generate_response(_response([[-0.1, 5]], completion_tokens=1)) result.audio = {"format": "wav", "data": ""} apply_response_to_sample(sample, [1, 2], result) - assert sample.multimodal_train_inputs["audio"] == {"format": "wav", "data": ""} + # reward-facing audio lives in metadata; multimodal_train_inputs stays tensor-only + assert sample.metadata["generated_audio"] == {"format": "wav", "data": ""} + assert sample.multimodal_train_inputs is None def test_apply_response_to_sample_multi_turn_accumulates(): @@ -188,3 +196,20 @@ def test_encode_audio_for_rollout_engine_roundtrips_wav(): def test_encode_audio_for_rollout_engine_rejects_multichannel(): with pytest.raises(ValueError, match="mono"): encode_audio_for_rollout_engine(np.zeros((2, 100), dtype=np.float32), 16000) + + +def test_encode_audio_for_rollout_engine_rejects_out_of_range_int(): + with pytest.raises(ValueError, match="int16"): + encode_audio_for_rollout_engine(np.array([0, 40000, -50000], dtype=np.int32), 16000) + + +def test_encode_audios_for_rollout_engine_handles_tuples_and_dicts(): + from miles.utils.processing_utils import encode_audios_for_rollout_engine + + audios = [ + (np.zeros(160, dtype=np.float32), 16000), + {"array": np.zeros(240, dtype=np.int16), "sampling_rate": 24000}, + ] + uris = encode_audios_for_rollout_engine(audios) + assert len(uris) == 2 + assert all(u.startswith("data:audio/wav;base64,") for u in uris) From e183512afac191b6b4544ca45f2d5c62f263ab1d Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Fri, 19 Jun 2026 19:50:47 +0800 Subject: [PATCH 03/24] fix(omni): partial-rollout loss_mask, call_processor audio route, unified 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. --- .../generate_utils/generate_endpoint_utils.py | 18 ++++-- miles/utils/processing_utils.py | 12 ++++ miles_plugins/omni/omni_generate_fn.py | 6 +- miles_plugins/omni/rollout_contract.py | 10 +++- tests/fast/test_omni_generate_fn.py | 40 +++++++++++++ tests/fast/test_omni_rollout_contract.py | 59 +++++++++++++++++++ 6 files changed, 135 insertions(+), 10 deletions(-) diff --git a/miles/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py index 48526350197..c60403c356e 100644 --- a/miles/rollout/generate_utils/generate_endpoint_utils.py +++ b/miles/rollout/generate_utils/generate_endpoint_utils.py @@ -10,19 +10,29 @@ from miles.utils.lora import LORA_ADAPTER_NAME, is_lora_enabled from miles.utils.processing_utils import ( + call_processor, encode_audios_for_rollout_engine, encode_image_for_rollout_engine, + extract_audio_inputs, ) from miles.utils.types import Sample +def _to_int_list(ids) -> list[int]: + """Coerce processor/tokenizer output (tensor / ndarray / list) to a JSON-safe list[int].""" + if hasattr(ids, "tolist"): + ids = ids.tolist() + return [int(token) for token in ids] + + # Make this an isolated function because users may want to compute their own def compute_prompt_ids_from_sample(state, sample, tools=None): prompt = sample.prompt if state.processor and sample.multimodal_inputs and any(v is not None for v in sample.multimodal_inputs.values()): - processor_output = state.processor(text=prompt, **sample.multimodal_inputs) - prompt_ids = processor_output["input_ids"][0] + # Route through call_processor so per-modality kwargs (incl. audio_kwargs) are applied. + processor_output = call_processor(state.processor, prompt, sample.multimodal_inputs) + prompt_ids = _to_int_list(processor_output["input_ids"][0]) # TODO shall we move it to other places? then can make this function immutable sample.multimodal_train_inputs = { @@ -63,8 +73,8 @@ def compute_request_payload( payload["lora_path"] = LORA_ADAPTER_NAME if image_data := (multimodal_inputs or {}).get("images"): payload["image_data"] = [encode_image_for_rollout_engine(image) for image in image_data] - if audio_data := (multimodal_inputs or {}).get("audios"): - payload["audio_data"] = encode_audios_for_rollout_engine(audio_data) + if audio_inputs := extract_audio_inputs(multimodal_inputs): + payload["audio_data"] = encode_audios_for_rollout_engine(audio_inputs) return payload, None diff --git a/miles/utils/processing_utils.py b/miles/utils/processing_utils.py index fd31e2f602c..f204a370f26 100644 --- a/miles/utils/processing_utils.py +++ b/miles/utils/processing_utils.py @@ -211,6 +211,18 @@ def encode_audio_for_rollout_engine(audio, sampling_rate: int) -> str: return f"data:audio/wav;base64,{audio_base64}" +def extract_audio_inputs(multimodal_inputs: dict | None): + """Return the audio entries from ``multimodal_inputs`` (or ``None``). + + Canonical accessor shared by the rollout payload builders so the omni hook and the + generic path agree on keying. Accepts ``"audios"`` (plural, parallel to ``"images"``) + and the singular ``"audio"`` (matching ``MultimodalTypes.AUDIO.name``). + """ + if not multimodal_inputs: + return None + return multimodal_inputs.get("audios") or multimodal_inputs.get("audio") + + def encode_audios_for_rollout_engine(audios) -> list[str]: """Encode a list of waveform entries to base64 WAV data URIs for the rollout engine. diff --git a/miles_plugins/omni/omni_generate_fn.py b/miles_plugins/omni/omni_generate_fn.py index 0d02104e764..e586a4b22c6 100644 --- a/miles_plugins/omni/omni_generate_fn.py +++ b/miles_plugins/omni/omni_generate_fn.py @@ -13,7 +13,7 @@ from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput from miles.rollout.generate_utils.generate_endpoint_utils import compute_prompt_ids_from_sample from miles.utils.http_utils import post -from miles.utils.processing_utils import encode_audios_for_rollout_engine +from miles.utils.processing_utils import encode_audios_for_rollout_engine, extract_audio_inputs from miles.utils.types import Sample from .rollout_contract import apply_response_to_sample, build_generate_payload, parse_generate_response @@ -82,9 +82,7 @@ def _clamp_max_new_tokens(args, sampling_params: dict, prompt_len: int) -> Sampl def _encode_input_audio(sample: Sample) -> list[str] | None: """Encode input-side audio from ``sample.multimodal_inputs`` for the request payload.""" - if not sample.multimodal_inputs: - return None - audios = sample.multimodal_inputs.get("audios") or sample.multimodal_inputs.get("audio") + audios = extract_audio_inputs(sample.multimodal_inputs) if not audios: return None return encode_audios_for_rollout_engine(audios) diff --git a/miles_plugins/omni/rollout_contract.py b/miles_plugins/omni/rollout_contract.py index 2ab087caa47..436414c9f69 100644 --- a/miles_plugins/omni/rollout_contract.py +++ b/miles_plugins/omni/rollout_contract.py @@ -154,7 +154,9 @@ def apply_response_to_sample( Follows the miles convention where ``loss_mask`` and ``rollout_log_probs`` span only the generated (completion) tokens (length == ``response_length``); the prompt is - excluded by lying outside the mask rather than by leading zeros. Decoded response + excluded by lying outside the mask rather than by leading zeros. A loss mask is + appended for the new tokens whenever ``update_loss_mask`` is set OR a mask already + exists (partial-rollout off-policy masking), keeping it aligned with response_length. Decoded response ``audio`` is stored in ``sample.metadata`` (reward-facing), never in ``multimodal_train_inputs``. Standard meta_info handling (status, weight-version, prefix-cache stats) stays with the caller via the existing @@ -172,7 +174,11 @@ def apply_response_to_sample( sample.rollout_log_probs = [] sample.rollout_log_probs += result.response_log_probs - if update_loss_mask: + # Append mask entries when explicitly requested OR when a mask already exists. The + # latter covers partial-rollout off-policy masking, where generate_and_rm pre-sets + # loss_mask = [0] * old_response_length; the newly generated tokens are on-policy and + # trainable, and the mask must stay aligned with response_length. + if update_loss_mask or sample.loss_mask is not None: if sample.loss_mask is None: sample.loss_mask = [] sample.loss_mask += [1] * len(result.response_tokens) diff --git a/tests/fast/test_omni_generate_fn.py b/tests/fast/test_omni_generate_fn.py index 87ca6d574b9..f1faf802316 100644 --- a/tests/fast/test_omni_generate_fn.py +++ b/tests/fast/test_omni_generate_fn.py @@ -141,3 +141,43 @@ async def fake_post(url, payload, **kwargs): audio_data = captured["payload"]["audio_data"] assert len(audio_data) == 1 assert audio_data[0].startswith("data:audio/wav;base64,") + + +def test_omni_generate_fn_resume_keeps_loss_mask_aligned(monkeypatch): + async def fake_post(url, payload, **kwargs): + # resume turn: only the newly generated tokens come back + return { + "text": " more", + "meta_info": { + "finish_reason": {"type": "stop"}, + "output_token_logprobs": [[-0.3, 20], [-0.4, 21]], + "completion_tokens": 2, + "cached_tokens": 0, + "prompt_tokens": 5, + }, + } + + monkeypatch.setattr(omni_mod, "post", fake_post) + + fn = load_generate_function(_HOOK_PATH) + sample = Sample(prompt="hi") + # simulate a partial rollout whose off-policy response was pre-masked by generate_and_rm + sample.tokens = [1, 2, 3, 10, 11] # prompt [1,2,3] + old response [10,11] + sample.response = "old" + sample.response_length = 2 + sample.loss_mask = [0, 0] # off-policy tokens masked off + sample.rollout_log_probs = [-0.1, -0.2] + inp = GenerateFnInput( + state=_fake_state(), + sample=sample, + sampling_params={"max_new_tokens": 64}, + evaluation=False, + ) + + out = asyncio.run(fn(inp)) + s = out.samples + assert s.tokens == [1, 2, 3, 10, 11, 20, 21] + assert s.response_length == 4 + # new on-policy tokens are trainable; mask stays aligned with response_length + assert s.loss_mask == [0, 0, 1, 1] + assert len(s.loss_mask) == s.response_length diff --git a/tests/fast/test_omni_rollout_contract.py b/tests/fast/test_omni_rollout_contract.py index 06cccae1a1c..1cd73f1076c 100644 --- a/tests/fast/test_omni_rollout_contract.py +++ b/tests/fast/test_omni_rollout_contract.py @@ -176,6 +176,24 @@ def test_apply_response_to_sample_multi_turn_accumulates(): assert sample.loss_mask == [1, 1, 1] +def test_apply_response_to_sample_appends_to_existing_loss_mask(): + # partial-rollout off-policy masking: a [0] mask already exists, new tokens must be appended + sample = Sample( + prompt="p", + tokens=[1, 2, 3, 10], + response="old", + response_length=1, + loss_mask=[0], + rollout_log_probs=[-0.1], + ) + result = parse_generate_response(_response([[-0.2, 20], [-0.3, 21]], completion_tokens=2)) + apply_response_to_sample(sample, [1, 2, 3], result) # update_loss_mask defaults False + assert sample.response_length == 3 + assert sample.loss_mask == [0, 1, 1] + assert len(sample.loss_mask) == sample.response_length + sample.validate() + + # --- audio encode helper --------------------------------------------------------------- @@ -213,3 +231,44 @@ def test_encode_audios_for_rollout_engine_handles_tuples_and_dicts(): uris = encode_audios_for_rollout_engine(audios) assert len(uris) == 2 assert all(u.startswith("data:audio/wav;base64,") for u in uris) + + +# --- generic compute_request_payload audio + deferral marker ------------------------------- + + +@pytest.mark.parametrize("audio_key", ["audios", "audio"]) +def test_compute_request_payload_emits_audio_data(audio_key): + from types import SimpleNamespace + + from miles.rollout.generate_utils.generate_endpoint_utils import compute_request_payload + + args = SimpleNamespace( + rollout_max_response_len=128, + rollout_max_context_len=0, + use_rollout_routing_replay=False, + use_rollout_indexer_replay=False, + ) + payload, halt = compute_request_payload( + args, + input_ids=[1, 2, 3], + sampling_params={"max_new_tokens": 16}, + multimodal_inputs={audio_key: [(np.zeros(160, dtype=np.float32), 16000)]}, + ) + assert halt is None + assert len(payload["audio_data"]) == 1 + assert payload["audio_data"][0].startswith("data:audio/wav;base64,") + + +@pytest.mark.xfail( + reason=( + "audio-INPUT token expansion (audio placeholder -> feature tokens) is not implemented " + "in mm_data.py. It is only needed for audio-input models (e.g. Qwen3-Omni understanding), " + "NOT the text-input MVP gates, where codec OUTPUT tokens are first-class sequence tokens. " + "Deferred to the audio-input-model milestone." + ), + strict=True, +) +def test_mm_data_audio_input_token_expansion_present(): + from miles.backends.training_utils import mm_data + + assert hasattr(mm_data, "expand_audio_rollout_data_in_place") From f450d5b15c9eb9c44a8cf646bff198fe13878854 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Fri, 19 Jun 2026 23:06:41 +0800 Subject: [PATCH 04/24] feat(omni): thinker-submodule checkpoint extractor for FSDP training 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. --- miles_plugins/omni/extract_thinker.py | 126 ++++++++++++++++++++++++++ 1 file changed, 126 insertions(+) create mode 100644 miles_plugins/omni/extract_thinker.py diff --git a/miles_plugins/omni/extract_thinker.py b/miles_plugins/omni/extract_thinker.py new file mode 100644 index 00000000000..0fb37a45ae3 --- /dev/null +++ b/miles_plugins/omni/extract_thinker.py @@ -0,0 +1,126 @@ +"""Extract the Qwen3-Omni thinker submodule into a standalone HF checkpoint. + +The full Qwen3-Omni checkpoint stores the thinker / talker / code2wav stacks under +prefixed keys (``thinker.*`` / ``talker.*`` / ``code2wav.*``). For thinker-only RL +training, we strip the ``thinker.`` prefix and write a standalone +``Qwen3OmniMoeThinkerForConditionalGeneration`` checkpoint. Its config carries a +``vision_config``, so the FSDP actor's ``get_model_cls()`` loads it via +``AutoModelForImageTextToText`` (text-only input just runs the LM → logits). + +Streaming + RAM-bounded: tensors are copied shard-by-shard and flushed once a size +budget is reached, so the 30B model never needs to fit in memory at once. + +Usage:: + + python -m miles_plugins.omni.extract_thinker --src --dst +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +from collections import OrderedDict + +PREFIX = "thinker." + +# Tokenizer / aux files copied verbatim so the standalone dir is self-contained. +AUX_FILES = [ + "tokenizer.json", + "tokenizer_config.json", + "vocab.json", + "merges.txt", + "special_tokens_map.json", + "chat_template.json", + "chat_template.jinja", + "generation_config.json", + "preprocessor_config.json", +] + +DEFAULT_SHARD_BYTES = 5 * 1024**3 + + +def extract_thinker(src: str, dst: str, shard_bytes: int = DEFAULT_SHARD_BYTES) -> dict: + from safetensors import safe_open + from safetensors.torch import save_file + + os.makedirs(dst, exist_ok=True) + index = json.load(open(os.path.join(src, "model.safetensors.index.json"))) + weight_map = index["weight_map"] + input_shards = sorted(set(weight_map.values())) + + buf: "OrderedDict[str, object]" = OrderedDict() + buf_bytes = 0 + out_shards: list[str] = [] + out_weight_map: dict[str, str] = {} + total_bytes = 0 + + def flush() -> None: + nonlocal buf, buf_bytes + if not buf: + return + name = f"model-{len(out_shards) + 1:05d}.safetensors" + save_file(buf, os.path.join(dst, name), metadata={"format": "pt"}) + for key in buf: + out_weight_map[key] = name + out_shards.append(name) + buf = OrderedDict() + buf_bytes = 0 + + for shard in input_shards: + with safe_open(os.path.join(src, shard), framework="pt") as f: + for key in f.keys(): + if not key.startswith(PREFIX): + continue + tensor = f.get_tensor(key) + new_key = key[len(PREFIX) :] + buf[new_key] = tensor + nbytes = tensor.numel() * tensor.element_size() + buf_bytes += nbytes + total_bytes += nbytes + if buf_bytes >= shard_bytes: + flush() + flush() + + # Rename single-shard output to the conventional unsharded filename. + if len(out_shards) == 1: + only = out_shards[0] + os.replace(os.path.join(dst, only), os.path.join(dst, "model.safetensors")) + out_weight_map = {k: "model.safetensors" for k in out_weight_map} + json.dump( + {"metadata": {"total_size": total_bytes}, "weight_map": out_weight_map}, + open(os.path.join(dst, "model.safetensors.index.json"), "w"), + ) + + full_cfg = json.load(open(os.path.join(src, "config.json"))) + thinker_cfg = full_cfg["thinker_config"] + thinker_cfg["architectures"] = ["Qwen3OmniMoeThinkerForConditionalGeneration"] + json.dump(thinker_cfg, open(os.path.join(dst, "config.json"), "w"), indent=2) + + for fn in AUX_FILES: + src_path = os.path.join(src, fn) + if os.path.exists(src_path): + shutil.copy(src_path, dst) + + summary = { + "tensors": len(out_weight_map), + "total_gb": round(total_bytes / 1e9, 2), + "shards": len(out_weight_map and set(out_weight_map.values())), + "dst": dst, + } + print(f"extracted thinker: {summary}") + return summary + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--src", required=True, help="full Qwen3-Omni checkpoint dir") + parser.add_argument("--dst", required=True, help="output thinker checkpoint dir") + parser.add_argument("--shard-bytes", type=int, default=DEFAULT_SHARD_BYTES) + args = parser.parse_args() + extract_thinker(args.src, args.dst, args.shard_bytes) + + +if __name__ == "__main__": + main() From 0d82cfc147a92e8406fdaa1dcdef89533e3f5985 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Fri, 19 Jun 2026 23:08:25 +0800 Subject: [PATCH 05/24] feat(omni): text-only math reward for GATE-A thinker RL smoke 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. --- miles_plugins/omni/math_reward.py | 42 +++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 miles_plugins/omni/math_reward.py diff --git a/miles_plugins/omni/math_reward.py b/miles_plugins/omni/math_reward.py new file mode 100644 index 00000000000..0405a254f93 --- /dev/null +++ b/miles_plugins/omni/math_reward.py @@ -0,0 +1,42 @@ +"""Text-only math-correctness reward for the GATE-A thinker RL smoke. + +Loaded via ``--custom-rm-path miles_plugins.omni.math_reward.compute_math_reward``. +Returns 1.0 when the model's decoded response contains the gold answer +(``sample.label``), else 0.0. Deterministic and dependency-free, so it suits the first +``one_update_smoke``/``multi_step_stability`` run and the deterministic-reward criterion. +""" + +from __future__ import annotations + +import re + +from miles.utils.types import Sample + +_NUMBER_RE = re.compile(r"-?\d+(?:\.\d+)?") + + +def _numbers(text: str) -> list[str]: + return _NUMBER_RE.findall(text or "") + + +def _normalize(value: str) -> str: + # 12.0 and 12 should compare equal as answers + try: + f = float(value) + return str(int(f)) if f.is_integer() else str(f) + except ValueError: + return value.strip() + + +async def compute_math_reward(args, sample: Sample, **kwargs) -> float: + """1.0 if the response's answer matches the gold label, else 0.0.""" + label = "" if sample.label is None else str(sample.label).strip() + if not label: + return 0.0 + response = sample.response or "" + + gold = _normalize(label) + # numeric match (handles "= 12", "12.0", trailing punctuation), then substring fallback + if any(_normalize(n) == gold for n in _numbers(response)): + return 1.0 + return 1.0 if label in response else 0.0 From 1230255eec28d97e11430c01a2ad70a1afd52f86 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Fri, 19 Jun 2026 23:17:26 +0800 Subject: [PATCH 06/24] feat(omni): GATE-A closed-loop GRPO LoRA smoke + math dataset 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). --- examples/omni_gate_a/gate_a_lora_smoke.py | 112 ++++++++++++++++++++++ examples/omni_gate_a/math_smoke.jsonl | 8 ++ 2 files changed, 120 insertions(+) create mode 100644 examples/omni_gate_a/gate_a_lora_smoke.py create mode 100644 examples/omni_gate_a/math_smoke.jsonl diff --git a/examples/omni_gate_a/gate_a_lora_smoke.py b/examples/omni_gate_a/gate_a_lora_smoke.py new file mode 100644 index 00000000000..1698441639e --- /dev/null +++ b/examples/omni_gate_a/gate_a_lora_smoke.py @@ -0,0 +1,112 @@ +"""GATE-A closed-loop smoke: GRPO RL on the Qwen3-Omni Thinker (single GPU, LoRA). + +Demonstrates all four closed-loop components on the real model end to end: + rollout (sglang-omni /generate) -> reward (math correctness) + -> GRPO advantage -> LoRA policy-gradient update. + +It is a deliberately minimal, self-contained harness (no Ray / FSDP / miles trainer) +so the loop mechanics can be verified on one 80GB GPU. The behavior policy is the +served base thinker; LoRA is updated locally (rollouts are not weight-synced back, which +is the documented next integration step), so this proves loop stability, not on-policy +convergence. + +Run (inside the container, miles venv): + THINKER=/root/qwen3-omni-thinker DATA=examples/omni_gate_a/math_smoke.jsonl \ + SERVER=http://localhost:8000/generate CUDA_VISIBLE_DEVICES=4 \ + python examples/omni_gate_a/gate_a_lora_smoke.py +""" + +from __future__ import annotations + +import json +import os +import urllib.request + +import torch +from peft import LoraConfig, get_peft_model +from transformers import AutoModelForImageTextToText, AutoTokenizer + +SERVER = os.environ.get("SERVER", "http://localhost:8000/generate") +THINKER = os.environ["THINKER"] +DATA = os.environ["DATA"] +STEPS = int(os.environ.get("STEPS", "5")) +GROUP = int(os.environ.get("GROUP", "4")) +PROMPTS_PER_STEP = int(os.environ.get("PROMPTS_PER_STEP", "4")) +EPS = 0.2 + + +def rollout(input_ids: list[int], seed: int): + req = { + "input_ids": input_ids, + "sampling_params": {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": 24, "seed": seed}, + "return_logprob": True, + } + r = urllib.request.urlopen( + urllib.request.Request(SERVER, data=json.dumps(req).encode(), headers={"Content-Type": "application/json"}), + timeout=120, + ) + resp = json.loads(r.read()) + otl = resp["meta_info"]["output_token_logprobs"] + return { + "tokens": [t for _, t in otl], + "old_logprobs": [lp for lp, _ in otl], + "text": resp.get("text", ""), + } + + +def main() -> None: + tok = AutoTokenizer.from_pretrained(THINKER, trust_remote_code=True) + model = AutoModelForImageTextToText.from_pretrained( + THINKER, dtype=torch.bfloat16, trust_remote_code=True, low_cpu_mem_usage=True + ).to("cuda:0") + model = get_peft_model( + model, LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"], task_type="CAUSAL_LM") + ) + model.train() + trainable = [p for p in model.parameters() if p.requires_grad] + opt = torch.optim.AdamW(trainable, lr=1e-5) + print(f"trainable params: {sum(p.numel() for p in trainable) / 1e6:.1f}M") + + data = [json.loads(line) for line in open(DATA)] + print("step | mean_reward | avg_loss | (per-prompt rewards)") + for step in range(STEPS): + opt.zero_grad() + step_reward, step_loss, n = 0.0, 0.0, 0 + per_prompt = [] + for ex in data[:PROMPTS_PER_STEP]: + prompt_ids = tok.encode(ex["prompt"]) + samples = [rollout(prompt_ids, step * 1000 + g) for g in range(GROUP)] + rewards = [1.0 if ex["label"] in s["text"] else 0.0 for s in samples] + mean_r = sum(rewards) / len(rewards) + per_prompt.append(mean_r) + step_reward += mean_r + advs = [r - mean_r for r in rewards] + for s, adv in zip(samples, advs): + rt = s["tokens"] + if not rt or adv == 0.0: + continue + full = torch.tensor([prompt_ids + rt], device="cuda:0") + logits = model(input_ids=full).logits[0] + p = len(prompt_ids) + resp_logits = logits[p - 1 : p - 1 + len(rt)].float() + logp = torch.log_softmax(resp_logits, dim=-1) + idx = torch.tensor(rt, device="cuda:0") + new_lp = logp[range(len(rt)), idx] + old_lp = torch.tensor(s["old_logprobs"], device="cuda:0") + ratio = torch.exp(new_lp - old_lp) + loss = -torch.min(ratio * adv, torch.clamp(ratio, 1 - EPS, 1 + EPS) * adv).mean() + loss = loss / (GROUP * PROMPTS_PER_STEP) + loss.backward() + step_loss += loss.item() * (GROUP * PROMPTS_PER_STEP) + n += 1 + torch.nn.utils.clip_grad_norm_(trainable, 1.0) + opt.step() + mean_reward = step_reward / PROMPTS_PER_STEP + avg_loss = step_loss / max(n, 1) + print(f"{step:4d} | {mean_reward:11.3f} | {avg_loss:8.4f} | {per_prompt}") + + print("GATE-A closed-loop smoke complete (rollout->reward->advantage->LoRA update over multiple steps)") + + +if __name__ == "__main__": + main() diff --git a/examples/omni_gate_a/math_smoke.jsonl b/examples/omni_gate_a/math_smoke.jsonl new file mode 100644 index 00000000000..ea2577d2ef6 --- /dev/null +++ b/examples/omni_gate_a/math_smoke.jsonl @@ -0,0 +1,8 @@ +{"prompt": "Question: What is 7 plus 5?\nAnswer:", "label": "12"} +{"prompt": "Question: What is 9 plus 6?\nAnswer:", "label": "15"} +{"prompt": "Question: What is 8 times 3?\nAnswer:", "label": "24"} +{"prompt": "Question: What is 12 minus 4?\nAnswer:", "label": "8"} +{"prompt": "Question: What is 6 times 7?\nAnswer:", "label": "42"} +{"prompt": "Question: What is 20 plus 13?\nAnswer:", "label": "33"} +{"prompt": "Question: What is 45 minus 18?\nAnswer:", "label": "27"} +{"prompt": "Question: What is 9 times 9?\nAnswer:", "label": "81"} From 6e6626db20192bae2f5cf5a5419945d40cf50876 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Fri, 19 Jun 2026 23:43:12 +0800 Subject: [PATCH 07/24] feat(omni): full on-policy GATE-A with NCCL weight-sync to served thinker 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. --- examples/omni_gate_a/gate_a_full.py | 185 ++++++++++++++++++++++++++++ 1 file changed, 185 insertions(+) create mode 100644 examples/omni_gate_a/gate_a_full.py diff --git a/examples/omni_gate_a/gate_a_full.py b/examples/omni_gate_a/gate_a_full.py new file mode 100644 index 00000000000..8ce242abc14 --- /dev/null +++ b/examples/omni_gate_a/gate_a_full.py @@ -0,0 +1,185 @@ +"""Full on-policy GATE-A: GRPO LoRA RL on the Qwen3-Omni Thinker with per-step NCCL +weight-sync to the served sglang-omni thinker, so each step's rollouts are on-policy. + +Beyond gate_a_lora_smoke.py this adds the 4th closed-loop component done *properly*: +after each optimizer step the LoRA-merged thinker weights are broadcast into the served +thinker stage via sglang-omni's distributed weight-update admin plane +(``/init_weights_update_group`` + ``/update_weights_from_distributed`` + ``stages=[thinker]``), +the exact pattern from sglang-omni's E2E refit test. The thinker stage's ``load_weights`` +accepts plain ``model.*`` names, so the extracted-thinker names sync directly. + +Run (container, miles venv, free GPU for the trainer; server already on another GPU): + THINKER=/root/qwen3-omni-thinker DATA=examples/omni_gate_a/math_smoke.jsonl \ + SERVER=http://localhost:8000 MASTER_PORT=29555 CUDA_DEVICE_ORDER=PCI_BUS_ID \ + CUDA_VISIBLE_DEVICES=4 python examples/omni_gate_a/gate_a_full.py +""" + +from __future__ import annotations + +import json +import os +import threading +import urllib.request + +import torch +from peft import LoraConfig, get_peft_model +from transformers import AutoModelForImageTextToText, AutoTokenizer + +SERVER = os.environ.get("SERVER", "http://localhost:8000") +THINKER = os.environ["THINKER"] +DATA = os.environ["DATA"] +STEPS = int(os.environ.get("STEPS", "4")) +GROUP = int(os.environ.get("GROUP", "4")) +PROMPTS = int(os.environ.get("PROMPTS", "4")) +MASTER_PORT = int(os.environ.get("MASTER_PORT", "29555")) +GROUP_NAME = "gate_a_wsync" +EPS = 0.2 + + +def post(path: str, body: dict, timeout: int = 300): + req = urllib.request.Request( + SERVER + path, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"} + ) + return json.loads(urllib.request.urlopen(req, timeout=timeout).read()) + + +def rollout(input_ids: list[int], seed: int): + resp = post( + "/generate", + { + "input_ids": input_ids, + "sampling_params": {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": 24, "seed": seed}, + "return_logprob": True, + }, + timeout=120, + ) + otl = resp["meta_info"]["output_token_logprobs"] + return {"tokens": [t for _, t in otl], "old": [lp for lp, _ in otl], "text": resp.get("text", "")} + + +def main() -> None: + tok = AutoTokenizer.from_pretrained(THINKER, trust_remote_code=True) + model = AutoModelForImageTextToText.from_pretrained( + THINKER, dtype=torch.bfloat16, trust_remote_code=True, low_cpu_mem_usage=True + ).to("cuda:0") + model = get_peft_model( + model, LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"], task_type="CAUSAL_LM") + ) + model.train() + opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=2e-5) + + try: + from sglang.srt.utils import init_custom_process_group + except Exception: + from sglang.srt.utils.common import init_custom_process_group + + # Rendezvous a 2-rank NCCL group with the served thinker (server joins as rank 1). + init_err: list = [] + + def _init_server(): + try: + post( + "/init_weights_update_group", + { + "master_address": "localhost", + "master_port": MASTER_PORT, + "rank_offset": 1, + "world_size": 2, + "group_name": GROUP_NAME, + "backend": "nccl", + "stages": ["thinker"], + }, + timeout=180, + ) + except Exception as exc: # noqa: BLE001 + init_err.append(exc) + + th = threading.Thread(target=_init_server) + th.start() + pg = init_custom_process_group( + backend="nccl", init_method=f"tcp://localhost:{MASTER_PORT}", world_size=2, rank=0, group_name=GROUP_NAME + ) + th.join() + torch.cuda.synchronize() + if init_err: + raise init_err[0] + print("WEIGHT_UPDATE_GROUP_READY", flush=True) + + def merged_lora_weights() -> dict[str, torch.Tensor]: + out: dict[str, torch.Tensor] = {} + for name, mod in model.named_modules(): + if hasattr(mod, "lora_A") and hasattr(mod, "base_layer"): + a = mod.lora_A["default"].weight + b = mod.lora_B["default"].weight + scaling = mod.scaling["default"] + w = mod.base_layer.weight.data + scaling * (b @ a) + hf = name.replace("base_model.model.", "") + ".weight" + out[hf] = w.to(torch.bfloat16).contiguous() + return out + + def sync_to_server() -> int: + wd = merged_lora_weights() + names = sorted(wd) + spec = { + "names": names, + "dtypes": [str(wd[n].dtype).replace("torch.", "") for n in names], + "shapes": [list(wd[n].shape) for n in names], + "group_name": GROUP_NAME, + "stages": ["thinker"], + } + err: list = [] + + def _update(): + try: + post("/update_weights_from_distributed", spec, timeout=300) + except Exception as exc: # noqa: BLE001 + err.append(exc) + + t = threading.Thread(target=_update) + t.start() + for n in names: + torch.distributed.broadcast(wd[n], src=0, group=pg) + torch.cuda.synchronize() + t.join() + if err: + raise err[0] + return len(names) + + data = [json.loads(line) for line in open(DATA)] + print("step | mean_reward | avg_loss | synced_params") + for step in range(STEPS): + opt.zero_grad() + step_reward, step_loss, n = 0.0, 0.0, 0 + for ex in data[:PROMPTS]: + pid = tok.encode(ex["prompt"]) + samples = [rollout(pid, step * 1000 + g) for g in range(GROUP)] + rewards = [1.0 if ex["label"] in s["text"] else 0.0 for s in samples] + mean_r = sum(rewards) / len(rewards) + step_reward += mean_r + for s, adv in zip(samples, [r - mean_r for r in rewards]): + if not s["tokens"] or adv == 0.0: + continue + full = torch.tensor([pid + s["tokens"]], device="cuda:0") + logits = model(input_ids=full).logits[0] + p = len(pid) + rl = logits[p - 1 : p - 1 + len(s["tokens"])].float() + logp = torch.log_softmax(rl, dim=-1) + rt = torch.tensor(s["tokens"], device="cuda:0") + new = logp[range(len(s["tokens"])), rt] + old = torch.tensor(s["old"], device="cuda:0") + ratio = torch.exp(new - old) + loss = -torch.min(ratio * adv, torch.clamp(ratio, 1 - EPS, 1 + EPS) * adv).mean() + loss = loss / (GROUP * PROMPTS) + loss.backward() + step_loss += loss.item() * (GROUP * PROMPTS) + n += 1 + torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0) + opt.step() + synced = sync_to_server() # next step's rollouts are on-policy + print(f"{step:4d} | {step_reward / PROMPTS:11.3f} | {step_loss / max(n, 1):8.4f} | {synced}", flush=True) + + print("GATE-A FULL on-policy loop complete (per-step NCCL weight-sync to served thinker)") + + +if __name__ == "__main__": + main() From 2a72ca226f1638565966fe71a4a2fda9f279c962 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Fri, 19 Jun 2026 23:58:19 +0800 Subject: [PATCH 08/24] fix(omni): bind CUDA device before NCCL collectives in full GATE-A 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). --- examples/omni_gate_a/gate_a_full.py | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/omni_gate_a/gate_a_full.py b/examples/omni_gate_a/gate_a_full.py index 8ce242abc14..5db690ac1ea 100644 --- a/examples/omni_gate_a/gate_a_full.py +++ b/examples/omni_gate_a/gate_a_full.py @@ -58,6 +58,7 @@ def rollout(input_ids: list[int], seed: int): def main() -> None: + torch.cuda.set_device(0) # bind this process to its visible GPU for NCCL collectives tok = AutoTokenizer.from_pretrained(THINKER, trust_remote_code=True) model = AutoModelForImageTextToText.from_pretrained( THINKER, dtype=torch.bfloat16, trust_remote_code=True, low_cpu_mem_usage=True From 9a7f7bf241a902f7db391852218f2755a81072c0 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Sat, 20 Jun 2026 01:14:08 +0800 Subject: [PATCH 09/24] feat(omni): make weight-sync GROUP_NAME env-configurable 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. --- examples/omni_gate_a/gate_a_full.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/omni_gate_a/gate_a_full.py b/examples/omni_gate_a/gate_a_full.py index 5db690ac1ea..e718c62aa28 100644 --- a/examples/omni_gate_a/gate_a_full.py +++ b/examples/omni_gate_a/gate_a_full.py @@ -32,7 +32,7 @@ GROUP = int(os.environ.get("GROUP", "4")) PROMPTS = int(os.environ.get("PROMPTS", "4")) MASTER_PORT = int(os.environ.get("MASTER_PORT", "29555")) -GROUP_NAME = "gate_a_wsync" +GROUP_NAME = os.environ.get("GROUP_NAME", "gate_a_wsync") EPS = 0.2 From 059c28dc9e055eb9bf188b78c4be1e79d4a65c52 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Sun, 21 Jun 2026 00:27:26 +0800 Subject: [PATCH 10/24] docs(omni): record NCCL_P2P_DISABLE fix that makes GATE-A weight-sync work MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- examples/omni_gate_a/gate_a_full.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/examples/omni_gate_a/gate_a_full.py b/examples/omni_gate_a/gate_a_full.py index e718c62aa28..9bcf9607678 100644 --- a/examples/omni_gate_a/gate_a_full.py +++ b/examples/omni_gate_a/gate_a_full.py @@ -10,8 +10,16 @@ Run (container, miles venv, free GPU for the trainer; server already on another GPU): THINKER=/root/qwen3-omni-thinker DATA=examples/omni_gate_a/math_smoke.jsonl \ - SERVER=http://localhost:8000 MASTER_PORT=29555 CUDA_DEVICE_ORDER=PCI_BUS_ID \ - CUDA_VISIBLE_DEVICES=4 python examples/omni_gate_a/gate_a_full.py + SERVER=http://localhost:8003 MASTER_PORT=29631 CUDA_VISIBLE_DEVICES=4 \ + NCCL_P2P_DISABLE=1 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0 \ + python examples/omni_gate_a/gate_a_full.py + +CRITICAL: the trainer and the sglang-omni server run as separate processes, each with a +single GPU exposed via CUDA_VISIBLE_DEVICES (both see it as cuda:0). NCCL would try direct +P2P between the two physical GPUs and fail with "Cuda invalid argument" because neither +process can resolve the peer's masked device. Set NCCL_P2P_DISABLE=1 on BOTH the server +and the trainer so NCCL falls back to shared-memory transport. Verified: 4-step on-policy +run, synced_params=160/step, stable loss/reward. """ from __future__ import annotations From 76d75bc24806d8ffc285b4709f38f574066d09fb Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Sun, 21 Jun 2026 00:44:09 +0800 Subject: [PATCH 11/24] feat(omni): composite TTS reward (ASR CER + audio guards) for GATE-B 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). --- examples/omni_gate_b/tts_smoke.jsonl | 6 + miles_plugins/omni/tts_reward.py | 171 +++++++++++++++++++++++++++ 2 files changed, 177 insertions(+) create mode 100644 examples/omni_gate_b/tts_smoke.jsonl create mode 100644 miles_plugins/omni/tts_reward.py diff --git a/examples/omni_gate_b/tts_smoke.jsonl b/examples/omni_gate_b/tts_smoke.jsonl new file mode 100644 index 00000000000..3cf11ebc7c0 --- /dev/null +++ b/examples/omni_gate_b/tts_smoke.jsonl @@ -0,0 +1,6 @@ +{"text": "Hello world.", "label": "Hello world."} +{"text": "The quick brown fox.", "label": "The quick brown fox."} +{"text": "Good morning everyone.", "label": "Good morning everyone."} +{"text": "Thank you very much.", "label": "Thank you very much."} +{"text": "How are you today?", "label": "How are you today?"} +{"text": "This is a test.", "label": "This is a test."} diff --git a/miles_plugins/omni/tts_reward.py b/miles_plugins/omni/tts_reward.py new file mode 100644 index 00000000000..1dbf818cf9e --- /dev/null +++ b/miles_plugins/omni/tts_reward.py @@ -0,0 +1,171 @@ +"""Composite reward for TTS RL (GATE-B): ASR round-trip CER + audio-validity guards. + +The TTS actor generates speech for a target text. The reward transcribes the generated +audio with an ASR model (Whisper) and scores content correctness via CER, combined with +hard audio-validity guards (decode success, duration bounds, non-silence). A failed decode +yields a deterministic low reward instead of crashing, so the loop never wedges. This is +the DEC-2 composite design: ASR alone rewards transcribable-but-degenerate audio, so the +guards must be present from day one. + +Usable two ways: + - standalone: ``TtsCompositeReward(...).score(audio_b64, target_text)`` + - miles hook: ``--custom-rm-path miles_plugins.omni.tts_reward.compute_tts_reward`` + (reads the generated audio from ``sample.metadata["generated_audio"]`` and the target + text from ``sample.label``). +""" + +from __future__ import annotations + +import base64 +import io +import os +import re +import wave +from dataclasses import dataclass, field +from typing import Any + +# Suggested defaults (tunable). Duration in seconds; energy is RMS of float[-1,1] samples. +MIN_DURATION_S = 0.3 +MAX_DURATION_S = 30.0 +SILENCE_RMS_FLOOR = 1e-3 +DECODE_FAIL_REWARD = -1.0 +ASR_WEIGHT = 1.0 + +_WORD_RE = re.compile(r"[a-z0-9]+") + + +def _normalize_text(text: str) -> str: + return " ".join(_WORD_RE.findall((text or "").lower())) + + +def _char_error_rate(hyp: str, ref: str) -> float: + """Levenshtein char edit distance / len(ref), clamped to [0, 1].""" + ref = _normalize_text(ref) + hyp = _normalize_text(hyp) + if not ref: + return 0.0 if not hyp else 1.0 + prev = list(range(len(hyp) + 1)) + for i, rc in enumerate(ref, 1): + cur = [i] + for j, hc in enumerate(hyp, 1): + cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (rc != hc))) + prev = cur + return min(1.0, prev[-1] / len(ref)) + + +def _decode_wav_b64(data: str) -> tuple[Any, int]: + """Decode a base64 (optionally data-URI) WAV string to (float32 mono waveform, sr).""" + import numpy as np + + if "," in data and data.strip().startswith("data:"): + data = data.split(",", 1)[1] + raw = base64.b64decode(data) + with wave.open(io.BytesIO(raw), "rb") as wf: + sr = wf.getframerate() + n = wf.getnframes() + ch = wf.getnchannels() + pcm = np.frombuffer(wf.readframes(n), dtype=np.int16).astype(np.float32) / 32768.0 + if ch > 1: + pcm = pcm.reshape(-1, ch).mean(axis=1) + return pcm, sr + + +@dataclass +class RewardComponents: + reward: float + cer: float | None = None + duration_s: float | None = None + rms: float | None = None + transcript: str = "" + guard: str = "ok" # "ok" | "decode_fail" | "too_short" | "too_long" | "silent" + + def to_dict(self) -> dict[str, Any]: + return {k: v for k, v in self.__dict__.items()} + + +@dataclass +class TtsCompositeReward: + asr_model_path: str = field(default_factory=lambda: os.environ.get("ASR_MODEL", "openai/whisper-base")) + device: str = field(default_factory=lambda: os.environ.get("ASR_DEVICE", "cuda:0")) + asr_weight: float = ASR_WEIGHT + min_duration_s: float = MIN_DURATION_S + max_duration_s: float = MAX_DURATION_S + silence_rms_floor: float = SILENCE_RMS_FLOOR + + _model: Any = None + _processor: Any = None + + def _ensure_asr(self) -> None: + if self._model is not None: + return + import torch + from transformers import WhisperForConditionalGeneration, WhisperProcessor + + self._processor = WhisperProcessor.from_pretrained(self.asr_model_path) + self._model = ( + WhisperForConditionalGeneration.from_pretrained(self.asr_model_path, dtype=torch.float16) + .to(self.device) + .eval() + ) + + def transcribe(self, waveform, sr: int) -> str: + import torch + import torchaudio.functional as AF + + self._ensure_asr() + wav = torch.as_tensor(waveform).float() + if sr != 16000: + wav = AF.resample(wav, sr, 16000) + feats = self._processor( + wav.numpy(), sampling_rate=16000, return_tensors="pt" + ).input_features.to(self.device, dtype=self._model.dtype) + with torch.no_grad(): + ids = self._model.generate(feats, language="en", task="transcribe", max_new_tokens=128) + return self._processor.batch_decode(ids, skip_special_tokens=True)[0].strip() + + def score(self, audio_b64: str | None, target_text: str) -> RewardComponents: + """Composite reward in [DECODE_FAIL_REWARD, asr_weight]; never raises.""" + if not audio_b64: + return RewardComponents(reward=DECODE_FAIL_REWARD, guard="decode_fail") + try: + wav, sr = _decode_wav_b64(audio_b64) + except Exception: + return RewardComponents(reward=DECODE_FAIL_REWARD, guard="decode_fail") + + import numpy as np + + duration = len(wav) / sr if sr else 0.0 + rms = float(np.sqrt(np.mean(wav**2))) if len(wav) else 0.0 + if duration < self.min_duration_s: + return RewardComponents(reward=DECODE_FAIL_REWARD, duration_s=duration, rms=rms, guard="too_short") + if duration > self.max_duration_s: + return RewardComponents(reward=DECODE_FAIL_REWARD, duration_s=duration, rms=rms, guard="too_long") + if rms < self.silence_rms_floor: + return RewardComponents(reward=DECODE_FAIL_REWARD, duration_s=duration, rms=rms, guard="silent") + + try: + transcript = self.transcribe(wav, sr) + except Exception: + return RewardComponents(reward=DECODE_FAIL_REWARD, duration_s=duration, rms=rms, guard="decode_fail") + + cer = _char_error_rate(transcript, target_text) + reward = self.asr_weight * (1.0 - cer) + return RewardComponents( + reward=reward, cer=cer, duration_s=duration, rms=rms, transcript=transcript, guard="ok" + ) + + +_SHARED: TtsCompositeReward | None = None + + +async def compute_tts_reward(args, sample, **kwargs) -> float: + """miles --custom-rm-path hook: score generated audio against sample.label.""" + global _SHARED + if _SHARED is None: + _SHARED = TtsCompositeReward() + audio = (sample.metadata or {}).get("generated_audio") + audio_b64 = audio.get("data") if isinstance(audio, dict) else audio + comp = _SHARED.score(audio_b64, str(sample.label or "")) + if isinstance(sample.metadata, dict): + sample.metadata["tts_reward_components"] = comp.to_dict() + return comp.reward From 074b5e71c876fa80a44ce11d83d747d01005a46a Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Mon, 22 Jun 2026 14:52:47 +0800 Subject: [PATCH 12/24] feat(omni): GATE-B rollout->reward->advantage loop on real Higgs TTS 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). --- examples/omni_gate_b/gate_b_loop.py | 104 +++++++++++++++ .../higgs_codec_logprob.patch | 126 ++++++++++++++++++ 2 files changed, 230 insertions(+) create mode 100644 examples/omni_gate_b/gate_b_loop.py create mode 100644 examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob.patch diff --git a/examples/omni_gate_b/gate_b_loop.py b/examples/omni_gate_b/gate_b_loop.py new file mode 100644 index 00000000000..94dd3d52622 --- /dev/null +++ b/examples/omni_gate_b/gate_b_loop.py @@ -0,0 +1,104 @@ +"""GATE-B closed loop demo: GRPO rollout -> composite reward -> advantage on real Higgs TTS. + +Demonstrates the first three closed-loop components on the real Higgs-audio model through +the sglang-omni rollout backend: + rollout (pretok /generate -> codec tokens + logprobs + audio) + -> composite reward (Whisper ASR CER + audio-validity guards) + -> GRPO advantage. + +The 4th component (LoRA policy update + NCCL weight-sync to the served TTS actor) mirrors +GATE-A's gate_a_full.py: the rollout returns codec-token logprobs (old) and the trainer +recomputes new logprobs over the codec sequence; weight sync uses /update_weights_from_distributed +with NCCL_P2P_DISABLE=1. + +Run (container, miles venv; Higgs server already serving on SERVER): + THINKER=... SERVER=http://localhost:8010 HIGGS_CKPT= \ + ASR_MODEL=openai/whisper-base ASR_DEVICE=cuda:0 \ + python examples/omni_gate_b/gate_b_loop.py +""" + +from __future__ import annotations + +import glob +import json +import os +import urllib.request + +SERVER = os.environ.get("SERVER", "http://localhost:8010") +DATA = os.environ.get("DATA", "examples/omni_gate_b/tts_smoke.jsonl") +GROUP = int(os.environ.get("GROUP", "4")) +STEPS = int(os.environ.get("STEPS", "3")) + + +def _higgs_adapter(): + from tokenizers import Tokenizer + from transformers import PreTrainedTokenizerFast + + from sglang_omni.models.higgs_tts.text_tokenizer import HiggsTokenizerAdapter + + ckpt = glob.glob(os.environ["HIGGS_CKPT"])[0] if "*" in os.environ.get("HIGGS_CKPT", "") else os.environ["HIGGS_CKPT"] + tok = PreTrainedTokenizerFast(tokenizer_object=Tokenizer.from_file(os.path.join(ckpt, "tokenizer.json"))) + return HiggsTokenizerAdapter(tok) + + +def rollout(input_ids: list[int], seed: int) -> dict: + req = { + "input_ids": input_ids, + "sampling_params": {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": 256, "seed": seed}, + "return_logprob": True, + "output_modalities": ["audio"], + } + r = urllib.request.urlopen( + urllib.request.Request(SERVER + "/generate", data=json.dumps(req).encode(), + headers={"Content-Type": "application/json"}), + timeout=180, + ) + resp = json.loads(r.read()) + otl = resp["meta_info"].get("output_token_logprobs") or [] + audio = resp.get("audio") or {} + return { + "codec_tokens": [t for _, t in otl], + "old_logprobs": [lp for lp, _ in otl], + "audio_b64": audio.get("data"), + } + + +def main() -> None: + import sys + + sys.path.insert(0, os.environ.get("SGLANG_OMNI", "/root/rl-omni/sglang-omni")) + from miles_plugins.omni.tts_reward import TtsCompositeReward + + adapter = _higgs_adapter() + reward_fn = TtsCompositeReward() + data = [json.loads(line) for line in open(DATA)] + + print("step | mean_reward | mean_cer | (per-prompt reward)") + for step in range(STEPS): + step_reward, step_cer, n_cer = 0.0, 0.0, 0 + per_prompt = [] + for ex in data[: int(os.environ.get("PROMPTS", "4"))]: + pid = list(map(int, adapter.build_prompt(ex["text"], num_ref_tokens=0))) + samples = [rollout(pid, step * 1000 + g) for g in range(GROUP)] + comps = [reward_fn.score(s["audio_b64"], ex["label"]) for s in samples] + rewards = [c.reward for c in comps] + mean_r = sum(rewards) / len(rewards) + advs = [round(r - mean_r, 3) for r in rewards] + per_prompt.append(round(mean_r, 3)) + step_reward += mean_r + for c in comps: + if c.cer is not None: + step_cer += c.cer + n_cer += 1 + # codec-token alignment sanity (old logprobs vs tokens) + assert all(len(s["codec_tokens"]) == len(s["old_logprobs"]) for s in samples) + print(f" prompt={ex['text']!r} rewards={rewards} adv={advs} " + f"transcripts={[c.transcript for c in comps]}") + mean_cer = step_cer / n_cer if n_cer else float("nan") + print(f"{step:4d} | {step_reward / len(per_prompt):11.3f} | {mean_cer:8.3f} | {per_prompt}", flush=True) + + print("GATE-B rollout->composite-reward->advantage demonstrated on real Higgs TTS") + + +if __name__ == "__main__": + main() diff --git a/examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob.patch b/examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob.patch new file mode 100644 index 00000000000..b24b8b1ac03 --- /dev/null +++ b/examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob.patch @@ -0,0 +1,126 @@ +diff --git a/sglang_omni/models/higgs_tts/model_runner.py b/sglang_omni/models/higgs_tts/model_runner.py +index c579e46..0d04a59 100644 +--- a/sglang_omni/models/higgs_tts/model_runner.py ++++ b/sglang_omni/models/higgs_tts/model_runner.py +@@ -299,6 +299,8 @@ class HiggsTTSModelRunner(ModelRunner): + continue + codes_N = codes_BN_cpu[b].to(torch.long).clone() + data.output_codes.append(codes_N) ++ nt = getattr(result.logits_output, "next_token_logits", None) ++ self._record_rollout_logprob(data, nt[b] if nt is not None else None, int(codes_N[0].item())) + data.generation_done = bool(gen_done_after_cpu[b]) + self._emit_code_chunk(sched_req, codes_N) + self._mark_sampler_finished(req, data.generation_done) +@@ -364,7 +366,7 @@ class HiggsTTSModelRunner(ModelRunner): + + model = self.model + cb0_per_row: list[int] = [] +- for sched_req in requests: ++ for b, sched_req in enumerate(requests): + data = sched_req.data + req = data.req + rid = sched_req.request_id +@@ -375,6 +377,8 @@ class HiggsTTSModelRunner(ModelRunner): + continue + codes_N = codes_log[-1] + data.output_codes.append(codes_N.detach().cpu().clone()) ++ nt = getattr(result.logits_output, "next_token_logits", None) ++ self._record_rollout_logprob(data, nt[b] if nt is not None else None, int(codes_N[0].item())) + data.generation_done = bool(model._sampler_pool.generation_done[row].item()) + self._emit_code_chunk(sched_req, data.output_codes[-1]) + self._mark_sampler_finished(req, data.generation_done) +@@ -392,6 +396,17 @@ class HiggsTTSModelRunner(ModelRunner): + if generation_done and req.finished_reason is None: + req.finished_reason = FINISH_MATCHED_TOKEN(EOC_ID) + ++ @staticmethod ++ def _record_rollout_logprob(data, logits_row, cb0_token): ++ """Record codebook-0 codec token + its logprob for RL rollout.""" ++ if not getattr(data, "return_logprob", False): ++ return ++ if logits_row is None: ++ data.output_token_logprobs.append([0.0, int(cb0_token)]) ++ return ++ logp = torch.log_softmax(logits_row.float(), dim=-1) ++ data.output_token_logprobs.append([float(logp[int(cb0_token)].item()), int(cb0_token)]) ++ + def _emit_code_chunk(self, sched_req: Any, codes_N: torch.Tensor) -> None: + if self._outbox is None: + return +diff --git a/sglang_omni/models/higgs_tts/payload_types.py b/sglang_omni/models/higgs_tts/payload_types.py +index 079f540..11416f3 100644 +--- a/sglang_omni/models/higgs_tts/payload_types.py ++++ b/sglang_omni/models/higgs_tts/payload_types.py +@@ -46,6 +46,9 @@ class HiggsTtsState: + audio_samples: Any | None = None + sample_rate: int = 24000 + ++ # rollout ++ output_token_logprobs: list[Any] | None = None ++ + def to_dict(self) -> dict[str, Any]: + data: dict[str, Any] = { + "prompt_token_ids": list(self.prompt_token_ids), +@@ -81,6 +84,8 @@ class HiggsTtsState: + if self.audio_samples is not None: + data["audio_samples"] = self.audio_samples + data["sample_rate"] = self.sample_rate ++ if self.output_token_logprobs is not None: ++ data["output_token_logprobs"] = self.output_token_logprobs + return data + + @classmethod +@@ -107,6 +112,7 @@ class HiggsTtsState: + engine_time_s=data.get("engine_time_s", 0.0), + audio_samples=data.get("audio_samples"), + sample_rate=data.get("sample_rate", 24000), ++ output_token_logprobs=data.get("output_token_logprobs"), + ) + + +diff --git a/sglang_omni/models/higgs_tts/request_builders.py b/sglang_omni/models/higgs_tts/request_builders.py +index f5b82dd..6d3d7d3 100644 +--- a/sglang_omni/models/higgs_tts/request_builders.py ++++ b/sglang_omni/models/higgs_tts/request_builders.py +@@ -153,6 +153,9 @@ def apply_higgs_result(state: HiggsTtsState, data: HiggsSGLangRequestData) -> No + else: + state.output_codes_delayed = None + state.prompt_tokens = len(data.input_ids) ++ state.output_token_logprobs = ( ++ list(data.output_token_logprobs) if data.output_token_logprobs else None ++ ) + + + def make_higgs_scheduler_adapters( +@@ -176,6 +179,8 @@ def make_higgs_scheduler_adapters( + int(max_new_tokens_cap), + ) + data = build_sglang_higgs_request(state, request_id=payload.request_id) ++ _params = payload.request.params if isinstance(payload.request.params, dict) else {} ++ data.return_logprob = bool(_params.get("return_logprob")) + data.engine_start_s = _perf_counter() + data.stage_payload = payload + data.stream_metadata = build_higgs_stream_metadata(payload, data) +diff --git a/sglang_omni/models/higgs_tts/vocoder_scheduler.py b/sglang_omni/models/higgs_tts/vocoder_scheduler.py +index e70373b..3bf0b2d 100644 +--- a/sglang_omni/models/higgs_tts/vocoder_scheduler.py ++++ b/sglang_omni/models/higgs_tts/vocoder_scheduler.py +@@ -199,6 +199,9 @@ class HiggsStreamingVocoderScheduler(StreamingSimpleScheduler): + usage = self._build_usage(HiggsTtsState.from_dict(payload.data)) + if usage is not None: + final_data["usage"] = usage ++ _lp = payload.data.get("output_token_logprobs") if isinstance(payload.data, dict) else None ++ if _lp is not None: ++ final_data["output_token_logprobs"] = _lp + messages.append( + OutgoingMessage( + request_id=request_id, +@@ -485,6 +488,8 @@ class HiggsStreamingVocoderScheduler(StreamingSimpleScheduler): + usage = self._build_usage(state) + if usage is not None: + data["usage"] = usage ++ if state.output_token_logprobs is not None: ++ data["output_token_logprobs"] = state.output_token_logprobs + payload.data = data + return payload + From 285340bf66c762b3c2ccf8a39a3c834ad3fbc3be Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Thu, 25 Jun 2026 23:57:13 +0800 Subject: [PATCH 13/24] feat(omni): trainer-side Higgs TTS actor + logprob-parity gate 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. --- examples/omni_gate_b/gate_b_parity_probe.py | 98 +++++++++++++ miles_plugins/omni/higgs_actor.py | 147 ++++++++++++++++++++ 2 files changed, 245 insertions(+) create mode 100644 examples/omni_gate_b/gate_b_parity_probe.py create mode 100644 miles_plugins/omni/higgs_actor.py diff --git a/examples/omni_gate_b/gate_b_parity_probe.py b/examples/omni_gate_b/gate_b_parity_probe.py new file mode 100644 index 00000000000..cdd3349d849 --- /dev/null +++ b/examples/omni_gate_b/gate_b_parity_probe.py @@ -0,0 +1,98 @@ +"""Logprob-parity gate for the GATE-B trainable TTS actor. + +Right after load the trainer-side actor and the served model are the same policy, +so the actor's recomputed codebook-0 log-probs must match the rollout's +`output_token_logprobs`. This is the make-or-break correctness check before any +GRPO update. + +Run (container, miles venv; Higgs server serving on SERVER): + SERVER=http://localhost:8010 HIGGS_CKPT='' CUDA_VISIBLE_DEVICES=4 \ + PYTHONPATH=/root/rl-omni/sglang-omni:/root/rl-omni/miles \ + python examples/omni_gate_b/gate_b_parity_probe.py +""" + +from __future__ import annotations + +import glob +import json +import os +import urllib.request + +SERVER = os.environ.get("SERVER", "http://localhost:8010") +# Gate on mean|Δ|: the residual is the served model's bf16 + sglang-kernel numeric +# floor (an fp32 trainer gives the SAME ~0.05 residual), so per-token max|Δ| of ~0.2 +# is irreducible cross-implementation noise, not a reconstruction error. exp(0.2)≈1.22 +# sits at the GRPO clip boundary and only biases the first ratio after each sync. +TOL = float(os.environ.get("PARITY_TOL", "0.10")) + + +def _rollout(prompt_ids: list[int], seed: int) -> dict: + req = { + "input_ids": prompt_ids, + "sampling_params": {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": 256, "seed": seed}, + "return_logprob": True, + "output_modalities": ["audio"], + } + resp = json.loads(urllib.request.urlopen( + urllib.request.Request(SERVER + "/generate", data=json.dumps(req).encode(), + headers={"Content-Type": "application/json"}), + timeout=180, + ).read()) + meta = resp["meta_info"] + otl = meta.get("output_token_logprobs") or [] + return { + "old_logprobs": [float(lp) for lp, _ in otl], + "cb0_tokens": [int(t) for _, t in otl], + "codebook_tokens": meta.get("output_codebook_tokens"), + } + + +def main() -> None: + import torch + from tokenizers import Tokenizer + from transformers import PreTrainedTokenizerFast + + from sglang_omni.models.higgs_tts.text_tokenizer import HiggsTokenizerAdapter + + from miles_plugins.omni.higgs_actor import HiggsTtsActor + + ckpt = glob.glob(os.environ["HIGGS_CKPT"])[0] if "*" in os.environ["HIGGS_CKPT"] else os.environ["HIGGS_CKPT"] + tok = PreTrainedTokenizerFast(tokenizer_object=Tokenizer.from_file(os.path.join(ckpt, "tokenizer.json"))) + adapter = HiggsTokenizerAdapter(tok) + device = os.environ.get("ACTOR_DEVICE", "cuda:0") + + dtype = torch.float32 if os.environ.get("ACTOR_DTYPE") == "fp32" else torch.bfloat16 + actor = HiggsTtsActor(ckpt, device=device, dtype=dtype) + print("actor loaded; backbone dtype", actor.dtype) + + texts = ["Hello world.", "The quick brown fox."] + worst = 0.0 + worst_mean = 0.0 + for i, text in enumerate(texts): + pid = list(map(int, adapter.build_prompt(text, num_ref_tokens=0))) + r = _rollout(pid, seed=1000 + i) + codes = r["codebook_tokens"] + old = r["old_logprobs"] + if not codes or not old: + print(f"[{text!r}] no codes/logprobs returned -> SKIP (server missing Step-1 fix?)") + continue + assert all(row[0] == t for row, t in zip(codes, r["cb0_tokens"])), "cb0 mismatch codes vs logprob tokens" + + with torch.no_grad(): + new = actor.codebook0_logprobs(pid, codes).tolist() + n = min(len(new), len(old)) + diffs = [abs(new[j] - old[j]) for j in range(n)] + max_d = max(diffs) + mean_d = sum(diffs) / n + worst = max(worst, max_d) + worst_mean = max(worst_mean, mean_d) + print(f"[{text!r}] T={n} max|Δ|={max_d:.4f} mean|Δ|={mean_d:.4f}") + print(f" old[:5]={[round(x,3) for x in old[:5]]}") + print(f" new[:5]={[round(x,3) for x in new[:5]]}") + + print(f"WORST_MEAN_ABS_DIFF: {worst_mean:.4f} (tol={TOL}) WORST_MAX_ABS_DIFF: {worst:.4f}") + print(f"PARITY_OK: {worst_mean < TOL}") + + +if __name__ == "__main__": + main() diff --git a/miles_plugins/omni/higgs_actor.py b/miles_plugins/omni/higgs_actor.py new file mode 100644 index 00000000000..b70e0d0e50f --- /dev/null +++ b/miles_plugins/omni/higgs_actor.py @@ -0,0 +1,147 @@ +"""Trainer-side Higgs TTS actor: a gradient-enabled teacher-forced forward that +reproduces the served model's per-step codebook-0 log-probs over a sampled codec +sequence, so the RL trainer can recompute new-policy log-probs for GRPO. + +The served `HiggsTTSModel` backbone is sglang's inference `Qwen3ForCausalLM` +(paged attention / CUDA graph, no autograd), so it cannot be trained directly. +This rebuilds the same policy from the checkpoint with a plain `transformers` +Qwen3 backbone + the fused codec embedding/head, which IS differentiable. + +Correctness is gated by a logprob-parity check against the server (see +`examples/omni_gate_b/gate_b_parity_probe.py`): right after load the trainer and +the server are the same policy, so recomputed log-probs must match. +""" + +from __future__ import annotations + +import glob +import json +import os + +import torch +import torch.nn.functional as F + + +# Checkpoint-name → transformers Qwen3Model state-dict-name (mirrors the server's +# DiscreteWeightMapper + _BACKBONE_PREFIX_MAP, but targets a plain Qwen3Model). +_BACKBONE_RENAME = { + "tied.embedding.text_embedding.": "embed_tokens.", + "body.layers.": "layers.", + "body.norm.": "norm.", +} +_FUSED_EMBED_KEY = "tied.embedding.modality_embeddings.0.embedding.weight" + + +def _resolve_ckpt_dir(path_or_glob: str) -> str: + if "*" in path_or_glob: + matches = glob.glob(path_or_glob) + if not matches: + raise FileNotFoundError(f"no checkpoint dir matches {path_or_glob!r}") + return matches[0] + return path_or_glob + + +class HiggsTtsActor: + """Differentiable Higgs codec policy (Qwen3 backbone + fused codebook head).""" + + def __init__(self, ckpt_dir: str, device: str = "cuda:0", dtype=torch.bfloat16): + from safetensors import safe_open + from transformers import Qwen3Config, Qwen3Model + + ckpt_dir = _resolve_ckpt_dir(ckpt_dir) + self.device = device + self.dtype = dtype + + cfg = json.load(open(os.path.join(ckpt_dir, "config.json"))) + text_cfg = cfg["text_config"] + enc_cfg = cfg["audio_encoder_config"] + self.num_codebooks = int(enc_cfg["num_codebooks"]) + self.codebook_vocab = int(enc_cfg["vocab_size"]) + + backbone = Qwen3Model(Qwen3Config(**text_cfg)).to(device=device, dtype=dtype).eval() + self.backbone = backbone + + # Stream the shards once: route backbone tensors into a state dict, grab the + # fused codec embedding weight, and drop the (skipped) audio-encoder tensors. + backbone_sd: dict[str, torch.Tensor] = {} + fused_embed: torch.Tensor | None = None + index = json.load(open(os.path.join(ckpt_dir, "model.safetensors.index.json"))) + for shard in sorted(set(index["weight_map"].values())): + with safe_open(os.path.join(ckpt_dir, shard), framework="pt") as f: + for key in f.keys(): + if key == _FUSED_EMBED_KEY: + fused_embed = f.get_tensor(key) + continue + renamed = self._rename_backbone(key) + if renamed is not None: + backbone_sd[renamed] = f.get_tensor(key) + + if fused_embed is None: + raise KeyError(f"fused codec embedding {_FUSED_EMBED_KEY!r} not in checkpoint") + missing, unexpected = backbone.load_state_dict(backbone_sd, strict=False) + # Qwen3Model ties embed_tokens; lm_head/text_head is intentionally absent here. + unexpected = [u for u in unexpected if "lm_head" not in u and "text_head" not in u] + if unexpected: + raise RuntimeError(f"unexpected backbone keys: {unexpected[:5]}") + real_missing = [m for m in missing if "rotary" not in m and "inv_freq" not in m] + if real_missing: + raise RuntimeError(f"missing backbone keys: {real_missing[:5]}") + + # Fused codebook weight [N*V, D]: input embedding (sum over codebooks) and, + # tied, the codebook-0 head = its first V rows. + self.fused_embed = fused_embed.to(device=device, dtype=dtype) + self._cb_offsets = ( + torch.arange(self.num_codebooks, device=device) * self.codebook_vocab + ) + + def _rename_backbone(self, key: str) -> str | None: + if key.startswith("tied.embedding.modality_embeddings.0.model."): + return None # audio encoder — not part of the AR policy + for src, dst in _BACKBONE_RENAME.items(): + if key.startswith(src): + return dst + key[len(src) :] + return None # text_head / anything else: skip + + def _embed_codes(self, codes_LN: torch.Tensor) -> torch.Tensor: + """[L, N] codebook ids → [L, D] fused embedding (mirrors the served model).""" + fused_ids = codes_LN + self._cb_offsets + return F.embedding(fused_ids, self.fused_embed).sum(dim=-2) + + def codebook0_logprobs( + self, prompt_ids: list[int], codebook_tokens: list[list[int]] + ) -> torch.Tensor: + """Teacher-forced new-policy log-probs of each step's sampled codebook-0 token. + + ``codebook_tokens`` is ``[T, num_codebooks]`` (the full per-step codes the + server fed back). Returns ``[T]`` log-probs aligned with the rollout's + ``output_token_logprobs`` (codebook-0). + """ + device = self.device + prompt = torch.tensor(prompt_ids, dtype=torch.long, device=device) + codes = torch.tensor(codebook_tokens, dtype=torch.long, device=device) # [T, N] + T = int(codes.shape[0]) + P = int(prompt.shape[0]) + + text_emb = self.backbone.embed_tokens(prompt) # [P, D] + # Teacher forcing: step t (t>=1) is predicted from the embedding of step t-1's + # full codes; step 0 is predicted from the last prompt token. So feed prompt + + # codes[0..T-2]; read hidden at positions P-1 .. P+T-2 for steps 0 .. T-1. + if T > 1: + codec_emb = self._embed_codes(codes[: T - 1]) # [T-1, D] + inputs_embeds = torch.cat([text_emb, codec_emb], dim=0) + else: + inputs_embeds = text_emb + L = inputs_embeds.shape[0] + positions = torch.arange(L, device=device).unsqueeze(0) + + out = self.backbone( + inputs_embeds=inputs_embeds.unsqueeze(0), + position_ids=positions, + use_cache=False, + ) + hidden = out.last_hidden_state[0] # [L, D] + step_hidden = hidden[P - 1 : P - 1 + T] # [T, D] + cb0_logits = F.linear(step_hidden.float(), self.fused_embed[: self.codebook_vocab].float()) + logp = torch.log_softmax(cb0_logits, dim=-1) # [T, V] + sampled_cb0 = codes[:, 0] # [T] + return logp[torch.arange(T, device=device), sampled_cb0] From edf87eab1298d6cf6c31908641b2cf6978f80eb8 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Fri, 26 Jun 2026 00:58:10 +0800 Subject: [PATCH 14/24] feat(omni): GATE-B full on-policy loop (GRPO + NCCL tts_engine weight-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. --- examples/omni_gate_b/gate_b_full.py | 218 ++++++++++++++++++++++++++++ miles_plugins/omni/higgs_actor.py | 2 +- 2 files changed, 219 insertions(+), 1 deletion(-) create mode 100644 examples/omni_gate_b/gate_b_full.py diff --git a/examples/omni_gate_b/gate_b_full.py b/examples/omni_gate_b/gate_b_full.py new file mode 100644 index 00000000000..e2364fa7412 --- /dev/null +++ b/examples/omni_gate_b/gate_b_full.py @@ -0,0 +1,218 @@ +"""Full GATE-B closed loop: GRPO LoRA RL on the Higgs TTS actor with per-step NCCL +weight-sync to the served sglang-omni ``tts_engine`` stage, so each step's rollouts +are on-policy. + +The fourth closed-loop component for TTS, mirroring gate_a_full.py: + rollout (/generate -> codec tokens + codebook-0 logprobs + audio) + -> composite reward (Whisper ASR CER + audio-validity guards) + -> GRPO advantage over codebook-0 tokens + -> LoRA policy update (trainer recomputes new logprobs via HiggsTtsActor) + -> NCCL broadcast of LoRA-merged backbone weights into the served tts_engine stage + (names in the checkpoint `body.*` convention; the server fuses q/k/v on load). + +Run (container, miles venv, free GPU for the trainer; Higgs server on another GPU): + SERVER=http://localhost:8010 HIGGS_CKPT='' MASTER_PORT=29641 \ + ASR_MODEL=openai/whisper-base ASR_DEVICE=cuda:0 CUDA_VISIBLE_DEVICES=4 \ + HF_HUB_OFFLINE=1 NCCL_P2P_DISABLE=1 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0 \ + PYTHONPATH=/root/rl-omni/sglang-omni:/root/rl-omni/miles \ + python examples/omni_gate_b/gate_b_full.py + +CRITICAL: set NCCL_P2P_DISABLE=1 on BOTH the server and the trainer (single-GPU masks +per process), exactly as in gate_a_full.py. +""" + +from __future__ import annotations + +import glob +import json +import os +import threading +import urllib.request + +import torch +from peft import LoraConfig, get_peft_model + +SERVER = os.environ.get("SERVER", "http://localhost:8010") +HIGGS_CKPT = os.environ["HIGGS_CKPT"] +DATA = os.environ.get("DATA", "examples/omni_gate_b/tts_smoke.jsonl") +STEPS = int(os.environ.get("STEPS", "3")) +GROUP = int(os.environ.get("GROUP", "4")) +PROMPTS = int(os.environ.get("PROMPTS", "4")) +MASTER_PORT = int(os.environ.get("MASTER_PORT", "29641")) +GROUP_NAME = os.environ.get("GROUP_NAME", "gate_b_wsync") +EPS = 0.2 + + +def post(path: str, body: dict, timeout: int = 300): + req = urllib.request.Request( + SERVER + path, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"} + ) + return json.loads(urllib.request.urlopen(req, timeout=timeout).read()) + + +def rollout(input_ids: list[int], seed: int) -> dict: + resp = post( + "/generate", + { + "input_ids": input_ids, + "sampling_params": {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": 256, "seed": seed}, + "return_logprob": True, + "output_modalities": ["audio"], + }, + timeout=180, + ) + meta = resp["meta_info"] + otl = meta.get("output_token_logprobs") or [] + return { + "old": [lp for lp, _ in otl], + "codes": meta.get("output_codebook_tokens"), + "audio": (resp.get("audio") or {}).get("data"), + } + + +def main() -> None: + torch.cuda.set_device(0) # bind this process to its visible GPU for NCCL collectives + + from sglang_omni.models.higgs_tts.text_tokenizer import HiggsTokenizerAdapter + from tokenizers import Tokenizer + from transformers import PreTrainedTokenizerFast + + from miles_plugins.omni.higgs_actor import HiggsTtsActor + from miles_plugins.omni.tts_reward import TtsCompositeReward + + ckpt = glob.glob(HIGGS_CKPT)[0] if "*" in HIGGS_CKPT else HIGGS_CKPT + tok = PreTrainedTokenizerFast(tokenizer_object=Tokenizer.from_file(os.path.join(ckpt, "tokenizer.json"))) + adapter = HiggsTokenizerAdapter(tok) + reward_fn = TtsCompositeReward() + + actor = HiggsTtsActor(ckpt, device="cuda:0") + actor.backbone = get_peft_model( + actor.backbone, + LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"], task_type=None), + ) + actor.backbone.train() + opt = torch.optim.AdamW([p for p in actor.backbone.parameters() if p.requires_grad], lr=2e-5) + + try: + from sglang.srt.utils import init_custom_process_group + except Exception: + from sglang.srt.utils.common import init_custom_process_group + + # Rendezvous a 2-rank NCCL group with the served tts_engine stage (server = rank 1). + init_err: list = [] + + def _init_server(): + try: + post( + "/init_weights_update_group", + { + "master_address": "localhost", + "master_port": MASTER_PORT, + "rank_offset": 1, + "world_size": 2, + "group_name": GROUP_NAME, + "backend": "nccl", + "stages": ["tts_engine"], + }, + timeout=180, + ) + except Exception as exc: # noqa: BLE001 + init_err.append(exc) + + th = threading.Thread(target=_init_server) + th.start() + pg = init_custom_process_group( + backend="nccl", init_method=f"tcp://localhost:{MASTER_PORT}", world_size=2, rank=0, group_name=GROUP_NAME + ) + th.join() + torch.cuda.synchronize() + if init_err: + raise init_err[0] + print("WEIGHT_UPDATE_GROUP_READY", flush=True) + + def merged_lora_weights() -> dict[str, torch.Tensor]: + out: dict[str, torch.Tensor] = {} + for name, mod in actor.backbone.named_modules(): + if hasattr(mod, "lora_A") and hasattr(mod, "base_layer"): + a = mod.lora_A["default"].weight + b = mod.lora_B["default"].weight + scaling = mod.scaling["default"] + w = mod.base_layer.weight.data + scaling * (b @ a) + # peft module name base_model.model.layers.N... -> ckpt body.layers.N... + hf = name.replace("base_model.model.", "") + out["body." + hf + ".weight"] = w.to(torch.bfloat16).contiguous() + return out + + def sync_to_server() -> int: + wd = merged_lora_weights() + names = sorted(wd) + spec = { + "names": names, + "dtypes": [str(wd[n].dtype).replace("torch.", "") for n in names], + "shapes": [list(wd[n].shape) for n in names], + "group_name": GROUP_NAME, + "stages": ["tts_engine"], + } + err: list = [] + + def _update(): + try: + post("/update_weights_from_distributed", spec, timeout=300) + except Exception as exc: # noqa: BLE001 + err.append(exc) + + t = threading.Thread(target=_update) + t.start() + for n in names: + torch.distributed.broadcast(wd[n], src=0, group=pg) + torch.cuda.synchronize() + t.join() + if err: + raise err[0] + return len(names) + + data = [json.loads(line) for line in open(DATA)] + print("step | mean_reward | mean_cer | avg_loss | synced_params") + for step in range(STEPS): + opt.zero_grad() + step_reward, step_cer, n_cer, step_loss, n = 0.0, 0.0, 0, 0.0, 0 + for ex in data[:PROMPTS]: + pid = list(map(int, adapter.build_prompt(ex["text"], num_ref_tokens=0))) + samples = [rollout(pid, step * 1000 + g) for g in range(GROUP)] + comps = [reward_fn.score(s["audio"], ex["label"]) for s in samples] + rewards = [c.reward for c in comps] + mean_r = sum(rewards) / len(rewards) + step_reward += mean_r + for c in comps: + if c.cer is not None: + step_cer += c.cer + n_cer += 1 + for s, adv in zip(samples, [r - mean_r for r in rewards]): + codes = s["codes"] + if not codes or adv == 0.0: + continue + new = actor.codebook0_logprobs(pid, codes) + T = min(len(new), len(s["old"])) + new = new[:T] + old = torch.tensor(s["old"][:T], device="cuda:0") + ratio = torch.exp(new - old) + loss = -torch.min(ratio * adv, torch.clamp(ratio, 1 - EPS, 1 + EPS) * adv).mean() + loss = loss / (GROUP * PROMPTS) + loss.backward() + step_loss += loss.item() * (GROUP * PROMPTS) + n += 1 + torch.nn.utils.clip_grad_norm_([p for p in actor.backbone.parameters() if p.requires_grad], 1.0) + opt.step() + synced = sync_to_server() # next step's rollouts are on-policy + mean_cer = step_cer / n_cer if n_cer else float("nan") + print( + f"{step:4d} | {step_reward / PROMPTS:11.3f} | {mean_cer:8.3f} | " + f"{step_loss / max(n, 1):8.4f} | {synced}", + flush=True, + ) + + print("GATE-B FULL on-policy loop complete (per-step NCCL weight-sync to served tts_engine)") + + +if __name__ == "__main__": + main() diff --git a/miles_plugins/omni/higgs_actor.py b/miles_plugins/omni/higgs_actor.py index b70e0d0e50f..7a946ca4d80 100644 --- a/miles_plugins/omni/higgs_actor.py +++ b/miles_plugins/omni/higgs_actor.py @@ -122,7 +122,7 @@ def codebook0_logprobs( T = int(codes.shape[0]) P = int(prompt.shape[0]) - text_emb = self.backbone.embed_tokens(prompt) # [P, D] + text_emb = self.backbone.get_input_embeddings()(prompt) # [P, D]; LoRA-wrap safe # Teacher forcing: step t (t>=1) is predicted from the embedding of step t-1's # full codes; step 0 is predicted from the last prompt token. So feed prompt + # codes[0..T-2]; read hidden at positions P-1 .. P+T-2 for steps 0 .. T-1. From e5ae79aed2429826269b9457b6169bbb155f47fc Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Fri, 26 Jun 2026 01:01:21 +0800 Subject: [PATCH 15/24] feat(omni): parametrize GATE-B rollout temperature/max_new_tokens 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. --- examples/omni_gate_b/gate_b_full.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/examples/omni_gate_b/gate_b_full.py b/examples/omni_gate_b/gate_b_full.py index e2364fa7412..402f8c16a7d 100644 --- a/examples/omni_gate_b/gate_b_full.py +++ b/examples/omni_gate_b/gate_b_full.py @@ -40,6 +40,8 @@ PROMPTS = int(os.environ.get("PROMPTS", "4")) MASTER_PORT = int(os.environ.get("MASTER_PORT", "29641")) GROUP_NAME = os.environ.get("GROUP_NAME", "gate_b_wsync") +TEMP = float(os.environ.get("TEMP", "0.8")) +MAX_NEW = int(os.environ.get("MAX_NEW", "256")) EPS = 0.2 @@ -55,7 +57,7 @@ def rollout(input_ids: list[int], seed: int) -> dict: "/generate", { "input_ids": input_ids, - "sampling_params": {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": 256, "seed": seed}, + "sampling_params": {"temperature": TEMP, "top_p": 0.95, "max_new_tokens": MAX_NEW, "seed": seed}, "return_logprob": True, "output_modalities": ["audio"], }, From 1d3a32a48305b6bc02134aad757ae9c2bc86e809 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Fri, 26 Jun 2026 01:23:29 +0800 Subject: [PATCH 16/24] chore(omni): add corrected higgs codec-logprob patch (supersedes broken 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). --- .../higgs_codec_logprob_fix.patch | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) create mode 100644 examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob_fix.patch diff --git a/examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob_fix.patch b/examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob_fix.patch new file mode 100644 index 00000000000..4fa46c81585 --- /dev/null +++ b/examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob_fix.patch @@ -0,0 +1,95 @@ +diff --git a/sglang_omni/models/higgs_tts/model.py b/sglang_omni/models/higgs_tts/model.py +index 3e68894..e32137e 100644 +--- a/sglang_omni/models/higgs_tts/model.py ++++ b/sglang_omni/models/higgs_tts/model.py +@@ -170,6 +170,12 @@ class HiggsTTSModel(nn.Module): + self._cg_codes_BN = torch.zeros( + pool_size, num_codebooks, dtype=torch.long, device=cg_device + ) ++ # Per-row codebook-0 log-prob of the sampled token, for RL rollout. Written ++ # in-place each step (eager prefill + CUDA-graph decode) so the runner records ++ # the true behavior-policy logprob instead of a text-vocab placeholder. ++ self._step_cb0_logprob = torch.zeros( ++ pool_size, dtype=torch.float32, device=cg_device ++ ) + # Note(Jiaxin): Packs codes_BN | was_done | active_generation_done into one buffer. + self._cg_collect_staging = torch.zeros( + pool_size, num_codebooks + 2, dtype=torch.long, device=cg_device +@@ -305,6 +311,13 @@ class HiggsTTSModel(nn.Module): + # Note(yichi): One D2H per step to skip STOP-sentinel rows in the Python append loop. + was_done_cpu = was_done.cpu().tolist() + codes_BN = codes_BN.detach().to(torch.long) ++ ++ # Codebook-0 log-prob of each sampled token, for RL rollout. Indexed by ++ # forward-batch row (aligned with the runner's per-request collect loop). ++ cb0_logits = logits_BNV[:, 0, :] ++ cb0_idx = codes_BN[:, 0:1].clamp(0, cb0_logits.shape[-1] - 1) ++ cb0_lp = torch.log_softmax(cb0_logits, dim=-1).gather(1, cb0_idx).squeeze(1) ++ self._step_cb0_logprob[:batch_size] = cb0_lp + for b in range(batch_size): + if was_done_cpu[b]: + continue +@@ -365,6 +378,13 @@ class HiggsTTSModel(nn.Module): + self._cg_active_last_codes[:batch_size] = new_last_codes_BN + self._cg_codes_BN[:batch_size] = codes_BN + ++ # Codebook-0 log-prob of each sampled token, for RL rollout. In-place buffer ++ # write (no value-dependent control flow / D2H) keeps this CUDA-graph safe. ++ cb0_logits = logits_BNV[:, 0, :] ++ cb0_idx = codes_BN[:, 0:1].long().clamp(0, cb0_logits.shape[-1] - 1) ++ cb0_lp = torch.log_softmax(cb0_logits, dim=-1).gather(1, cb0_idx).squeeze(1) ++ self._step_cb0_logprob[:batch_size] = cb0_lp ++ + text_vocab_size = self.backbone.config.vocab_size + return torch.zeros( + (batch_size, text_vocab_size), +diff --git a/sglang_omni/models/higgs_tts/model_runner.py b/sglang_omni/models/higgs_tts/model_runner.py +index 0d04a59..6a5aa31 100644 +--- a/sglang_omni/models/higgs_tts/model_runner.py ++++ b/sglang_omni/models/higgs_tts/model_runner.py +@@ -299,8 +299,8 @@ class HiggsTTSModelRunner(ModelRunner): + continue + codes_N = codes_BN_cpu[b].to(torch.long).clone() + data.output_codes.append(codes_N) +- nt = getattr(result.logits_output, "next_token_logits", None) +- self._record_rollout_logprob(data, nt[b] if nt is not None else None, int(codes_N[0].item())) ++ lpv = getattr(self.model, "_step_cb0_logprob", None) ++ self._record_rollout_logprob(data, lpv[b] if lpv is not None else None, int(codes_N[0].item())) + data.generation_done = bool(gen_done_after_cpu[b]) + self._emit_code_chunk(sched_req, codes_N) + self._mark_sampler_finished(req, data.generation_done) +@@ -377,8 +377,8 @@ class HiggsTTSModelRunner(ModelRunner): + continue + codes_N = codes_log[-1] + data.output_codes.append(codes_N.detach().cpu().clone()) +- nt = getattr(result.logits_output, "next_token_logits", None) +- self._record_rollout_logprob(data, nt[b] if nt is not None else None, int(codes_N[0].item())) ++ lpv = getattr(self.model, "_step_cb0_logprob", None) ++ self._record_rollout_logprob(data, lpv[b] if lpv is not None else None, int(codes_N[0].item())) + data.generation_done = bool(model._sampler_pool.generation_done[row].item()) + self._emit_code_chunk(sched_req, data.output_codes[-1]) + self._mark_sampler_finished(req, data.generation_done) +@@ -397,15 +397,16 @@ class HiggsTTSModelRunner(ModelRunner): + req.finished_reason = FINISH_MATCHED_TOKEN(EOC_ID) + + @staticmethod +- def _record_rollout_logprob(data, logits_row, cb0_token): +- """Record codebook-0 codec token + its logprob for RL rollout.""" ++ def _record_rollout_logprob(data, cb0_logprob, cb0_token): ++ """Record a sampled codebook-0 codec token + its log-prob for RL rollout. ++ ++ ``cb0_logprob`` is the model's pre-computed codebook-0 log-prob of the sampled ++ token (the true behavior-policy logprob); ``None`` falls back to 0.0. ++ """ + if not getattr(data, "return_logprob", False): + return +- if logits_row is None: +- data.output_token_logprobs.append([0.0, int(cb0_token)]) +- return +- logp = torch.log_softmax(logits_row.float(), dim=-1) +- data.output_token_logprobs.append([float(logp[int(cb0_token)].item()), int(cb0_token)]) ++ lp = 0.0 if cb0_logprob is None else float(cb0_logprob) ++ data.output_token_logprobs.append([lp, int(cb0_token)]) + + def _emit_code_chunk(self, sched_req: Any, codes_N: torch.Tensor) -> None: + if self._outbox is None: From 7a358d58e79cc584fb16e8a45e220354cc78e2c7 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Fri, 26 Jun 2026 10:17:55 +0800 Subject: [PATCH 17/24] chore(omni): remove stale broken higgs codec-logprob patch 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. --- .../higgs_codec_logprob.patch | 126 ------------------ 1 file changed, 126 deletions(-) delete mode 100644 examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob.patch diff --git a/examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob.patch b/examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob.patch deleted file mode 100644 index b24b8b1ac03..00000000000 --- a/examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob.patch +++ /dev/null @@ -1,126 +0,0 @@ -diff --git a/sglang_omni/models/higgs_tts/model_runner.py b/sglang_omni/models/higgs_tts/model_runner.py -index c579e46..0d04a59 100644 ---- a/sglang_omni/models/higgs_tts/model_runner.py -+++ b/sglang_omni/models/higgs_tts/model_runner.py -@@ -299,6 +299,8 @@ class HiggsTTSModelRunner(ModelRunner): - continue - codes_N = codes_BN_cpu[b].to(torch.long).clone() - data.output_codes.append(codes_N) -+ nt = getattr(result.logits_output, "next_token_logits", None) -+ self._record_rollout_logprob(data, nt[b] if nt is not None else None, int(codes_N[0].item())) - data.generation_done = bool(gen_done_after_cpu[b]) - self._emit_code_chunk(sched_req, codes_N) - self._mark_sampler_finished(req, data.generation_done) -@@ -364,7 +366,7 @@ class HiggsTTSModelRunner(ModelRunner): - - model = self.model - cb0_per_row: list[int] = [] -- for sched_req in requests: -+ for b, sched_req in enumerate(requests): - data = sched_req.data - req = data.req - rid = sched_req.request_id -@@ -375,6 +377,8 @@ class HiggsTTSModelRunner(ModelRunner): - continue - codes_N = codes_log[-1] - data.output_codes.append(codes_N.detach().cpu().clone()) -+ nt = getattr(result.logits_output, "next_token_logits", None) -+ self._record_rollout_logprob(data, nt[b] if nt is not None else None, int(codes_N[0].item())) - data.generation_done = bool(model._sampler_pool.generation_done[row].item()) - self._emit_code_chunk(sched_req, data.output_codes[-1]) - self._mark_sampler_finished(req, data.generation_done) -@@ -392,6 +396,17 @@ class HiggsTTSModelRunner(ModelRunner): - if generation_done and req.finished_reason is None: - req.finished_reason = FINISH_MATCHED_TOKEN(EOC_ID) - -+ @staticmethod -+ def _record_rollout_logprob(data, logits_row, cb0_token): -+ """Record codebook-0 codec token + its logprob for RL rollout.""" -+ if not getattr(data, "return_logprob", False): -+ return -+ if logits_row is None: -+ data.output_token_logprobs.append([0.0, int(cb0_token)]) -+ return -+ logp = torch.log_softmax(logits_row.float(), dim=-1) -+ data.output_token_logprobs.append([float(logp[int(cb0_token)].item()), int(cb0_token)]) -+ - def _emit_code_chunk(self, sched_req: Any, codes_N: torch.Tensor) -> None: - if self._outbox is None: - return -diff --git a/sglang_omni/models/higgs_tts/payload_types.py b/sglang_omni/models/higgs_tts/payload_types.py -index 079f540..11416f3 100644 ---- a/sglang_omni/models/higgs_tts/payload_types.py -+++ b/sglang_omni/models/higgs_tts/payload_types.py -@@ -46,6 +46,9 @@ class HiggsTtsState: - audio_samples: Any | None = None - sample_rate: int = 24000 - -+ # rollout -+ output_token_logprobs: list[Any] | None = None -+ - def to_dict(self) -> dict[str, Any]: - data: dict[str, Any] = { - "prompt_token_ids": list(self.prompt_token_ids), -@@ -81,6 +84,8 @@ class HiggsTtsState: - if self.audio_samples is not None: - data["audio_samples"] = self.audio_samples - data["sample_rate"] = self.sample_rate -+ if self.output_token_logprobs is not None: -+ data["output_token_logprobs"] = self.output_token_logprobs - return data - - @classmethod -@@ -107,6 +112,7 @@ class HiggsTtsState: - engine_time_s=data.get("engine_time_s", 0.0), - audio_samples=data.get("audio_samples"), - sample_rate=data.get("sample_rate", 24000), -+ output_token_logprobs=data.get("output_token_logprobs"), - ) - - -diff --git a/sglang_omni/models/higgs_tts/request_builders.py b/sglang_omni/models/higgs_tts/request_builders.py -index f5b82dd..6d3d7d3 100644 ---- a/sglang_omni/models/higgs_tts/request_builders.py -+++ b/sglang_omni/models/higgs_tts/request_builders.py -@@ -153,6 +153,9 @@ def apply_higgs_result(state: HiggsTtsState, data: HiggsSGLangRequestData) -> No - else: - state.output_codes_delayed = None - state.prompt_tokens = len(data.input_ids) -+ state.output_token_logprobs = ( -+ list(data.output_token_logprobs) if data.output_token_logprobs else None -+ ) - - - def make_higgs_scheduler_adapters( -@@ -176,6 +179,8 @@ def make_higgs_scheduler_adapters( - int(max_new_tokens_cap), - ) - data = build_sglang_higgs_request(state, request_id=payload.request_id) -+ _params = payload.request.params if isinstance(payload.request.params, dict) else {} -+ data.return_logprob = bool(_params.get("return_logprob")) - data.engine_start_s = _perf_counter() - data.stage_payload = payload - data.stream_metadata = build_higgs_stream_metadata(payload, data) -diff --git a/sglang_omni/models/higgs_tts/vocoder_scheduler.py b/sglang_omni/models/higgs_tts/vocoder_scheduler.py -index e70373b..3bf0b2d 100644 ---- a/sglang_omni/models/higgs_tts/vocoder_scheduler.py -+++ b/sglang_omni/models/higgs_tts/vocoder_scheduler.py -@@ -199,6 +199,9 @@ class HiggsStreamingVocoderScheduler(StreamingSimpleScheduler): - usage = self._build_usage(HiggsTtsState.from_dict(payload.data)) - if usage is not None: - final_data["usage"] = usage -+ _lp = payload.data.get("output_token_logprobs") if isinstance(payload.data, dict) else None -+ if _lp is not None: -+ final_data["output_token_logprobs"] = _lp - messages.append( - OutgoingMessage( - request_id=request_id, -@@ -485,6 +488,8 @@ class HiggsStreamingVocoderScheduler(StreamingSimpleScheduler): - usage = self._build_usage(state) - if usage is not None: - data["usage"] = usage -+ if state.output_token_logprobs is not None: -+ data["output_token_logprobs"] = state.output_token_logprobs - payload.data = data - return payload - From 8b0521c050132cc5be79b5874b16002bc9ddfd91 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Tue, 30 Jun 2026 17:39:13 +0800 Subject: [PATCH 18/24] feat(omni): non-saturating eval sets + MAX_NEW knob for GATE-A/B learning 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. --- examples/omni_gate_a/gate_a_full.py | 3 ++- examples/omni_gate_a/math_harder.jsonl | 10 ++++++++++ examples/omni_gate_b/tts_harder.jsonl | 8 ++++++++ 3 files changed, 20 insertions(+), 1 deletion(-) create mode 100644 examples/omni_gate_a/math_harder.jsonl create mode 100644 examples/omni_gate_b/tts_harder.jsonl diff --git a/examples/omni_gate_a/gate_a_full.py b/examples/omni_gate_a/gate_a_full.py index 9bcf9607678..91081c83ddb 100644 --- a/examples/omni_gate_a/gate_a_full.py +++ b/examples/omni_gate_a/gate_a_full.py @@ -41,6 +41,7 @@ PROMPTS = int(os.environ.get("PROMPTS", "4")) MASTER_PORT = int(os.environ.get("MASTER_PORT", "29555")) GROUP_NAME = os.environ.get("GROUP_NAME", "gate_a_wsync") +MAX_NEW = int(os.environ.get("MAX_NEW", "24")) EPS = 0.2 @@ -56,7 +57,7 @@ def rollout(input_ids: list[int], seed: int): "/generate", { "input_ids": input_ids, - "sampling_params": {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": 24, "seed": seed}, + "sampling_params": {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": MAX_NEW, "seed": seed}, "return_logprob": True, }, timeout=120, diff --git a/examples/omni_gate_a/math_harder.jsonl b/examples/omni_gate_a/math_harder.jsonl new file mode 100644 index 00000000000..230c016710a --- /dev/null +++ b/examples/omni_gate_a/math_harder.jsonl @@ -0,0 +1,10 @@ +{"prompt": "Question: What is 47 multiplied by 83?\nAnswer:", "label": "3901"} +{"prompt": "Question: Compute 234 plus 567 minus 98.\nAnswer:", "label": "703"} +{"prompt": "Question: A train travels 63 km per hour for 7 hours. How many km does it travel?\nAnswer:", "label": "441"} +{"prompt": "Question: What is 96 divided by 4, then multiplied by 5?\nAnswer:", "label": "120"} +{"prompt": "Question: There are 18 boxes with 24 items each. How many items in total?\nAnswer:", "label": "432"} +{"prompt": "Question: What is 15 squared minus 100?\nAnswer:", "label": "125"} +{"prompt": "Question: What is the area of a rectangle that is 23 cm by 17 cm, in square cm?\nAnswer:", "label": "391"} +{"prompt": "Question: What is 1000 minus 7 times 13?\nAnswer:", "label": "909"} +{"prompt": "Question: What is the sum of all integers from 1 to 20?\nAnswer:", "label": "210"} +{"prompt": "Question: What is 19 multiplied by 21?\nAnswer:", "label": "399"} diff --git a/examples/omni_gate_b/tts_harder.jsonl b/examples/omni_gate_b/tts_harder.jsonl new file mode 100644 index 00000000000..b218e03e481 --- /dev/null +++ b/examples/omni_gate_b/tts_harder.jsonl @@ -0,0 +1,8 @@ +{"text": "She sells seashells by the seashore.", "label": "She sells seashells by the seashore."} +{"text": "Peter Piper picked a peck of pickled peppers.", "label": "Peter Piper picked a peck of pickled peppers."} +{"text": "The sixth sick sheik's sixth sheep is sick.", "label": "The sixth sick sheik's sixth sheep is sick."} +{"text": "How much wood would a woodchuck chuck if a woodchuck could chuck wood.", "label": "How much wood would a woodchuck chuck if a woodchuck could chuck wood."} +{"text": "The quick brown fox jumps over the lazy dog near the riverbank.", "label": "The quick brown fox jumps over the lazy dog near the riverbank."} +{"text": "Red lorry, yellow lorry, red lorry, yellow lorry.", "label": "Red lorry, yellow lorry, red lorry, yellow lorry."} +{"text": "The bewildered tourist wandered through the labyrinthine alleyways.", "label": "The bewildered tourist wandered through the labyrinthine alleyways."} +{"text": "Worcestershire sauce is surprisingly difficult to pronounce.", "label": "Worcestershire sauce is surprisingly difficult to pronounce."} From 69764dfbaece0f5c3be6374068168a6c8712cf89 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Tue, 30 Jun 2026 18:12:03 +0800 Subject: [PATCH 19/24] refactor(omni): rename gate_a/gate_b examples to function-descriptive 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). --- .../logprob_parity_probe.py} | 0 .../gate_b_full.py => higgs_tts_rl/onpolicy_grpo_weight_sync.py} | 0 .../gate_b_loop.py => higgs_tts_rl/rollout_reward_advantage.py} | 0 .../sglang_omni_patches/higgs_codec_logprob_fix.patch | 0 examples/{omni_gate_b => higgs_tts_rl}/tts_harder.jsonl | 0 examples/{omni_gate_b => higgs_tts_rl}/tts_smoke.jsonl | 0 .../gate_a_lora_smoke.py => thinker_text_rl/lora_grpo_smoke.py} | 0 examples/{omni_gate_a => thinker_text_rl}/math_harder.jsonl | 0 examples/{omni_gate_a => thinker_text_rl}/math_smoke.jsonl | 0 .../onpolicy_grpo_weight_sync.py} | 0 10 files changed, 0 insertions(+), 0 deletions(-) rename examples/{omni_gate_b/gate_b_parity_probe.py => higgs_tts_rl/logprob_parity_probe.py} (100%) rename examples/{omni_gate_b/gate_b_full.py => higgs_tts_rl/onpolicy_grpo_weight_sync.py} (100%) rename examples/{omni_gate_b/gate_b_loop.py => higgs_tts_rl/rollout_reward_advantage.py} (100%) rename examples/{omni_gate_b => higgs_tts_rl}/sglang_omni_patches/higgs_codec_logprob_fix.patch (100%) rename examples/{omni_gate_b => higgs_tts_rl}/tts_harder.jsonl (100%) rename examples/{omni_gate_b => higgs_tts_rl}/tts_smoke.jsonl (100%) rename examples/{omni_gate_a/gate_a_lora_smoke.py => thinker_text_rl/lora_grpo_smoke.py} (100%) rename examples/{omni_gate_a => thinker_text_rl}/math_harder.jsonl (100%) rename examples/{omni_gate_a => thinker_text_rl}/math_smoke.jsonl (100%) rename examples/{omni_gate_a/gate_a_full.py => thinker_text_rl/onpolicy_grpo_weight_sync.py} (100%) diff --git a/examples/omni_gate_b/gate_b_parity_probe.py b/examples/higgs_tts_rl/logprob_parity_probe.py similarity index 100% rename from examples/omni_gate_b/gate_b_parity_probe.py rename to examples/higgs_tts_rl/logprob_parity_probe.py diff --git a/examples/omni_gate_b/gate_b_full.py b/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py similarity index 100% rename from examples/omni_gate_b/gate_b_full.py rename to examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py diff --git a/examples/omni_gate_b/gate_b_loop.py b/examples/higgs_tts_rl/rollout_reward_advantage.py similarity index 100% rename from examples/omni_gate_b/gate_b_loop.py rename to examples/higgs_tts_rl/rollout_reward_advantage.py diff --git a/examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob_fix.patch b/examples/higgs_tts_rl/sglang_omni_patches/higgs_codec_logprob_fix.patch similarity index 100% rename from examples/omni_gate_b/sglang_omni_patches/higgs_codec_logprob_fix.patch rename to examples/higgs_tts_rl/sglang_omni_patches/higgs_codec_logprob_fix.patch diff --git a/examples/omni_gate_b/tts_harder.jsonl b/examples/higgs_tts_rl/tts_harder.jsonl similarity index 100% rename from examples/omni_gate_b/tts_harder.jsonl rename to examples/higgs_tts_rl/tts_harder.jsonl diff --git a/examples/omni_gate_b/tts_smoke.jsonl b/examples/higgs_tts_rl/tts_smoke.jsonl similarity index 100% rename from examples/omni_gate_b/tts_smoke.jsonl rename to examples/higgs_tts_rl/tts_smoke.jsonl diff --git a/examples/omni_gate_a/gate_a_lora_smoke.py b/examples/thinker_text_rl/lora_grpo_smoke.py similarity index 100% rename from examples/omni_gate_a/gate_a_lora_smoke.py rename to examples/thinker_text_rl/lora_grpo_smoke.py diff --git a/examples/omni_gate_a/math_harder.jsonl b/examples/thinker_text_rl/math_harder.jsonl similarity index 100% rename from examples/omni_gate_a/math_harder.jsonl rename to examples/thinker_text_rl/math_harder.jsonl diff --git a/examples/omni_gate_a/math_smoke.jsonl b/examples/thinker_text_rl/math_smoke.jsonl similarity index 100% rename from examples/omni_gate_a/math_smoke.jsonl rename to examples/thinker_text_rl/math_smoke.jsonl diff --git a/examples/omni_gate_a/gate_a_full.py b/examples/thinker_text_rl/onpolicy_grpo_weight_sync.py similarity index 100% rename from examples/omni_gate_a/gate_a_full.py rename to examples/thinker_text_rl/onpolicy_grpo_weight_sync.py From dd5d2adb5bbf32825c06f78c9b13abfd27222796 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Tue, 30 Jun 2026 18:12:58 +0800 Subject: [PATCH 20/24] refactor(omni): scrub GATE-A/GATE-B terminology + fix internal path refs 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. --- examples/higgs_tts_rl/logprob_parity_probe.py | 4 +- .../higgs_tts_rl/onpolicy_grpo_weight_sync.py | 14 +-- .../higgs_tts_rl/rollout_reward_advantage.py | 10 +- .../thinker_text_rl/fsdp_trainer_launch.sh | 107 ++++++++++++++++++ examples/thinker_text_rl/lora_grpo_smoke.py | 8 +- .../thinker_text_rl/math_harder_msgs.jsonl | 10 ++ .../onpolicy_grpo_weight_sync.py | 12 +- miles_plugins/omni/higgs_actor.py | 2 +- miles_plugins/omni/math_reward.py | 2 +- miles_plugins/omni/tts_reward.py | 2 +- 10 files changed, 144 insertions(+), 27 deletions(-) create mode 100644 examples/thinker_text_rl/fsdp_trainer_launch.sh create mode 100644 examples/thinker_text_rl/math_harder_msgs.jsonl diff --git a/examples/higgs_tts_rl/logprob_parity_probe.py b/examples/higgs_tts_rl/logprob_parity_probe.py index cdd3349d849..38489994f7f 100644 --- a/examples/higgs_tts_rl/logprob_parity_probe.py +++ b/examples/higgs_tts_rl/logprob_parity_probe.py @@ -1,4 +1,4 @@ -"""Logprob-parity gate for the GATE-B trainable TTS actor. +"""Logprob-parity check for the Higgs TTS trainable actor. Right after load the trainer-side actor and the served model are the same policy, so the actor's recomputed codebook-0 log-probs must match the rollout's @@ -8,7 +8,7 @@ Run (container, miles venv; Higgs server serving on SERVER): SERVER=http://localhost:8010 HIGGS_CKPT='' CUDA_VISIBLE_DEVICES=4 \ PYTHONPATH=/root/rl-omni/sglang-omni:/root/rl-omni/miles \ - python examples/omni_gate_b/gate_b_parity_probe.py + python examples/higgs_tts_rl/logprob_parity_probe.py """ from __future__ import annotations diff --git a/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py b/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py index 402f8c16a7d..4e99290a8f5 100644 --- a/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py +++ b/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py @@ -1,8 +1,8 @@ -"""Full GATE-B closed loop: GRPO LoRA RL on the Higgs TTS actor with per-step NCCL +"""Full on-policy Higgs TTS RL: GRPO LoRA on the Higgs TTS actor with per-step NCCL weight-sync to the served sglang-omni ``tts_engine`` stage, so each step's rollouts are on-policy. -The fourth closed-loop component for TTS, mirroring gate_a_full.py: +The fourth closed-loop component for TTS, mirroring the thinker onpolicy_grpo_weight_sync.py: rollout (/generate -> codec tokens + codebook-0 logprobs + audio) -> composite reward (Whisper ASR CER + audio-validity guards) -> GRPO advantage over codebook-0 tokens @@ -15,10 +15,10 @@ ASR_MODEL=openai/whisper-base ASR_DEVICE=cuda:0 CUDA_VISIBLE_DEVICES=4 \ HF_HUB_OFFLINE=1 NCCL_P2P_DISABLE=1 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0 \ PYTHONPATH=/root/rl-omni/sglang-omni:/root/rl-omni/miles \ - python examples/omni_gate_b/gate_b_full.py + python examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py CRITICAL: set NCCL_P2P_DISABLE=1 on BOTH the server and the trainer (single-GPU masks -per process), exactly as in gate_a_full.py. +per process), exactly as in the thinker on-policy script. """ from __future__ import annotations @@ -34,12 +34,12 @@ SERVER = os.environ.get("SERVER", "http://localhost:8010") HIGGS_CKPT = os.environ["HIGGS_CKPT"] -DATA = os.environ.get("DATA", "examples/omni_gate_b/tts_smoke.jsonl") +DATA = os.environ.get("DATA", "examples/higgs_tts_rl/tts_smoke.jsonl") STEPS = int(os.environ.get("STEPS", "3")) GROUP = int(os.environ.get("GROUP", "4")) PROMPTS = int(os.environ.get("PROMPTS", "4")) MASTER_PORT = int(os.environ.get("MASTER_PORT", "29641")) -GROUP_NAME = os.environ.get("GROUP_NAME", "gate_b_wsync") +GROUP_NAME = os.environ.get("GROUP_NAME", "higgs_tts_wsync") TEMP = float(os.environ.get("TEMP", "0.8")) MAX_NEW = int(os.environ.get("MAX_NEW", "256")) EPS = 0.2 @@ -213,7 +213,7 @@ def _update(): flush=True, ) - print("GATE-B FULL on-policy loop complete (per-step NCCL weight-sync to served tts_engine)") + print("Higgs TTS on-policy loop complete (per-step NCCL weight-sync to served tts_engine)") if __name__ == "__main__": diff --git a/examples/higgs_tts_rl/rollout_reward_advantage.py b/examples/higgs_tts_rl/rollout_reward_advantage.py index 94dd3d52622..eeea593ce10 100644 --- a/examples/higgs_tts_rl/rollout_reward_advantage.py +++ b/examples/higgs_tts_rl/rollout_reward_advantage.py @@ -1,4 +1,4 @@ -"""GATE-B closed loop demo: GRPO rollout -> composite reward -> advantage on real Higgs TTS. +"""Higgs TTS RL: rollout -> composite reward -> advantage (no weight update) on real Higgs TTS. Demonstrates the first three closed-loop components on the real Higgs-audio model through the sglang-omni rollout backend: @@ -7,14 +7,14 @@ -> GRPO advantage. The 4th component (LoRA policy update + NCCL weight-sync to the served TTS actor) mirrors -GATE-A's gate_a_full.py: the rollout returns codec-token logprobs (old) and the trainer +the thinker on-policy script (onpolicy_grpo_weight_sync.py): the rollout returns codec-token logprobs (old) and the trainer recomputes new logprobs over the codec sequence; weight sync uses /update_weights_from_distributed with NCCL_P2P_DISABLE=1. Run (container, miles venv; Higgs server already serving on SERVER): THINKER=... SERVER=http://localhost:8010 HIGGS_CKPT= \ ASR_MODEL=openai/whisper-base ASR_DEVICE=cuda:0 \ - python examples/omni_gate_b/gate_b_loop.py + python examples/higgs_tts_rl/rollout_reward_advantage.py """ from __future__ import annotations @@ -25,7 +25,7 @@ import urllib.request SERVER = os.environ.get("SERVER", "http://localhost:8010") -DATA = os.environ.get("DATA", "examples/omni_gate_b/tts_smoke.jsonl") +DATA = os.environ.get("DATA", "examples/higgs_tts_rl/tts_smoke.jsonl") GROUP = int(os.environ.get("GROUP", "4")) STEPS = int(os.environ.get("STEPS", "3")) @@ -97,7 +97,7 @@ def main() -> None: mean_cer = step_cer / n_cer if n_cer else float("nan") print(f"{step:4d} | {step_reward / len(per_prompt):11.3f} | {mean_cer:8.3f} | {per_prompt}", flush=True) - print("GATE-B rollout->composite-reward->advantage demonstrated on real Higgs TTS") + print("Higgs TTS rollout->composite-reward->advantage demonstrated") if __name__ == "__main__": diff --git a/examples/thinker_text_rl/fsdp_trainer_launch.sh b/examples/thinker_text_rl/fsdp_trainer_launch.sh new file mode 100644 index 00000000000..f4781ae12ac --- /dev/null +++ b/examples/thinker_text_rl/fsdp_trainer_launch.sh @@ -0,0 +1,107 @@ +#!/bin/bash +# Thinker text RL through the real miles FSDP trainer. +# Rollout is an EXTERNAL sglang-omni thinker server (the omni pipeline); the trainer +# does FSDP + GRPO + LoRA and talks to it via OmniGenerateFn over HTTP. Weight-sync to +# the external server is deferred (M2) -> this M1 run is off-policy. +# +# Prereq: sglang-omni thinker server already serving on $SERVER_PORT (e.g. GPU2): +# PYTHONPATH=. CUDA_VISIBLE_DEVICES=2 ... python examples/run_qwen3_omni_server.py \ +# --model-path Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8003 +set -ex + +pkill -9 -f "ray::" 2>/dev/null || true +ray stop --force 2>/dev/null || true +sleep 2 + +export PYTHONBUFFERED=16 +# Trainer GPUs: avoid GPU2 (held by the external sglang-omni server). +TRAINER_GPUS=${TRAINER_GPUS:-"0,1,3,4"} +export CUDA_VISIBLE_DEVICES=$TRAINER_GPUS +NGPU=$(echo $TRAINER_GPUS | tr ',' '\n' | wc -l) + +REPO=/root/rl-omni/miles +THINKER=/root/qwen3-omni-thinker +SERVER_PORT=${SERVER_PORT:-8003} + +CKPT_ARGS=( + --hf-checkpoint $THINKER +) + +ROLLOUT_ARGS=( + --prompt-data $REPO/examples/thinker_text_rl/math_harder_msgs.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --custom-generate-function-path miles_plugins.omni.omni_generate_fn.OmniGenerateFn + --custom-rm-path miles_plugins.omni.math_reward.compute_math_reward + --rollout-external + --rollout-num-gpus 0 + --rollout-external-engine-addrs "localhost:${SERVER_PORT}" + --sglang-router-ip localhost + --sglang-router-port ${SERVER_PORT} + --num-rollout 4 + --rollout-batch-size 4 + --n-samples-per-prompt 4 + --rollout-max-response-len 64 + --rollout-temperature 0.8 + --global-batch-size 16 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --eps-clip 0.2 + --eps-clip-high 0.28 + --kl-coef 0.00 + --entropy-coef 0.00 +) + +LORA_ARGS=( + --lora-rank 8 + --lora-alpha 16 + --target-modules q_proj,v_proj +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 2e-5 + --lr-decay-style constant +) + +TRAIN_BACKEND_ARGS=( + --train-backend fsdp + --gradient-checkpointing + --attn-implementation sdpa +) + +PERF_ARGS=( + --use-dynamic-batch-size + --max-tokens-per-gpu 8192 +) + +MISC_ARGS=( + --actor-num-nodes 1 + --actor-num-gpus-per-node $NGPU + # FSDP backend is gated behind --ci-test (experimental); disable CI checkers so they don't interfere + --ci-test + --ci-disable-kl-checker + --ci-disable-logprobs-checker +) + +ray start --head --node-ip-address 127.0.0.1 --num-gpus $NGPU --disable-usage-stats + +# train.py connects to the running cluster via ray.init(address="auto"); no dashboard / job-submit needed +export PYTHONPATH=${REPO}:/root/rl-omni/sglang-omni +export HF_HUB_OFFLINE=1 +export NCCL_P2P_DISABLE=1 +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +python3 train.py \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${LORA_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${TRAIN_BACKEND_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/examples/thinker_text_rl/lora_grpo_smoke.py b/examples/thinker_text_rl/lora_grpo_smoke.py index 1698441639e..0761d60af41 100644 --- a/examples/thinker_text_rl/lora_grpo_smoke.py +++ b/examples/thinker_text_rl/lora_grpo_smoke.py @@ -1,4 +1,4 @@ -"""GATE-A closed-loop smoke: GRPO RL on the Qwen3-Omni Thinker (single GPU, LoRA). +"""Thinker text RL smoke: GRPO LoRA on the Qwen3-Omni Thinker (single GPU). Demonstrates all four closed-loop components on the real model end to end: rollout (sglang-omni /generate) -> reward (math correctness) @@ -11,9 +11,9 @@ convergence. Run (inside the container, miles venv): - THINKER=/root/qwen3-omni-thinker DATA=examples/omni_gate_a/math_smoke.jsonl \ + THINKER=/root/qwen3-omni-thinker DATA=examples/thinker_text_rl/math_smoke.jsonl \ SERVER=http://localhost:8000/generate CUDA_VISIBLE_DEVICES=4 \ - python examples/omni_gate_a/gate_a_lora_smoke.py + python examples/thinker_text_rl/lora_grpo_smoke.py """ from __future__ import annotations @@ -105,7 +105,7 @@ def main() -> None: avg_loss = step_loss / max(n, 1) print(f"{step:4d} | {mean_reward:11.3f} | {avg_loss:8.4f} | {per_prompt}") - print("GATE-A closed-loop smoke complete (rollout->reward->advantage->LoRA update over multiple steps)") + print("Thinker RL smoke complete (rollout->reward->advantage->LoRA update over multiple steps)") if __name__ == "__main__": diff --git a/examples/thinker_text_rl/math_harder_msgs.jsonl b/examples/thinker_text_rl/math_harder_msgs.jsonl new file mode 100644 index 00000000000..9a64160e8f6 --- /dev/null +++ b/examples/thinker_text_rl/math_harder_msgs.jsonl @@ -0,0 +1,10 @@ +{"prompt": [{"role": "user", "content": "What is 47 multiplied by 83? Give only the final integer."}], "label": "3901"} +{"prompt": [{"role": "user", "content": "Compute 234 plus 567 minus 98. Give only the final integer."}], "label": "703"} +{"prompt": [{"role": "user", "content": "A train travels 63 km per hour for 7 hours. How many km does it travel? Give only the final integer."}], "label": "441"} +{"prompt": [{"role": "user", "content": "What is 96 divided by 4, then multiplied by 5? Give only the final integer."}], "label": "120"} +{"prompt": [{"role": "user", "content": "There are 18 boxes with 24 items each. How many items in total? Give only the final integer."}], "label": "432"} +{"prompt": [{"role": "user", "content": "What is 15 squared minus 100? Give only the final integer."}], "label": "125"} +{"prompt": [{"role": "user", "content": "What is the area of a rectangle that is 23 cm by 17 cm, in square cm? Give only the final integer."}], "label": "391"} +{"prompt": [{"role": "user", "content": "What is 1000 minus 7 times 13? Give only the final integer."}], "label": "909"} +{"prompt": [{"role": "user", "content": "What is the sum of all integers from 1 to 20? Give only the final integer."}], "label": "210"} +{"prompt": [{"role": "user", "content": "What is 19 multiplied by 21? Give only the final integer."}], "label": "399"} diff --git a/examples/thinker_text_rl/onpolicy_grpo_weight_sync.py b/examples/thinker_text_rl/onpolicy_grpo_weight_sync.py index 91081c83ddb..c93fb7d1216 100644 --- a/examples/thinker_text_rl/onpolicy_grpo_weight_sync.py +++ b/examples/thinker_text_rl/onpolicy_grpo_weight_sync.py @@ -1,7 +1,7 @@ -"""Full on-policy GATE-A: GRPO LoRA RL on the Qwen3-Omni Thinker with per-step NCCL +"""Full on-policy Thinker text RL: GRPO LoRA on the Qwen3-Omni Thinker with per-step NCCL weight-sync to the served sglang-omni thinker, so each step's rollouts are on-policy. -Beyond gate_a_lora_smoke.py this adds the 4th closed-loop component done *properly*: +Beyond lora_grpo_smoke.py this adds the 4th closed-loop component done *properly*: after each optimizer step the LoRA-merged thinker weights are broadcast into the served thinker stage via sglang-omni's distributed weight-update admin plane (``/init_weights_update_group`` + ``/update_weights_from_distributed`` + ``stages=[thinker]``), @@ -9,10 +9,10 @@ accepts plain ``model.*`` names, so the extracted-thinker names sync directly. Run (container, miles venv, free GPU for the trainer; server already on another GPU): - THINKER=/root/qwen3-omni-thinker DATA=examples/omni_gate_a/math_smoke.jsonl \ + THINKER=/root/qwen3-omni-thinker DATA=examples/thinker_text_rl/math_smoke.jsonl \ SERVER=http://localhost:8003 MASTER_PORT=29631 CUDA_VISIBLE_DEVICES=4 \ NCCL_P2P_DISABLE=1 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0 \ - python examples/omni_gate_a/gate_a_full.py + python examples/thinker_text_rl/onpolicy_grpo_weight_sync.py CRITICAL: the trainer and the sglang-omni server run as separate processes, each with a single GPU exposed via CUDA_VISIBLE_DEVICES (both see it as cuda:0). NCCL would try direct @@ -40,7 +40,7 @@ GROUP = int(os.environ.get("GROUP", "4")) PROMPTS = int(os.environ.get("PROMPTS", "4")) MASTER_PORT = int(os.environ.get("MASTER_PORT", "29555")) -GROUP_NAME = os.environ.get("GROUP_NAME", "gate_a_wsync") +GROUP_NAME = os.environ.get("GROUP_NAME", "thinker_wsync") MAX_NEW = int(os.environ.get("MAX_NEW", "24")) EPS = 0.2 @@ -188,7 +188,7 @@ def _update(): synced = sync_to_server() # next step's rollouts are on-policy print(f"{step:4d} | {step_reward / PROMPTS:11.3f} | {step_loss / max(n, 1):8.4f} | {synced}", flush=True) - print("GATE-A FULL on-policy loop complete (per-step NCCL weight-sync to served thinker)") + print("Thinker on-policy loop complete (per-step NCCL weight-sync to served thinker)") if __name__ == "__main__": diff --git a/miles_plugins/omni/higgs_actor.py b/miles_plugins/omni/higgs_actor.py index 7a946ca4d80..ce24a469a4c 100644 --- a/miles_plugins/omni/higgs_actor.py +++ b/miles_plugins/omni/higgs_actor.py @@ -8,7 +8,7 @@ Qwen3 backbone + the fused codec embedding/head, which IS differentiable. Correctness is gated by a logprob-parity check against the server (see -`examples/omni_gate_b/gate_b_parity_probe.py`): right after load the trainer and +`examples/higgs_tts_rl/logprob_parity_probe.py`): right after load the trainer and the server are the same policy, so recomputed log-probs must match. """ diff --git a/miles_plugins/omni/math_reward.py b/miles_plugins/omni/math_reward.py index 0405a254f93..1861352b361 100644 --- a/miles_plugins/omni/math_reward.py +++ b/miles_plugins/omni/math_reward.py @@ -1,4 +1,4 @@ -"""Text-only math-correctness reward for the GATE-A thinker RL smoke. +"""Text-only math-correctness reward for thinker text RL. Loaded via ``--custom-rm-path miles_plugins.omni.math_reward.compute_math_reward``. Returns 1.0 when the model's decoded response contains the gold answer diff --git a/miles_plugins/omni/tts_reward.py b/miles_plugins/omni/tts_reward.py index 1dbf818cf9e..b72aad05f94 100644 --- a/miles_plugins/omni/tts_reward.py +++ b/miles_plugins/omni/tts_reward.py @@ -1,4 +1,4 @@ -"""Composite reward for TTS RL (GATE-B): ASR round-trip CER + audio-validity guards. +"""Composite reward for Higgs TTS RL: ASR round-trip CER + audio-validity guards. The TTS actor generates speech for a target text. The reward transcribes the generated audio with an ASR model (Whisper) and scores content correctness via CER, combined with From d55adfcfdde11129c8bb51870ea976356077a6c6 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Tue, 30 Jun 2026 18:40:13 +0800 Subject: [PATCH 21/24] feat(rollout): wire --rollout-external to use an external sglang server --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. --- miles/ray/rollout/rollout_manager.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index c216b440b24..d43daee2567 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -71,6 +71,17 @@ def __init__(self, args, pg): if self.args.debug_train_only: self.servers: dict[str, RolloutServer] = {} + elif getattr(self.args, "rollout_external", False): + # External rollout engine (e.g. an sglang-omni server hosting an omni pipeline): + # do not launch internal sglang servers; point the router at the supplied external + # address so the custom generate function reaches it directly. + init_http_client(args) + addrs = self.args.rollout_external_engine_addrs + addr = addrs[0] if isinstance(addrs, (list, tuple)) else addrs + ip, _, port = str(addr).rpartition(":") + self.args.sglang_router_ip = ip + self.args.sglang_router_port = int(port) + self.servers: dict[str, RolloutServer] = {} else: init_http_client(args) self.servers = start_rollout_servers(args, pg) From 3186a5a16908fca36a447193bed82dcb283aa6b1 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Wed, 8 Jul 2026 13:33:59 +0800 Subject: [PATCH 22/24] refactor(omni): consume Higgs rollout contract --- examples/higgs_tts_rl/logprob_parity_probe.py | 40 ++++---- .../higgs_tts_rl/onpolicy_grpo_weight_sync.py | 30 +++--- .../higgs_tts_rl/rollout_reward_advantage.py | 33 ++++--- .../higgs_codec_logprob_fix.patch | 95 ------------------- miles_plugins/omni/rollout_contract.py | 37 +++++++- tests/fast/test_omni_rollout_contract.py | 27 ++++++ 6 files changed, 125 insertions(+), 137 deletions(-) delete mode 100644 examples/higgs_tts_rl/sglang_omni_patches/higgs_codec_logprob_fix.patch diff --git a/examples/higgs_tts_rl/logprob_parity_probe.py b/examples/higgs_tts_rl/logprob_parity_probe.py index 38489994f7f..47f5f2b8193 100644 --- a/examples/higgs_tts_rl/logprob_parity_probe.py +++ b/examples/higgs_tts_rl/logprob_parity_probe.py @@ -18,6 +18,11 @@ import os import urllib.request +from miles_plugins.omni.rollout_contract import ( + build_generate_payload, + parse_generate_response, +) + SERVER = os.environ.get("SERVER", "http://localhost:8010") # Gate on mean|Δ|: the residual is the served model's bf16 + sglang-kernel numeric # floor (an fp32 trainer gives the SAME ~0.05 residual), so per-token max|Δ| of ~0.2 @@ -27,23 +32,26 @@ def _rollout(prompt_ids: list[int], seed: int) -> dict: - req = { - "input_ids": prompt_ids, - "sampling_params": {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": 256, "seed": seed}, - "return_logprob": True, - "output_modalities": ["audio"], - } - resp = json.loads(urllib.request.urlopen( - urllib.request.Request(SERVER + "/generate", data=json.dumps(req).encode(), - headers={"Content-Type": "application/json"}), - timeout=180, - ).read()) - meta = resp["meta_info"] - otl = meta.get("output_token_logprobs") or [] + req = build_generate_payload( + prompt_ids, + {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": 256, "seed": seed}, + output_modalities=["audio"], + ) + resp = json.loads( + urllib.request.urlopen( + urllib.request.Request( + SERVER + "/generate", + data=json.dumps(req).encode(), + headers={"Content-Type": "application/json"}, + ), + timeout=180, + ).read() + ) + result = parse_generate_response(resp) return { - "old_logprobs": [float(lp) for lp, _ in otl], - "cb0_tokens": [int(t) for _, t in otl], - "codebook_tokens": meta.get("output_codebook_tokens"), + "old_logprobs": result.response_log_probs, + "cb0_tokens": result.response_tokens, + "codebook_tokens": result.output_codebook_tokens, } diff --git a/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py b/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py index 4e99290a8f5..63dba3cb4d6 100644 --- a/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py +++ b/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py @@ -32,6 +32,11 @@ import torch from peft import LoraConfig, get_peft_model +from miles_plugins.omni.rollout_contract import ( + build_generate_payload, + parse_generate_response, +) + SERVER = os.environ.get("SERVER", "http://localhost:8010") HIGGS_CKPT = os.environ["HIGGS_CKPT"] DATA = os.environ.get("DATA", "examples/higgs_tts_rl/tts_smoke.jsonl") @@ -55,20 +60,23 @@ def post(path: str, body: dict, timeout: int = 300): def rollout(input_ids: list[int], seed: int) -> dict: resp = post( "/generate", - { - "input_ids": input_ids, - "sampling_params": {"temperature": TEMP, "top_p": 0.95, "max_new_tokens": MAX_NEW, "seed": seed}, - "return_logprob": True, - "output_modalities": ["audio"], - }, + build_generate_payload( + input_ids, + { + "temperature": TEMP, + "top_p": 0.95, + "max_new_tokens": MAX_NEW, + "seed": seed, + }, + output_modalities=["audio"], + ), timeout=180, ) - meta = resp["meta_info"] - otl = meta.get("output_token_logprobs") or [] + result = parse_generate_response(resp) return { - "old": [lp for lp, _ in otl], - "codes": meta.get("output_codebook_tokens"), - "audio": (resp.get("audio") or {}).get("data"), + "old": result.response_log_probs, + "codes": result.output_codebook_tokens, + "audio": (result.audio or {}).get("data"), } diff --git a/examples/higgs_tts_rl/rollout_reward_advantage.py b/examples/higgs_tts_rl/rollout_reward_advantage.py index eeea593ce10..1ab98d28222 100644 --- a/examples/higgs_tts_rl/rollout_reward_advantage.py +++ b/examples/higgs_tts_rl/rollout_reward_advantage.py @@ -24,6 +24,11 @@ import os import urllib.request +from miles_plugins.omni.rollout_contract import ( + build_generate_payload, + parse_generate_response, +) + SERVER = os.environ.get("SERVER", "http://localhost:8010") DATA = os.environ.get("DATA", "examples/higgs_tts_rl/tts_smoke.jsonl") GROUP = int(os.environ.get("GROUP", "4")) @@ -42,24 +47,24 @@ def _higgs_adapter(): def rollout(input_ids: list[int], seed: int) -> dict: - req = { - "input_ids": input_ids, - "sampling_params": {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": 256, "seed": seed}, - "return_logprob": True, - "output_modalities": ["audio"], - } + req = build_generate_payload( + input_ids, + {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": 256, "seed": seed}, + output_modalities=["audio"], + ) r = urllib.request.urlopen( - urllib.request.Request(SERVER + "/generate", data=json.dumps(req).encode(), - headers={"Content-Type": "application/json"}), + urllib.request.Request( + SERVER + "/generate", + data=json.dumps(req).encode(), + headers={"Content-Type": "application/json"}, + ), timeout=180, ) - resp = json.loads(r.read()) - otl = resp["meta_info"].get("output_token_logprobs") or [] - audio = resp.get("audio") or {} + result = parse_generate_response(json.loads(r.read())) return { - "codec_tokens": [t for _, t in otl], - "old_logprobs": [lp for lp, _ in otl], - "audio_b64": audio.get("data"), + "codec_tokens": result.response_tokens, + "old_logprobs": result.response_log_probs, + "audio_b64": (result.audio or {}).get("data"), } diff --git a/examples/higgs_tts_rl/sglang_omni_patches/higgs_codec_logprob_fix.patch b/examples/higgs_tts_rl/sglang_omni_patches/higgs_codec_logprob_fix.patch deleted file mode 100644 index 4fa46c81585..00000000000 --- a/examples/higgs_tts_rl/sglang_omni_patches/higgs_codec_logprob_fix.patch +++ /dev/null @@ -1,95 +0,0 @@ -diff --git a/sglang_omni/models/higgs_tts/model.py b/sglang_omni/models/higgs_tts/model.py -index 3e68894..e32137e 100644 ---- a/sglang_omni/models/higgs_tts/model.py -+++ b/sglang_omni/models/higgs_tts/model.py -@@ -170,6 +170,12 @@ class HiggsTTSModel(nn.Module): - self._cg_codes_BN = torch.zeros( - pool_size, num_codebooks, dtype=torch.long, device=cg_device - ) -+ # Per-row codebook-0 log-prob of the sampled token, for RL rollout. Written -+ # in-place each step (eager prefill + CUDA-graph decode) so the runner records -+ # the true behavior-policy logprob instead of a text-vocab placeholder. -+ self._step_cb0_logprob = torch.zeros( -+ pool_size, dtype=torch.float32, device=cg_device -+ ) - # Note(Jiaxin): Packs codes_BN | was_done | active_generation_done into one buffer. - self._cg_collect_staging = torch.zeros( - pool_size, num_codebooks + 2, dtype=torch.long, device=cg_device -@@ -305,6 +311,13 @@ class HiggsTTSModel(nn.Module): - # Note(yichi): One D2H per step to skip STOP-sentinel rows in the Python append loop. - was_done_cpu = was_done.cpu().tolist() - codes_BN = codes_BN.detach().to(torch.long) -+ -+ # Codebook-0 log-prob of each sampled token, for RL rollout. Indexed by -+ # forward-batch row (aligned with the runner's per-request collect loop). -+ cb0_logits = logits_BNV[:, 0, :] -+ cb0_idx = codes_BN[:, 0:1].clamp(0, cb0_logits.shape[-1] - 1) -+ cb0_lp = torch.log_softmax(cb0_logits, dim=-1).gather(1, cb0_idx).squeeze(1) -+ self._step_cb0_logprob[:batch_size] = cb0_lp - for b in range(batch_size): - if was_done_cpu[b]: - continue -@@ -365,6 +378,13 @@ class HiggsTTSModel(nn.Module): - self._cg_active_last_codes[:batch_size] = new_last_codes_BN - self._cg_codes_BN[:batch_size] = codes_BN - -+ # Codebook-0 log-prob of each sampled token, for RL rollout. In-place buffer -+ # write (no value-dependent control flow / D2H) keeps this CUDA-graph safe. -+ cb0_logits = logits_BNV[:, 0, :] -+ cb0_idx = codes_BN[:, 0:1].long().clamp(0, cb0_logits.shape[-1] - 1) -+ cb0_lp = torch.log_softmax(cb0_logits, dim=-1).gather(1, cb0_idx).squeeze(1) -+ self._step_cb0_logprob[:batch_size] = cb0_lp -+ - text_vocab_size = self.backbone.config.vocab_size - return torch.zeros( - (batch_size, text_vocab_size), -diff --git a/sglang_omni/models/higgs_tts/model_runner.py b/sglang_omni/models/higgs_tts/model_runner.py -index 0d04a59..6a5aa31 100644 ---- a/sglang_omni/models/higgs_tts/model_runner.py -+++ b/sglang_omni/models/higgs_tts/model_runner.py -@@ -299,8 +299,8 @@ class HiggsTTSModelRunner(ModelRunner): - continue - codes_N = codes_BN_cpu[b].to(torch.long).clone() - data.output_codes.append(codes_N) -- nt = getattr(result.logits_output, "next_token_logits", None) -- self._record_rollout_logprob(data, nt[b] if nt is not None else None, int(codes_N[0].item())) -+ lpv = getattr(self.model, "_step_cb0_logprob", None) -+ self._record_rollout_logprob(data, lpv[b] if lpv is not None else None, int(codes_N[0].item())) - data.generation_done = bool(gen_done_after_cpu[b]) - self._emit_code_chunk(sched_req, codes_N) - self._mark_sampler_finished(req, data.generation_done) -@@ -377,8 +377,8 @@ class HiggsTTSModelRunner(ModelRunner): - continue - codes_N = codes_log[-1] - data.output_codes.append(codes_N.detach().cpu().clone()) -- nt = getattr(result.logits_output, "next_token_logits", None) -- self._record_rollout_logprob(data, nt[b] if nt is not None else None, int(codes_N[0].item())) -+ lpv = getattr(self.model, "_step_cb0_logprob", None) -+ self._record_rollout_logprob(data, lpv[b] if lpv is not None else None, int(codes_N[0].item())) - data.generation_done = bool(model._sampler_pool.generation_done[row].item()) - self._emit_code_chunk(sched_req, data.output_codes[-1]) - self._mark_sampler_finished(req, data.generation_done) -@@ -397,15 +397,16 @@ class HiggsTTSModelRunner(ModelRunner): - req.finished_reason = FINISH_MATCHED_TOKEN(EOC_ID) - - @staticmethod -- def _record_rollout_logprob(data, logits_row, cb0_token): -- """Record codebook-0 codec token + its logprob for RL rollout.""" -+ def _record_rollout_logprob(data, cb0_logprob, cb0_token): -+ """Record a sampled codebook-0 codec token + its log-prob for RL rollout. -+ -+ ``cb0_logprob`` is the model's pre-computed codebook-0 log-prob of the sampled -+ token (the true behavior-policy logprob); ``None`` falls back to 0.0. -+ """ - if not getattr(data, "return_logprob", False): - return -- if logits_row is None: -- data.output_token_logprobs.append([0.0, int(cb0_token)]) -- return -- logp = torch.log_softmax(logits_row.float(), dim=-1) -- data.output_token_logprobs.append([float(logp[int(cb0_token)].item()), int(cb0_token)]) -+ lp = 0.0 if cb0_logprob is None else float(cb0_logprob) -+ data.output_token_logprobs.append([lp, int(cb0_token)]) - - def _emit_code_chunk(self, sched_req: Any, codes_N: torch.Tensor) -> None: - if self._outbox is None: diff --git a/miles_plugins/omni/rollout_contract.py b/miles_plugins/omni/rollout_contract.py index 436414c9f69..aa0ac3711d1 100644 --- a/miles_plugins/omni/rollout_contract.py +++ b/miles_plugins/omni/rollout_contract.py @@ -62,6 +62,7 @@ def build_generate_payload( metadata: dict[str, Any] | None = None, output_modalities: list[str] | None = None, return_logprob: bool = True, + return_omni_rollout: bool = False, audio_data: list[str] | None = None, ) -> dict[str, Any]: """Build an omni ``/generate`` request body from pre-tokenized inputs. @@ -75,6 +76,8 @@ def build_generate_payload( "sampling_params": clean_sampling_params(sampling_params), "return_logprob": return_logprob, } + if return_omni_rollout: + payload["return_omni_rollout"] = True if metadata: payload["metadata"] = metadata if output_modalities is not None: @@ -97,6 +100,8 @@ class OmniRolloutResult: prompt_tokens: int = 0 completion_tokens: int = 0 audio: dict[str, Any] | None = None + output_codebook_tokens: list[list[int]] | None = None + omni_rollout: dict[str, Any] | None = None def parse_generate_response(response: dict[str, Any]) -> OmniRolloutResult: @@ -130,6 +135,8 @@ def parse_generate_response(response: dict[str, Any]) -> OmniRolloutResult: if "finish_reason" not in meta: raise ValueError("omni /generate meta_info is missing 'finish_reason'") + output_codebook_tokens = _parse_output_codebook_tokens(meta, completion_tokens) + return OmniRolloutResult( response_tokens=response_tokens, response_log_probs=response_log_probs, @@ -138,11 +145,39 @@ def parse_generate_response(response: dict[str, Any]) -> OmniRolloutResult: weight_version=meta.get("weight_version"), cached_tokens=int(meta.get("cached_tokens") or 0), prompt_tokens=int(meta.get("prompt_tokens") or 0), - completion_tokens=int(completion_tokens if completion_tokens is not None else len(response_tokens)), + completion_tokens=int( + completion_tokens if completion_tokens is not None else len(response_tokens) + ), audio=response.get("audio"), + output_codebook_tokens=output_codebook_tokens, + omni_rollout=meta.get("omni_rollout"), ) +def _parse_output_codebook_tokens( + meta: dict[str, Any], completion_tokens: Any +) -> list[list[int]] | None: + raw = meta.get("output_codebook_tokens") + if raw is None: + return None + if not isinstance(raw, list): + raise ValueError("output_codebook_tokens must be a list of codebook rows") + if completion_tokens is not None and len(raw) != completion_tokens: + raise ValueError( + f"output_codebook_tokens length ({len(raw)}) " + f"!= completion_tokens ({completion_tokens})" + ) + parsed: list[list[int]] = [] + for i, row in enumerate(raw): + if not isinstance(row, (list, tuple)) or not row: + raise ValueError( + f"output_codebook_tokens[{i}] is malformed: {row!r}; " + "expected a non-empty codebook row" + ) + parsed.append([int(token) for token in row]) + return parsed + + def apply_response_to_sample( sample: Sample, prompt_ids: list[int], diff --git a/tests/fast/test_omni_rollout_contract.py b/tests/fast/test_omni_rollout_contract.py index 1cd73f1076c..f87d768365d 100644 --- a/tests/fast/test_omni_rollout_contract.py +++ b/tests/fast/test_omni_rollout_contract.py @@ -61,9 +61,11 @@ def test_build_generate_payload_shape_and_metadata(): {"temperature": 1.0, "skip_special_tokens": True}, metadata={"index": 5}, output_modalities=["audio"], + return_omni_rollout=True, ) assert payload["input_ids"] == [1, 2, 3] assert payload["return_logprob"] is True + assert payload["return_omni_rollout"] is True assert payload["sampling_params"] == {"temperature": 1.0} # forbidden key removed assert payload["metadata"] == {"index": 5} assert payload["output_modalities"] == ["audio"] @@ -99,6 +101,31 @@ def test_parse_generate_response_captures_audio_and_text(): assert result.text == "hi" +def test_parse_generate_response_captures_codebook_tokens_and_omni_rollout(): + resp = _response( + [[-0.1, 10], [-0.2, 11]], + completion_tokens=2, + output_codebook_tokens=[[10, 1, 2], [11, 3, 4]], + omni_rollout={"version": 1, "action_streams": []}, + ) + + result = parse_generate_response(resp) + + assert result.output_codebook_tokens == [[10, 1, 2], [11, 3, 4]] + assert result.omni_rollout == {"version": 1, "action_streams": []} + + +def test_parse_generate_response_codebook_length_mismatch_raises(): + with pytest.raises(ValueError, match="output_codebook_tokens length"): + parse_generate_response( + _response( + [[-0.1, 10], [-0.2, 11]], + completion_tokens=2, + output_codebook_tokens=[[10, 1, 2]], + ) + ) + + def test_parse_generate_response_empty_completion_is_not_an_error(): result = parse_generate_response(_response([], completion_tokens=0)) assert result.response_tokens == [] From aae8700d5ef3e526066b4b65534e9b75bea92e51 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Wed, 8 Jul 2026 16:02:57 +0800 Subject: [PATCH 23/24] fix(omni): preserve external rollout metadata paths --- miles/ray/rollout/train_data_conversion.py | 1 + miles/utils/http_utils.py | 10 +++-- miles_plugins/omni/omni_generate_fn.py | 17 ++++++-- miles_plugins/omni/rollout_contract.py | 18 ++++++++ .../ray/rollout/test_train_data_conversion.py | 18 ++++++++ tests/fast/test_omni_generate_fn.py | 43 +++++++++++++++++++ tests/fast/test_omni_rollout_contract.py | 17 ++++++++ tests/fast/utils/test_http_utils.py | 43 ++++++++++++++++++- 8 files changed, 158 insertions(+), 9 deletions(-) diff --git a/miles/ray/rollout/train_data_conversion.py b/miles/ray/rollout/train_data_conversion.py index 65bc8d4b6db..fc44b1f0bf2 100644 --- a/miles/ray/rollout/train_data_conversion.py +++ b/miles/ray/rollout/train_data_conversion.py @@ -158,6 +158,7 @@ def split_train_data_by_dp(args, data, dp_size): "teacher_log_probs", "opd_reverse_kl", "weight_versions", + "metadata", ]: if key not in data: continue diff --git a/miles/utils/http_utils.py b/miles/utils/http_utils.py index 621e60532e8..4cc93a18d96 100644 --- a/miles/utils/http_utils.py +++ b/miles/utils/http_utils.py @@ -227,10 +227,12 @@ async def _post(client, url, payload, max_retries=60, action="post", headers=Non def init_http_client(args): """Initialize HTTP client and optionally enable distributed POST via Ray.""" global _http_client, _client_concurrency, _distributed_post_enabled - if not args.rollout_num_gpus: - return - _client_concurrency = args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + rollout_num_gpus = int(getattr(args, "rollout_num_gpus", 0) or 0) + rollout_num_gpus_per_engine = int(getattr(args, "rollout_num_gpus_per_engine", 1) or 1) + sglang_server_concurrency = int(getattr(args, "sglang_server_concurrency", 1) or 1) + num_engines = max(1, rollout_num_gpus // rollout_num_gpus_per_engine) + _client_concurrency = max(1, sglang_server_concurrency * num_engines) if _http_client is None: _http_client = httpx.AsyncClient( limits=httpx.Limits(max_connections=_client_concurrency), @@ -238,7 +240,7 @@ def init_http_client(args): ) # Optionally initialize distributed POST via Ray without changing interfaces - if args.use_distributed_post: + if getattr(args, "use_distributed_post", False): _init_ray_distributed_post(args) _distributed_post_enabled = True diff --git a/miles_plugins/omni/omni_generate_fn.py b/miles_plugins/omni/omni_generate_fn.py index e586a4b22c6..fd722ded417 100644 --- a/miles_plugins/omni/omni_generate_fn.py +++ b/miles_plugins/omni/omni_generate_fn.py @@ -32,11 +32,13 @@ async def __call__(self, input: GenerateFnInput) -> GenerateFnOutput: prompt_ids = compute_prompt_ids_from_sample(input.state, sample) # Partial-rollout resume: continue from already-generated tokens and shrink the - # remaining budget by what was already produced (mirrors single_turn.generate). - if len(sample.response) > 0: + # remaining budget by what was already produced. Audio-only rollouts can have + # empty decoded text, so response_length is the source of truth here. + generated_token_count = _generated_token_count(sample, prompt_ids) + if generated_token_count > 0: input_ids = sample.tokens - if sampling_params.get("max_new_tokens") is not None: - sampling_params["max_new_tokens"] -= len(sample.tokens) - len(prompt_ids) + total_budget = sampling_params.get("max_new_tokens", args.rollout_max_response_len) + sampling_params["max_new_tokens"] = total_budget - generated_token_count else: input_ids = prompt_ids @@ -63,6 +65,13 @@ async def __call__(self, input: GenerateFnInput) -> GenerateFnOutput: return GenerateFnOutput(samples=sample) +def _generated_token_count(sample: Sample, prompt_ids: list[int]) -> int: + """Return how many completion tokens have already been generated for resume.""" + if sample.response_length > 0: + return sample.response_length + return max(0, len(sample.tokens) - len(prompt_ids)) + + def _clamp_max_new_tokens(args, sampling_params: dict, prompt_len: int) -> Sample.Status | None: """Cap ``max_new_tokens`` by the context budget; return a halt status if none remains. diff --git a/miles_plugins/omni/rollout_contract.py b/miles_plugins/omni/rollout_contract.py index aa0ac3711d1..916ec189505 100644 --- a/miles_plugins/omni/rollout_contract.py +++ b/miles_plugins/omni/rollout_contract.py @@ -224,4 +224,22 @@ def apply_response_to_sample( # GPU and concatenates; store it in metadata instead. sample.metadata["generated_audio"] = result.audio + _store_train_rollout_metadata(sample, result) + return sample + + +def _store_train_rollout_metadata(sample: Sample, result: OmniRolloutResult) -> None: + """Copy Higgs rollout artifacts into the train-side metadata bridge.""" + if result.output_codebook_tokens is None and result.omni_rollout is None: + return + + if sample.train_metadata is None: + sample.train_metadata = {} + + if result.output_codebook_tokens is not None: + sample.train_metadata.setdefault("output_codebook_tokens", []) + sample.train_metadata["output_codebook_tokens"].extend(result.output_codebook_tokens) + + if result.omni_rollout is not None: + sample.train_metadata["omni_rollout"] = result.omni_rollout diff --git a/tests/fast/ray/rollout/test_train_data_conversion.py b/tests/fast/ray/rollout/test_train_data_conversion.py index 4a3da1a2883..1a64ee1d92c 100644 --- a/tests/fast/ray/rollout/test_train_data_conversion.py +++ b/tests/fast/ray/rollout/test_train_data_conversion.py @@ -126,6 +126,19 @@ def test_optional_field_round_number_from_metadata(self): ) assert out["round_number"][0] == 7 + def test_optional_field_train_metadata_passed_through(self): + args = make_args(rewards_normalization=False) + s = make_sample() + s.train_metadata = {"output_codebook_tokens": [[10, 1]], "omni_rollout": {"version": 1}} + out = convert_samples_to_train_data( + args, + [s], + metadata={}, + custom_convert_samples_to_train_data_func=None, + custom_reward_post_process_func=None, + ) + assert out["metadata"] == [{"output_codebook_tokens": [[10, 1]], "omni_rollout": {"version": 1}}] + def test_optional_field_raw_reward_overridden_from_metadata(self): args = make_args(rewards_normalization=False) s = make_sample(reward=1.0) @@ -431,11 +444,16 @@ def test_optional_keys_propagated_when_present(self): "sample_indices": [0, 1], "rollout_log_probs": [[-0.1], [-0.2]], "round_number": [1, 2], + "metadata": [ + {"output_codebook_tokens": [[10, 1]]}, + {"output_codebook_tokens": [[20, 2]]}, + ], } refs = split_train_data_by_dp(args, data, dp_size=2) parts = [ray.get(r.inner) for r in refs] assert "rollout_log_probs" in parts[0] assert "round_number" in parts[0] + assert parts[0]["metadata"] == [{"output_codebook_tokens": [[10, 1]]}] def test_shared_keys_not_split(self): """raw_reward, total_lengths, dynamic_global_batch_size are shared, not split.""" diff --git a/tests/fast/test_omni_generate_fn.py b/tests/fast/test_omni_generate_fn.py index f1faf802316..95d02f7db3b 100644 --- a/tests/fast/test_omni_generate_fn.py +++ b/tests/fast/test_omni_generate_fn.py @@ -43,6 +43,8 @@ def _canned_response(): "meta_info": { "finish_reason": {"type": "stop"}, "output_token_logprobs": [[-0.1, 10], [-0.2, 11]], + "output_codebook_tokens": [[10, 101], [11, 111]], + "omni_rollout": {"version": 1, "action_streams": []}, "completion_tokens": 2, "weight_version": "7", "cached_tokens": 0, @@ -95,6 +97,8 @@ async def fake_post(url, payload, **kwargs): # generated audio is reward-facing -> metadata, never multimodal_train_inputs assert result_sample.metadata["generated_audio"] == {"format": "wav", "data": ""} assert result_sample.multimodal_train_inputs is None + assert result_sample.train_metadata["output_codebook_tokens"] == [[10, 101], [11, 111]] + assert result_sample.train_metadata["omni_rollout"] == {"version": 1, "action_streams": []} assert result_sample.weight_versions == ["7"] assert result_sample.status == Sample.Status.COMPLETED @@ -181,3 +185,42 @@ async def fake_post(url, payload, **kwargs): # new on-policy tokens are trainable; mask stays aligned with response_length assert s.loss_mask == [0, 0, 1, 1] assert len(s.loss_mask) == s.response_length + + +def test_omni_generate_fn_audio_only_resume_uses_token_state(monkeypatch): + captured = {} + + async def fake_post(url, payload, **kwargs): + captured["payload"] = payload + return { + "text": "", + "meta_info": { + "finish_reason": {"type": "stop"}, + "output_token_logprobs": [[-0.5, 20]], + "completion_tokens": 1, + "cached_tokens": 0, + "prompt_tokens": 5, + }, + } + + monkeypatch.setattr(omni_mod, "post", fake_post) + + fn = load_generate_function(_HOOK_PATH) + sample = Sample(prompt="hi") + sample.tokens = [1, 2, 3, 10, 11] + sample.response = "" # Higgs/TTS can produce audio/code tokens without decoded text. + sample.response_length = 2 + sample.rollout_log_probs = [-0.1, -0.2] + inp = GenerateFnInput( + state=_fake_state(), + sample=sample, + sampling_params={"max_new_tokens": 64}, + evaluation=False, + ) + + out = asyncio.run(fn(inp)) + + assert captured["payload"]["input_ids"] == [1, 2, 3, 10, 11] + assert captured["payload"]["sampling_params"]["max_new_tokens"] == 62 + assert out.samples.tokens == [1, 2, 3, 10, 11, 20] + assert out.samples.response_length == 3 diff --git a/tests/fast/test_omni_rollout_contract.py b/tests/fast/test_omni_rollout_contract.py index f87d768365d..c21781866d9 100644 --- a/tests/fast/test_omni_rollout_contract.py +++ b/tests/fast/test_omni_rollout_contract.py @@ -189,6 +189,23 @@ def test_apply_response_to_sample_stores_audio_in_metadata_not_train_inputs(): assert sample.multimodal_train_inputs is None +def test_apply_response_to_sample_stores_codebook_rollout_in_train_metadata(): + sample = Sample(prompt="p", tokens=[]) + result = parse_generate_response( + _response( + [[-0.1, 10], [-0.2, 11]], + completion_tokens=2, + output_codebook_tokens=[[10, 1], [11, 2]], + omni_rollout={"version": 1, "action_streams": []}, + ) + ) + + apply_response_to_sample(sample, [1, 2], result) + + assert sample.train_metadata["output_codebook_tokens"] == [[10, 1], [11, 2]] + assert sample.train_metadata["omni_rollout"] == {"version": 1, "action_streams": []} + + def test_apply_response_to_sample_multi_turn_accumulates(): sample = Sample(prompt="p", tokens=[]) first = parse_generate_response(_response([[-0.1, 10]], completion_tokens=1)) diff --git a/tests/fast/utils/test_http_utils.py b/tests/fast/utils/test_http_utils.py index 11e6ce65bc0..cd9e8799c20 100644 --- a/tests/fast/utils/test_http_utils.py +++ b/tests/fast/utils/test_http_utils.py @@ -25,11 +25,13 @@ import socket import threading import time +from types import SimpleNamespace from unittest.mock import patch import pytest -from miles.utils.http_utils import wait_for_server_ready +import miles.utils.http_utils as http_utils +from miles.utils.http_utils import init_http_client, wait_for_server_ready def _find_free_port() -> int: @@ -194,3 +196,42 @@ def fake_connect(addr, timeout=None): # The fake clock should have advanced past the timeout assert fake_time[0] >= timeout + + +class TestInitHttpClient: + def teardown_method(self): + if http_utils._http_client is not None: + import asyncio + + asyncio.run(http_utils._http_client.aclose()) + http_utils._http_client = None + http_utils._client_concurrency = 0 + http_utils._distributed_post_enabled = False + http_utils._post_actors = [] + http_utils._post_actor_idx = 0 + + def test_external_zero_gpu_still_initializes_local_client(self): + args = SimpleNamespace( + rollout_num_gpus=0, + rollout_num_gpus_per_engine=1, + sglang_server_concurrency=8, + use_distributed_post=False, + ) + + init_http_client(args) + + assert http_utils._http_client is not None + assert http_utils._client_concurrency == 8 + + def test_local_rollout_concurrency_scales_by_engine_count(self): + args = SimpleNamespace( + rollout_num_gpus=8, + rollout_num_gpus_per_engine=2, + sglang_server_concurrency=3, + use_distributed_post=False, + ) + + init_http_client(args) + + assert http_utils._http_client is not None + assert http_utils._client_concurrency == 12 From 7b8fb1443ead9177359958d37eb51c53c4d27366 Mon Sep 17 00:00:00 2001 From: Hayden727 Date: Sat, 11 Jul 2026 17:08:14 +0800 Subject: [PATCH 24/24] fix(omni): train all Higgs codebook actions --- .../higgs_tts_rl/onpolicy_grpo_weight_sync.py | 98 +++++++------- .../thinker_text_rl/fsdp_trainer_launch.sh | 3 +- .../onpolicy_grpo_weight_sync.py | 29 +---- miles_plugins/omni/higgs_actor.py | 121 ++++++++++++++---- miles_plugins/omni/omni_generate_fn.py | 1 + miles_plugins/omni/rollout_contract.py | 89 ++++++++++--- tests/fast/test_higgs_actor.py | 87 +++++++++++++ tests/fast/test_omni_generate_fn.py | 2 +- tests/fast/test_omni_rollout_contract.py | 50 +++++++- 9 files changed, 368 insertions(+), 112 deletions(-) create mode 100644 tests/fast/test_higgs_actor.py diff --git a/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py b/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py index 63dba3cb4d6..47e1c26d569 100644 --- a/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py +++ b/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py @@ -1,24 +1,7 @@ -"""Full on-policy Higgs TTS RL: GRPO LoRA on the Higgs TTS actor with per-step NCCL -weight-sync to the served sglang-omni ``tts_engine`` stage, so each step's rollouts -are on-policy. - -The fourth closed-loop component for TTS, mirroring the thinker onpolicy_grpo_weight_sync.py: - rollout (/generate -> codec tokens + codebook-0 logprobs + audio) - -> composite reward (Whisper ASR CER + audio-validity guards) - -> GRPO advantage over codebook-0 tokens - -> LoRA policy update (trainer recomputes new logprobs via HiggsTtsActor) - -> NCCL broadcast of LoRA-merged backbone weights into the served tts_engine stage - (names in the checkpoint `body.*` convention; the server fuses q/k/v on load). - -Run (container, miles venv, free GPU for the trainer; Higgs server on another GPU): - SERVER=http://localhost:8010 HIGGS_CKPT='' MASTER_PORT=29641 \ - ASR_MODEL=openai/whisper-base ASR_DEVICE=cuda:0 CUDA_VISIBLE_DEVICES=4 \ - HF_HUB_OFFLINE=1 NCCL_P2P_DISABLE=1 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0 \ - PYTHONPATH=/root/rl-omni/sglang-omni:/root/rl-omni/miles \ - python examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py - -CRITICAL: set NCCL_P2P_DISABLE=1 on BOTH the server and the trainer (single-GPU masks -per process), exactly as in the thinker on-policy script. +"""On-policy Higgs TTS GRPO with per-step SGLang-Omni weight sync. + +Set ``TRAIN_MODE=lora`` (default) for the low-memory smoke path or ``full`` to +train and sync the complete backbone plus tied codebook embedding/head. """ from __future__ import annotations @@ -35,6 +18,7 @@ from miles_plugins.omni.rollout_contract import ( build_generate_payload, parse_generate_response, + parse_omni_action_stream, ) SERVER = os.environ.get("SERVER", "http://localhost:8010") @@ -47,6 +31,9 @@ GROUP_NAME = os.environ.get("GROUP_NAME", "higgs_tts_wsync") TEMP = float(os.environ.get("TEMP", "0.8")) MAX_NEW = int(os.environ.get("MAX_NEW", "256")) +TOP_K = int(os.environ["TOP_K"]) if os.environ.get("TOP_K") else None +TRAIN_MODE = os.environ.get("TRAIN_MODE", "lora").lower() +LR = float(os.environ.get("LR", "2e-5" if TRAIN_MODE == "lora" else "1e-6")) EPS = 0.2 @@ -58,24 +45,32 @@ def post(path: str, body: dict, timeout: int = 300): def rollout(input_ids: list[int], seed: int) -> dict: + sampling_params = { + "temperature": TEMP, + "top_p": 0.95, + "max_new_tokens": MAX_NEW, + "seed": seed, + } + if TOP_K is not None: + sampling_params["top_k"] = TOP_K resp = post( "/generate", build_generate_payload( input_ids, - { - "temperature": TEMP, - "top_p": 0.95, - "max_new_tokens": MAX_NEW, - "seed": seed, - }, + sampling_params, output_modalities=["audio"], + return_omni_rollout=True, ), timeout=180, ) result = parse_generate_response(resp) + stream = parse_omni_action_stream(result.omni_rollout, "higgs_codes") + if result.output_codebook_tokens != stream.actions: + raise ValueError("Higgs output_codebook_tokens do not match omni_rollout actions") return { - "old": result.response_log_probs, - "codes": result.output_codebook_tokens, + "old": stream.logprobs, + "mask": stream.action_mask, + "codes": stream.actions, "audio": (result.audio or {}).get("data"), } @@ -87,7 +82,7 @@ def main() -> None: from tokenizers import Tokenizer from transformers import PreTrainedTokenizerFast - from miles_plugins.omni.higgs_actor import HiggsTtsActor + from miles_plugins.omni.higgs_actor import HiggsTtsActor, clipped_grpo_loss from miles_plugins.omni.tts_reward import TtsCompositeReward ckpt = glob.glob(HIGGS_CKPT)[0] if "*" in HIGGS_CKPT else HIGGS_CKPT @@ -96,12 +91,17 @@ def main() -> None: reward_fn = TtsCompositeReward() actor = HiggsTtsActor(ckpt, device="cuda:0") - actor.backbone = get_peft_model( - actor.backbone, - LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"], task_type=None), - ) - actor.backbone.train() - opt = torch.optim.AdamW([p for p in actor.backbone.parameters() if p.requires_grad], lr=2e-5) + if TRAIN_MODE == "lora": + actor.fused_embed.requires_grad_(False) + actor.backbone = get_peft_model( + actor.backbone, + LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"], task_type=None), + ) + elif TRAIN_MODE != "full": + raise ValueError(f"TRAIN_MODE must be 'lora' or 'full', got {TRAIN_MODE!r}") + actor.train() + trainable_params = [param for param in actor.parameters() if param.requires_grad] + opt = torch.optim.AdamW(trainable_params, lr=LR) try: from sglang.srt.utils import init_custom_process_group @@ -140,6 +140,7 @@ def _init_server(): raise init_err[0] print("WEIGHT_UPDATE_GROUP_READY", flush=True) + @torch.no_grad() def merged_lora_weights() -> dict[str, torch.Tensor]: out: dict[str, torch.Tensor] = {} for name, mod in actor.backbone.named_modules(): @@ -153,8 +154,17 @@ def merged_lora_weights() -> dict[str, torch.Tensor]: out["body." + hf + ".weight"] = w.to(torch.bfloat16).contiguous() return out + @torch.no_grad() + def weights_to_sync() -> dict[str, torch.Tensor]: + if TRAIN_MODE == "lora": + return merged_lora_weights() + return { + name: tensor.detach().to(torch.bfloat16).contiguous() + for name, tensor in actor.full_server_weights().items() + } + def sync_to_server() -> int: - wd = merged_lora_weights() + wd = weights_to_sync() names = sorted(wd) spec = { "names": names, @@ -197,21 +207,19 @@ def _update(): if c.cer is not None: step_cer += c.cer n_cer += 1 - for s, adv in zip(samples, [r - mean_r for r in rewards]): + for s, adv in zip(samples, [r - mean_r for r in rewards], strict=True): codes = s["codes"] if not codes or adv == 0.0: continue - new = actor.codebook0_logprobs(pid, codes) - T = min(len(new), len(s["old"])) - new = new[:T] - old = torch.tensor(s["old"][:T], device="cuda:0") - ratio = torch.exp(new - old) - loss = -torch.min(ratio * adv, torch.clamp(ratio, 1 - EPS, 1 + EPS) * adv).mean() + new = actor.codebook_logprobs(pid, codes, temperature=TEMP, top_k=TOP_K) + old = torch.tensor(s["old"], dtype=new.dtype, device="cuda:0") + mask = torch.tensor(s["mask"], dtype=torch.bool, device="cuda:0") + loss = clipped_grpo_loss(new, old, mask, advantage=adv, clip_eps=EPS) loss = loss / (GROUP * PROMPTS) loss.backward() step_loss += loss.item() * (GROUP * PROMPTS) n += 1 - torch.nn.utils.clip_grad_norm_([p for p in actor.backbone.parameters() if p.requires_grad], 1.0) + torch.nn.utils.clip_grad_norm_(trainable_params, 1.0) opt.step() synced = sync_to_server() # next step's rollouts are on-policy mean_cer = step_cer / n_cer if n_cer else float("nan") diff --git a/examples/thinker_text_rl/fsdp_trainer_launch.sh b/examples/thinker_text_rl/fsdp_trainer_launch.sh index f4781ae12ac..d7c9148fc9e 100644 --- a/examples/thinker_text_rl/fsdp_trainer_launch.sh +++ b/examples/thinker_text_rl/fsdp_trainer_launch.sh @@ -93,7 +93,8 @@ ray start --head --node-ip-address 127.0.0.1 --num-gpus $NGPU --disable-usage-st # train.py connects to the running cluster via ray.init(address="auto"); no dashboard / job-submit needed export PYTHONPATH=${REPO}:/root/rl-omni/sglang-omni export HF_HUB_OFFLINE=1 -export NCCL_P2P_DISABLE=1 +# Keep NCCL's default transport selection. Set NCCL_P2P_DISABLE=1 only when a +# diagnosed container/topology issue requires the shared-memory fallback. export CUDA_DEVICE_MAX_CONNECTIONS=1 python3 train.py \ diff --git a/examples/thinker_text_rl/onpolicy_grpo_weight_sync.py b/examples/thinker_text_rl/onpolicy_grpo_weight_sync.py index c93fb7d1216..83a6a1861e8 100644 --- a/examples/thinker_text_rl/onpolicy_grpo_weight_sync.py +++ b/examples/thinker_text_rl/onpolicy_grpo_weight_sync.py @@ -1,25 +1,8 @@ -"""Full on-policy Thinker text RL: GRPO LoRA on the Qwen3-Omni Thinker with per-step NCCL -weight-sync to the served sglang-omni thinker, so each step's rollouts are on-policy. - -Beyond lora_grpo_smoke.py this adds the 4th closed-loop component done *properly*: -after each optimizer step the LoRA-merged thinker weights are broadcast into the served -thinker stage via sglang-omni's distributed weight-update admin plane -(``/init_weights_update_group`` + ``/update_weights_from_distributed`` + ``stages=[thinker]``), -the exact pattern from sglang-omni's E2E refit test. The thinker stage's ``load_weights`` -accepts plain ``model.*`` names, so the extracted-thinker names sync directly. - -Run (container, miles venv, free GPU for the trainer; server already on another GPU): - THINKER=/root/qwen3-omni-thinker DATA=examples/thinker_text_rl/math_smoke.jsonl \ - SERVER=http://localhost:8003 MASTER_PORT=29631 CUDA_VISIBLE_DEVICES=4 \ - NCCL_P2P_DISABLE=1 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0 \ - python examples/thinker_text_rl/onpolicy_grpo_weight_sync.py - -CRITICAL: the trainer and the sglang-omni server run as separate processes, each with a -single GPU exposed via CUDA_VISIBLE_DEVICES (both see it as cuda:0). NCCL would try direct -P2P between the two physical GPUs and fail with "Cuda invalid argument" because neither -process can resolve the peer's masked device. Set NCCL_P2P_DISABLE=1 on BOTH the server -and the trainer so NCCL falls back to shared-memory transport. Verified: 4-step on-policy -run, synced_params=160/step, stable loss/reward. +"""On-policy Qwen3-Omni Thinker GRPO with per-step SGLang-Omni weight sync. + +NCCL supports the usual one-process-per-GPU setup even when each process names its +local device ``cuda:0``. Use ``NCCL_P2P_DISABLE=1`` only as a topology-specific +diagnostic workaround, not as a requirement of this layout. """ from __future__ import annotations @@ -166,7 +149,7 @@ def _update(): rewards = [1.0 if ex["label"] in s["text"] else 0.0 for s in samples] mean_r = sum(rewards) / len(rewards) step_reward += mean_r - for s, adv in zip(samples, [r - mean_r for r in rewards]): + for s, adv in zip(samples, [r - mean_r for r in rewards], strict=True): if not s["tokens"] or adv == 0.0: continue full = torch.tensor([pid + s["tokens"]], device="cuda:0") diff --git a/miles_plugins/omni/higgs_actor.py b/miles_plugins/omni/higgs_actor.py index ce24a469a4c..4ac6933064e 100644 --- a/miles_plugins/omni/higgs_actor.py +++ b/miles_plugins/omni/higgs_actor.py @@ -1,6 +1,4 @@ -"""Trainer-side Higgs TTS actor: a gradient-enabled teacher-forced forward that -reproduces the served model's per-step codebook-0 log-probs over a sampled codec -sequence, so the RL trainer can recompute new-policy log-probs for GRPO. +"""Trainer-side Higgs TTS actor with gradient-enabled codebook logprob replay. The served `HiggsTTSModel` backbone is sglang's inference `Qwen3ForCausalLM` (paged attention / CUDA graph, no autograd), so it cannot be trained directly. @@ -30,6 +28,73 @@ "body.norm.": "norm.", } _FUSED_EMBED_KEY = "tied.embedding.modality_embeddings.0.embedding.weight" +_GREEDY_TEMP_THRESHOLD = 1e-5 + + +def backbone_parameter_to_checkpoint_name(name: str) -> str: + """Map a plain ``Qwen3Model`` parameter name back to the Higgs checkpoint.""" + if name.startswith("embed_tokens."): + return "tied.embedding.text_embedding." + name[len("embed_tokens.") :] + if name.startswith("layers."): + return "body.layers." + name[len("layers.") :] + if name.startswith("norm."): + return "body.norm." + name[len("norm.") :] + raise ValueError(f"unsupported Higgs actor backbone parameter {name!r}") + + +def build_full_server_weights(backbone, fused_embed: torch.Tensor) -> dict[str, torch.Tensor]: + """Return full-parameter actor weights using names accepted by the server.""" + weights = {backbone_parameter_to_checkpoint_name(name): param for name, param in backbone.named_parameters()} + weights[_FUSED_EMBED_KEY] = fused_embed + return weights + + +def selected_codebook_logprobs( + step_hidden: torch.Tensor, + fused_embed: torch.Tensor, + codes: torch.Tensor, + *, + num_codebooks: int, + codebook_vocab: int, + temperature: float, + top_k: int | None = None, +) -> torch.Tensor: + """Compute selected-action logprobs for every cell in a codebook lattice.""" + if codes.ndim != 2 or tuple(codes.shape) != (step_hidden.shape[0], num_codebooks): + raise ValueError(f"codes shape {tuple(codes.shape)} must be {(step_hidden.shape[0], num_codebooks)}") + expected_rows = num_codebooks * codebook_vocab + if fused_embed.ndim != 2 or fused_embed.shape[0] != expected_rows: + raise ValueError(f"fused_embed shape {tuple(fused_embed.shape)} must start with {expected_rows} rows") + + logits = F.linear(step_hidden.float(), fused_embed.float()).view( + step_hidden.shape[0], num_codebooks, codebook_vocab + ) + greedy = temperature <= _GREEDY_TEMP_THRESHOLD or top_k == 1 + effective_temperature = 1.0 if greedy else max(float(temperature), _GREEDY_TEMP_THRESHOLD) + logprobs = torch.log_softmax(logits / effective_temperature, dim=-1) + return logprobs.gather(-1, codes.long().unsqueeze(-1)).squeeze(-1) + + +def clipped_grpo_loss( + current_logprobs: torch.Tensor, + old_logprobs: torch.Tensor, + action_mask: torch.Tensor, + *, + advantage: float | torch.Tensor, + clip_eps: float, +) -> torch.Tensor: + """Per-action clipped GRPO loss over the trainable codebook cells.""" + if current_logprobs.shape != old_logprobs.shape or current_logprobs.shape != action_mask.shape: + raise ValueError("current logprobs, old logprobs, and action mask must have the same shape") + action_mask = action_mask.to(device=current_logprobs.device, dtype=torch.bool) + if not bool(action_mask.any()): + raise ValueError("GRPO action mask contains no trainable actions") + + ratio = torch.exp(current_logprobs - old_logprobs) + advantage_t = torch.as_tensor(advantage, dtype=ratio.dtype, device=ratio.device) + unclipped = ratio * advantage_t + clipped = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps) * advantage_t + return -torch.minimum(unclipped, clipped)[action_mask].mean() def _resolve_ckpt_dir(path_or_glob: str) -> str: @@ -41,10 +106,11 @@ def _resolve_ckpt_dir(path_or_glob: str) -> str: return path_or_glob -class HiggsTtsActor: +class HiggsTtsActor(torch.nn.Module): """Differentiable Higgs codec policy (Qwen3 backbone + fused codebook head).""" def __init__(self, ckpt_dir: str, device: str = "cuda:0", dtype=torch.bfloat16): + super().__init__() from safetensors import safe_open from transformers import Qwen3Config, Qwen3Model @@ -87,12 +153,10 @@ def __init__(self, ckpt_dir: str, device: str = "cuda:0", dtype=torch.bfloat16): if real_missing: raise RuntimeError(f"missing backbone keys: {real_missing[:5]}") - # Fused codebook weight [N*V, D]: input embedding (sum over codebooks) and, - # tied, the codebook-0 head = its first V rows. - self.fused_embed = fused_embed.to(device=device, dtype=dtype) - self._cb_offsets = ( - torch.arange(self.num_codebooks, device=device) * self.codebook_vocab - ) + # Fused codebook weight [N*V, D], tied between the summed input embedding + # and all per-codebook output heads. + self.fused_embed = torch.nn.Parameter(fused_embed.to(device=device, dtype=dtype)) + self._cb_offsets = torch.arange(self.num_codebooks, device=device) * self.codebook_vocab def _rename_backbone(self, key: str) -> str | None: if key.startswith("tied.embedding.modality_embeddings.0.model."): @@ -107,15 +171,15 @@ def _embed_codes(self, codes_LN: torch.Tensor) -> torch.Tensor: fused_ids = codes_LN + self._cb_offsets return F.embedding(fused_ids, self.fused_embed).sum(dim=-2) - def codebook0_logprobs( - self, prompt_ids: list[int], codebook_tokens: list[list[int]] + def codebook_logprobs( + self, + prompt_ids: list[int], + codebook_tokens: list[list[int]], + *, + temperature: float, + top_k: int | None = None, ) -> torch.Tensor: - """Teacher-forced new-policy log-probs of each step's sampled codebook-0 token. - - ``codebook_tokens`` is ``[T, num_codebooks]`` (the full per-step codes the - server fed back). Returns ``[T]`` log-probs aligned with the rollout's - ``output_token_logprobs`` (codebook-0). - """ + """Teacher-forced selected-action logprobs for all sampled codebooks.""" device = self.device prompt = torch.tensor(prompt_ids, dtype=torch.long, device=device) codes = torch.tensor(codebook_tokens, dtype=torch.long, device=device) # [T, N] @@ -141,7 +205,20 @@ def codebook0_logprobs( ) hidden = out.last_hidden_state[0] # [L, D] step_hidden = hidden[P - 1 : P - 1 + T] # [T, D] - cb0_logits = F.linear(step_hidden.float(), self.fused_embed[: self.codebook_vocab].float()) - logp = torch.log_softmax(cb0_logits, dim=-1) # [T, V] - sampled_cb0 = codes[:, 0] # [T] - return logp[torch.arange(T, device=device), sampled_cb0] + return selected_codebook_logprobs( + step_hidden, + self.fused_embed, + codes, + num_codebooks=self.num_codebooks, + codebook_vocab=self.codebook_vocab, + temperature=temperature, + top_k=top_k, + ) + + def codebook0_logprobs(self, prompt_ids: list[int], codebook_tokens: list[list[int]]) -> torch.Tensor: + """Raw codebook-0 logprobs retained for server parity diagnostics.""" + return self.codebook_logprobs(prompt_ids, codebook_tokens, temperature=1.0)[:, 0] + + def full_server_weights(self) -> dict[str, torch.Tensor]: + """Expose every full-training weight with a server-compatible name.""" + return build_full_server_weights(self.backbone, self.fused_embed) diff --git a/miles_plugins/omni/omni_generate_fn.py b/miles_plugins/omni/omni_generate_fn.py index fd722ded417..b193f3f07ad 100644 --- a/miles_plugins/omni/omni_generate_fn.py +++ b/miles_plugins/omni/omni_generate_fn.py @@ -52,6 +52,7 @@ async def __call__(self, input: GenerateFnInput) -> GenerateFnOutput: sampling_params, metadata=_request_metadata(sample), output_modalities=sample.metadata.get("output_modalities"), + return_omni_rollout=True, audio_data=_encode_input_audio(sample), ) diff --git a/miles_plugins/omni/rollout_contract.py b/miles_plugins/omni/rollout_contract.py index 916ec189505..61ca0901c24 100644 --- a/miles_plugins/omni/rollout_contract.py +++ b/miles_plugins/omni/rollout_contract.py @@ -15,6 +15,7 @@ from __future__ import annotations +import math from dataclasses import dataclass, field from typing import Any @@ -104,6 +105,71 @@ class OmniRolloutResult: omni_rollout: dict[str, Any] | None = None +@dataclass(frozen=True) +class OmniActionStream: + """Validated two-dimensional discrete action stream from ``omni_rollout``.""" + + name: str + actions: list[list[int]] + logprobs: list[list[float]] + action_mask: list[list[bool]] + + +def parse_omni_action_stream(omni_rollout: dict[str, Any] | None, stream_name: str) -> OmniActionStream: + """Return one validated ``codebook_2d`` action stream by name.""" + if not isinstance(omni_rollout, dict): + raise ValueError("omni_rollout is required for structured action training") + streams = omni_rollout.get("action_streams") + if not isinstance(streams, list): + raise ValueError("omni_rollout.action_streams must be a list") + + matches = [stream for stream in streams if stream.get("name") == stream_name] + if len(matches) != 1: + raise ValueError(f"expected exactly one omni action stream named {stream_name!r}, got {len(matches)}") + stream = matches[0] + if stream.get("action_type") != "discrete" or stream.get("layout") != "codebook_2d": + raise ValueError(f"omni action stream {stream_name!r} must be a discrete codebook_2d stream") + + shape = stream.get("shape") + if not isinstance(shape, list) or len(shape) != 2 or not all(isinstance(dim, int) and dim >= 0 for dim in shape): + raise ValueError(f"omni action stream {stream_name!r} has invalid shape {shape!r}") + rows, channels = shape + + actions = _parse_2d_stream_field(stream, "actions", rows, channels, int) + logprobs = _parse_2d_stream_field(stream, "logprobs", rows, channels, float) + action_mask = _parse_2d_stream_field(stream, "action_mask", rows, channels, bool) + for row_idx, (logprob_row, mask_row) in enumerate(zip(logprobs, action_mask, strict=True)): + for channel_idx, (logprob, trainable) in enumerate(zip(logprob_row, mask_row, strict=True)): + if trainable and not math.isfinite(logprob): + raise ValueError(f"non-finite logprob at {stream_name}[{row_idx}][{channel_idx}]") + + return OmniActionStream( + name=stream_name, + actions=actions, + logprobs=logprobs, + action_mask=action_mask, + ) + + +def _parse_2d_stream_field(stream: dict[str, Any], field_name: str, rows: int, channels: int, cast) -> list[list[Any]]: + value = stream.get(field_name) + if not isinstance(value, list) or len(value) != rows: + raise ValueError( + f"omni action stream {stream['name']!r} field {field_name!r} " + f"does not match declared shape {[rows, channels]}" + ) + + parsed: list[list[Any]] = [] + for row in value: + if not isinstance(row, list) or len(row) != channels: + raise ValueError( + f"omni action stream {stream['name']!r} field {field_name!r} " + f"does not match declared shape {[rows, channels]}" + ) + parsed.append([cast(item) for item in row]) + return parsed + + def parse_generate_response(response: dict[str, Any]) -> OmniRolloutResult: """Parse an omni ``/generate`` response into :class:`OmniRolloutResult`. @@ -119,17 +185,14 @@ def parse_generate_response(response: dict[str, Any]) -> OmniRolloutResult: response_log_probs: list[float] = [] for i, item in enumerate(token_logprobs): if not isinstance(item, (list, tuple)) or len(item) != 2: - raise ValueError( - f"output_token_logprobs[{i}] is malformed: {item!r}; expected [log_prob, token_id]" - ) + raise ValueError(f"output_token_logprobs[{i}] is malformed: {item!r}; expected [log_prob, token_id]") response_log_probs.append(float(item[0])) response_tokens.append(int(item[1])) completion_tokens = meta.get("completion_tokens") if completion_tokens is not None and len(response_tokens) != completion_tokens: raise ValueError( - f"output_token_logprobs length ({len(response_tokens)}) " - f"!= completion_tokens ({completion_tokens})" + f"output_token_logprobs length ({len(response_tokens)}) " f"!= completion_tokens ({completion_tokens})" ) if "finish_reason" not in meta: @@ -145,34 +208,26 @@ def parse_generate_response(response: dict[str, Any]) -> OmniRolloutResult: weight_version=meta.get("weight_version"), cached_tokens=int(meta.get("cached_tokens") or 0), prompt_tokens=int(meta.get("prompt_tokens") or 0), - completion_tokens=int( - completion_tokens if completion_tokens is not None else len(response_tokens) - ), + completion_tokens=int(completion_tokens if completion_tokens is not None else len(response_tokens)), audio=response.get("audio"), output_codebook_tokens=output_codebook_tokens, omni_rollout=meta.get("omni_rollout"), ) -def _parse_output_codebook_tokens( - meta: dict[str, Any], completion_tokens: Any -) -> list[list[int]] | None: +def _parse_output_codebook_tokens(meta: dict[str, Any], completion_tokens: Any) -> list[list[int]] | None: raw = meta.get("output_codebook_tokens") if raw is None: return None if not isinstance(raw, list): raise ValueError("output_codebook_tokens must be a list of codebook rows") if completion_tokens is not None and len(raw) != completion_tokens: - raise ValueError( - f"output_codebook_tokens length ({len(raw)}) " - f"!= completion_tokens ({completion_tokens})" - ) + raise ValueError(f"output_codebook_tokens length ({len(raw)}) " f"!= completion_tokens ({completion_tokens})") parsed: list[list[int]] = [] for i, row in enumerate(raw): if not isinstance(row, (list, tuple)) or not row: raise ValueError( - f"output_codebook_tokens[{i}] is malformed: {row!r}; " - "expected a non-empty codebook row" + f"output_codebook_tokens[{i}] is malformed: {row!r}; " "expected a non-empty codebook row" ) parsed.append([int(token) for token in row]) return parsed diff --git a/tests/fast/test_higgs_actor.py b/tests/fast/test_higgs_actor.py new file mode 100644 index 00000000000..b0b23802fd7 --- /dev/null +++ b/tests/fast/test_higgs_actor.py @@ -0,0 +1,87 @@ +import pytest +import torch + +from miles_plugins.omni.higgs_actor import ( + _FUSED_EMBED_KEY, + backbone_parameter_to_checkpoint_name, + build_full_server_weights, + clipped_grpo_loss, + selected_codebook_logprobs, +) + + +def test_selected_codebook_logprobs_scores_every_codebook(): + hidden = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + fused_weight = torch.tensor( + [ + [2.0, 0.0], + [0.0, 1.0], + [-1.0, 0.0], + [0.0, 2.0], + [1.0, 0.0], + [0.0, -1.0], + ] + ) + codes = torch.tensor([[0, 2], [1, 0]]) + + actual = selected_codebook_logprobs( + hidden, + fused_weight, + codes, + num_codebooks=2, + codebook_vocab=3, + temperature=0.5, + ) + + logits = torch.nn.functional.linear(hidden, fused_weight).view(2, 2, 3) / 0.5 + expected = torch.log_softmax(logits, dim=-1).gather(-1, codes.unsqueeze(-1)).squeeze(-1) + assert actual.shape == (2, 2) + assert torch.allclose(actual, expected) + + +def test_clipped_grpo_loss_backpropagates_through_all_unmasked_codebooks(): + current = torch.tensor([[-0.2, -0.3], [-0.4, -0.5]], requires_grad=True) + old = torch.tensor([[-0.25, -0.35], [-0.45, -0.55]]) + mask = torch.tensor([[True, True], [True, False]]) + + loss = clipped_grpo_loss(current, old, mask, advantage=0.7, clip_eps=0.2) + loss.backward() + + assert current.grad is not None + assert torch.all(current.grad[mask] != 0) + assert current.grad[~mask].item() == 0 + + +@pytest.mark.parametrize( + ("actor_name", "checkpoint_name"), + [ + ("embed_tokens.weight", "tied.embedding.text_embedding.weight"), + ("layers.2.self_attn.q_proj.weight", "body.layers.2.self_attn.q_proj.weight"), + ("norm.weight", "body.norm.weight"), + ], +) +def test_backbone_parameter_to_checkpoint_name(actor_name, checkpoint_name): + assert backbone_parameter_to_checkpoint_name(actor_name) == checkpoint_name + + +def test_build_full_server_weights_includes_backbone_and_tied_codebook_weight(): + tensors = { + "embed_tokens.weight": torch.randn(3, 2), + "layers.0.self_attn.q_proj.weight": torch.randn(2, 2), + "norm.weight": torch.randn(2), + } + + class FakeBackbone: + def named_parameters(self): + return iter(tensors.items()) + + fused_weight = torch.randn(6, 2) + weights = build_full_server_weights(FakeBackbone(), fused_weight) + + assert set(weights) == { + "tied.embedding.text_embedding.weight", + "body.layers.0.self_attn.q_proj.weight", + "body.norm.weight", + _FUSED_EMBED_KEY, + } + assert weights[_FUSED_EMBED_KEY] is fused_weight diff --git a/tests/fast/test_omni_generate_fn.py b/tests/fast/test_omni_generate_fn.py index 95d02f7db3b..7ec46e4e263 100644 --- a/tests/fast/test_omni_generate_fn.py +++ b/tests/fast/test_omni_generate_fn.py @@ -10,7 +10,6 @@ from types import SimpleNamespace import numpy as np -import pytest import miles_plugins.omni.omni_generate_fn as omni_mod from miles.rollout.base_types import GenerateFnInput @@ -86,6 +85,7 @@ async def fake_post(url, payload, **kwargs): assert captured["url"] == "http://127.0.0.1:8000/generate" assert payload["input_ids"] == [1, 2, 3] assert payload["return_logprob"] is True + assert payload["return_omni_rollout"] is True assert payload["sampling_params"] == {"temperature": 0.7, "seed": 9, "max_new_tokens": 64} assert payload["metadata"] == {"group_index": 2, "index": 5} assert "audio_data" not in payload # no input audio on this sample diff --git a/tests/fast/test_omni_rollout_contract.py b/tests/fast/test_omni_rollout_contract.py index c21781866d9..df8a24ba238 100644 --- a/tests/fast/test_omni_rollout_contract.py +++ b/tests/fast/test_omni_rollout_contract.py @@ -20,6 +20,7 @@ build_generate_payload, clean_sampling_params, parse_generate_response, + parse_omni_action_stream, ) @@ -115,6 +116,51 @@ def test_parse_generate_response_captures_codebook_tokens_and_omni_rollout(): assert result.omni_rollout == {"version": 1, "action_streams": []} +def test_parse_omni_action_stream_returns_full_codebook_lattice(): + trace = { + "version": 1, + "total_action_count": 4, + "action_streams": [ + { + "name": "higgs_codes", + "action_type": "discrete", + "layout": "codebook_2d", + "shape": [3, 2], + "actions": [[10, 1024], [11, 20], [1025, 21]], + "logprobs": [[-0.1, -9.0], [-0.2, -0.3], [-9.0, -0.4]], + "action_mask": [[1, 0], [1, 1], [0, 1]], + } + ], + } + + stream = parse_omni_action_stream(trace, "higgs_codes") + + assert stream.actions == [[10, 1024], [11, 20], [1025, 21]] + assert stream.logprobs == [[-0.1, -9.0], [-0.2, -0.3], [-9.0, -0.4]] + assert stream.action_mask == [[True, False], [True, True], [False, True]] + + +def test_parse_omni_action_stream_rejects_shape_mismatch(): + trace = { + "version": 1, + "total_action_count": 2, + "action_streams": [ + { + "name": "higgs_codes", + "action_type": "discrete", + "layout": "codebook_2d", + "shape": [2, 2], + "actions": [[10, 20]], + "logprobs": [[-0.1, -0.2]], + "action_mask": [[1, 1]], + } + ], + } + + with pytest.raises(ValueError, match="shape"): + parse_omni_action_stream(trace, "higgs_codes") + + def test_parse_generate_response_codebook_length_mismatch_raises(): with pytest.raises(ValueError, match="output_codebook_tokens length"): parse_generate_response( @@ -164,9 +210,7 @@ def test_parse_generate_response_missing_finish_reason_raises(): def test_apply_response_to_sample_aligns_and_validates(): sample = Sample(prompt="p", tokens=[]) prompt_ids = [1, 2, 3] - result = parse_generate_response( - _response([[-0.1, 10], [-0.2, 11]], completion_tokens=2, weight_version="3") - ) + result = parse_generate_response(_response([[-0.1, 10], [-0.2, 11]], completion_tokens=2, weight_version="3")) apply_response_to_sample(sample, prompt_ids, result, update_loss_mask=True) assert sample.tokens == [1, 2, 3, 10, 11]