diff --git a/examples/higgs_tts_rl/logprob_parity_probe.py b/examples/higgs_tts_rl/logprob_parity_probe.py new file mode 100644 index 00000000000..47f5f2b8193 --- /dev/null +++ b/examples/higgs_tts_rl/logprob_parity_probe.py @@ -0,0 +1,106 @@ +"""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 +`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/higgs_tts_rl/logprob_parity_probe.py +""" + +from __future__ import annotations + +import glob +import json +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 +# 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 = 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": result.response_log_probs, + "cb0_tokens": result.response_tokens, + "codebook_tokens": result.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/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py b/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py new file mode 100644 index 00000000000..47e1c26d569 --- /dev/null +++ b/examples/higgs_tts_rl/onpolicy_grpo_weight_sync.py @@ -0,0 +1,236 @@ +"""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 + +import glob +import json +import os +import threading +import urllib.request + +import torch +from peft import LoraConfig, get_peft_model + +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") +HIGGS_CKPT = os.environ["HIGGS_CKPT"] +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", "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 + + +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: + 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, + 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": stream.logprobs, + "mask": stream.action_mask, + "codes": stream.actions, + "audio": (result.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, clipped_grpo_loss + 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") + 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 + 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) + + @torch.no_grad() + 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 + + @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 = weights_to_sync() + 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], strict=True): + codes = s["codes"] + if not codes or adv == 0.0: + continue + 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_(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") + 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("Higgs TTS on-policy loop complete (per-step NCCL weight-sync to served tts_engine)") + + +if __name__ == "__main__": + main() diff --git a/examples/higgs_tts_rl/rollout_reward_advantage.py b/examples/higgs_tts_rl/rollout_reward_advantage.py new file mode 100644 index 00000000000..1ab98d28222 --- /dev/null +++ b/examples/higgs_tts_rl/rollout_reward_advantage.py @@ -0,0 +1,109 @@ +"""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: + 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 +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/higgs_tts_rl/rollout_reward_advantage.py +""" + +from __future__ import annotations + +import glob +import json +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")) +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 = 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"}, + ), + timeout=180, + ) + result = parse_generate_response(json.loads(r.read())) + return { + "codec_tokens": result.response_tokens, + "old_logprobs": result.response_log_probs, + "audio_b64": (result.audio or {}).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("Higgs TTS rollout->composite-reward->advantage demonstrated") + + +if __name__ == "__main__": + main() diff --git a/examples/higgs_tts_rl/tts_harder.jsonl b/examples/higgs_tts_rl/tts_harder.jsonl new file mode 100644 index 00000000000..b218e03e481 --- /dev/null +++ b/examples/higgs_tts_rl/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."} diff --git a/examples/higgs_tts_rl/tts_smoke.jsonl b/examples/higgs_tts_rl/tts_smoke.jsonl new file mode 100644 index 00000000000..3cf11ebc7c0 --- /dev/null +++ b/examples/higgs_tts_rl/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/examples/thinker_text_rl/fsdp_trainer_launch.sh b/examples/thinker_text_rl/fsdp_trainer_launch.sh new file mode 100644 index 00000000000..d7c9148fc9e --- /dev/null +++ b/examples/thinker_text_rl/fsdp_trainer_launch.sh @@ -0,0 +1,108 @@ +#!/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 +# 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 \ + ${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 new file mode 100644 index 00000000000..0761d60af41 --- /dev/null +++ b/examples/thinker_text_rl/lora_grpo_smoke.py @@ -0,0 +1,112 @@ +"""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) + -> 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/thinker_text_rl/math_smoke.jsonl \ + SERVER=http://localhost:8000/generate CUDA_VISIBLE_DEVICES=4 \ + python examples/thinker_text_rl/lora_grpo_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("Thinker RL smoke complete (rollout->reward->advantage->LoRA update over multiple steps)") + + +if __name__ == "__main__": + main() diff --git a/examples/thinker_text_rl/math_harder.jsonl b/examples/thinker_text_rl/math_harder.jsonl new file mode 100644 index 00000000000..230c016710a --- /dev/null +++ b/examples/thinker_text_rl/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/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/math_smoke.jsonl b/examples/thinker_text_rl/math_smoke.jsonl new file mode 100644 index 00000000000..ea2577d2ef6 --- /dev/null +++ b/examples/thinker_text_rl/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"} diff --git a/examples/thinker_text_rl/onpolicy_grpo_weight_sync.py b/examples/thinker_text_rl/onpolicy_grpo_weight_sync.py new file mode 100644 index 00000000000..83a6a1861e8 --- /dev/null +++ b/examples/thinker_text_rl/onpolicy_grpo_weight_sync.py @@ -0,0 +1,178 @@ +"""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 + +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 = os.environ.get("GROUP_NAME", "thinker_wsync") +MAX_NEW = int(os.environ.get("MAX_NEW", "24")) +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": MAX_NEW, "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: + 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 + ).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], strict=True): + 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("Thinker on-policy loop complete (per-step NCCL weight-sync to served thinker)") + + +if __name__ == "__main__": + main() 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) 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/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py index d50098e686b..c60403c356e 100644 --- a/miles/rollout/generate_utils/generate_endpoint_utils.py +++ b/miles/rollout/generate_utils/generate_endpoint_utils.py @@ -9,17 +9,30 @@ 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 ( + 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 = { @@ -60,6 +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_inputs := extract_audio_inputs(multimodal_inputs): + payload["audio_data"] = encode_audios_for_rollout_engine(audio_inputs) return payload, None 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/utils/processing_utils.py b/miles/utils/processing_utils.py index 855a06a8fb2..f204a370f26 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,66 @@ 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: + 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: + 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}" + + +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. + + 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/__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/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() diff --git a/miles_plugins/omni/higgs_actor.py b/miles_plugins/omni/higgs_actor.py new file mode 100644 index 00000000000..4ac6933064e --- /dev/null +++ b/miles_plugins/omni/higgs_actor.py @@ -0,0 +1,224 @@ +"""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. +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/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. +""" + +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" +_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: + 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(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 + + 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], 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."): + 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 codebook_logprobs( + self, + prompt_ids: list[int], + codebook_tokens: list[list[int]], + *, + temperature: float, + top_k: int | None = None, + ) -> torch.Tensor: + """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] + T = int(codes.shape[0]) + P = int(prompt.shape[0]) + + 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. + 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] + 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/math_reward.py b/miles_plugins/omni/math_reward.py new file mode 100644 index 00000000000..1861352b361 --- /dev/null +++ b/miles_plugins/omni/math_reward.py @@ -0,0 +1,42 @@ +"""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 +(``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 diff --git a/miles_plugins/omni/omni_generate_fn.py b/miles_plugins/omni/omni_generate_fn.py new file mode 100644 index 00000000000..b193f3f07ad --- /dev/null +++ b/miles_plugins/omni/omni_generate_fn.py @@ -0,0 +1,108 @@ +"""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 (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 + +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, extract_audio_inputs +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 = 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 and shrink the + # 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 + 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 + + 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"), + return_omni_rollout=True, + audio_data=_encode_input_audio(sample), + ) + + 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 _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. + + 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.""" + audios = extract_audio_inputs(sample.multimodal_inputs) + 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 = { + "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..61ca0901c24 --- /dev/null +++ b/miles_plugins/omni/rollout_contract.py @@ -0,0 +1,300 @@ +"""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 + +import math +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, + return_omni_rollout: bool = False, + 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 return_omni_rollout: + payload["return_omni_rollout"] = True + 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 + output_codebook_tokens: list[list[int]] | None = None + 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`. + + 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'") + + output_codebook_tokens = _parse_output_codebook_tokens(meta, completion_tokens) + + 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"), + 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], + 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. 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 + ``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 + + # 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) + + if result.audio is not None: + # 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 + + _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/miles_plugins/omni/tts_reward.py b/miles_plugins/omni/tts_reward.py new file mode 100644 index 00000000000..b72aad05f94 --- /dev/null +++ b/miles_plugins/omni/tts_reward.py @@ -0,0 +1,171 @@ +"""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 +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 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_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 new file mode 100644 index 00000000000..7ec46e4e263 --- /dev/null +++ b/tests/fast/test_omni_generate_fn.py @@ -0,0 +1,226 @@ +"""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 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]], + "output_codebook_tokens": [[10, 101], [11, 111]], + "omni_rollout": {"version": 1, "action_streams": []}, + "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["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 + + 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.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 + + +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,") + + +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 + + +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 new file mode 100644 index 00000000000..df8a24ba238 --- /dev/null +++ b/tests/fast/test_omni_rollout_contract.py @@ -0,0 +1,362 @@ +"""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, + parse_omni_action_stream, +) + + +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"], + 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"] + # 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_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_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( + _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 == [] + 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_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": ""}) + + +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_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) + # 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_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)) + 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] + + +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 --------------------------------------------------------------- + + +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) + + +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) + + +# --- 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") 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