From c8a1084f6219f59c91da25aa0ee13f4170980bf3 Mon Sep 17 00:00:00 2001
From: Jingwen Gu <75733630+JingwenGu0829@users.noreply.github.com>
Date: Tue, 21 Jul 2026 22:40:02 +0000
Subject: [PATCH 01/32] Add SGLang Omni multimodal input adapter

---
 miles/rollout/generate_hub/sglang_omni.py     |  68 ++++++++++
 .../generate_utils/generate_endpoint_utils.py | 114 ++++++++++++++--
 miles/utils/processing_utils.py               |  28 ++++
 requirements.txt                              |   1 +
 .../rollout/generate_hub/test_sglang_omni.py  | 122 ++++++++++++++++++
 5 files changed, 324 insertions(+), 9 deletions(-)
 create mode 100644 miles/rollout/generate_hub/sglang_omni.py
 create mode 100644 tests/fast/rollout/generate_hub/test_sglang_omni.py

diff --git a/miles/rollout/generate_hub/sglang_omni.py b/miles/rollout/generate_hub/sglang_omni.py
new file mode 100644
index 00000000000..ecb8c46e93b
--- /dev/null
+++ b/miles/rollout/generate_hub/sglang_omni.py
@@ -0,0 +1,68 @@
+"""Single-turn rollout using SGLang Omni's processed multimodal contract.
+
+Select with ``--custom-generate-function-path
+miles.rollout.generate_hub.sglang_omni.generate``.
+"""
+
+from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput
+from miles.rollout.generate_utils.generate_endpoint_utils import (
+    compute_prompt_ids_from_sample,
+    compute_request_payload,
+    compute_routing_headers,
+    multimodal_route_headers,
+    update_sample_from_response,
+)
+from miles.utils.http_utils import post
+from miles.utils.types import Sample
+
+
+async def generate(input: GenerateFnInput) -> GenerateFnOutput:
+    args = input.args
+    sample = input.sample
+    sampling_params = input.sampling_params
+    assert sample.status in {
+        Sample.Status.PENDING,
+        Sample.Status.ABORTED,
+    }, f"{sample.status=}"
+    url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate"
+
+    prompt_ids = compute_prompt_ids_from_sample(input.state, sample)
+    has_multimodal_inputs = sample.multimodal_inputs and any(
+        value is not None for value in sample.multimodal_inputs.values()
+    )
+    if has_multimodal_inputs and sample.multimodal_train_inputs is None:
+        raise ValueError(
+            "SGLang Omni multimodal rollout requires processor-produced " "sample.multimodal_train_inputs"
+        )
+
+    if sample.response:
+        input_ids = sample.tokens
+        sampling_params["max_new_tokens"] -= len(sample.tokens) - len(prompt_ids)
+        assert sampling_params["max_new_tokens"] >= 0
+        if sampling_params["max_new_tokens"] == 0:
+            sample.status = Sample.Status.TRUNCATED
+            return GenerateFnOutput(samples=sample)
+    else:
+        input_ids = prompt_ids
+
+    payload, halt_status = compute_request_payload(
+        args,
+        input_ids=input_ids,
+        sampling_params=sampling_params,
+        multimodal_inputs=sample.multimodal_inputs,
+        multimodal_train_inputs=sample.multimodal_train_inputs,
+    )
+    if payload is None:
+        sample.status = halt_status
+        return GenerateFnOutput(samples=sample)
+
+    headers = compute_routing_headers(args, sample) or {}
+    headers.update(multimodal_route_headers(payload.get("multimodal_train_inputs")) or {})
+    output = await post(url, payload, headers=headers or None)
+    await update_sample_from_response(
+        args,
+        sample,
+        payload=payload,
+        output=output,
+    )
+    return GenerateFnOutput(samples=sample)
diff --git a/miles/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py
index 1de0e66d079..e5b210a51c4 100644
--- a/miles/rollout/generate_utils/generate_endpoint_utils.py
+++ b/miles/rollout/generate_utils/generate_endpoint_utils.py
@@ -2,31 +2,124 @@
 Utils to integrate SGLang's `/generate` endpoint with RL things like Sample.
 """
 
+import base64
 from copy import deepcopy
 from typing import Any
 
 import numpy as np
 import pybase64
+import torch
 
 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_image_for_rollout_engine,
+    extract_multimodal_train_inputs,
+)
 from miles.utils.types import Sample
 
 
+_MULTIMODAL_TENSOR_DTYPES = {
+    "bool",
+    "uint8",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "float16",
+    "bfloat16",
+    "float32",
+    "float64",
+}
+_SGLANG_OMNI_MULTIMODAL_TENSOR_NAMES = frozenset(
+    {
+        "pixel_values",
+        "image_grid_thw",
+        "input_features",
+        "feature_attention_mask",
+        "audio_feature_lengths",
+        "pixel_values_videos",
+        "video_grid_thw",
+        "video_second_per_grid",
+    }
+)
+
+
+def _multimodal_modalities(tensors: dict[str, torch.Tensor]) -> list[str]:
+    modalities = []
+    if "pixel_values" in tensors:
+        modalities.append("image")
+    if "input_features" in tensors:
+        modalities.append("audio")
+    if "pixel_values_videos" in tensors:
+        modalities.append("video")
+    return modalities
+
+
+def serialize_multimodal_train_inputs(
+    multimodal_train_inputs: dict[str, Any],
+) -> dict[str, Any]:
+    """Encode the processor tensor kwargs shared with SGLang Omni."""
+    if not multimodal_train_inputs:
+        raise ValueError("multimodal_train_inputs must not be empty")
+    unknown = sorted(set(multimodal_train_inputs) - _SGLANG_OMNI_MULTIMODAL_TENSOR_NAMES)
+    if unknown:
+        raise ValueError("unsupported SGLang Omni processor tensor fields: " + ", ".join(unknown))
+
+    tensors: dict[str, dict[str, Any]] = {}
+    tensor_values: dict[str, torch.Tensor] = {}
+    for name, value in multimodal_train_inputs.items():
+        if not isinstance(value, torch.Tensor):
+            raise TypeError(
+                "multimodal_train_inputs must contain only tensors; " f"{name!r} has type {type(value).__name__}"
+            )
+        tensor = value.detach().to(device="cpu").contiguous()
+        dtype = str(tensor.dtype).removeprefix("torch.")
+        if dtype not in _MULTIMODAL_TENSOR_DTYPES:
+            raise TypeError(f"unsupported multimodal tensor dtype for {name!r}: {dtype}")
+        raw = tensor.reshape(-1).view(torch.uint8).numpy().tobytes()
+        tensors[name] = {
+            "dtype": dtype,
+            "shape": list(tensor.shape),
+            "data": base64.b64encode(raw).decode("ascii"),
+        }
+        tensor_values[name] = tensor
+
+    modalities = _multimodal_modalities(tensor_values)
+    if not modalities:
+        raise ValueError("multimodal_train_inputs does not contain an image, audio, or video " "encoder tensor")
+    return {"version": 1, "modalities": modalities, "tensors": tensors}
+
+
+def multimodal_route_headers(
+    serialized_inputs: dict[str, Any] | None,
+) -> dict[str, str] | None:
+    """Route large JSON bundles without asking the router to scan tensor data."""
+    if not serialized_inputs:
+        return None
+    capabilities = [f"{modality}_input" for modality in serialized_inputs.get("modalities", [])]
+    if not capabilities:
+        return None
+    return {"x-sglang-omni-route-capabilities": ",".join(capabilities)}
+
+
 # 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)
+    if (
+        state.processor
+        and sample.multimodal_inputs
+        and any(value is not None for value in sample.multimodal_inputs.values())
+    ):
+        processor_output = call_processor(state.processor, prompt, sample.multimodal_inputs)
         prompt_ids = processor_output["input_ids"][0]
 
-        # TODO shall we move it to other places? then can make this function immutable
-        sample.multimodal_train_inputs = {
-            k: v for k, v in processor_output.items() if k not in ["input_ids", "attention_mask"]
-        } or None
+        sample.multimodal_train_inputs = extract_multimodal_train_inputs(processor_output)
 
-        return prompt_ids
+        if hasattr(prompt_ids, "tolist"):
+            prompt_ids = prompt_ids.tolist()
+        return [int(token_id) for token_id in prompt_ids]
     else:
         if not isinstance(prompt, str):
             prompt = state.tokenizer.apply_chat_template(
@@ -56,6 +149,7 @@ def compute_request_payload(
     input_ids: list[int],
     sampling_params: dict,
     multimodal_inputs: dict | None = None,
+    multimodal_train_inputs: dict[str, Any] | None = None,
 ) -> tuple[dict[str, Any] | None, Sample.Status | None]:
     sampling_params = deepcopy(sampling_params)
     max_new_tokens = sampling_params.pop("max_new_tokens", args.rollout_max_response_len)
@@ -73,7 +167,9 @@ def compute_request_payload(
     }
     if is_lora_enabled(args):
         payload["lora_path"] = LORA_ADAPTER_NAME
-    if image_data := (multimodal_inputs or {}).get("images"):
+    if multimodal_train_inputs:
+        payload["multimodal_train_inputs"] = serialize_multimodal_train_inputs(multimodal_train_inputs)
+    elif image_data := (multimodal_inputs or {}).get("images"):
         payload["image_data"] = [encode_image_for_rollout_engine(image) for image in image_data]
 
     return payload, None
diff --git a/miles/utils/processing_utils.py b/miles/utils/processing_utils.py
index 855a06a8fb2..fbcd8695ac1 100644
--- a/miles/utils/processing_utils.py
+++ b/miles/utils/processing_utils.py
@@ -137,6 +137,23 @@ def call_processor(processor, text, multimodal_inputs: dict | None = None):
     return processor(text=text, **kwargs)
 
 
+def extract_multimodal_train_inputs(processor_output):
+    """Normalize processor kwargs to the tensor-only Megatron contract."""
+    import torch
+
+    result = {}
+    for name, value in processor_output.items():
+        if name in ("input_ids", "attention_mask"):
+            continue
+        if not isinstance(value, torch.Tensor):
+            try:
+                value = torch.as_tensor(value)
+            except (TypeError, ValueError, RuntimeError) as exc:
+                raise TypeError(f"processor output {name!r} cannot be converted to a tensor") from exc
+        result[name] = value
+    return result or None
+
+
 def load_processor(name_or_path: str, **kwargs):
     try:
         proc = AutoProcessor.from_pretrained(name_or_path, **kwargs)
@@ -153,6 +170,17 @@ def load_processor(name_or_path: str, **kwargs):
 
 def process_vision_info(prompt, processor):
     # TODO: temporary solution, will write image utils for miles later
+    if getattr(processor, "audio_token", None) is not None:
+        from qwen_omni_utils import process_mm_info
+
+        image_patch_size = getattr(processor.image_processor, "patch_size", DEFAULT_PATCH_SIZE)
+        audios, images, videos = process_mm_info(
+            prompt,
+            use_audio_in_video=False,
+            image_patch_size=image_patch_size,
+        )
+        return {"audio": audios, "images": images, "videos": videos}
+
     from qwen_vl_utils import process_vision_info as qwen_process_vision_info
 
     if hasattr(processor.image_processor, "patch_size"):
diff --git a/requirements.txt b/requirements.txt
index 74b84a2977f..523877bc4d8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -16,6 +16,7 @@ pybase64
 pylatexenc
 pytest-asyncio
 pyyaml
+qwen_omni_utils==0.0.9 # for Qwen Omni audio/video inputs
 qwen_vl_utils # for VLM
 ray[default]
 ring_flash_attn; platform_system == "Linux"
diff --git a/tests/fast/rollout/generate_hub/test_sglang_omni.py b/tests/fast/rollout/generate_hub/test_sglang_omni.py
new file mode 100644
index 00000000000..4a08d9d31ee
--- /dev/null
+++ b/tests/fast/rollout/generate_hub/test_sglang_omni.py
@@ -0,0 +1,122 @@
+import asyncio
+import base64
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from miles.rollout.generate_utils.generate_endpoint_utils import serialize_multimodal_train_inputs
+from miles.utils.types import Sample
+
+
+def test_serialize_audio_video_processor_tensors():
+    inputs = {
+        "input_features": torch.arange(6, dtype=torch.float32).reshape(1, 2, 3),
+        "feature_attention_mask": torch.ones((1, 2), dtype=torch.long),
+        "pixel_values_videos": torch.arange(12, dtype=torch.bfloat16).reshape(2, 2, 3),
+        "video_grid_thw": torch.tensor([[1, 2, 3]], dtype=torch.long),
+        "video_second_per_grid": torch.tensor([0.5], dtype=torch.float32),
+    }
+
+    bundle = serialize_multimodal_train_inputs(inputs)
+
+    assert bundle["version"] == 1
+    assert bundle["modalities"] == ["audio", "video"]
+    assert set(bundle["tensors"]) == set(inputs)
+    for name, tensor in inputs.items():
+        encoded = bundle["tensors"][name]
+        assert encoded["dtype"] == str(tensor.dtype).removeprefix("torch.")
+        assert encoded["shape"] == list(tensor.shape)
+        assert base64.b64decode(encoded["data"]) == (
+            tensor.contiguous().reshape(-1).view(torch.uint8).numpy().tobytes()
+        )
+
+
+def test_qwen_omni_media_extraction_and_tensor_normalization(monkeypatch):
+    from miles.utils import processing_utils
+
+    qwen_omni_utils = pytest.importorskip("qwen_omni_utils")
+    prompt = [{"role": "user", "content": [{"type": "audio", "audio": "a.wav"}]}]
+    captured = {}
+
+    def fake_process_mm_info(conversations, **kwargs):
+        captured["conversations"], captured["kwargs"] = conversations, kwargs
+        return ["audio samples"], ["image"], ["video frames"]
+
+    monkeypatch.setattr(qwen_omni_utils, "process_mm_info", fake_process_mm_info)
+    processor = SimpleNamespace(audio_token="<audio>", image_processor=SimpleNamespace(patch_size=16))
+
+    media = processing_utils.process_vision_info(prompt, processor)
+    train_inputs = processing_utils.extract_multimodal_train_inputs(
+        {
+            "input_ids": [[1, 2]],
+            "attention_mask": [[1, 1]],
+            "pixel_values_videos": torch.ones((1, 2)),
+            "video_second_per_grid": [0.5],
+        }
+    )
+
+    assert media == {
+        "audio": ["audio samples"],
+        "images": ["image"],
+        "videos": ["video frames"],
+    }
+    assert captured == {
+        "conversations": prompt,
+        "kwargs": {"use_audio_in_video": False, "image_patch_size": 16},
+    }
+    assert torch.equal(train_inputs["video_second_per_grid"], torch.tensor([0.5]))
+
+
+def test_sglang_omni_adapter_sends_processed_audio_video(monkeypatch):
+    from miles.rollout.generate_hub import sglang_omni
+
+    captured = {}
+
+    async def fake_post(url, payload, headers=None):
+        captured.update(url=url, payload=payload, headers=headers)
+        return {}
+
+    async def fake_update(*args, **kwargs):
+        return None
+
+    monkeypatch.setattr(sglang_omni, "post", fake_post)
+    monkeypatch.setattr(sglang_omni, "update_sample_from_response", fake_update)
+    monkeypatch.setattr(
+        sglang_omni,
+        "compute_prompt_ids_from_sample",
+        lambda state, sample: [1, 2, 3],
+    )
+
+    args = SimpleNamespace(
+        sglang_router_ip="127.0.0.1",
+        sglang_router_port=30000,
+        sglang_router_policy="round_robin",
+        rollout_max_response_len=128,
+        rollout_max_context_len=0,
+        use_rollout_routing_replay=False,
+        use_rollout_indexer_replay=False,
+    )
+    sample = Sample(
+        multimodal_inputs={"audio": ["audio"], "videos": ["video"]},
+        multimodal_train_inputs={
+            "input_features": torch.ones((1, 2, 3)),
+            "pixel_values_videos": torch.ones((2, 2, 3)),
+        },
+    )
+    generate_input = SimpleNamespace(
+        args=args,
+        sample=sample,
+        sampling_params={"temperature": 1.0, "max_new_tokens": 16},
+        state=None,
+    )
+
+    asyncio.run(sglang_omni.generate(generate_input))
+
+    assert captured["url"].endswith("/generate")
+    assert captured["payload"]["input_ids"] == [1, 2, 3]
+    assert captured["payload"]["multimodal_train_inputs"]["modalities"] == [
+        "audio",
+        "video",
+    ]
+    assert captured["headers"] == {"x-sglang-omni-route-capabilities": "audio_input,video_input"}

From 1d2d67906eb140c92fee7787c8e8acbfeee0d474 Mon Sep 17 00:00:00 2001
From: Jingwen Gu <75733630+JingwenGu0829@users.noreply.github.com>
Date: Wed, 22 Jul 2026 01:16:28 +0000
Subject: [PATCH 02/32] Remove multimodal capability route headers

---
 miles/rollout/generate_hub/sglang_omni.py     |  5 +---
 .../generate_utils/generate_endpoint_utils.py | 30 +------------------
 .../rollout/generate_hub/test_sglang_omni.py  | 12 ++++----
 3 files changed, 8 insertions(+), 39 deletions(-)

diff --git a/miles/rollout/generate_hub/sglang_omni.py b/miles/rollout/generate_hub/sglang_omni.py
index ecb8c46e93b..6dd96e0ea26 100644
--- a/miles/rollout/generate_hub/sglang_omni.py
+++ b/miles/rollout/generate_hub/sglang_omni.py
@@ -9,7 +9,6 @@
     compute_prompt_ids_from_sample,
     compute_request_payload,
     compute_routing_headers,
-    multimodal_route_headers,
     update_sample_from_response,
 )
 from miles.utils.http_utils import post
@@ -56,9 +55,7 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
         sample.status = halt_status
         return GenerateFnOutput(samples=sample)
 
-    headers = compute_routing_headers(args, sample) or {}
-    headers.update(multimodal_route_headers(payload.get("multimodal_train_inputs")) or {})
-    output = await post(url, payload, headers=headers or None)
+    output = await post(url, payload, headers=compute_routing_headers(args, sample))
     await update_sample_from_response(
         args,
         sample,
diff --git a/miles/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py
index e5b210a51c4..01b18b17624 100644
--- a/miles/rollout/generate_utils/generate_endpoint_utils.py
+++ b/miles/rollout/generate_utils/generate_endpoint_utils.py
@@ -45,17 +45,6 @@
 )
 
 
-def _multimodal_modalities(tensors: dict[str, torch.Tensor]) -> list[str]:
-    modalities = []
-    if "pixel_values" in tensors:
-        modalities.append("image")
-    if "input_features" in tensors:
-        modalities.append("audio")
-    if "pixel_values_videos" in tensors:
-        modalities.append("video")
-    return modalities
-
-
 def serialize_multimodal_train_inputs(
     multimodal_train_inputs: dict[str, Any],
 ) -> dict[str, Any]:
@@ -67,7 +56,6 @@ def serialize_multimodal_train_inputs(
         raise ValueError("unsupported SGLang Omni processor tensor fields: " + ", ".join(unknown))
 
     tensors: dict[str, dict[str, Any]] = {}
-    tensor_values: dict[str, torch.Tensor] = {}
     for name, value in multimodal_train_inputs.items():
         if not isinstance(value, torch.Tensor):
             raise TypeError(
@@ -83,24 +71,8 @@ def serialize_multimodal_train_inputs(
             "shape": list(tensor.shape),
             "data": base64.b64encode(raw).decode("ascii"),
         }
-        tensor_values[name] = tensor
-
-    modalities = _multimodal_modalities(tensor_values)
-    if not modalities:
-        raise ValueError("multimodal_train_inputs does not contain an image, audio, or video " "encoder tensor")
-    return {"version": 1, "modalities": modalities, "tensors": tensors}
 
-
-def multimodal_route_headers(
-    serialized_inputs: dict[str, Any] | None,
-) -> dict[str, str] | None:
-    """Route large JSON bundles without asking the router to scan tensor data."""
-    if not serialized_inputs:
-        return None
-    capabilities = [f"{modality}_input" for modality in serialized_inputs.get("modalities", [])]
-    if not capabilities:
-        return None
-    return {"x-sglang-omni-route-capabilities": ",".join(capabilities)}
+    return {"version": 1, "tensors": tensors}
 
 
 # Make this an isolated function because users may want to compute their own
diff --git a/tests/fast/rollout/generate_hub/test_sglang_omni.py b/tests/fast/rollout/generate_hub/test_sglang_omni.py
index 4a08d9d31ee..f4cc0d25f2a 100644
--- a/tests/fast/rollout/generate_hub/test_sglang_omni.py
+++ b/tests/fast/rollout/generate_hub/test_sglang_omni.py
@@ -21,7 +21,7 @@ def test_serialize_audio_video_processor_tensors():
     bundle = serialize_multimodal_train_inputs(inputs)
 
     assert bundle["version"] == 1
-    assert bundle["modalities"] == ["audio", "video"]
+    assert set(bundle) == {"version", "tensors"}
     assert set(bundle["tensors"]) == set(inputs)
     for name, tensor in inputs.items():
         encoded = bundle["tensors"][name]
@@ -115,8 +115,8 @@ async def fake_update(*args, **kwargs):
 
     assert captured["url"].endswith("/generate")
     assert captured["payload"]["input_ids"] == [1, 2, 3]
-    assert captured["payload"]["multimodal_train_inputs"]["modalities"] == [
-        "audio",
-        "video",
-    ]
-    assert captured["headers"] == {"x-sglang-omni-route-capabilities": "audio_input,video_input"}
+    assert set(captured["payload"]["multimodal_train_inputs"]["tensors"]) == {
+        "input_features",
+        "pixel_values_videos",
+    }
+    assert captured["headers"] is None

From 3dc88854a53e676c3b39d41abd986a36aacfd69e Mon Sep 17 00:00:00 2001
From: Jingwen Gu <75733630+JingwenGu0829@users.noreply.github.com>
Date: Wed, 22 Jul 2026 01:32:11 +0000
Subject: [PATCH 03/32] Simplify multimodal processor contract

---
 .../generate_utils/generate_endpoint_utils.py | 46 ++-----------------
 miles/utils/processing_utils.py               | 29 ++++--------
 .../rollout/generate_hub/test_sglang_omni.py  |  7 ++-
 3 files changed, 17 insertions(+), 65 deletions(-)

diff --git a/miles/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py
index 01b18b17624..b772830a7e6 100644
--- a/miles/rollout/generate_utils/generate_endpoint_utils.py
+++ b/miles/rollout/generate_utils/generate_endpoint_utils.py
@@ -19,52 +19,14 @@
 from miles.utils.types import Sample
 
 
-_MULTIMODAL_TENSOR_DTYPES = {
-    "bool",
-    "uint8",
-    "int8",
-    "int16",
-    "int32",
-    "int64",
-    "float16",
-    "bfloat16",
-    "float32",
-    "float64",
-}
-_SGLANG_OMNI_MULTIMODAL_TENSOR_NAMES = frozenset(
-    {
-        "pixel_values",
-        "image_grid_thw",
-        "input_features",
-        "feature_attention_mask",
-        "audio_feature_lengths",
-        "pixel_values_videos",
-        "video_grid_thw",
-        "video_second_per_grid",
-    }
-)
-
-
 def serialize_multimodal_train_inputs(
-    multimodal_train_inputs: dict[str, Any],
+    multimodal_train_inputs: dict[str, torch.Tensor],
 ) -> dict[str, Any]:
     """Encode the processor tensor kwargs shared with SGLang Omni."""
-    if not multimodal_train_inputs:
-        raise ValueError("multimodal_train_inputs must not be empty")
-    unknown = sorted(set(multimodal_train_inputs) - _SGLANG_OMNI_MULTIMODAL_TENSOR_NAMES)
-    if unknown:
-        raise ValueError("unsupported SGLang Omni processor tensor fields: " + ", ".join(unknown))
-
     tensors: dict[str, dict[str, Any]] = {}
-    for name, value in multimodal_train_inputs.items():
-        if not isinstance(value, torch.Tensor):
-            raise TypeError(
-                "multimodal_train_inputs must contain only tensors; " f"{name!r} has type {type(value).__name__}"
-            )
-        tensor = value.detach().to(device="cpu").contiguous()
+    for name, tensor in multimodal_train_inputs.items():
+        tensor = tensor.detach().cpu().contiguous()
         dtype = str(tensor.dtype).removeprefix("torch.")
-        if dtype not in _MULTIMODAL_TENSOR_DTYPES:
-            raise TypeError(f"unsupported multimodal tensor dtype for {name!r}: {dtype}")
         raw = tensor.reshape(-1).view(torch.uint8).numpy().tobytes()
         tensors[name] = {
             "dtype": dtype,
@@ -121,7 +83,7 @@ def compute_request_payload(
     input_ids: list[int],
     sampling_params: dict,
     multimodal_inputs: dict | None = None,
-    multimodal_train_inputs: dict[str, Any] | None = None,
+    multimodal_train_inputs: dict[str, torch.Tensor] | None = None,
 ) -> tuple[dict[str, Any] | None, Sample.Status | None]:
     sampling_params = deepcopy(sampling_params)
     max_new_tokens = sampling_params.pop("max_new_tokens", args.rollout_max_response_len)
diff --git a/miles/utils/processing_utils.py b/miles/utils/processing_utils.py
index fbcd8695ac1..9b24d99e729 100644
--- a/miles/utils/processing_utils.py
+++ b/miles/utils/processing_utils.py
@@ -5,9 +5,11 @@
 import os
 from pathlib import Path
 
+import torch
 from huggingface_hub import hf_hub_download
+from qwen_omni_utils import process_mm_info
 from tokenizers import Tokenizer as RawTokenizer
-from transformers import AutoProcessor, AutoTokenizer, PreTrainedTokenizerBase, ProcessorMixin
+from transformers import AutoProcessor, AutoTokenizer, PreTrainedTokenizerBase, ProcessorMixin, Qwen3OmniMoeProcessor
 
 from miles.utils.hf_config import register_hf_config_aliases
 
@@ -139,19 +141,11 @@ def call_processor(processor, text, multimodal_inputs: dict | None = None):
 
 def extract_multimodal_train_inputs(processor_output):
     """Normalize processor kwargs to the tensor-only Megatron contract."""
-    import torch
-
-    result = {}
-    for name, value in processor_output.items():
-        if name in ("input_ids", "attention_mask"):
-            continue
-        if not isinstance(value, torch.Tensor):
-            try:
-                value = torch.as_tensor(value)
-            except (TypeError, ValueError, RuntimeError) as exc:
-                raise TypeError(f"processor output {name!r} cannot be converted to a tensor") from exc
-        result[name] = value
-    return result or None
+    return {
+        name: torch.as_tensor(value)
+        for name, value in processor_output.items()
+        if name not in ("input_ids", "attention_mask")
+    } or None
 
 
 def load_processor(name_or_path: str, **kwargs):
@@ -170,14 +164,11 @@ def load_processor(name_or_path: str, **kwargs):
 
 def process_vision_info(prompt, processor):
     # TODO: temporary solution, will write image utils for miles later
-    if getattr(processor, "audio_token", None) is not None:
-        from qwen_omni_utils import process_mm_info
-
-        image_patch_size = getattr(processor.image_processor, "patch_size", DEFAULT_PATCH_SIZE)
+    if isinstance(processor, Qwen3OmniMoeProcessor):
         audios, images, videos = process_mm_info(
             prompt,
             use_audio_in_video=False,
-            image_patch_size=image_patch_size,
+            image_patch_size=processor.image_processor.patch_size,
         )
         return {"audio": audios, "images": images, "videos": videos}
 
diff --git a/tests/fast/rollout/generate_hub/test_sglang_omni.py b/tests/fast/rollout/generate_hub/test_sglang_omni.py
index f4cc0d25f2a..d5e0e1a2bd9 100644
--- a/tests/fast/rollout/generate_hub/test_sglang_omni.py
+++ b/tests/fast/rollout/generate_hub/test_sglang_omni.py
@@ -2,7 +2,6 @@
 import base64
 from types import SimpleNamespace
 
-import pytest
 import torch
 
 from miles.rollout.generate_utils.generate_endpoint_utils import serialize_multimodal_train_inputs
@@ -35,7 +34,6 @@ def test_serialize_audio_video_processor_tensors():
 def test_qwen_omni_media_extraction_and_tensor_normalization(monkeypatch):
     from miles.utils import processing_utils
 
-    qwen_omni_utils = pytest.importorskip("qwen_omni_utils")
     prompt = [{"role": "user", "content": [{"type": "audio", "audio": "a.wav"}]}]
     captured = {}
 
@@ -43,8 +41,9 @@ def fake_process_mm_info(conversations, **kwargs):
         captured["conversations"], captured["kwargs"] = conversations, kwargs
         return ["audio samples"], ["image"], ["video frames"]
 
-    monkeypatch.setattr(qwen_omni_utils, "process_mm_info", fake_process_mm_info)
-    processor = SimpleNamespace(audio_token="<audio>", image_processor=SimpleNamespace(patch_size=16))
+    monkeypatch.setattr(processing_utils, "process_mm_info", fake_process_mm_info)
+    processor = object.__new__(processing_utils.Qwen3OmniMoeProcessor)
+    processor.image_processor = SimpleNamespace(patch_size=16)
 
     media = processing_utils.process_vision_info(prompt, processor)
     train_inputs = processing_utils.extract_multimodal_train_inputs(

From 0fe40f693f8ca4b570e813d985dbb26c237f94cc Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Thu, 23 Jul 2026 12:38:42 +0800
Subject: [PATCH 04/32] perf(rollout): use pybase64 for multimodal tensor
 serialization

---
 miles/rollout/generate_utils/generate_endpoint_utils.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/miles/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py
index b772830a7e6..37472d553b5 100644
--- a/miles/rollout/generate_utils/generate_endpoint_utils.py
+++ b/miles/rollout/generate_utils/generate_endpoint_utils.py
@@ -2,7 +2,6 @@
 Utils to integrate SGLang's `/generate` endpoint with RL things like Sample.
 """
 
-import base64
 from copy import deepcopy
 from typing import Any
 
@@ -31,7 +30,7 @@ def serialize_multimodal_train_inputs(
         tensors[name] = {
             "dtype": dtype,
             "shape": list(tensor.shape),
-            "data": base64.b64encode(raw).decode("ascii"),
+            "data": pybase64.b64encode(raw).decode("ascii"),
         }
 
     return {"version": 1, "tensors": tensors}

From 088a44a78e45d8a2c6ef8e0cea6d4ef829751989 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Thu, 23 Jul 2026 12:39:46 +0800
Subject: [PATCH 05/32] fix(utils): lazily import qwen_omni_utils in
 process_vision_info

---
 miles/utils/processing_utils.py                     | 3 ++-
 requirements.txt                                    | 2 +-
 tests/fast/rollout/generate_hub/test_sglang_omni.py | 7 ++++++-
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/miles/utils/processing_utils.py b/miles/utils/processing_utils.py
index 9b24d99e729..d417f77e1a3 100644
--- a/miles/utils/processing_utils.py
+++ b/miles/utils/processing_utils.py
@@ -7,7 +7,6 @@
 
 import torch
 from huggingface_hub import hf_hub_download
-from qwen_omni_utils import process_mm_info
 from tokenizers import Tokenizer as RawTokenizer
 from transformers import AutoProcessor, AutoTokenizer, PreTrainedTokenizerBase, ProcessorMixin, Qwen3OmniMoeProcessor
 
@@ -165,6 +164,8 @@ def load_processor(name_or_path: str, **kwargs):
 def process_vision_info(prompt, processor):
     # TODO: temporary solution, will write image utils for miles later
     if isinstance(processor, Qwen3OmniMoeProcessor):
+        from qwen_omni_utils import process_mm_info
+
         audios, images, videos = process_mm_info(
             prompt,
             use_audio_in_video=False,
diff --git a/requirements.txt b/requirements.txt
index 523877bc4d8..1820f01b3ce 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -16,7 +16,7 @@ pybase64
 pylatexenc
 pytest-asyncio
 pyyaml
-qwen_omni_utils==0.0.9 # for Qwen Omni audio/video inputs
+qwen_omni_utils==0.0.9 # for Qwen Omni audio/video inputs; requires torchvision from the base image
 qwen_vl_utils # for VLM
 ray[default]
 ring_flash_attn; platform_system == "Linux"
diff --git a/tests/fast/rollout/generate_hub/test_sglang_omni.py b/tests/fast/rollout/generate_hub/test_sglang_omni.py
index d5e0e1a2bd9..83310d094fa 100644
--- a/tests/fast/rollout/generate_hub/test_sglang_omni.py
+++ b/tests/fast/rollout/generate_hub/test_sglang_omni.py
@@ -1,5 +1,6 @@
 import asyncio
 import base64
+import sys
 from types import SimpleNamespace
 
 import torch
@@ -41,7 +42,11 @@ def fake_process_mm_info(conversations, **kwargs):
         captured["conversations"], captured["kwargs"] = conversations, kwargs
         return ["audio samples"], ["image"], ["video frames"]
 
-    monkeypatch.setattr(processing_utils, "process_mm_info", fake_process_mm_info)
+    monkeypatch.setitem(
+        sys.modules,
+        "qwen_omni_utils",
+        SimpleNamespace(process_mm_info=fake_process_mm_info),
+    )
     processor = object.__new__(processing_utils.Qwen3OmniMoeProcessor)
     processor.image_processor = SimpleNamespace(patch_size=16)
 

From 17c4aef5a984fa05c97458d31ce59d6deadd5c00 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Thu, 23 Jul 2026 12:40:06 +0800
Subject: [PATCH 06/32] refactor(rollout): reuse
 extract_multimodal_train_inputs in sglang_rollout

---
 miles/rollout/sglang_rollout.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py
index 274ccecb1be..1517f061192 100644
--- a/miles/rollout/sglang_rollout.py
+++ b/miles/rollout/sglang_rollout.py
@@ -27,6 +27,7 @@
 from miles.utils.processing_utils import (
     call_processor,
     encode_image_for_rollout_engine,
+    extract_multimodal_train_inputs,
     load_processor,
     load_tokenizer,
 )
@@ -151,9 +152,7 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A
         processor_output = call_processor(state.processor, sample.prompt, sample.multimodal_inputs)
         prompt_ids = processor_output["input_ids"][0]
         prompt_ids = prompt_ids.tolist() if hasattr(prompt_ids, "tolist") else list(prompt_ids)
-        sample.multimodal_train_inputs = {
-            k: v for k, v in processor_output.items() if k not in ["input_ids", "attention_mask"]
-        } or None
+        sample.multimodal_train_inputs = extract_multimodal_train_inputs(processor_output)
     else:
         prompt_ids = state.tokenizer.encode(sample.prompt, add_special_tokens=False)
 

From 5833a6db346cab04c340897d746675480c181d9a Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Thu, 23 Jul 2026 12:40:56 +0800
Subject: [PATCH 07/32] test(rollout): cover sglang-omni adapter guard and
 resume branches

---
 .../rollout/generate_hub/test_sglang_omni.py  | 100 +++++++++++++++++-
 1 file changed, 97 insertions(+), 3 deletions(-)

diff --git a/tests/fast/rollout/generate_hub/test_sglang_omni.py b/tests/fast/rollout/generate_hub/test_sglang_omni.py
index 83310d094fa..a405db9900f 100644
--- a/tests/fast/rollout/generate_hub/test_sglang_omni.py
+++ b/tests/fast/rollout/generate_hub/test_sglang_omni.py
@@ -3,6 +3,7 @@
 import sys
 from types import SimpleNamespace
 
+import pytest
 import torch
 
 from miles.rollout.generate_utils.generate_endpoint_utils import serialize_multimodal_train_inputs
@@ -27,9 +28,11 @@ def test_serialize_audio_video_processor_tensors():
         encoded = bundle["tensors"][name]
         assert encoded["dtype"] == str(tensor.dtype).removeprefix("torch.")
         assert encoded["shape"] == list(tensor.shape)
-        assert base64.b64decode(encoded["data"]) == (
-            tensor.contiguous().reshape(-1).view(torch.uint8).numpy().tobytes()
-        )
+        decoded = torch.frombuffer(
+            bytearray(base64.b64decode(encoded["data"])),
+            dtype=getattr(torch, encoded["dtype"]),
+        ).reshape(encoded["shape"])
+        assert torch.equal(decoded, tensor)
 
 
 def test_qwen_omni_media_extraction_and_tensor_normalization(monkeypatch):
@@ -124,3 +127,94 @@ async def fake_update(*args, **kwargs):
         "pixel_values_videos",
     }
     assert captured["headers"] is None
+
+
+def _adapter_args() -> SimpleNamespace:
+    return SimpleNamespace(
+        sglang_router_ip="127.0.0.1",
+        sglang_router_port=30000,
+        sglang_router_policy="round_robin",
+        rollout_max_response_len=128,
+        rollout_max_context_len=0,
+        use_rollout_routing_replay=False,
+        use_rollout_indexer_replay=False,
+    )
+
+
+def test_sglang_omni_adapter_rejects_multimodal_without_train_inputs(monkeypatch):
+    from miles.rollout.generate_hub import sglang_omni
+
+    monkeypatch.setattr(
+        sglang_omni,
+        "compute_prompt_ids_from_sample",
+        lambda state, sample: [1, 2, 3],
+    )
+    generate_input = SimpleNamespace(
+        args=_adapter_args(),
+        sample=Sample(multimodal_inputs={"audio": ["audio"]}),
+        sampling_params={"temperature": 1.0, "max_new_tokens": 16},
+        state=None,
+    )
+
+    with pytest.raises(ValueError, match="requires processor-produced"):
+        asyncio.run(sglang_omni.generate(generate_input))
+
+
+def test_sglang_omni_adapter_truncates_exhausted_resume(monkeypatch):
+    from miles.rollout.generate_hub import sglang_omni
+
+    async def fail_post(url, payload, headers=None):
+        raise AssertionError("post must not run for an exhausted resume")
+
+    monkeypatch.setattr(sglang_omni, "post", fail_post)
+    monkeypatch.setattr(
+        sglang_omni,
+        "compute_prompt_ids_from_sample",
+        lambda state, sample: [1, 2, 3],
+    )
+    sample = Sample(response="hi", tokens=[1, 2, 3, 4, 5])
+    generate_input = SimpleNamespace(
+        args=_adapter_args(),
+        sample=sample,
+        sampling_params={"temperature": 1.0, "max_new_tokens": 2},
+        state=None,
+    )
+
+    output = asyncio.run(sglang_omni.generate(generate_input))
+
+    assert sample.status == Sample.Status.TRUNCATED
+    assert output.samples is sample
+
+
+def test_sglang_omni_adapter_resumes_partial_rollout(monkeypatch):
+    from miles.rollout.generate_hub import sglang_omni
+
+    captured = {}
+
+    async def fake_post(url, payload, headers=None):
+        captured["payload"] = payload
+        return {}
+
+    async def fake_update(*args, **kwargs):
+        return None
+
+    monkeypatch.setattr(sglang_omni, "post", fake_post)
+    monkeypatch.setattr(sglang_omni, "update_sample_from_response", fake_update)
+    monkeypatch.setattr(
+        sglang_omni,
+        "compute_prompt_ids_from_sample",
+        lambda state, sample: [1, 2, 3],
+    )
+    sample = Sample(response="hi", tokens=[1, 2, 3, 4, 5])
+    sampling_params = {"temperature": 1.0, "max_new_tokens": 16}
+    generate_input = SimpleNamespace(
+        args=_adapter_args(),
+        sample=sample,
+        sampling_params=sampling_params,
+        state=None,
+    )
+
+    asyncio.run(sglang_omni.generate(generate_input))
+
+    assert captured["payload"]["input_ids"] == [1, 2, 3, 4, 5]
+    assert sampling_params["max_new_tokens"] == 14

From 83872b3262d89edf5a60b9c01a7a02c4215b8396 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Mon, 15 Jun 2026 00:59:40 +0800
Subject: [PATCH 08/32] [omni] Qwen3-Omni thinker RL bridge + thin rollout glue
 (first omni<->miles loop)

First omni<->miles RL loop on the Qwen3-Omni-30B-A3B thinker (text MoE).
---
 .../run_qwen3_omni_thinker_grpo.py            | 171 +++++++++++++++
 .../megatron_utils/megatron_to_hf/__init__.py |   3 +
 .../megatron_to_hf/qwen3omni_moe.py           |  14 ++
 miles/rollout/generate_hub/omni_thinker.py    |  57 +++++
 scripts/models/qwen3-omni-30B-A3B-thinker.sh  |  52 +++++
 tests/fast/test_qwen3_omni_thinker.py         | 205 ++++++++++++++++++
 tools/extract_qwen3_omni_thinker.py           | 120 ++++++++++
 7 files changed, 622 insertions(+)
 create mode 100644 examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
 create mode 100644 miles/backends/megatron_utils/megatron_to_hf/qwen3omni_moe.py
 create mode 100644 miles/rollout/generate_hub/omni_thinker.py
 create mode 100644 scripts/models/qwen3-omni-30B-A3B-thinker.sh
 create mode 100644 tests/fast/test_qwen3_omni_thinker.py
 create mode 100644 tools/extract_qwen3_omni_thinker.py

diff --git a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
new file mode 100644
index 00000000000..9023bc12d3f
--- /dev/null
+++ b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
@@ -0,0 +1,171 @@
+"""GRPO on the Qwen3-Omni-30B-A3B thinker (text MoE) with DAPO-math reward.
+
+Off-policy until live weight-sync lands: point --sglang-router-ip/port at a standalone omni
+server; the rollout half serves the frozen base while the trainer updates its own copy, with
+TIS + --get-mismatch-metrics absorbing/measuring the gap.
+"""
+
+import os
+from dataclasses import dataclass
+from typing import Literal
+
+import typer
+
+import miles.utils.external_utils.command_utils as U
+
+OMNI_MODEL = "Qwen3-Omni-30B-A3B-Instruct"
+THINKER_MODEL = "Qwen3-Omni-30B-A3B-Thinker"
+MEGATRON_MODEL_TYPE = "qwen3-omni-30B-A3B-thinker"
+
+
+@dataclass
+class ScriptArgs(U.ExecuteTrainConfig):
+    mode: Literal["normal", "debug_minimal"] = "normal"
+    run_id: str = U.create_run_id()
+    num_gpus_per_node: int = 8
+    data_dir: str = "/root/datasets"
+    model_dir: str = "/root/models"
+    megatron_path: str = "/root/Megatron-LM"
+    omni_router_ip: str = "127.0.0.1"
+    omni_router_port: int = 30000
+    extra_args: str = ""
+
+
+def prepare(args: ScriptArgs):
+    U.exec_command(f"mkdir -p {args.model_dir} {args.data_dir}")
+    U.exec_command(f"hf download Qwen/{OMNI_MODEL} --local-dir {args.model_dir}/{OMNI_MODEL}")
+    repo = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    U.exec_command(
+        f"python {repo}/tools/extract_qwen3_omni_thinker.py "
+        f"--src {args.model_dir}/{OMNI_MODEL} --dst {args.model_dir}/{THINKER_MODEL}"
+    )
+    U.convert_checkpoint(
+        model_name=THINKER_MODEL,
+        megatron_model_type=MEGATRON_MODEL_TYPE,
+        num_gpus_per_node=args.num_gpus_per_node,
+        dir_dst=args.model_dir,
+        hf_checkpoint=f"{args.model_dir}/{THINKER_MODEL}",
+        megatron_path=args.megatron_path,
+    )
+    U.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=args.data_dir)
+
+
+def execute(args: ScriptArgs):
+    ref_load_path = f"{args.model_dir}/{THINKER_MODEL}_torch_dist"
+    load_save_path = f"{args.output_dir}/{args.run_id}/checkpoints"
+
+    ckpt_args = (
+        f"--hf-checkpoint {args.model_dir}/{THINKER_MODEL}/ "
+        f"--ref-load {ref_load_path} "
+        f"--load {load_save_path} "
+        "--model-name qwen3omni_moe "  # body.* broadcast naming
+    )
+
+    rollout_args = (
+        "--custom-generate-function-path miles.rollout.generate_hub.omni_thinker.generate "
+        f"--prompt-data {args.data_dir}/dapo-math-17k/dapo-math-17k.jsonl "
+        "--input-key prompt "
+        "--label-key label "
+        "--apply-chat-template "
+        "--rollout-shuffle "
+        "--rm-type dapo "
+        "--reward-key score "
+        "--num-rollout 3000 "
+        "--rollout-batch-size 32 "
+        "--n-samples-per-prompt 8 "
+        f"--rollout-max-response-len {100 if args.mode == 'debug_minimal' else 8192} "
+        "--rollout-temperature 1 "
+        "--global-batch-size 256 "
+        "--balance-data "
+        f"--sglang-router-ip {args.omni_router_ip} "
+        f"--sglang-router-port {args.omni_router_port} "
+    )
+
+    consistency_args = (
+        "--use-rollout-logprobs "
+        "--get-mismatch-metrics "
+        "--use-tis "
+        "--tis-clip 2.0 "
+    )
+
+    perf_args = (
+        "--tensor-model-parallel-size 8 "
+        "--sequence-parallel "
+        "--pipeline-model-parallel-size 1 "
+        "--context-parallel-size 1 "
+        "--expert-model-parallel-size 8 "
+        "--expert-tensor-parallel-size 1 "
+        "--recompute-granularity full "
+        "--recompute-method uniform "
+        "--recompute-num-layers 1 "
+        "--use-dynamic-batch-size "
+        "--max-tokens-per-gpu 9216 "
+        "--optimizer-cpu-offload "
+        "--overlap-cpu-optimizer-d2h-h2d "
+        "--use-precision-aware-optimizer "
+    )
+
+    grpo_args = (
+        "--advantage-estimator grpo "
+        "--use-kl-loss "
+        "--kl-loss-coef 0.00 "
+        "--kl-loss-type low_var_kl "
+        "--entropy-coef 0.00 "
+        "--eps-clip 0.2 "
+        "--eps-clip-high 0.28 "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+    )
+
+    misc_args = (
+        "--attention-dropout 0.0 "
+        "--hidden-dropout 0.0 "
+        "--accumulate-allreduce-grads-in-fp32 "
+        "--attention-softmax-in-fp32 "
+        "--attention-backend flash "
+        "--actor-num-nodes 1 "
+        f"--actor-num-gpus-per-node {args.num_gpus_per_node} "
+        f"--num-gpus-per-node {args.num_gpus_per_node} "
+        f"--rollout-num-gpus {args.num_gpus_per_node} "
+    )
+
+    train_args = (
+        f"{ckpt_args} "
+        f"{rollout_args} "
+        f"{consistency_args} "
+        f"{optimizer_args} "
+        f"{grpo_args} "
+        f"{U.get_default_wandb_args(__file__, run_id=args.run_id)} "
+        f"{perf_args} "
+        f"{misc_args} "
+        f"{args.extra_args} "
+    )
+
+    U.execute_train(
+        train_args=train_args,
+        num_gpus_per_node=args.num_gpus_per_node,
+        megatron_model_type=MEGATRON_MODEL_TYPE,
+        train_script="train.py",
+        megatron_path=args.megatron_path,
+        extra_env_vars={
+            "FLASHINFER_DISABLE_VERSION_CHECK": "1",
+            "PYTHONPATH": f"{args.megatron_path}",
+        },
+    )
+
+
+@U.dataclass_cli
+def main(args: ScriptArgs):
+    prepare(args)
+    execute(args)
+
+
+if __name__ == "__main__":
+    typer.run(main)
diff --git a/miles/backends/megatron_utils/megatron_to_hf/__init__.py b/miles/backends/megatron_utils/megatron_to_hf/__init__.py
index 0873e6f8e32..ba5d67d76d6 100644
--- a/miles/backends/megatron_utils/megatron_to_hf/__init__.py
+++ b/miles/backends/megatron_utils/megatron_to_hf/__init__.py
@@ -10,6 +10,7 @@
 from .qwen3_5 import convert_qwen3_5_to_hf
 from .qwen3_next import convert_qwen3_next_to_hf
 from .qwen3moe import convert_qwen3moe_to_hf
+from .qwen3omni_moe import convert_qwen3omni_moe_to_hf
 
 
 # TODO unify w/ `convert_to_hf`
@@ -42,6 +43,8 @@ def _convert_to_hf_core(args, model_name, name, param):
         converted_named_tensors = convert_glm4moe_to_hf(args, name, param)
     elif "glm4" in model_name:
         converted_named_tensors = convert_glm4_to_hf(args, name, param)
+    elif "qwen3omni" in model_name:  # body.* names; keep before generic qwen3
+        converted_named_tensors = convert_qwen3omni_moe_to_hf(args, name, param)
     elif "qwen3moe" in model_name:
         converted_named_tensors = convert_qwen3moe_to_hf(args, name, param)
     elif "qwen3next" in model_name:
diff --git a/miles/backends/megatron_utils/megatron_to_hf/qwen3omni_moe.py b/miles/backends/megatron_utils/megatron_to_hf/qwen3omni_moe.py
new file mode 100644
index 00000000000..b33b52b8a18
--- /dev/null
+++ b/miles/backends/megatron_utils/megatron_to_hf/qwen3omni_moe.py
@@ -0,0 +1,14 @@
+"""Megatron->HF broadcast converter for the Qwen3-Omni thinker (Qwen3-MoE text).
+
+Same conversion as qwen3moe, but prefixes every HF name with `body.` (the namespace
+the sglang-omni weight-receive side demuxes on). Selected via `--model-name
+qwen3omni_moe`. Nothing is dropped: the thinker is untied; `tied.*` is talker-side.
+"""
+
+from .qwen3moe import convert_qwen3moe_to_hf
+
+BODY_PREFIX = "body."
+
+
+def convert_qwen3omni_moe_to_hf(args, name, param):
+    return [(f"{BODY_PREFIX}{n}", t) for n, t in convert_qwen3moe_to_hf(args, name, param)]
diff --git a/miles/rollout/generate_hub/omni_thinker.py b/miles/rollout/generate_hub/omni_thinker.py
new file mode 100644
index 00000000000..5ef8325260d
--- /dev/null
+++ b/miles/rollout/generate_hub/omni_thinker.py
@@ -0,0 +1,57 @@
+"""Rollout adapter for the sglang-omni `/generate` endpoint (thinker / text).
+
+Thin wrapper over the default single-turn generate: reuses miles' payload builder and
+response parser, adds only the omni-specific fields. Wire via
+`--custom-generate-function-path miles.rollout.generate_hub.omni_thinker.generate`.
+"""
+
+from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput
+from miles.rollout.generate_utils.generate_endpoint_utils import (
+    compute_prompt_ids_from_sample,
+    compute_request_payload,
+    update_sample_from_response,
+)
+from miles.utils.http_utils import post
+from miles.utils.types import Sample
+
+
+async def generate(input: GenerateFnInput) -> GenerateFnOutput:
+    args = input.args
+    sample = input.sample
+    sampling_params = input.sampling_params
+    assert sample.status in {Sample.Status.PENDING, Sample.Status.ABORTED}, f"{sample.status=}"
+    # omni thinker path is text-only with no MoE/indexer replay yet (server forward-declares both)
+    assert not (args.use_rollout_routing_replay or args.use_rollout_indexer_replay), (
+        "omni rollout has no routing/indexer replay; unset --use-rollout-routing-replay / --use-rollout-indexer-replay"
+    )
+    assert not (sample.multimodal_inputs or {}).get("images"), "omni thinker rollout is text-only"
+    url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate"
+
+    prompt_ids = compute_prompt_ids_from_sample(input.state, sample)
+    if len(sample.response) > 0:  # partial rollout resume
+        input_ids = sample.tokens
+        sampling_params["max_new_tokens"] -= len(sample.tokens) - len(prompt_ids)
+        assert sampling_params["max_new_tokens"] >= 0
+        if sampling_params["max_new_tokens"] == 0:
+            sample.status = Sample.Status.TRUNCATED
+            return GenerateFnOutput(samples=sample)
+    else:
+        input_ids = prompt_ids
+
+    payload, halt_status = compute_request_payload(
+        args, input_ids=input_ids, sampling_params=sampling_params, multimodal_inputs=sample.multimodal_inputs
+    )
+    if payload is None:
+        sample.status = halt_status
+        return GenerateFnOutput(samples=sample)
+
+    payload["output_modalities"] = ["text"]
+    payload["return_omni_rollout"] = False
+    # rep_penalty=1: only sampler warper the trainer recompute can't replay (temp it divides out symmetrically)
+    payload["sampling_params"]["repetition_penalty"] = 1.0
+    if sample.metadata:
+        payload["metadata"] = sample.metadata
+
+    output = await post(url, payload)
+    await update_sample_from_response(args, sample, payload=payload, output=output)
+    return GenerateFnOutput(samples=sample)
diff --git a/scripts/models/qwen3-omni-30B-A3B-thinker.sh b/scripts/models/qwen3-omni-30B-A3B-thinker.sh
new file mode 100644
index 00000000000..a071a2e6e32
--- /dev/null
+++ b/scripts/models/qwen3-omni-30B-A3B-thinker.sh
@@ -0,0 +1,52 @@
+# Megatron model args for the Qwen3-Omni-30B-A3B thinker (text MoE).
+# Architecturally Qwen3-30B-A3B; only vocab differs (152064 vs 151936, from the real config.json).
+
+NLAYERS="${MODEL_ARGS_NUM_LAYERS:-48}"
+FIRST_K_DENSE_REPLACE=0
+
+arr=()
+for ((i=0; i<NLAYERS; i++)); do
+  if (( i < FIRST_K_DENSE_REPLACE )); then
+    arr+=(0)
+  else
+    arr+=(1)
+  fi
+done
+
+printf -v MOE_LAYER_FREQ "[%s]" "$(IFS=', '; echo "${arr[*]}")"
+
+
+MODEL_ARGS=(
+   --disable-bias-linear
+   --qk-layernorm
+   --group-query-attention
+   --num-attention-heads 32
+   --num-query-groups 4
+   --kv-channels 128
+   --num-layers $NLAYERS
+   --hidden-size 2048
+   --ffn-hidden-size 6144
+
+   --normalization RMSNorm
+   --position-embedding-type rope
+   --norm-epsilon 1e-6
+   --rotary-percent 1.0
+   --swiglu
+   --untie-embeddings-and-output-weights
+   --vocab-size "${MODEL_ARGS_VOCAB_SIZE:-152064}"
+
+   --rotary-base "${MODEL_ARGS_ROTARY_BASE:-1000000}"
+
+   # moe
+   --moe-ffn-hidden-size 768
+   --moe-router-score-function softmax
+   --moe-token-dispatcher-type alltoall
+   --moe-router-topk 8
+   --moe-layer-freq "$MOE_LAYER_FREQ"
+   --num-experts 128
+   --moe-grouped-gemm
+   --moe-token-drop-policy probs
+   --moe-router-dtype fp32
+   --moe-permute-fusion
+   --moe-aux-loss-coeff 0
+)
diff --git a/tests/fast/test_qwen3_omni_thinker.py b/tests/fast/test_qwen3_omni_thinker.py
new file mode 100644
index 00000000000..d8af6364436
--- /dev/null
+++ b/tests/fast/test_qwen3_omni_thinker.py
@@ -0,0 +1,205 @@
+"""Qwen3-Omni thinker bridge: extraction name/config logic + the body.* converter contract."""
+
+import asyncio
+import importlib.util
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+_REPO = Path(__file__).resolve().parents[2]
+
+
+def _load_extract_tool():
+    path = _REPO / "tools" / "extract_qwen3_omni_thinker.py"
+    spec = importlib.util.spec_from_file_location("extract_qwen3_omni_thinker", path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def test_map_thinker_param_name_keeps_text_backbone():
+    m = _load_extract_tool().map_thinker_param_name
+    assert m("thinker.model.embed_tokens.weight") == "model.embed_tokens.weight"
+    assert m("thinker.lm_head.weight") == "lm_head.weight"
+    assert m("thinker.model.norm.weight") == "model.norm.weight"
+    assert m("thinker.model.layers.7.mlp.experts.3.down_proj.weight") == "model.layers.7.mlp.experts.3.down_proj.weight"
+    assert m("thinker.model.layers.0.self_attn.q_norm.weight") == "model.layers.0.self_attn.q_norm.weight"
+
+
+def test_map_thinker_param_name_normalizes_language_model_infix():
+    m = _load_extract_tool().map_thinker_param_name
+    assert m("thinker.model.language_model.layers.0.mlp.gate.weight") == "model.layers.0.mlp.gate.weight"
+
+
+def test_map_thinker_param_name_drops_non_thinker_and_non_text():
+    m = _load_extract_tool().map_thinker_param_name
+    assert m("talker.model.layers.0.self_attn.q_proj.weight") is None
+    assert m("code2wav.decoder.weight") is None
+    assert m("thinker.audio_tower.layers.0.fc1.weight") is None
+    assert m("thinker.visual.blocks.0.attn.qkv.weight") is None
+
+
+def test_synthesize_thinker_config_stamps_qwen3_moe():
+    synth = _load_extract_tool().synthesize_thinker_config
+    omni = {
+        "model_type": "qwen3_omni_moe",
+        "thinker_config": {
+            "eos_token_id": 151645,
+            "text_config": {"vocab_size": 152064, "num_experts": 128, "tie_word_embeddings": False},
+        },
+    }
+    cfg = synth(omni)
+    assert cfg["model_type"] == "qwen3_moe"
+    assert cfg["architectures"] == ["Qwen3MoeForCausalLM"]
+    assert cfg["num_experts"] == 128
+    assert cfg["vocab_size"] == 152064
+    assert cfg["tie_word_embeddings"] is False
+    assert cfg["eos_token_id"] == 151645
+
+
+def test_synthesize_thinker_config_flat_fallback():
+    cfg = _load_extract_tool().synthesize_thinker_config({"vocab_size": 100, "num_experts": 8})
+    assert cfg["model_type"] == "qwen3_moe"
+    assert cfg["vocab_size"] == 100
+    assert cfg["tie_word_embeddings"] is False
+
+
+def _args():
+    # head_dim=kv_channels=2, value_num_per_group=heads//groups=2
+    return SimpleNamespace(hidden_size=8, num_attention_heads=4, num_query_groups=2, kv_channels=2)
+
+
+def _conv():
+    return pytest.importorskip(
+        "miles.backends.megatron_utils.megatron_to_hf.qwen3omni_moe"
+    ).convert_qwen3omni_moe_to_hf
+
+
+def test_body_converter_simple_params():
+    torch = pytest.importorskip("torch")
+    conv, args = _conv(), _args()
+    names = lambda n, p: [x for x, _ in conv(args, n, p)]
+    assert names("module.module.embedding.word_embeddings.weight", torch.zeros(10, 8)) == ["body.model.embed_tokens.weight"]
+    assert names("module.module.output_layer.weight", torch.zeros(10, 8)) == ["body.lm_head.weight"]
+    assert names("module.module.decoder.final_layernorm.weight", torch.zeros(8)) == ["body.model.norm.weight"]
+    assert names("module.module.decoder.layers.3.mlp.router.weight", torch.zeros(128, 8)) == ["body.model.layers.3.mlp.gate.weight"]
+    assert names("module.module.decoder.layers.3.self_attention.q_layernorm.weight", torch.zeros(2)) == ["body.model.layers.3.self_attn.q_norm.weight"]
+    assert names("module.module.decoder.layers.3.self_attention.k_layernorm.weight", torch.zeros(2)) == ["body.model.layers.3.self_attn.k_norm.weight"]
+
+
+def test_body_converter_fused_splits():
+    torch = pytest.importorskip("torch")
+    conv, args = _conv(), _args()
+
+    # qkv: (groups*(vpg+2)*head_dim, hidden) = (16, 8) -> q(8,8) k(4,8) v(4,8)
+    out = conv(args, "module.module.decoder.layers.1.self_attention.linear_qkv.weight", torch.zeros(16, 8))
+    assert {n: tuple(t.shape) for n, t in out} == {
+        "body.model.layers.1.self_attn.q_proj.weight": (8, 8),
+        "body.model.layers.1.self_attn.k_proj.weight": (4, 8),
+        "body.model.layers.1.self_attn.v_proj.weight": (4, 8),
+    }
+    # expert fc1: gate+up fused on dim 0 -> chunk(2)
+    out = conv(args, "module.module.decoder.layers.1.mlp.experts.linear_fc1.weight5", torch.zeros(8, 8))
+    assert {n: tuple(t.shape) for n, t in out} == {
+        "body.model.layers.1.mlp.experts.5.gate_proj.weight": (4, 8),
+        "body.model.layers.1.mlp.experts.5.up_proj.weight": (4, 8),
+    }
+    out = conv(args, "module.module.decoder.layers.1.mlp.experts.linear_fc2.weight5", torch.zeros(8, 4))
+    assert [n for n, _ in out] == ["body.model.layers.1.mlp.experts.5.down_proj.weight"]
+
+
+def test_body_converter_dispatch_routes_qwen3omni():
+    torch = pytest.importorskip("torch")
+    mod = pytest.importorskip("miles.backends.megatron_utils.megatron_to_hf")
+    out = mod._convert_to_hf_core(
+        _args(), "qwen3omni_moe", "module.module.decoder.layers.0.mlp.router.weight", torch.zeros(128, 8)
+    )
+    assert [n for n, _ in out] == ["body.model.layers.0.mlp.gate.weight"]
+
+
+# --- rollout adapter (miles.rollout.generate_hub.omni_thinker.generate) ---
+
+
+def _omni_mod():
+    return pytest.importorskip("miles.rollout.generate_hub.omni_thinker")
+
+
+def _omni_input(monkeypatch, omni, *, replay=False, indexer_replay=False, multimodal=None, metadata=None):
+    """A generate() input with post/prompt-id mocked; returns (input, sample, captured)."""
+    Sample = pytest.importorskip("miles.utils.types").Sample
+    captured = {}
+
+    async def fake_post(url, payload):
+        captured["url"], captured["payload"] = url, payload
+        return {
+            "text": "hello",
+            "meta_info": {
+                "finish_reason": {"type": "stop"},
+                "output_token_logprobs": [[-0.1, 11], [-0.2, 22]],
+                "prompt_tokens": 3,
+                "completion_tokens": 2,
+                "cached_tokens": 0,
+            },
+        }
+
+    monkeypatch.setattr(omni, "post", fake_post)
+    monkeypatch.setattr(omni, "compute_prompt_ids_from_sample", lambda state, sample: [1, 2, 3])
+
+    args = SimpleNamespace(
+        sglang_router_ip="127.0.0.1",
+        sglang_router_port=30000,
+        use_rollout_routing_replay=replay,
+        use_rollout_indexer_replay=indexer_replay,
+        rollout_max_response_len=128,
+        rollout_max_context_len=0,
+        sglang_speculative_algorithm=None,
+    )
+    sample = Sample(status=Sample.Status.PENDING, metadata=metadata or {}, multimodal_inputs=multimodal)
+    inp = SimpleNamespace(
+        args=args, sample=sample, sampling_params={"temperature": 1.0, "max_new_tokens": 64}, state=None
+    )
+    return inp, sample, captured
+
+
+def test_generate_builds_text_only_omni_payload_and_parses_response(monkeypatch):
+    omni = _omni_mod()
+    inp, sample, captured = _omni_input(monkeypatch, omni, metadata={"task": "math"})
+
+    asyncio.run(omni.generate(inp))
+
+    p = captured["payload"]
+    assert captured["url"].endswith("/generate")
+    assert p["input_ids"] == [1, 2, 3]
+    assert p["output_modalities"] == ["text"]
+    assert p["return_omni_rollout"] is False
+    assert p["return_routed_experts"] is False
+    assert p["return_indexer_topk"] is False
+    assert p["sampling_params"]["repetition_penalty"] == 1.0
+    assert p["metadata"] == {"task": "math"}
+    # response parse: tokens = input_ids + decoded ids, logprobs aligned, status from finish_reason
+    assert sample.tokens == [1, 2, 3, 11, 22]
+    assert sample.rollout_log_probs == [-0.1, -0.2]
+    assert sample.response == "hello"
+    assert sample.status.name == "COMPLETED"
+
+
+def test_generate_rejects_moe_replay_for_omni(monkeypatch):
+    omni = _omni_mod()
+    inp, _, _ = _omni_input(monkeypatch, omni, replay=True)
+    with pytest.raises(AssertionError, match="replay"):
+        asyncio.run(omni.generate(inp))
+
+
+def test_generate_rejects_indexer_replay_for_omni(monkeypatch):
+    omni = _omni_mod()
+    inp, _, _ = _omni_input(monkeypatch, omni, indexer_replay=True)
+    with pytest.raises(AssertionError, match="replay"):
+        asyncio.run(omni.generate(inp))
+
+
+def test_generate_rejects_multimodal_on_text_only_path(monkeypatch):
+    omni = _omni_mod()
+    inp, _, _ = _omni_input(monkeypatch, omni, multimodal={"images": ["img"]})
+    with pytest.raises(AssertionError, match="text-only"):
+        asyncio.run(omni.generate(inp))
diff --git a/tools/extract_qwen3_omni_thinker.py b/tools/extract_qwen3_omni_thinker.py
new file mode 100644
index 00000000000..71c1d91a49a
--- /dev/null
+++ b/tools/extract_qwen3_omni_thinker.py
@@ -0,0 +1,120 @@
+"""Extract a standalone Qwen3-MoE thinker checkpoint from a full Qwen3-Omni model.
+
+miles loads HF via AutoBridge keyed on model_type; the composite omni checkpoint has
+no bridge, but the thinker text backbone is a plain Qwen3-MoE. This writes a
+self-contained HF dir (thinker text + lm_head, renamed, model_type=qwen3_moe) for the
+existing Qwen3MoEBridge.
+
+    python tools/extract_qwen3_omni_thinker.py --src <omni> --dst <thinker>
+"""
+
+from __future__ import annotations
+
+_NON_TEXT_THINKER_SUBMODULES = ("audio_tower.", "visual.", "model.audio_tower.", "model.visual.")
+
+
+def map_thinker_param_name(name: str) -> str | None:
+    """Full omni param name -> standalone Qwen3-MoE name, or None to drop."""
+    if not name.startswith("thinker."):
+        return None
+    rest = name[len("thinker.") :]
+    if rest.startswith(_NON_TEXT_THINKER_SUBMODULES):
+        return None
+    return rest.replace("model.language_model.", "model.")
+
+
+def synthesize_thinker_config(omni_config: dict) -> dict:
+    """Plain Qwen3-MoE config dict from the full omni config."""
+    thinker = omni_config.get("thinker_config", omni_config)
+    cfg = dict(thinker.get("text_config", thinker))
+    cfg["model_type"] = "qwen3_moe"
+    cfg["architectures"] = ["Qwen3MoeForCausalLM"]
+    cfg.setdefault("tie_word_embeddings", False)
+    for k in ("bos_token_id", "eos_token_id", "pad_token_id"):
+        if k not in cfg and k in thinker:
+            cfg[k] = thinker[k]
+    return cfg
+
+
+def main() -> None:
+    import argparse
+    import json
+    import shutil
+    from pathlib import Path
+
+    from safetensors import safe_open
+    from safetensors.torch import save_file
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--src", required=True)
+    parser.add_argument("--dst", required=True)
+    parser.add_argument("--shard-size-gb", type=float, default=5.0)
+    args = parser.parse_args()
+
+    src, dst = Path(args.src), Path(args.dst)
+    dst.mkdir(parents=True, exist_ok=True)
+
+    with open(src / "config.json") as f:
+        omni_config = json.load(f)
+    with open(dst / "config.json", "w") as f:
+        json.dump(synthesize_thinker_config(omni_config), f, indent=2)
+
+    for fname in (
+        "tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt",
+        "special_tokens_map.json", "generation_config.json", "chat_template.json", "chat_template.jinja",
+    ):
+        if (src / fname).exists():
+            shutil.copy2(src / fname, dst / fname)
+
+    index_path = src / "model.safetensors.index.json"
+    if index_path.exists():
+        with open(index_path) as f:
+            shard_files = sorted(set(json.load(f)["weight_map"].values()))
+    else:
+        shard_files = ["model.safetensors"]
+
+    shard_size_bytes = int(args.shard_size_gb * (1024 ** 3))
+    out_index: dict[str, str] = {}
+    out_shards: list[tuple[str, dict]] = []
+    buf: dict = {}
+    buf_bytes = 0
+    total_kept = 0
+
+    def flush():
+        nonlocal buf, buf_bytes
+        if not buf:
+            return
+        shard_name = f"model-{len(out_shards) + 1:05d}.safetensors"
+        out_shards.append((shard_name, buf))
+        for k in buf:
+            out_index[k] = shard_name
+        buf, buf_bytes = {}, 0
+
+    for shard_file in shard_files:
+        with safe_open(src / shard_file, framework="pt") as reader:
+            for key in reader.keys():
+                new_key = map_thinker_param_name(key)
+                if new_key is None:
+                    continue
+                tensor = reader.get_tensor(key)
+                buf[new_key] = tensor
+                buf_bytes += tensor.numel() * tensor.element_size()
+                total_kept += 1
+                if buf_bytes >= shard_size_bytes:
+                    flush()
+    flush()
+
+    if len(out_shards) == 1:
+        save_file(out_shards[0][1], dst / "model.safetensors", metadata={"format": "pt"})
+    else:
+        total_size = sum(t.numel() * t.element_size() for _, ts in out_shards for t in ts.values())
+        for shard_name, tensors in out_shards:
+            save_file(tensors, dst / shard_name, metadata={"format": "pt"})
+        with open(dst / "model.safetensors.index.json", "w") as f:
+            json.dump({"metadata": {"total_size": total_size}, "weight_map": out_index}, f, indent=2)
+
+    print(f"[done] {total_kept} thinker tensors -> {dst}")
+
+
+if __name__ == "__main__":
+    main()

From 7e10fe62b791326272700a9f69c7b29eb9d253eb Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Wed, 17 Jun 2026 22:32:11 +0800
Subject: [PATCH 09/32] [omni] enforce rollout_temperature==1 (omni emits
 temp-1 logprobs)

---
 miles/rollout/generate_hub/omni_thinker.py |  9 +++++++--
 tests/fast/test_qwen3_omni_thinker.py      | 11 ++++++++++-
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/miles/rollout/generate_hub/omni_thinker.py b/miles/rollout/generate_hub/omni_thinker.py
index 5ef8325260d..f49ec5ea5d0 100644
--- a/miles/rollout/generate_hub/omni_thinker.py
+++ b/miles/rollout/generate_hub/omni_thinker.py
@@ -20,7 +20,12 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
     sample = input.sample
     sampling_params = input.sampling_params
     assert sample.status in {Sample.Status.PENDING, Sample.Status.ABORTED}, f"{sample.status=}"
-    # omni thinker path is text-only with no MoE/indexer replay yet (server forward-declares both)
+    # omni /generate emits temp-1 (pre-temperature) full-vocab logprobs; the trainer recompute
+    # divides logits by rollout_temperature, so they agree only at temp=1.
+    assert args.rollout_temperature == 1.0, (
+        f"omni rollout logprob is temp-1; rollout_temperature must be 1.0, got {args.rollout_temperature}"
+    )
+    # text-only path with no MoE/indexer replay yet (server forward-declares both)
     assert not (args.use_rollout_routing_replay or args.use_rollout_indexer_replay), (
         "omni rollout has no routing/indexer replay; unset --use-rollout-routing-replay / --use-rollout-indexer-replay"
     )
@@ -47,7 +52,7 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
 
     payload["output_modalities"] = ["text"]
     payload["return_omni_rollout"] = False
-    # rep_penalty=1: only sampler warper the trainer recompute can't replay (temp it divides out symmetrically)
+    # rep_penalty=1: the trainer recompute can't replay a repetition penalty (logprobs would diverge)
     payload["sampling_params"]["repetition_penalty"] = 1.0
     if sample.metadata:
         payload["metadata"] = sample.metadata
diff --git a/tests/fast/test_qwen3_omni_thinker.py b/tests/fast/test_qwen3_omni_thinker.py
index d8af6364436..b9773e3565b 100644
--- a/tests/fast/test_qwen3_omni_thinker.py
+++ b/tests/fast/test_qwen3_omni_thinker.py
@@ -125,7 +125,7 @@ def _omni_mod():
     return pytest.importorskip("miles.rollout.generate_hub.omni_thinker")
 
 
-def _omni_input(monkeypatch, omni, *, replay=False, indexer_replay=False, multimodal=None, metadata=None):
+def _omni_input(monkeypatch, omni, *, replay=False, indexer_replay=False, multimodal=None, metadata=None, temperature=1.0):
     """A generate() input with post/prompt-id mocked; returns (input, sample, captured)."""
     Sample = pytest.importorskip("miles.utils.types").Sample
     captured = {}
@@ -149,6 +149,7 @@ async def fake_post(url, payload):
     args = SimpleNamespace(
         sglang_router_ip="127.0.0.1",
         sglang_router_port=30000,
+        rollout_temperature=temperature,
         use_rollout_routing_replay=replay,
         use_rollout_indexer_replay=indexer_replay,
         rollout_max_response_len=128,
@@ -203,3 +204,11 @@ def test_generate_rejects_multimodal_on_text_only_path(monkeypatch):
     inp, _, _ = _omni_input(monkeypatch, omni, multimodal={"images": ["img"]})
     with pytest.raises(AssertionError, match="text-only"):
         asyncio.run(omni.generate(inp))
+
+
+def test_generate_rejects_nonunit_temperature(monkeypatch):
+    # omni emits temp-1 logprobs; the trainer recompute divides by rollout_temperature, so temp!=1 desyncs
+    omni = _omni_mod()
+    inp, _, _ = _omni_input(monkeypatch, omni, temperature=2.0)
+    with pytest.raises(AssertionError, match="temp-1"):
+        asyncio.run(omni.generate(inp))

From 37ba2906809e60567ec5af69146f138e1e0980b8 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Fri, 24 Jul 2026 01:58:14 +0800
Subject: [PATCH 10/32] [omni] retire text-only adapter for sglang_omni
 contract; thinker.* weight naming

- omni_thinker.py folded into generate_hub/sglang_omni.py: output_modalities/return_omni_rollout,
  pinned repetition_penalty=1.0, metadata passthrough, replay guard; temp-1 assert dropped
  (server sampler now returns post-temperature logprobs, matching trainer recompute)
- megatron_to_hf/qwen3omni_moe.py: body.* -> thinker.* (body.* is the Higgs namespace;
  sglang-omni thinker load_weights strips thinker. and ignores tower/talker names)
- extract tool: stream shards from flush() (peak RAM = one shard), fail loud on zero tensors
- tests: converter/extract moved under tests/fast/backends/megatron_utils with explicit
  register_cpu_ci; sglang_omni variant registered in generation_fixtures, harness basic+resume
---
 .../run_qwen3_omni_thinker_grpo.py            |   2 +-
 .../megatron_utils/megatron_to_hf/__init__.py |   2 +-
 .../megatron_to_hf/qwen3omni_moe.py           |  11 +-
 miles/rollout/generate_hub/omni_thinker.py    |  62 -----
 miles/rollout/generate_hub/sglang_omni.py     |  19 ++
 .../test_extract_qwen3_omni_thinker.py        | 152 +++++++++++++
 .../test_qwen3omni_moe_converter.py           |  80 +++++++
 tests/fast/fixtures/generation_fixtures.py    |   1 +
 .../rollout/generate_hub/test_sglang_omni.py  | 131 +++++++++++
 tests/fast/test_qwen3_omni_thinker.py         | 214 ------------------
 tools/extract_qwen3_omni_thinker.py           |  58 +++--
 11 files changed, 428 insertions(+), 304 deletions(-)
 delete mode 100644 miles/rollout/generate_hub/omni_thinker.py
 create mode 100644 tests/fast/backends/megatron_utils/test_extract_qwen3_omni_thinker.py
 create mode 100644 tests/fast/backends/megatron_utils/test_qwen3omni_moe_converter.py
 delete mode 100644 tests/fast/test_qwen3_omni_thinker.py

diff --git a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
index 9023bc12d3f..c2ec90bd191 100644
--- a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
+++ b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
@@ -62,7 +62,7 @@ def execute(args: ScriptArgs):
     )
 
     rollout_args = (
-        "--custom-generate-function-path miles.rollout.generate_hub.omni_thinker.generate "
+        "--custom-generate-function-path miles.rollout.generate_hub.sglang_omni.generate "
         f"--prompt-data {args.data_dir}/dapo-math-17k/dapo-math-17k.jsonl "
         "--input-key prompt "
         "--label-key label "
diff --git a/miles/backends/megatron_utils/megatron_to_hf/__init__.py b/miles/backends/megatron_utils/megatron_to_hf/__init__.py
index ba5d67d76d6..8fb2fa428ba 100644
--- a/miles/backends/megatron_utils/megatron_to_hf/__init__.py
+++ b/miles/backends/megatron_utils/megatron_to_hf/__init__.py
@@ -43,7 +43,7 @@ def _convert_to_hf_core(args, model_name, name, param):
         converted_named_tensors = convert_glm4moe_to_hf(args, name, param)
     elif "glm4" in model_name:
         converted_named_tensors = convert_glm4_to_hf(args, name, param)
-    elif "qwen3omni" in model_name:  # body.* names; keep before generic qwen3
+    elif "qwen3omni" in model_name:  # thinker.* names; keep before generic qwen3
         converted_named_tensors = convert_qwen3omni_moe_to_hf(args, name, param)
     elif "qwen3moe" in model_name:
         converted_named_tensors = convert_qwen3moe_to_hf(args, name, param)
diff --git a/miles/backends/megatron_utils/megatron_to_hf/qwen3omni_moe.py b/miles/backends/megatron_utils/megatron_to_hf/qwen3omni_moe.py
index b33b52b8a18..1051a0471ff 100644
--- a/miles/backends/megatron_utils/megatron_to_hf/qwen3omni_moe.py
+++ b/miles/backends/megatron_utils/megatron_to_hf/qwen3omni_moe.py
@@ -1,14 +1,15 @@
 """Megatron->HF broadcast converter for the Qwen3-Omni thinker (Qwen3-MoE text).
 
-Same conversion as qwen3moe, but prefixes every HF name with `body.` (the namespace
-the sglang-omni weight-receive side demuxes on). Selected via `--model-name
-qwen3omni_moe`. Nothing is dropped: the thinker is untied; `tied.*` is talker-side.
+Same conversion as qwen3moe, but prefixes every HF name with `thinker.` — the namespace
+the sglang-omni thinker stage strips in `load_weights` (names under `audio_tower.` /
+`visual.` / `talker.` / `code2wav.` are ignored there, so pushing only the text backbone
+is safe). Selected via `--model-name qwen3omni_moe`. The thinker is untied.
 """
 
 from .qwen3moe import convert_qwen3moe_to_hf
 
-BODY_PREFIX = "body."
+THINKER_PREFIX = "thinker."
 
 
 def convert_qwen3omni_moe_to_hf(args, name, param):
-    return [(f"{BODY_PREFIX}{n}", t) for n, t in convert_qwen3moe_to_hf(args, name, param)]
+    return [(f"{THINKER_PREFIX}{n}", t) for n, t in convert_qwen3moe_to_hf(args, name, param)]
diff --git a/miles/rollout/generate_hub/omni_thinker.py b/miles/rollout/generate_hub/omni_thinker.py
deleted file mode 100644
index f49ec5ea5d0..00000000000
--- a/miles/rollout/generate_hub/omni_thinker.py
+++ /dev/null
@@ -1,62 +0,0 @@
-"""Rollout adapter for the sglang-omni `/generate` endpoint (thinker / text).
-
-Thin wrapper over the default single-turn generate: reuses miles' payload builder and
-response parser, adds only the omni-specific fields. Wire via
-`--custom-generate-function-path miles.rollout.generate_hub.omni_thinker.generate`.
-"""
-
-from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput
-from miles.rollout.generate_utils.generate_endpoint_utils import (
-    compute_prompt_ids_from_sample,
-    compute_request_payload,
-    update_sample_from_response,
-)
-from miles.utils.http_utils import post
-from miles.utils.types import Sample
-
-
-async def generate(input: GenerateFnInput) -> GenerateFnOutput:
-    args = input.args
-    sample = input.sample
-    sampling_params = input.sampling_params
-    assert sample.status in {Sample.Status.PENDING, Sample.Status.ABORTED}, f"{sample.status=}"
-    # omni /generate emits temp-1 (pre-temperature) full-vocab logprobs; the trainer recompute
-    # divides logits by rollout_temperature, so they agree only at temp=1.
-    assert args.rollout_temperature == 1.0, (
-        f"omni rollout logprob is temp-1; rollout_temperature must be 1.0, got {args.rollout_temperature}"
-    )
-    # text-only path with no MoE/indexer replay yet (server forward-declares both)
-    assert not (args.use_rollout_routing_replay or args.use_rollout_indexer_replay), (
-        "omni rollout has no routing/indexer replay; unset --use-rollout-routing-replay / --use-rollout-indexer-replay"
-    )
-    assert not (sample.multimodal_inputs or {}).get("images"), "omni thinker rollout is text-only"
-    url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate"
-
-    prompt_ids = compute_prompt_ids_from_sample(input.state, sample)
-    if len(sample.response) > 0:  # partial rollout resume
-        input_ids = sample.tokens
-        sampling_params["max_new_tokens"] -= len(sample.tokens) - len(prompt_ids)
-        assert sampling_params["max_new_tokens"] >= 0
-        if sampling_params["max_new_tokens"] == 0:
-            sample.status = Sample.Status.TRUNCATED
-            return GenerateFnOutput(samples=sample)
-    else:
-        input_ids = prompt_ids
-
-    payload, halt_status = compute_request_payload(
-        args, input_ids=input_ids, sampling_params=sampling_params, multimodal_inputs=sample.multimodal_inputs
-    )
-    if payload is None:
-        sample.status = halt_status
-        return GenerateFnOutput(samples=sample)
-
-    payload["output_modalities"] = ["text"]
-    payload["return_omni_rollout"] = False
-    # rep_penalty=1: the trainer recompute can't replay a repetition penalty (logprobs would diverge)
-    payload["sampling_params"]["repetition_penalty"] = 1.0
-    if sample.metadata:
-        payload["metadata"] = sample.metadata
-
-    output = await post(url, payload)
-    await update_sample_from_response(args, sample, payload=payload, output=output)
-    return GenerateFnOutput(samples=sample)
diff --git a/miles/rollout/generate_hub/sglang_omni.py b/miles/rollout/generate_hub/sglang_omni.py
index 6dd96e0ea26..4552985ce25 100644
--- a/miles/rollout/generate_hub/sglang_omni.py
+++ b/miles/rollout/generate_hub/sglang_omni.py
@@ -2,6 +2,13 @@
 
 Select with ``--custom-generate-function-path
 miles.rollout.generate_hub.sglang_omni.generate``.
+
+Multimodal samples ship the processor-expanded ``input_ids`` plus the serialized
+processor tensors (``multimodal_train_inputs``); the server trusts those ids, feeds the
+tensors straight into its audio/vision towers, and never re-runs media processing, so
+rollout and training share one canonical token sequence. Logprobs come from the vanilla
+sglang sampler (post-temperature unless the server sets SGLANG_RETURN_ORIGINAL_LOGPROB),
+matching the trainer recompute convention.
 """
 
 from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput
@@ -23,6 +30,11 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
         Sample.Status.PENDING,
         Sample.Status.ABORTED,
     }, f"{sample.status=}"
+    # the omni server declares return_routed_experts/return_indexer_topk in its protocol but
+    # implements neither replay; fail loud instead of training on silently missing traces
+    assert not (args.use_rollout_routing_replay or args.use_rollout_indexer_replay), (
+        "sglang-omni rollout has no routing/indexer replay; unset --use-rollout-routing-replay / --use-rollout-indexer-replay"
+    )
     url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate"
 
     prompt_ids = compute_prompt_ids_from_sample(input.state, sample)
@@ -55,6 +67,13 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
         sample.status = halt_status
         return GenerateFnOutput(samples=sample)
 
+    payload["output_modalities"] = ["text"]
+    payload["return_omni_rollout"] = False
+    # the trainer recompute cannot replay a repetition penalty (logprobs would diverge)
+    payload["sampling_params"]["repetition_penalty"] = 1.0
+    if sample.metadata:
+        payload["metadata"] = sample.metadata
+
     output = await post(url, payload, headers=compute_routing_headers(args, sample))
     await update_sample_from_response(
         args,
diff --git a/tests/fast/backends/megatron_utils/test_extract_qwen3_omni_thinker.py b/tests/fast/backends/megatron_utils/test_extract_qwen3_omni_thinker.py
new file mode 100644
index 00000000000..47d477c2678
--- /dev/null
+++ b/tests/fast/backends/megatron_utils/test_extract_qwen3_omni_thinker.py
@@ -0,0 +1,152 @@
+"""Qwen3-Omni thinker extraction tool: name/config mapping + streaming shard writes."""
+
+import importlib.util
+import json
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[])
+
+_REPO = Path(__file__).resolve().parents[4]
+
+
+def _load_extract_tool():
+    path = _REPO / "tools" / "extract_qwen3_omni_thinker.py"
+    spec = importlib.util.spec_from_file_location("extract_qwen3_omni_thinker", path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def test_map_thinker_param_name_keeps_text_backbone():
+    m = _load_extract_tool().map_thinker_param_name
+    assert m("thinker.model.embed_tokens.weight") == "model.embed_tokens.weight"
+    assert m("thinker.lm_head.weight") == "lm_head.weight"
+    assert m("thinker.model.norm.weight") == "model.norm.weight"
+    assert (
+        m("thinker.model.layers.7.mlp.experts.3.down_proj.weight") == "model.layers.7.mlp.experts.3.down_proj.weight"
+    )
+    assert m("thinker.model.layers.0.self_attn.q_norm.weight") == "model.layers.0.self_attn.q_norm.weight"
+
+
+def test_map_thinker_param_name_normalizes_language_model_infix():
+    m = _load_extract_tool().map_thinker_param_name
+    assert m("thinker.model.language_model.layers.0.mlp.gate.weight") == "model.layers.0.mlp.gate.weight"
+
+
+def test_map_thinker_param_name_drops_non_thinker_and_non_text():
+    m = _load_extract_tool().map_thinker_param_name
+    assert m("talker.model.layers.0.self_attn.q_proj.weight") is None
+    assert m("code2wav.decoder.weight") is None
+    assert m("thinker.audio_tower.layers.0.fc1.weight") is None
+    assert m("thinker.visual.blocks.0.attn.qkv.weight") is None
+
+
+def test_synthesize_thinker_config_stamps_qwen3_moe():
+    synth = _load_extract_tool().synthesize_thinker_config
+    omni = {
+        "model_type": "qwen3_omni_moe",
+        "thinker_config": {
+            "eos_token_id": 151645,
+            "text_config": {"vocab_size": 152064, "num_experts": 128, "tie_word_embeddings": False},
+        },
+    }
+    cfg = synth(omni)
+    assert cfg["model_type"] == "qwen3_moe"
+    assert cfg["architectures"] == ["Qwen3MoeForCausalLM"]
+    assert cfg["num_experts"] == 128
+    assert cfg["vocab_size"] == 152064
+    assert cfg["tie_word_embeddings"] is False
+    assert cfg["eos_token_id"] == 151645
+
+
+def test_synthesize_thinker_config_flat_fallback():
+    cfg = _load_extract_tool().synthesize_thinker_config({"vocab_size": 100, "num_experts": 8})
+    assert cfg["model_type"] == "qwen3_moe"
+    assert cfg["vocab_size"] == 100
+    assert cfg["tie_word_embeddings"] is False
+
+
+def _write_omni_src(src: Path, tensors: dict[str, torch.Tensor]):
+    src.mkdir(parents=True, exist_ok=True)
+    with open(src / "config.json", "w") as f:
+        json.dump({"thinker_config": {"text_config": {"vocab_size": 32}}}, f)
+    save_file(tensors, src / "model.safetensors", metadata={"format": "pt"})
+
+
+def _read_all(dst: Path) -> dict[str, torch.Tensor]:
+    index_path = dst / "model.safetensors.index.json"
+    if index_path.exists():
+        with open(index_path) as f:
+            weight_map = json.load(f)["weight_map"]
+        shards = sorted(set(weight_map.values()))
+    else:
+        weight_map, shards = None, ["model.safetensors"]
+    out = {}
+    for shard in shards:
+        with safe_open(dst / shard, framework="pt") as r:
+            for k in r.keys():
+                out[k] = r.get_tensor(k)
+                if weight_map is not None:
+                    assert weight_map[k] == shard
+    return out
+
+
+def test_extract_single_shard_roundtrip(tmp_path):
+    tool = _load_extract_tool()
+    src, dst = tmp_path / "src", tmp_path / "dst"
+    _write_omni_src(
+        src,
+        {
+            "thinker.model.embed_tokens.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3),
+            "thinker.audio_tower.layers.0.fc1.weight": torch.zeros(2, 2),
+            "talker.model.norm.weight": torch.zeros(2),
+        },
+    )
+
+    total = tool.extract(src, dst, shard_size_gb=1.0)
+
+    assert total == 1
+    assert (dst / "model.safetensors").exists()
+    assert not (dst / "model.safetensors.index.json").exists()
+    out = _read_all(dst)
+    assert set(out) == {"model.embed_tokens.weight"}
+    assert torch.equal(out["model.embed_tokens.weight"], torch.arange(6, dtype=torch.float32).reshape(2, 3))
+    with open(dst / "config.json") as f:
+        assert json.load(f)["model_type"] == "qwen3_moe"
+
+
+def test_extract_multi_shard_writes_streaming_index(tmp_path):
+    tool = _load_extract_tool()
+    src, dst = tmp_path / "src", tmp_path / "dst"
+    tensors = {f"thinker.model.layers.{i}.mlp.gate.weight": torch.full((64, 64), float(i)) for i in range(4)}
+    _write_omni_src(src, tensors)
+
+    # 64*64*4B = 16KiB per tensor; 20KiB shards force one tensor per shard (flush at >=1 tensor)
+    total = tool.extract(src, dst, shard_size_gb=20 * 1024 / (1024**3))
+
+    assert total == 4
+    index_path = dst / "model.safetensors.index.json"
+    assert index_path.exists()
+    with open(index_path) as f:
+        index = json.load(f)
+    assert index["metadata"]["total_size"] == sum(t.numel() * t.element_size() for t in tensors.values())
+    assert len(set(index["weight_map"].values())) > 1, "expected multiple shards"
+    out = _read_all(dst)
+    for i in range(4):
+        assert torch.equal(out[f"model.layers.{i}.mlp.gate.weight"], torch.full((64, 64), float(i)))
+
+
+def test_extract_fails_loud_when_no_thinker_tensors(tmp_path):
+    tool = _load_extract_tool()
+    src, dst = tmp_path / "src", tmp_path / "dst"
+    _write_omni_src(src, {"talker.model.norm.weight": torch.zeros(2)})
+
+    with pytest.raises(ValueError, match="no thinker tensors"):
+        tool.extract(src, dst, shard_size_gb=1.0)
diff --git a/tests/fast/backends/megatron_utils/test_qwen3omni_moe_converter.py b/tests/fast/backends/megatron_utils/test_qwen3omni_moe_converter.py
new file mode 100644
index 00000000000..94f0d006e20
--- /dev/null
+++ b/tests/fast/backends/megatron_utils/test_qwen3omni_moe_converter.py
@@ -0,0 +1,80 @@
+"""Qwen3-Omni thinker Megatron->HF converter: thinker.* naming + dispatch contract."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="stage-a-cpu", labels=[])
+
+from miles.backends.megatron_utils.megatron_to_hf import _convert_to_hf_core
+from miles.backends.megatron_utils.megatron_to_hf.qwen3omni_moe import convert_qwen3omni_moe_to_hf
+
+
+def _args():
+    # head_dim=kv_channels=2, value_num_per_group=heads//groups=2
+    return SimpleNamespace(hidden_size=8, num_attention_heads=4, num_query_groups=2, kv_channels=2)
+
+
+def _names(name, param):
+    return [x for x, _ in convert_qwen3omni_moe_to_hf(_args(), name, param)]
+
+
+def test_thinker_converter_simple_params():
+    assert _names("module.module.embedding.word_embeddings.weight", torch.zeros(10, 8)) == [
+        "thinker.model.embed_tokens.weight"
+    ]
+    assert _names("module.module.output_layer.weight", torch.zeros(10, 8)) == ["thinker.lm_head.weight"]
+    assert _names("module.module.decoder.final_layernorm.weight", torch.zeros(8)) == ["thinker.model.norm.weight"]
+    assert _names("module.module.decoder.layers.3.mlp.router.weight", torch.zeros(128, 8)) == [
+        "thinker.model.layers.3.mlp.gate.weight"
+    ]
+    assert _names("module.module.decoder.layers.3.self_attention.q_layernorm.weight", torch.zeros(2)) == [
+        "thinker.model.layers.3.self_attn.q_norm.weight"
+    ]
+
+
+def test_thinker_converter_fused_splits():
+    conv, args = convert_qwen3omni_moe_to_hf, _args()
+
+    # qkv: (groups*(vpg+2)*head_dim, hidden) = (16, 8) -> q(8,8) k(4,8) v(4,8)
+    out = conv(args, "module.module.decoder.layers.1.self_attention.linear_qkv.weight", torch.zeros(16, 8))
+    assert {n: tuple(t.shape) for n, t in out} == {
+        "thinker.model.layers.1.self_attn.q_proj.weight": (8, 8),
+        "thinker.model.layers.1.self_attn.k_proj.weight": (4, 8),
+        "thinker.model.layers.1.self_attn.v_proj.weight": (4, 8),
+    }
+
+    # expert fc1: chunk into gate/up
+    out = conv(args, "module.module.decoder.layers.2.mlp.experts.linear_fc1.weight5", torch.zeros(12, 8))
+    assert {n: tuple(t.shape) for n, t in out} == {
+        "thinker.model.layers.2.mlp.experts.5.gate_proj.weight": (6, 8),
+        "thinker.model.layers.2.mlp.experts.5.up_proj.weight": (6, 8),
+    }
+
+
+def test_thinker_converter_never_emits_body_prefix():
+    # body.* is the Higgs-TTS namespace; Qwen3-Omni thinker load_weights expects thinker.*/HF names
+    for name, param in [
+        ("module.module.embedding.word_embeddings.weight", torch.zeros(10, 8)),
+        ("module.module.decoder.layers.0.mlp.experts.linear_fc2.weight0", torch.zeros(8, 6)),
+    ]:
+        for hf_name, _ in convert_qwen3omni_moe_to_hf(_args(), name, param):
+            assert hf_name.startswith("thinker."), hf_name
+            assert not hf_name.startswith("body."), hf_name
+
+
+def test_dispatch_routes_qwen3omni_before_generic_qwen3():
+    # "qwen3omni_moe" does not contain "qwen3moe"; without the dedicated branch it would fall
+    # through to the dense qwen2 converter and crash on MoE params
+    out = _convert_to_hf_core(
+        _args(), "qwen3omni_moe", "module.module.embedding.word_embeddings.weight", torch.zeros(10, 8)
+    )
+    assert [n for n, _ in out] == ["thinker.model.embed_tokens.weight"]
+
+
+def test_dispatch_rejects_unknown_model():
+    with pytest.raises(ValueError, match="Unsupported model"):
+        _convert_to_hf_core(_args(), "unknown_model_xyz", "module.module.x", torch.zeros(1))
diff --git a/tests/fast/fixtures/generation_fixtures.py b/tests/fast/fixtures/generation_fixtures.py
index 95838f5c748..3d7e1f3ba4a 100644
--- a/tests/fast/fixtures/generation_fixtures.py
+++ b/tests/fast/fixtures/generation_fixtures.py
@@ -31,6 +31,7 @@
 VARIANT_TO_GENERATE_FN_PATH = {
     "old_sglang_rollout": "miles.rollout.sglang_rollout.generate",
     "single_turn": "miles.rollout.generate_hub.single_turn.generate",
+    "sglang_omni": "miles.rollout.generate_hub.sglang_omni.generate",
     "multi_turn_single_sample": "miles.rollout.generate_hub.multi_turn.generate",
     "multi_turn_multi_samples": "miles.rollout.generate_hub.multi_turn.generate",
     "agentic_tool_call_single_sample": "miles.rollout.generate_hub.agentic_tool_call.generate",
diff --git a/tests/fast/rollout/generate_hub/test_sglang_omni.py b/tests/fast/rollout/generate_hub/test_sglang_omni.py
index a405db9900f..87cbf7fbc11 100644
--- a/tests/fast/rollout/generate_hub/test_sglang_omni.py
+++ b/tests/fast/rollout/generate_hub/test_sglang_omni.py
@@ -1,3 +1,7 @@
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=120, suite="stage-b-cpu", labels=[])
+
 import asyncio
 import base64
 import sys
@@ -5,10 +9,24 @@
 
 import pytest
 import torch
+from tests.fast.fixtures.generation_fixtures import generation_env, make_sample, run_generate
 
 from miles.rollout.generate_utils.generate_endpoint_utils import serialize_multimodal_train_inputs
+from miles.utils.test_utils.mock_sglang_server import ProcessResult
 from miles.utils.types import Sample
 
+_ = generation_env
+
+PROMPT_TOKENS = [3838, 374, 220, 16, 10, 22, 30]  # Qwen3-0.6B "What is 1+7?"
+RESPONSE_TOKENS = [59, 79075, 90, 23, 92]  # "\\boxed{8}"
+RESPONSE_LOG_PROBS = [-0.0, -0.0078125, -0.015625, -0.0234375, -0.03125]
+SAMPLING_PARAMS = {"max_new_tokens": 16, "temperature": 0.7}
+
+
+@pytest.fixture
+def variant():
+    return "sglang_omni"
+
 
 def test_serialize_audio_video_processor_tensors():
     inputs = {
@@ -186,6 +204,119 @@ async def fail_post(url, payload, headers=None):
     assert output.samples is sample
 
 
+class TestHarness:
+    """Through the real parse_args + mock sglang server harness (generation_fixtures)."""
+
+    def test_basic_generation_request_shape_and_sample(self, generation_env):
+        result = run_generate(generation_env, make_sample(), dict(SAMPLING_PARAMS), variant="sglang_omni")
+
+        [request] = result.requests
+        assert request["input_ids"] == PROMPT_TOKENS
+        assert request["output_modalities"] == ["text"]
+        assert request["return_omni_rollout"] is False
+        assert request["return_logprob"] is True
+        assert request["sampling_params"]["repetition_penalty"] == 1.0
+        assert request["sampling_params"]["max_new_tokens"] == SAMPLING_PARAMS["max_new_tokens"]
+        sample = result.sample
+        assert sample.status == Sample.Status.COMPLETED
+        assert sample.tokens == PROMPT_TOKENS + RESPONSE_TOKENS
+        assert sample.rollout_log_probs == RESPONSE_LOG_PROBS
+        assert sample.response == "\\boxed{8}"
+
+    def test_resumed_partial_rollout_decrements_budget(self, generation_env):
+        partial_tokens = [59, 79075]  # "\\boxed"
+        generation_env.mock_server.process_fn = lambda _: ProcessResult(text="\\boxed", finish_reason="abort")
+        sample = make_sample()
+        result1 = run_generate(generation_env, sample, dict(SAMPLING_PARAMS), variant="sglang_omni")
+        assert result1.sample.status == Sample.Status.ABORTED
+        assert result1.sample.tokens == PROMPT_TOKENS + partial_tokens
+
+        generation_env.mock_server.process_fn = lambda _: ProcessResult(text="{8}", finish_reason="stop")
+        result2 = run_generate(generation_env, result1.sample, dict(SAMPLING_PARAMS), variant="sglang_omni")
+        [request] = result2.requests
+        assert request["input_ids"] == PROMPT_TOKENS + partial_tokens
+        assert request["sampling_params"]["max_new_tokens"] == SAMPLING_PARAMS["max_new_tokens"] - len(partial_tokens)
+        assert result2.sample.status == Sample.Status.COMPLETED
+        assert result2.sample.tokens == PROMPT_TOKENS + partial_tokens + [90, 23, 92]
+        assert result2.sample.rollout_log_probs == [-0.0, -0.0078125, -0.0, -0.0078125, -0.015625]
+
+
+def test_sglang_omni_adapter_sets_omni_contract_fields(monkeypatch):
+    from miles.rollout.generate_hub import sglang_omni
+
+    captured = {}
+
+    async def fake_post(url, payload, headers=None):
+        captured["payload"] = payload
+        return {}
+
+    async def fake_update(*args, **kwargs):
+        return None
+
+    monkeypatch.setattr(sglang_omni, "post", fake_post)
+    monkeypatch.setattr(sglang_omni, "update_sample_from_response", fake_update)
+    monkeypatch.setattr(sglang_omni, "compute_prompt_ids_from_sample", lambda state, sample: [1, 2, 3])
+    generate_input = SimpleNamespace(
+        args=_adapter_args(),
+        sample=Sample(),
+        sampling_params={"temperature": 1.0, "max_new_tokens": 16, "repetition_penalty": 1.4},
+        state=None,
+    )
+
+    asyncio.run(sglang_omni.generate(generate_input))
+
+    payload = captured["payload"]
+    assert payload["output_modalities"] == ["text"]
+    assert payload["return_omni_rollout"] is False
+    # trainer recompute cannot replay a repetition penalty; the adapter pins it to 1.0
+    assert payload["sampling_params"]["repetition_penalty"] == 1.0
+    assert "metadata" not in payload
+
+
+def test_sglang_omni_adapter_forwards_sample_metadata(monkeypatch):
+    from miles.rollout.generate_hub import sglang_omni
+
+    captured = {}
+
+    async def fake_post(url, payload, headers=None):
+        captured["payload"] = payload
+        return {}
+
+    async def fake_update(*args, **kwargs):
+        return None
+
+    monkeypatch.setattr(sglang_omni, "post", fake_post)
+    monkeypatch.setattr(sglang_omni, "update_sample_from_response", fake_update)
+    monkeypatch.setattr(sglang_omni, "compute_prompt_ids_from_sample", lambda state, sample: [1, 2, 3])
+    generate_input = SimpleNamespace(
+        args=_adapter_args(),
+        sample=Sample(metadata={"task": "avqa"}),
+        sampling_params={"temperature": 1.0, "max_new_tokens": 16},
+        state=None,
+    )
+
+    asyncio.run(sglang_omni.generate(generate_input))
+
+    assert captured["payload"]["metadata"] == {"task": "avqa"}
+
+
+def test_sglang_omni_adapter_rejects_rollout_replay(monkeypatch):
+    from miles.rollout.generate_hub import sglang_omni
+
+    monkeypatch.setattr(sglang_omni, "compute_prompt_ids_from_sample", lambda state, sample: [1, 2, 3])
+    for flag in ("use_rollout_routing_replay", "use_rollout_indexer_replay"):
+        args = _adapter_args()
+        setattr(args, flag, True)
+        generate_input = SimpleNamespace(
+            args=args,
+            sample=Sample(),
+            sampling_params={"temperature": 1.0, "max_new_tokens": 16},
+            state=None,
+        )
+        with pytest.raises(AssertionError, match="replay"):
+            asyncio.run(sglang_omni.generate(generate_input))
+
+
 def test_sglang_omni_adapter_resumes_partial_rollout(monkeypatch):
     from miles.rollout.generate_hub import sglang_omni
 
diff --git a/tests/fast/test_qwen3_omni_thinker.py b/tests/fast/test_qwen3_omni_thinker.py
deleted file mode 100644
index b9773e3565b..00000000000
--- a/tests/fast/test_qwen3_omni_thinker.py
+++ /dev/null
@@ -1,214 +0,0 @@
-"""Qwen3-Omni thinker bridge: extraction name/config logic + the body.* converter contract."""
-
-import asyncio
-import importlib.util
-from pathlib import Path
-from types import SimpleNamespace
-
-import pytest
-
-_REPO = Path(__file__).resolve().parents[2]
-
-
-def _load_extract_tool():
-    path = _REPO / "tools" / "extract_qwen3_omni_thinker.py"
-    spec = importlib.util.spec_from_file_location("extract_qwen3_omni_thinker", path)
-    mod = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(mod)
-    return mod
-
-
-def test_map_thinker_param_name_keeps_text_backbone():
-    m = _load_extract_tool().map_thinker_param_name
-    assert m("thinker.model.embed_tokens.weight") == "model.embed_tokens.weight"
-    assert m("thinker.lm_head.weight") == "lm_head.weight"
-    assert m("thinker.model.norm.weight") == "model.norm.weight"
-    assert m("thinker.model.layers.7.mlp.experts.3.down_proj.weight") == "model.layers.7.mlp.experts.3.down_proj.weight"
-    assert m("thinker.model.layers.0.self_attn.q_norm.weight") == "model.layers.0.self_attn.q_norm.weight"
-
-
-def test_map_thinker_param_name_normalizes_language_model_infix():
-    m = _load_extract_tool().map_thinker_param_name
-    assert m("thinker.model.language_model.layers.0.mlp.gate.weight") == "model.layers.0.mlp.gate.weight"
-
-
-def test_map_thinker_param_name_drops_non_thinker_and_non_text():
-    m = _load_extract_tool().map_thinker_param_name
-    assert m("talker.model.layers.0.self_attn.q_proj.weight") is None
-    assert m("code2wav.decoder.weight") is None
-    assert m("thinker.audio_tower.layers.0.fc1.weight") is None
-    assert m("thinker.visual.blocks.0.attn.qkv.weight") is None
-
-
-def test_synthesize_thinker_config_stamps_qwen3_moe():
-    synth = _load_extract_tool().synthesize_thinker_config
-    omni = {
-        "model_type": "qwen3_omni_moe",
-        "thinker_config": {
-            "eos_token_id": 151645,
-            "text_config": {"vocab_size": 152064, "num_experts": 128, "tie_word_embeddings": False},
-        },
-    }
-    cfg = synth(omni)
-    assert cfg["model_type"] == "qwen3_moe"
-    assert cfg["architectures"] == ["Qwen3MoeForCausalLM"]
-    assert cfg["num_experts"] == 128
-    assert cfg["vocab_size"] == 152064
-    assert cfg["tie_word_embeddings"] is False
-    assert cfg["eos_token_id"] == 151645
-
-
-def test_synthesize_thinker_config_flat_fallback():
-    cfg = _load_extract_tool().synthesize_thinker_config({"vocab_size": 100, "num_experts": 8})
-    assert cfg["model_type"] == "qwen3_moe"
-    assert cfg["vocab_size"] == 100
-    assert cfg["tie_word_embeddings"] is False
-
-
-def _args():
-    # head_dim=kv_channels=2, value_num_per_group=heads//groups=2
-    return SimpleNamespace(hidden_size=8, num_attention_heads=4, num_query_groups=2, kv_channels=2)
-
-
-def _conv():
-    return pytest.importorskip(
-        "miles.backends.megatron_utils.megatron_to_hf.qwen3omni_moe"
-    ).convert_qwen3omni_moe_to_hf
-
-
-def test_body_converter_simple_params():
-    torch = pytest.importorskip("torch")
-    conv, args = _conv(), _args()
-    names = lambda n, p: [x for x, _ in conv(args, n, p)]
-    assert names("module.module.embedding.word_embeddings.weight", torch.zeros(10, 8)) == ["body.model.embed_tokens.weight"]
-    assert names("module.module.output_layer.weight", torch.zeros(10, 8)) == ["body.lm_head.weight"]
-    assert names("module.module.decoder.final_layernorm.weight", torch.zeros(8)) == ["body.model.norm.weight"]
-    assert names("module.module.decoder.layers.3.mlp.router.weight", torch.zeros(128, 8)) == ["body.model.layers.3.mlp.gate.weight"]
-    assert names("module.module.decoder.layers.3.self_attention.q_layernorm.weight", torch.zeros(2)) == ["body.model.layers.3.self_attn.q_norm.weight"]
-    assert names("module.module.decoder.layers.3.self_attention.k_layernorm.weight", torch.zeros(2)) == ["body.model.layers.3.self_attn.k_norm.weight"]
-
-
-def test_body_converter_fused_splits():
-    torch = pytest.importorskip("torch")
-    conv, args = _conv(), _args()
-
-    # qkv: (groups*(vpg+2)*head_dim, hidden) = (16, 8) -> q(8,8) k(4,8) v(4,8)
-    out = conv(args, "module.module.decoder.layers.1.self_attention.linear_qkv.weight", torch.zeros(16, 8))
-    assert {n: tuple(t.shape) for n, t in out} == {
-        "body.model.layers.1.self_attn.q_proj.weight": (8, 8),
-        "body.model.layers.1.self_attn.k_proj.weight": (4, 8),
-        "body.model.layers.1.self_attn.v_proj.weight": (4, 8),
-    }
-    # expert fc1: gate+up fused on dim 0 -> chunk(2)
-    out = conv(args, "module.module.decoder.layers.1.mlp.experts.linear_fc1.weight5", torch.zeros(8, 8))
-    assert {n: tuple(t.shape) for n, t in out} == {
-        "body.model.layers.1.mlp.experts.5.gate_proj.weight": (4, 8),
-        "body.model.layers.1.mlp.experts.5.up_proj.weight": (4, 8),
-    }
-    out = conv(args, "module.module.decoder.layers.1.mlp.experts.linear_fc2.weight5", torch.zeros(8, 4))
-    assert [n for n, _ in out] == ["body.model.layers.1.mlp.experts.5.down_proj.weight"]
-
-
-def test_body_converter_dispatch_routes_qwen3omni():
-    torch = pytest.importorskip("torch")
-    mod = pytest.importorskip("miles.backends.megatron_utils.megatron_to_hf")
-    out = mod._convert_to_hf_core(
-        _args(), "qwen3omni_moe", "module.module.decoder.layers.0.mlp.router.weight", torch.zeros(128, 8)
-    )
-    assert [n for n, _ in out] == ["body.model.layers.0.mlp.gate.weight"]
-
-
-# --- rollout adapter (miles.rollout.generate_hub.omni_thinker.generate) ---
-
-
-def _omni_mod():
-    return pytest.importorskip("miles.rollout.generate_hub.omni_thinker")
-
-
-def _omni_input(monkeypatch, omni, *, replay=False, indexer_replay=False, multimodal=None, metadata=None, temperature=1.0):
-    """A generate() input with post/prompt-id mocked; returns (input, sample, captured)."""
-    Sample = pytest.importorskip("miles.utils.types").Sample
-    captured = {}
-
-    async def fake_post(url, payload):
-        captured["url"], captured["payload"] = url, payload
-        return {
-            "text": "hello",
-            "meta_info": {
-                "finish_reason": {"type": "stop"},
-                "output_token_logprobs": [[-0.1, 11], [-0.2, 22]],
-                "prompt_tokens": 3,
-                "completion_tokens": 2,
-                "cached_tokens": 0,
-            },
-        }
-
-    monkeypatch.setattr(omni, "post", fake_post)
-    monkeypatch.setattr(omni, "compute_prompt_ids_from_sample", lambda state, sample: [1, 2, 3])
-
-    args = SimpleNamespace(
-        sglang_router_ip="127.0.0.1",
-        sglang_router_port=30000,
-        rollout_temperature=temperature,
-        use_rollout_routing_replay=replay,
-        use_rollout_indexer_replay=indexer_replay,
-        rollout_max_response_len=128,
-        rollout_max_context_len=0,
-        sglang_speculative_algorithm=None,
-    )
-    sample = Sample(status=Sample.Status.PENDING, metadata=metadata or {}, multimodal_inputs=multimodal)
-    inp = SimpleNamespace(
-        args=args, sample=sample, sampling_params={"temperature": 1.0, "max_new_tokens": 64}, state=None
-    )
-    return inp, sample, captured
-
-
-def test_generate_builds_text_only_omni_payload_and_parses_response(monkeypatch):
-    omni = _omni_mod()
-    inp, sample, captured = _omni_input(monkeypatch, omni, metadata={"task": "math"})
-
-    asyncio.run(omni.generate(inp))
-
-    p = captured["payload"]
-    assert captured["url"].endswith("/generate")
-    assert p["input_ids"] == [1, 2, 3]
-    assert p["output_modalities"] == ["text"]
-    assert p["return_omni_rollout"] is False
-    assert p["return_routed_experts"] is False
-    assert p["return_indexer_topk"] is False
-    assert p["sampling_params"]["repetition_penalty"] == 1.0
-    assert p["metadata"] == {"task": "math"}
-    # response parse: tokens = input_ids + decoded ids, logprobs aligned, status from finish_reason
-    assert sample.tokens == [1, 2, 3, 11, 22]
-    assert sample.rollout_log_probs == [-0.1, -0.2]
-    assert sample.response == "hello"
-    assert sample.status.name == "COMPLETED"
-
-
-def test_generate_rejects_moe_replay_for_omni(monkeypatch):
-    omni = _omni_mod()
-    inp, _, _ = _omni_input(monkeypatch, omni, replay=True)
-    with pytest.raises(AssertionError, match="replay"):
-        asyncio.run(omni.generate(inp))
-
-
-def test_generate_rejects_indexer_replay_for_omni(monkeypatch):
-    omni = _omni_mod()
-    inp, _, _ = _omni_input(monkeypatch, omni, indexer_replay=True)
-    with pytest.raises(AssertionError, match="replay"):
-        asyncio.run(omni.generate(inp))
-
-
-def test_generate_rejects_multimodal_on_text_only_path(monkeypatch):
-    omni = _omni_mod()
-    inp, _, _ = _omni_input(monkeypatch, omni, multimodal={"images": ["img"]})
-    with pytest.raises(AssertionError, match="text-only"):
-        asyncio.run(omni.generate(inp))
-
-
-def test_generate_rejects_nonunit_temperature(monkeypatch):
-    # omni emits temp-1 logprobs; the trainer recompute divides by rollout_temperature, so temp!=1 desyncs
-    omni = _omni_mod()
-    inp, _, _ = _omni_input(monkeypatch, omni, temperature=2.0)
-    with pytest.raises(AssertionError, match="temp-1"):
-        asyncio.run(omni.generate(inp))
diff --git a/tools/extract_qwen3_omni_thinker.py b/tools/extract_qwen3_omni_thinker.py
index 71c1d91a49a..575cc53d9e8 100644
--- a/tools/extract_qwen3_omni_thinker.py
+++ b/tools/extract_qwen3_omni_thinker.py
@@ -3,7 +3,9 @@
 miles loads HF via AutoBridge keyed on model_type; the composite omni checkpoint has
 no bridge, but the thinker text backbone is a plain Qwen3-MoE. This writes a
 self-contained HF dir (thinker text + lm_head, renamed, model_type=qwen3_moe) for the
-existing Qwen3MoEBridge.
+existing Qwen3MoEBridge. Audio/visual towers are dropped here on purpose: the trainer
+loads them frozen from the original omni checkpoint (see the audio-injection wrapper),
+and the rollout server keeps its own copies.
 
     python tools/extract_qwen3_omni_thinker.py --src <omni> --dst <thinker>
 """
@@ -36,22 +38,21 @@ def synthesize_thinker_config(omni_config: dict) -> dict:
     return cfg
 
 
-def main() -> None:
-    import argparse
+def extract(src, dst, shard_size_gb: float = 5.0) -> int:
+    """Stream thinker tensors from src into sharded safetensors under dst.
+
+    Each shard is written (and its buffer released) as soon as it fills, so peak RAM is
+    one shard, not the whole ~60 GB thinker. Returns the number of kept tensors.
+    """
     import json
+    import os
     import shutil
     from pathlib import Path
 
     from safetensors import safe_open
     from safetensors.torch import save_file
 
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--src", required=True)
-    parser.add_argument("--dst", required=True)
-    parser.add_argument("--shard-size-gb", type=float, default=5.0)
-    args = parser.parse_args()
-
-    src, dst = Path(args.src), Path(args.dst)
+    src, dst = Path(src), Path(dst)
     dst.mkdir(parents=True, exist_ok=True)
 
     with open(src / "config.json") as f:
@@ -73,21 +74,24 @@ def main() -> None:
     else:
         shard_files = ["model.safetensors"]
 
-    shard_size_bytes = int(args.shard_size_gb * (1024 ** 3))
+    shard_size_bytes = int(shard_size_gb * (1024**3))
     out_index: dict[str, str] = {}
-    out_shards: list[tuple[str, dict]] = []
+    shard_names: list[str] = []
     buf: dict = {}
     buf_bytes = 0
+    total_size = 0
     total_kept = 0
 
     def flush():
-        nonlocal buf, buf_bytes
+        nonlocal buf, buf_bytes, total_size
         if not buf:
             return
-        shard_name = f"model-{len(out_shards) + 1:05d}.safetensors"
-        out_shards.append((shard_name, buf))
-        for k in buf:
+        shard_name = f"model-{len(shard_names) + 1:05d}.safetensors"
+        save_file(buf, dst / shard_name, metadata={"format": "pt"})
+        shard_names.append(shard_name)
+        for k, t in buf.items():
             out_index[k] = shard_name
+            total_size += t.numel() * t.element_size()
         buf, buf_bytes = {}, 0
 
     for shard_file in shard_files:
@@ -104,16 +108,28 @@ def flush():
                     flush()
     flush()
 
-    if len(out_shards) == 1:
-        save_file(out_shards[0][1], dst / "model.safetensors", metadata={"format": "pt"})
+    if total_kept == 0:
+        raise ValueError(f"no thinker tensors found under {src} (wrong --src?)")
+
+    if len(shard_names) == 1:
+        os.rename(dst / shard_names[0], dst / "model.safetensors")
     else:
-        total_size = sum(t.numel() * t.element_size() for _, ts in out_shards for t in ts.values())
-        for shard_name, tensors in out_shards:
-            save_file(tensors, dst / shard_name, metadata={"format": "pt"})
         with open(dst / "model.safetensors.index.json", "w") as f:
             json.dump({"metadata": {"total_size": total_size}, "weight_map": out_index}, f, indent=2)
 
     print(f"[done] {total_kept} thinker tensors -> {dst}")
+    return total_kept
+
+
+def main() -> None:
+    import argparse
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--src", required=True)
+    parser.add_argument("--dst", required=True)
+    parser.add_argument("--shard-size-gb", type=float, default=5.0)
+    args = parser.parse_args()
+    extract(args.src, args.dst, shard_size_gb=args.shard_size_gb)
 
 
 if __name__ == "__main__":

From 766278611cabf86439a4481236a1524a94ad3da0 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Fri, 24 Jul 2026 02:15:43 +0800
Subject: [PATCH 11/32] [omni] frozen-audio-tower embedding injection for the
 thinker trainer

- miles_plugins/models/qwen3_omni_thinker.py: patch GPTModel.forward to fill
  input_ids==audio_token_id rows from the omni checkpoint's frozen audio tower
  (decoder_input path); audio-only prompts keep sequential position ids, so RoPE
  is untouched; image/video keys rejected loudly; tower loaded lazily from the
  omni safetensors (audio_tower.* only) and kept out of DDP/optimizer/checkpoints
- model_provider: install hook behind --qwen3-omni-audio-encoder-path (classic + bridge)
- training_utils/data.get_batch: right-pad variable-length audio tensors
  (input_features/feature_attention_mask) to the micro-batch max before dim-0 concat
---
 .../backends/megatron_utils/model_provider.py |   9 +
 miles/backends/training_utils/data.py         |  28 +-
 miles/utils/arguments.py                      |  10 +
 miles_plugins/models/qwen3_omni_thinker.py    | 174 +++++++++++
 .../test_get_batch_multimodal_audio.py        | 108 +++++++
 tests/fast/test_qwen3_omni_audio_injection.py | 276 ++++++++++++++++++
 6 files changed, 597 insertions(+), 8 deletions(-)
 create mode 100644 miles_plugins/models/qwen3_omni_thinker.py
 create mode 100644 tests/fast/backends/training_utils/test_get_batch_multimodal_audio.py
 create mode 100644 tests/fast/test_qwen3_omni_audio_injection.py

diff --git a/miles/backends/megatron_utils/model_provider.py b/miles/backends/megatron_utils/model_provider.py
index 145a486ba83..5f06fd60932 100644
--- a/miles/backends/megatron_utils/model_provider.py
+++ b/miles/backends/megatron_utils/model_provider.py
@@ -191,6 +191,7 @@ def _logits_only_forward(*args, **kwargs):
                 return out[0] if isinstance(out, tuple) else out
 
             model.forward = _logits_only_forward
+            _maybe_install_omni_audio_injection(args, model)
             return model
 
         return wrapped_bridge_provider
@@ -322,12 +323,20 @@ def model_provider(
             model.output_layer = LinearForLastLayer(input_size=config.hidden_size, output_size=1, config=config)
 
         _maybe_install_witness(args, model)
+        _maybe_install_omni_audio_injection(args, model)
 
         return model
 
     return model_provider
 
 
+def _maybe_install_omni_audio_injection(args: argparse.Namespace, model) -> None:
+    if getattr(args, "qwen3_omni_audio_encoder_path", None):
+        from miles_plugins.models.qwen3_omni_thinker import install_audio_injection
+
+        install_audio_injection(model, args)
+
+
 def _maybe_install_witness(
     args: argparse.Namespace,
     model: GPTModel,
diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py
index d29cb165750..3f4cb23df2b 100644
--- a/miles/backends/training_utils/data.py
+++ b/miles/backends/training_utils/data.py
@@ -323,23 +323,35 @@ def _compute_transform_like_token_ids(ids_list: list):
     # Process multimodal training tensors if present
     multimodal_train_inputs = batch.get("multimodal_train_inputs", None)
     if multimodal_train_inputs is not None:
-        multimodal_data = {}  # key -> concatenated tensor
+        multimodal_lists = {}  # key -> per-sequence tensors
         multimodal_num_items = {}  # key -> list of item counts per sequence
         for mm_input_dict in multimodal_train_inputs:
             if mm_input_dict is not None:
                 for key, mm_tensor in mm_input_dict.items():
-                    if key not in multimodal_data:
-                        multimodal_data[key] = mm_tensor
-                        multimodal_num_items[key] = [mm_tensor.size(0)]
-                    else:
-                        multimodal_data[key] = torch.cat([multimodal_data[key], mm_tensor], dim=0)
-                        multimodal_num_items[key].append(mm_tensor.size(0))
-        batch["multimodal_train_inputs"] = multimodal_data
+                    multimodal_lists.setdefault(key, []).append(mm_tensor)
+                    multimodal_num_items.setdefault(key, []).append(mm_tensor.size(0))
+        batch["multimodal_train_inputs"] = {
+            key: _concat_multimodal_tensors(key, tensors) for key, tensors in multimodal_lists.items()
+        }
         batch["multimodal_num_items"] = multimodal_num_items
 
     return batch
 
 
+# Audio processor tensors have a per-sample mel-frame time axis (last dim); samples in one
+# micro-batch differ, so right-pad to the max before dim-0 concat. Zero padding is exactly the
+# HF processor's own within-sample padding scheme: feature_lens = feature_attention_mask.sum(-1)
+# stays correct, and the encoder never reads the padded frames.
+_PAD_LAST_DIM_MM_KEYS = ("input_features", "feature_attention_mask")
+
+
+def _concat_multimodal_tensors(key: str, tensors: list[torch.Tensor]) -> torch.Tensor:
+    if key in _PAD_LAST_DIM_MM_KEYS:
+        max_len = max(t.size(-1) for t in tensors)
+        tensors = [F.pad(t, (0, max_len - t.size(-1)), value=0) for t in tensors]
+    return torch.cat(tensors, dim=0)
+
+
 class DataIterator:
     """Micro-batch iterator over rollout dicts.
 
diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py
index 0cabc2f20f1..d4469f6de03 100644
--- a/miles/utils/arguments.py
+++ b/miles/utils/arguments.py
@@ -823,6 +823,16 @@ def add_data_arguments(parser):
                     'JSON string for multimodal data mapping media types to data keys. Example: \'{"image": "image_file"}\''
                 ),
             )
+            parser.add_argument(
+                "--qwen3-omni-audio-encoder-path",
+                type=str,
+                default=None,
+                help=(
+                    "Path to the full Qwen3-Omni HF checkpoint; enables frozen-audio-tower embedding "
+                    "injection into the extracted text-backbone trainer (audio-input RL). "
+                    "See miles_plugins/models/qwen3_omni_thinker.py."
+                ),
+            )
             parser.add_argument("--metadata-key", type=str, default="metadata", help="JSON dataset key")
             parser.add_argument(
                 "--tool-key",
diff --git a/miles_plugins/models/qwen3_omni_thinker.py b/miles_plugins/models/qwen3_omni_thinker.py
new file mode 100644
index 00000000000..9f2abe75630
--- /dev/null
+++ b/miles_plugins/models/qwen3_omni_thinker.py
@@ -0,0 +1,174 @@
+"""Frozen-audio-encoder embedding injection for the Qwen3-Omni thinker text backbone.
+
+The trainer holds only the extracted Qwen3-MoE text model; audio placeholder tokens in a
+multimodal prompt would otherwise embed as ordinary vocab rows and the recomputed
+logprobs would be garbage. This patch fills those positions from the omni checkpoint's
+frozen audio tower at forward time, mirroring HF `get_audio_features` +
+`masked_scatter`. Audio-only prompts keep plain sequential position ids (HF
+`get_rope_index` only takes the mrope branch when image/video grids are present), so the
+text backbone's RoPE needs no change — images/videos are NOT supported here and are
+rejected loudly.
+
+The tower is deliberately not a registered submodule: it stays out of DDP buckets, the
+optimizer, and (dist-)checkpoints; it is lazily loaded once per process and shared.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+AUDIO_KWARG_KEYS = ("input_features", "feature_attention_mask", "audio_feature_lengths")
+_UNSUPPORTED_MM_KEYS = (
+    "pixel_values",
+    "pixel_values_videos",
+    "image_grid_thw",
+    "video_grid_thw",
+    "video_second_per_grid",
+)
+
+_ENCODER_CACHE: dict[tuple[str, str, torch.dtype], torch.nn.Module] = {}
+
+
+def resolve_audio_token_id(omni_config: dict) -> int:
+    thinker = omni_config.get("thinker_config", omni_config)
+    audio_token_id = thinker.get("audio_token_id")
+    assert audio_token_id is not None, "audio_token_id not found in omni config (thinker_config)"
+    return int(audio_token_id)
+
+
+def load_frozen_audio_encoder(omni_checkpoint: str | Path, device, dtype) -> torch.nn.Module:
+    """Build the audio tower from `thinker.audio_tower.*` tensors of the full omni checkpoint.
+
+    Reads only the tower's tensors (safetensors are lazy), so the 60 GB text/talker
+    weights never touch RAM.
+    """
+    from safetensors import safe_open
+    from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeAudioEncoderConfig
+    from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeAudioEncoder
+
+    src = Path(omni_checkpoint)
+    with open(src / "config.json") as f:
+        omni_config = json.load(f)
+    thinker = omni_config.get("thinker_config", omni_config)
+    audio_config = Qwen3OmniMoeAudioEncoderConfig(**thinker["audio_config"])
+
+    index_path = src / "model.safetensors.index.json"
+    if index_path.exists():
+        with open(index_path) as f:
+            weight_map = json.load(f)["weight_map"]
+        shard_files = sorted({shard for name, shard in weight_map.items() if name.startswith("thinker.audio_tower.")})
+    else:
+        shard_files = ["model.safetensors"]
+
+    prefix = "thinker.audio_tower."
+    state_dict = {}
+    for shard_file in shard_files:
+        with safe_open(src / shard_file, framework="pt") as reader:
+            for key in reader.keys():
+                if key.startswith(prefix):
+                    state_dict[key[len(prefix) :]] = reader.get_tensor(key)
+    assert state_dict, f"no {prefix}* tensors found under {src}"
+
+    encoder = Qwen3OmniMoeAudioEncoder(audio_config)
+    encoder.load_state_dict(state_dict, strict=True)
+    encoder = encoder.to(device=device, dtype=dtype)
+    encoder.requires_grad_(False)
+    encoder.eval()
+    logger.info(f"loaded frozen Qwen3-Omni audio tower from {src} ({sum(p.numel() for p in encoder.parameters()):,} params)")
+    return encoder
+
+
+def compute_audio_embeddings(encoder, input_features, feature_attention_mask, audio_feature_lengths) -> torch.Tensor:
+    """Mirror HF Qwen3OmniMoeThinker.get_audio_features: [num_audios, mel, T] -> [total_audio_tokens, output_dim]."""
+    with torch.no_grad():
+        if feature_attention_mask is not None:
+            feature_lens = feature_attention_mask.sum(dim=1)
+            input_features = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
+        else:
+            assert audio_feature_lengths is not None, "need feature_attention_mask or audio_feature_lengths"
+            feature_lens = audio_feature_lengths
+            input_features = input_features.permute(0, 2, 1).reshape(-1, input_features.shape[1]).permute(1, 0)
+        encoder_param = next(encoder.parameters())
+        input_features = input_features.to(device=encoder_param.device, dtype=encoder_param.dtype)
+        feature_lens = feature_lens.to(encoder_param.device)
+        return encoder(input_features, feature_lens=feature_lens).last_hidden_state
+
+
+def scatter_audio_embeddings(hidden, input_ids, audio_embeds, audio_token_id) -> torch.Tensor:
+    """Replace hidden rows at audio placeholder positions, out of place.
+
+    hidden: [s, b, h] (mcore embedding layout); input_ids: [b, s]; audio_embeds: [n, h].
+    Requires the packed b == 1 layout: masked_scatter fills s-major, which would
+    interleave samples for b > 1.
+    """
+    assert hidden.dim() == 3 and input_ids.dim() == 2, f"{hidden.shape=} {input_ids.shape=}"
+    assert hidden.size(1) == 1 and input_ids.size(0) == 1, f"audio injection requires the packed [1, s] layout, got {input_ids.shape}"
+    mask = input_ids[0] == audio_token_id  # [s]
+    num_positions = int(mask.sum())
+    assert num_positions == audio_embeds.size(0), (
+        f"audio placeholder/embedding mismatch: {num_positions} audio tokens in input_ids "
+        f"vs {audio_embeds.size(0)} encoder outputs"
+    )
+    audio_embeds = audio_embeds.to(device=hidden.device, dtype=hidden.dtype)
+    return hidden.masked_scatter(mask.view(-1, 1, 1).to(hidden.device), audio_embeds)
+
+
+def install_audio_injection(model, args, encoder_loader=None, audio_token_id: int | None = None):
+    """Patch `model.forward` to inject frozen-audio-tower embeddings via decoder_input.
+
+    Text-only batches pass through untouched (the encoder is not even loaded). Non-first
+    PP stages swallow the audio kwargs and pass through. Keeps the module identity
+    unchanged, so DDP buckets, the optimizer, and checkpoints see the same model.
+    """
+    checkpoint_path = getattr(args, "qwen3_omni_audio_encoder_path", None)
+    if audio_token_id is None:
+        with open(Path(checkpoint_path) / "config.json") as f:
+            audio_token_id = resolve_audio_token_id(json.load(f))
+    if encoder_loader is None:
+
+        def encoder_loader(device, dtype):
+            key = (str(checkpoint_path), str(device), dtype)
+            if key not in _ENCODER_CACHE:
+                _ENCODER_CACHE[key] = load_frozen_audio_encoder(checkpoint_path, device=device, dtype=dtype)
+            return _ENCODER_CACHE[key]
+
+    orig_forward = model.forward
+
+    def forward(*fargs, **kwargs):
+        audio_kwargs = {k: kwargs.pop(k) for k in AUDIO_KWARG_KEYS if k in kwargs}
+        unsupported = [k for k in _UNSUPPORTED_MM_KEYS if k in kwargs]
+        assert not unsupported, (
+            f"qwen3_omni_thinker audio injection is audio-only; got unsupported multimodal keys {unsupported} "
+            "(image/video need the mrope + deepstack path)"
+        )
+        input_features = audio_kwargs.get("input_features")
+        if input_features is None or not getattr(model, "pre_process", True):
+            return orig_forward(*fargs, **kwargs)
+
+        assert not fargs, "audio injection expects keyword-only forward calls"
+        assert "decoder_input" not in kwargs, "decoder_input already set upstream"
+        assert not getattr(args, "sequence_parallel", False), (
+            "audio injection scatters full-sequence embeddings; run with sequence-parallel off"
+        )
+        assert getattr(args, "context_parallel_size", 1) == 1, "audio injection requires context_parallel_size == 1"
+
+        input_ids = kwargs["input_ids"]
+        hidden = model.embedding(input_ids=input_ids, position_ids=kwargs.get("position_ids"))
+        encoder = encoder_loader(device=hidden.device, dtype=hidden.dtype)
+        audio_embeds = compute_audio_embeddings(
+            encoder,
+            input_features,
+            audio_kwargs.get("feature_attention_mask"),
+            audio_kwargs.get("audio_feature_lengths"),
+        )
+        kwargs["decoder_input"] = scatter_audio_embeddings(hidden, input_ids, audio_embeds, audio_token_id)
+        return orig_forward(**kwargs)
+
+    model.forward = forward
+    return model
diff --git a/tests/fast/backends/training_utils/test_get_batch_multimodal_audio.py b/tests/fast/backends/training_utils/test_get_batch_multimodal_audio.py
new file mode 100644
index 00000000000..1a3d588a6de
--- /dev/null
+++ b/tests/fast/backends/training_utils/test_get_batch_multimodal_audio.py
@@ -0,0 +1,108 @@
+"""get_batch: variable-length audio processor tensors must pad-and-concat, not crash."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import miles.backends.training_utils.cp_utils as cp_utils_mod
+import miles.backends.training_utils.data as data_mod
+from miles.backends.training_utils.data import get_batch
+
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="stage-a-cpu", labels=[])
+
+MEL = 8
+
+
+class _FakeIterator:
+    def __init__(self, batch: dict):
+        self._batch = batch
+        self.rollout_data = {}
+
+    def get_next(self, keys):
+        return {key: self._batch[key] for key in keys if key in self._batch}
+
+
+@pytest.fixture(autouse=True)
+def _stub_cuda(monkeypatch):
+    monkeypatch.setattr(torch.Tensor, "cuda", lambda self, *args, **kwargs: self, raising=False)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: "cpu", raising=False)
+    state = SimpleNamespace(cp=SimpleNamespace(rank=0, size=1), tp=SimpleNamespace(size=1))
+    monkeypatch.setattr(data_mod, "get_parallel_state", lambda: state)
+    monkeypatch.setattr(cp_utils_mod, "get_parallel_state", lambda: state)
+
+
+KEYS = ["tokens", "loss_masks", "total_lengths", "response_lengths", "multimodal_train_inputs"]
+
+
+def _make_batch(mm_dicts: list[dict | None]) -> _FakeIterator:
+    lengths = [6, 4][: len(mm_dicts)]
+    return _FakeIterator(
+        {
+            "tokens": [torch.arange(1, n + 1, dtype=torch.long) for n in lengths],
+            "loss_masks": [torch.ones(n // 2, dtype=torch.int) for n in lengths],
+            "total_lengths": lengths,
+            "response_lengths": [n // 2 for n in lengths],
+            "multimodal_train_inputs": mm_dicts,
+        }
+    )
+
+
+def test_variable_length_audio_features_padded_and_concatenated():
+    a_feat = torch.randn(1, MEL, 50)
+    a_mask = torch.ones(1, 50, dtype=torch.long)
+    b_feat = torch.randn(2, MEL, 30)
+    b_mask = torch.ones(2, 30, dtype=torch.long)
+    b_mask[1, 20:] = 0
+    iterator = _make_batch(
+        [
+            {"input_features": a_feat, "feature_attention_mask": a_mask},
+            {"input_features": b_feat, "feature_attention_mask": b_mask},
+        ]
+    )
+
+    batch = get_batch(iterator, KEYS, pad_multiplier=1, qkv_format="thd")
+
+    mm = batch["multimodal_train_inputs"]
+    assert mm["input_features"].shape == (3, MEL, 50)
+    assert torch.equal(mm["input_features"][0], a_feat[0])
+    assert torch.equal(mm["input_features"][1:, :, :30], b_feat)
+    assert torch.all(mm["input_features"][1:, :, 30:] == 0)
+    assert mm["feature_attention_mask"].shape == (3, 50)
+    # per-audio frame counts survive the padding: this is what feature_lens is computed from
+    assert mm["feature_attention_mask"].sum(dim=1).tolist() == [50, 30, 20]
+
+
+def test_same_length_audio_and_image_keys_concat_plainly():
+    mm_a = {
+        "input_features": torch.randn(1, MEL, 30),
+        "feature_attention_mask": torch.ones(1, 30, dtype=torch.long),
+        "pixel_values": torch.randn(4, 10),
+    }
+    mm_b = {
+        "input_features": torch.randn(1, MEL, 30),
+        "feature_attention_mask": torch.ones(1, 30, dtype=torch.long),
+        "pixel_values": torch.randn(2, 10),
+    }
+    iterator = _make_batch([mm_a, mm_b])
+
+    batch = get_batch(iterator, KEYS, pad_multiplier=1, qkv_format="thd")
+
+    mm = batch["multimodal_train_inputs"]
+    assert mm["input_features"].shape == (2, MEL, 30)
+    assert mm["pixel_values"].shape == (6, 10)
+    assert torch.equal(mm["pixel_values"], torch.cat([mm_a["pixel_values"], mm_b["pixel_values"]], dim=0))
+
+
+def test_none_mm_entries_are_skipped():
+    feat = torch.randn(1, MEL, 20)
+    mask = torch.ones(1, 20, dtype=torch.long)
+    iterator = _make_batch([None, {"input_features": feat, "feature_attention_mask": mask}])
+
+    batch = get_batch(iterator, KEYS, pad_multiplier=1, qkv_format="thd")
+
+    mm = batch["multimodal_train_inputs"]
+    assert mm["input_features"].shape == (1, MEL, 20)
+    assert torch.equal(mm["input_features"], feat)
diff --git a/tests/fast/test_qwen3_omni_audio_injection.py b/tests/fast/test_qwen3_omni_audio_injection.py
new file mode 100644
index 00000000000..17ea03d11e6
--- /dev/null
+++ b/tests/fast/test_qwen3_omni_audio_injection.py
@@ -0,0 +1,276 @@
+"""Frozen-audio-encoder injection for the Qwen3-Omni thinker text backbone."""
+
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=60, suite="stage-a-cpu", labels=[])
+
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeAudioEncoderConfig
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+    Qwen3OmniMoeAudioEncoder,
+    _get_feat_extract_output_lengths,
+)
+
+from miles_plugins.models.qwen3_omni_thinker import (
+    compute_audio_embeddings,
+    install_audio_injection,
+    load_frozen_audio_encoder,
+    resolve_audio_token_id,
+    scatter_audio_embeddings,
+)
+
+AUDIO_TOKEN_ID = 151675
+HIDDEN = 16
+MEL = 8
+
+
+def _tiny_encoder_config() -> Qwen3OmniMoeAudioEncoderConfig:
+    return Qwen3OmniMoeAudioEncoderConfig(
+        d_model=32,
+        encoder_layers=2,
+        encoder_attention_heads=4,
+        encoder_ffn_dim=64,
+        num_mel_bins=MEL,
+        output_dim=HIDDEN,
+        downsample_hidden_size=24,
+    )
+
+
+@pytest.fixture(scope="module")
+def tiny_encoder() -> Qwen3OmniMoeAudioEncoder:
+    torch.manual_seed(0)
+    return Qwen3OmniMoeAudioEncoder(_tiny_encoder_config()).eval()
+
+
+def _audio_batch(tiny_encoder, num_frames: int = 40):
+    """input_features [1, mel, T] + full attention mask, and the token count it expands to."""
+    torch.manual_seed(1)
+    input_features = torch.randn(1, MEL, num_frames)
+    feature_attention_mask = torch.ones(1, num_frames, dtype=torch.long)
+    # the module-level formula is what the processor uses for placeholder counts
+    out_lens = _get_feat_extract_output_lengths(torch.tensor([num_frames]))
+    return input_features, feature_attention_mask, int(out_lens.item())
+
+
+class _StubGPT(torch.nn.Module):
+    """Duck-typed mcore GPT chunk: embedding + forward, [s, b, h] layout."""
+
+    def __init__(self, pre_process: bool = True):
+        super().__init__()
+        self.pre_process = pre_process
+        torch.manual_seed(2)
+        self.word_embeddings = torch.nn.Embedding(152064, HIDDEN)
+        self.captured = None
+
+    def embedding(self, input_ids, position_ids=None):
+        return self.word_embeddings(input_ids).transpose(0, 1)  # [s, b, h]
+
+    def forward(self, **kwargs):
+        self.captured = kwargs
+        return torch.zeros(1)
+
+
+def _args(**overrides):
+    defaults = dict(sequence_parallel=False, context_parallel_size=1, qwen3_omni_audio_encoder_path="/nonexistent")
+    defaults.update(overrides)
+    return SimpleNamespace(**defaults)
+
+
+def _install(model, tiny_encoder, **arg_overrides):
+    return install_audio_injection(
+        model,
+        _args(**arg_overrides),
+        encoder_loader=lambda device, dtype: tiny_encoder,
+        audio_token_id=AUDIO_TOKEN_ID,
+    )
+
+
+def _packed_input_ids(audio_token_counts: list[int], text_len: int = 3) -> torch.Tensor:
+    """[1, s] packed sequence: text_len text tokens then N audio tokens, per sample."""
+    ids = []
+    for n in audio_token_counts:
+        ids += [7] * text_len + [AUDIO_TOKEN_ID] * n
+    return torch.tensor([ids], dtype=torch.long)
+
+
+# ------------------------------- scatter -------------------------------
+
+
+def test_scatter_replaces_audio_positions_in_order():
+    input_ids = _packed_input_ids([2, 3])
+    s = input_ids.size(1)
+    hidden = torch.arange(s * HIDDEN, dtype=torch.float32).reshape(s, 1, HIDDEN)
+    audio_embeds = -torch.arange(1, 5 * HIDDEN + 1, dtype=torch.float32).reshape(5, HIDDEN)
+
+    out = scatter_audio_embeddings(hidden, input_ids, audio_embeds, AUDIO_TOKEN_ID)
+
+    mask = (input_ids[0] == AUDIO_TOKEN_ID).tolist()
+    audio_row = 0
+    for pos, is_audio in enumerate(mask):
+        if is_audio:
+            assert torch.equal(out[pos, 0], audio_embeds[audio_row]), pos
+            audio_row += 1
+        else:
+            assert torch.equal(out[pos, 0], hidden[pos, 0]), pos
+    assert audio_row == 5
+    assert torch.equal(hidden[3, 0], torch.arange(3 * HIDDEN, 4 * HIDDEN, dtype=torch.float32)), "input mutated"
+
+
+def test_scatter_count_mismatch_fails_loud():
+    input_ids = _packed_input_ids([2])
+    hidden = torch.zeros(input_ids.size(1), 1, HIDDEN)
+    with pytest.raises(AssertionError, match="audio"):
+        scatter_audio_embeddings(hidden, input_ids, torch.zeros(3, HIDDEN), AUDIO_TOKEN_ID)
+
+
+def test_scatter_requires_packed_batch_dim_one():
+    input_ids = torch.full((2, 4), AUDIO_TOKEN_ID)
+    hidden = torch.zeros(4, 2, HIDDEN)
+    with pytest.raises(AssertionError, match="packed"):
+        scatter_audio_embeddings(hidden, input_ids, torch.zeros(8, HIDDEN), AUDIO_TOKEN_ID)
+
+
+# ------------------------------- encoder plumbing -------------------------------
+
+
+def test_compute_audio_embeddings_matches_direct_encoder_call(tiny_encoder):
+    input_features, feature_attention_mask, n_tokens = _audio_batch(tiny_encoder)
+
+    embeds = compute_audio_embeddings(tiny_encoder, input_features, feature_attention_mask, None)
+
+    assert embeds.shape == (n_tokens, HIDDEN)
+    lengths = feature_attention_mask.sum(dim=1)
+    flat = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
+    expected = tiny_encoder(flat, feature_lens=lengths).last_hidden_state
+    assert torch.allclose(embeds, expected)
+
+
+def test_compute_audio_embeddings_concatenates_multiple_audios(tiny_encoder):
+    torch.manual_seed(3)
+    input_features = torch.randn(2, MEL, 40)
+    feature_attention_mask = torch.ones(2, 40, dtype=torch.long)
+    feature_attention_mask[1, 24:] = 0  # second audio shorter
+
+    embeds = compute_audio_embeddings(tiny_encoder, input_features, feature_attention_mask, None)
+
+    out_lens = _get_feat_extract_output_lengths(torch.tensor([40, 24]))
+    assert embeds.shape == (int(out_lens.sum()), HIDDEN)
+
+
+# ------------------------------- forward patch -------------------------------
+
+
+def test_install_injects_decoder_input(tiny_encoder):
+    model = _StubGPT()
+    _install(model, tiny_encoder)
+    input_features, feature_attention_mask, n_tokens = _audio_batch(tiny_encoder)
+    input_ids = _packed_input_ids([n_tokens])
+
+    model.forward(
+        input_ids=input_ids,
+        position_ids=None,
+        attention_mask=None,
+        input_features=input_features,
+        feature_attention_mask=feature_attention_mask,
+    )
+
+    captured = model.captured
+    assert "input_features" not in captured and "feature_attention_mask" not in captured
+    decoder_input = captured["decoder_input"]
+    assert decoder_input.shape == (input_ids.size(1), 1, HIDDEN)
+    expected_embeds = compute_audio_embeddings(tiny_encoder, input_features, feature_attention_mask, None)
+    audio_positions = (input_ids[0] == AUDIO_TOKEN_ID).nonzero().squeeze(1)
+    assert torch.allclose(decoder_input[audio_positions, 0], expected_embeds.to(decoder_input.dtype))
+    text_positions = (input_ids[0] != AUDIO_TOKEN_ID).nonzero().squeeze(1)
+    ref = model.embedding(input_ids)
+    assert torch.equal(decoder_input[text_positions, 0], ref[text_positions, 0])
+
+
+def test_install_passthrough_without_audio(tiny_encoder):
+    model = _StubGPT()
+
+    def _fail_loader(device, dtype):
+        raise AssertionError("encoder must not load on text-only batches")
+
+    install_audio_injection(model, _args(), encoder_loader=_fail_loader, audio_token_id=AUDIO_TOKEN_ID)
+    model.forward(input_ids=_packed_input_ids([0]), position_ids=None, attention_mask=None)
+    assert "decoder_input" not in model.captured
+
+    model.forward(input_ids=_packed_input_ids([0]), input_features=None)
+    assert "decoder_input" not in model.captured
+    assert "input_features" not in model.captured
+
+
+def test_install_rejects_vision_keys(tiny_encoder):
+    model = _StubGPT()
+    _install(model, tiny_encoder)
+    with pytest.raises(AssertionError, match="pixel_values"):
+        model.forward(input_ids=_packed_input_ids([1]), pixel_values=torch.zeros(1, 3))
+
+
+def test_install_non_pre_process_swallows_audio_kwargs(tiny_encoder):
+    model = _StubGPT(pre_process=False)
+    _install(model, tiny_encoder)
+    input_features, feature_attention_mask, n_tokens = _audio_batch(tiny_encoder)
+
+    model.forward(
+        input_ids=_packed_input_ids([n_tokens]),
+        input_features=input_features,
+        feature_attention_mask=feature_attention_mask,
+    )
+
+    assert "input_features" not in model.captured
+    assert "decoder_input" not in model.captured
+
+
+def test_install_rejects_sequence_parallel_with_audio(tiny_encoder):
+    model = _StubGPT()
+    _install(model, tiny_encoder, sequence_parallel=True)
+    input_features, feature_attention_mask, n_tokens = _audio_batch(tiny_encoder)
+    with pytest.raises(AssertionError, match="sequence.parallel"):
+        model.forward(
+            input_ids=_packed_input_ids([n_tokens]),
+            input_features=input_features,
+            feature_attention_mask=feature_attention_mask,
+        )
+    # text-only batches stay allowed under SP
+    model.forward(input_ids=_packed_input_ids([0]))
+    assert model.captured is not None
+
+
+# ------------------------------- checkpoint loader -------------------------------
+
+
+def test_load_frozen_audio_encoder_from_omni_checkpoint(tmp_path, tiny_encoder):
+    ckpt = tmp_path / "omni"
+    ckpt.mkdir()
+    config = {
+        "model_type": "qwen3_omni_moe",
+        "thinker_config": {
+            "audio_token_id": AUDIO_TOKEN_ID,
+            "audio_config": _tiny_encoder_config().to_dict(),
+        },
+    }
+    with open(ckpt / "config.json", "w") as f:
+        json.dump(config, f)
+    tensors = {f"thinker.audio_tower.{k}": v for k, v in tiny_encoder.state_dict().items()}
+    tensors["thinker.model.embed_tokens.weight"] = torch.zeros(4, 4)  # decoys the loader must skip
+    tensors["talker.model.norm.weight"] = torch.zeros(4)
+    save_file(tensors, ckpt / "model.safetensors", metadata={"format": "pt"})
+
+    encoder = load_frozen_audio_encoder(ckpt, device=torch.device("cpu"), dtype=torch.float32)
+
+    assert not encoder.training
+    assert all(not p.requires_grad for p in encoder.parameters())
+    input_features, feature_attention_mask, _ = _audio_batch(tiny_encoder)
+    expected = compute_audio_embeddings(tiny_encoder, input_features, feature_attention_mask, None)
+    actual = compute_audio_embeddings(encoder, input_features, feature_attention_mask, None)
+    assert torch.allclose(actual, expected, atol=1e-5)
+
+    assert resolve_audio_token_id(config) == AUDIO_TOKEN_ID

From 11f5c5da435d4fe128270b8261cb2a874ae1912a Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Fri, 24 Jul 2026 02:32:08 +0800
Subject: [PATCH 12/32] [omni] AVQA audio-input example + omni-server
 weight-sync admin dialect

- examples/omni_thinker: prepare_avqa.py converts Joysw909/AVQA (r1aqa jsonl) into the
  miles <audio>-placeholder format, letter label + choices metadata scored by --rm-type gpqa;
  run script reworked to the external-server topology (--rollout-external + addr, server
  doubles as router), audio-tower injection flag, no SP, and --sync-mode skip|distributed
- sglang_engine: --rollout-external-admin-api sglang-omni no-ops /flush_cache and
  /begin|end_weight_update (routes the omni server does not expose; it quiesces+flushes
  internally) and delegates the post-update flush to the update request;
  --rollout-weight-update-stages scopes the three weight-group requests so a multi-stage
  server does not fan group init out to every stage and break NCCL world-size accounting
---
 examples/omni_thinker/prepare_avqa.py         |  67 +++++++++++
 .../run_qwen3_omni_thinker_grpo.py            |  76 ++++++++-----
 miles/backends/sglang_utils/sglang_engine.py  |  55 ++++++---
 miles/rollout/generate_hub/sglang_omni.py     |   6 +-
 miles/utils/arguments.py                      |  21 ++++
 miles_plugins/models/qwen3_omni_thinker.py    |  14 ++-
 .../sglang_utils/test_omni_admin_api.py       | 104 ++++++++++++++++++
 .../test_get_batch_multimodal_audio.py        |   3 +-
 tests/fast/test_prepare_avqa.py               |  79 +++++++++++++
 tools/extract_qwen3_omni_thinker.py           |  10 +-
 10 files changed, 377 insertions(+), 58 deletions(-)
 create mode 100644 examples/omni_thinker/prepare_avqa.py
 create mode 100644 tests/fast/backends/sglang_utils/test_omni_admin_api.py
 create mode 100644 tests/fast/test_prepare_avqa.py

diff --git a/examples/omni_thinker/prepare_avqa.py b/examples/omni_thinker/prepare_avqa.py
new file mode 100644
index 00000000000..129061bff0b
--- /dev/null
+++ b/examples/omni_thinker/prepare_avqa.py
@@ -0,0 +1,67 @@
+"""Convert AVQA (Joysw909/AVQA, r1aqa line-json) into a miles jsonl for audio-input GRPO.
+
+Each output row is the miles multimodal convention: a `<audio>` placeholder prompt (paired
+with --multimodal-keys '{"audio": "audios"}'), the option letter as label, and the raw
+choices in metadata so `--rm-type gpqa` scores letter or option-text answers.
+
+    python examples/omni_thinker/prepare_avqa.py \
+        --src <dataset_dir>/train_r1aqa_line.json --dst <data_dir>/avqa.jsonl \
+        --audio-root <dataset_dir> [--max-samples 5000]
+"""
+
+import argparse
+import json
+import string
+from pathlib import Path
+
+# AVQA questions say "in the video", but the r1aqa recipe feeds the audio track only; the
+# questions are answerable from sound (VGGSound clips).
+_PROMPT_TEMPLATE = (
+    "<audio>{question}\n{options}\n"
+    "Listen to the audio and choose the best option. Respond with the option letter in the form 'Answer: <letter>'."
+)
+_DATASET_PATH_PREFIX = "./Joysw909/AVQA/"
+
+
+def convert_row(row: dict, audio_root: str) -> dict:
+    choices = list(row["multi_choice"])
+    answer_index = int(row["answer"])
+    assert 0 <= answer_index < len(choices) <= len(string.ascii_uppercase), f"bad row: {row}"
+    letters = string.ascii_uppercase[: len(choices)]
+    options = "\n".join(f"{letter}. {choice}" for letter, choice in zip(letters, choices, strict=True))
+    audio_rel = row["audio_path"].removeprefix(_DATASET_PATH_PREFIX)
+    return {
+        "prompt": _PROMPT_TEMPLATE.format(question=row["question_text"].strip(), options=options),
+        "audios": [str(Path(audio_root) / audio_rel)],
+        "label": letters[answer_index],
+        "metadata": {"choices": choices},
+    }
+
+
+def convert_file(src, dst, audio_root: str, max_samples: int | None = None) -> int:
+    n = 0
+    with open(src) as fin, open(dst, "w") as fout:
+        for line in fin:
+            if not line.strip():
+                continue
+            if max_samples is not None and n >= max_samples:
+                break
+            fout.write(json.dumps(convert_row(json.loads(line), audio_root)) + "\n")
+            n += 1
+    assert n > 0, f"no rows converted from {src}"
+    return n
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--src", required=True, help="train_r1aqa_line.json from Joysw909/AVQA")
+    parser.add_argument("--dst", required=True, help="output jsonl path")
+    parser.add_argument("--audio-root", required=True, help="dataset dir containing VGG*/*.wav")
+    parser.add_argument("--max-samples", type=int, default=None)
+    args = parser.parse_args()
+    n = convert_file(args.src, args.dst, audio_root=args.audio_root, max_samples=args.max_samples)
+    print(f"[done] {n} samples -> {args.dst}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
index c2ec90bd191..e10beac52b8 100644
--- a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
+++ b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
@@ -1,33 +1,39 @@
-"""GRPO on the Qwen3-Omni-30B-A3B thinker (text MoE) with DAPO-math reward.
+"""GRPO on the Qwen3-Omni-30B-A3B thinker with audio-input AVQA (MCQ reward).
 
-Off-policy until live weight-sync lands: point --sglang-router-ip/port at a standalone omni
-server; the rollout half serves the frozen base while the trainer updates its own copy, with
-TIS + --get-mismatch-metrics absorbing/measuring the gap.
+Topology: rollout runs on a standalone sglang-omni text server (external engines; miles
+launches nothing locally); the trainer holds the extracted text backbone and injects
+frozen-audio-tower embeddings at placeholder positions (--qwen3-omni-audio-encoder-path).
+
+Weight sync: `--sync-mode skip` freezes the server (off-policy debug; TIS absorbs the gap),
+`--sync-mode distributed` pushes thinker.* weights over NCCL each step (on-policy).
 """
 
 import os
 from dataclasses import dataclass
 from typing import Literal
 
-import typer
 
 import miles.utils.external_utils.command_utils as U
 
 OMNI_MODEL = "Qwen3-Omni-30B-A3B-Instruct"
 THINKER_MODEL = "Qwen3-Omni-30B-A3B-Thinker"
 MEGATRON_MODEL_TYPE = "qwen3-omni-30B-A3B-thinker"
+AVQA_DATASET = "Joysw909/AVQA"
 
 
 @dataclass
 class ScriptArgs(U.ExecuteTrainConfig):
     mode: Literal["normal", "debug_minimal"] = "normal"
+    sync_mode: Literal["skip", "distributed"] = "distributed"
     run_id: str = U.create_run_id()
     num_gpus_per_node: int = 8
     data_dir: str = "/root/datasets"
     model_dir: str = "/root/models"
     megatron_path: str = "/root/Megatron-LM"
-    omni_router_ip: str = "127.0.0.1"
-    omni_router_port: int = 30000
+    omni_server_ip: str = "127.0.0.1"
+    omni_server_port: int = 30000
+    omni_server_tp: int = 4  # TP size of the external omni server (NCCL group world_size = tp + 1)
+    avqa_max_samples: int = 5120
     extra_args: str = ""
 
 
@@ -47,7 +53,13 @@ def prepare(args: ScriptArgs):
         hf_checkpoint=f"{args.model_dir}/{THINKER_MODEL}",
         megatron_path=args.megatron_path,
     )
-    U.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=args.data_dir)
+    avqa_dir = f"{args.data_dir}/{AVQA_DATASET.split('/')[-1]}"
+    U.exec_command(f"hf download {AVQA_DATASET} --repo-type dataset --local-dir {avqa_dir}")
+    U.exec_command(
+        f"python {repo}/examples/omni_thinker/prepare_avqa.py "
+        f"--src {avqa_dir}/train_r1aqa_line.json --dst {args.data_dir}/avqa.jsonl "
+        f"--audio-root {avqa_dir} --max-samples {args.avqa_max_samples}"
+    )
 
 
 def execute(args: ScriptArgs):
@@ -58,39 +70,47 @@ def execute(args: ScriptArgs):
         f"--hf-checkpoint {args.model_dir}/{THINKER_MODEL}/ "
         f"--ref-load {ref_load_path} "
         f"--load {load_save_path} "
-        "--model-name qwen3omni_moe "  # body.* broadcast naming
+        "--model-name qwen3omni_moe "  # thinker.* broadcast naming
     )
 
+    debug_minimal = args.mode == "debug_minimal"
     rollout_args = (
         "--custom-generate-function-path miles.rollout.generate_hub.sglang_omni.generate "
-        f"--prompt-data {args.data_dir}/dapo-math-17k/dapo-math-17k.jsonl "
+        f"--prompt-data {args.data_dir}/avqa.jsonl "
         "--input-key prompt "
         "--label-key label "
+        '--multimodal-keys \'{"audio": "audios"}\' '
         "--apply-chat-template "
         "--rollout-shuffle "
-        "--rm-type dapo "
-        "--reward-key score "
-        "--num-rollout 3000 "
-        "--rollout-batch-size 32 "
-        "--n-samples-per-prompt 8 "
-        f"--rollout-max-response-len {100 if args.mode == 'debug_minimal' else 8192} "
+        "--rm-type gpqa "
+        f"--num-rollout {6 if debug_minimal else args.avqa_max_samples // 32} "
+        f"--rollout-batch-size {8 if debug_minimal else 32} "
+        f"--n-samples-per-prompt {4 if debug_minimal else 8} "
+        f"--rollout-max-response-len {100 if debug_minimal else 2048} "
         "--rollout-temperature 1 "
-        "--global-batch-size 256 "
+        f"--global-batch-size {32 if debug_minimal else 256} "
         "--balance-data "
-        f"--sglang-router-ip {args.omni_router_ip} "
-        f"--sglang-router-port {args.omni_router_port} "
+        # the standalone omni server doubles as the router: the adapter posts straight to it
+        f"--sglang-router-ip {args.omni_server_ip} "
+        f"--sglang-router-port {args.omni_server_port} "
+        "--rollout-external "
+        f"--rollout-external-engine-addrs {args.omni_server_ip}:{args.omni_server_port} "
+        "--rollout-external-admin-api sglang-omni "
+        "--rollout-weight-update-stages thinker "
+        f"--rollout-num-gpus {args.omni_server_tp} "
+        f"--rollout-num-gpus-per-engine {args.omni_server_tp} "
     )
 
-    consistency_args = (
-        "--use-rollout-logprobs "
-        "--get-mismatch-metrics "
-        "--use-tis "
-        "--tis-clip 2.0 "
-    )
+    # frozen audio tower for placeholder-embedding injection (audio-input training)
+    mm_args = f"--qwen3-omni-audio-encoder-path {args.model_dir}/{OMNI_MODEL} "
+
+    consistency_args = "--use-rollout-logprobs " "--get-mismatch-metrics " "--use-tis " "--tis-clip 2.0 "
+    if args.sync_mode == "skip":
+        consistency_args += "--debug-skip-weight-update "
 
     perf_args = (
         "--tensor-model-parallel-size 8 "
-        "--sequence-parallel "
+        # no --sequence-parallel: audio injection scatters full-sequence embeddings
         "--pipeline-model-parallel-size 1 "
         "--context-parallel-size 1 "
         "--expert-model-parallel-size 8 "
@@ -133,12 +153,12 @@ def execute(args: ScriptArgs):
         "--actor-num-nodes 1 "
         f"--actor-num-gpus-per-node {args.num_gpus_per_node} "
         f"--num-gpus-per-node {args.num_gpus_per_node} "
-        f"--rollout-num-gpus {args.num_gpus_per_node} "
     )
 
     train_args = (
         f"{ckpt_args} "
         f"{rollout_args} "
+        f"{mm_args} "
         f"{consistency_args} "
         f"{optimizer_args} "
         f"{grpo_args} "
@@ -168,4 +188,4 @@ def main(args: ScriptArgs):
 
 
 if __name__ == "__main__":
-    typer.run(main)
+    main()
diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py
index 095ca2cd677..a175794035e 100644
--- a/miles/backends/sglang_utils/sglang_engine.py
+++ b/miles/backends/sglang_utils/sglang_engine.py
@@ -430,10 +430,23 @@ def load_lora_adapter_from_distributed(
             payload,
         )
 
+    @property
+    def _omni_admin_api(self) -> bool:
+        """External sglang-omni server: no /flush_cache and no weight-update session routes;
+        the server quiesces and flushes internally around /update_weights_from_distributed."""
+        return getattr(getattr(self, "args", None), "rollout_external_admin_api", "sglang") == "sglang-omni"
+
+    def _weight_update_stages(self) -> list[str] | None:
+        """Stage scoping for multi-stage (omni) servers; unset means the server fans the
+        op out to every registered stage, breaking the NCCL world-size accounting."""
+        return getattr(getattr(self, "args", None), "rollout_weight_update_stages", None)
+
     def flush_cache(self):
         """Flush the cache of the server."""
         if self.node_rank != 0:
             return
+        if self._omni_admin_api:
+            return  # no such route; the omni server flushes inside its update lifecycle
         last_message = None
         for _ in range(60):
             try:
@@ -559,26 +572,24 @@ def update_weights_from_disk(
         return self._make_request("update_weights_from_disk", payload)
 
     def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name, backend):
-        return self._make_request(
-            "init_weights_update_group",
-            {
-                "master_address": master_address,
-                "master_port": master_port,
-                "rank_offset": rank_offset,
-                "world_size": world_size,
-                "group_name": group_name,
-                "backend": backend,
-            },
-        )
+        payload = {
+            "master_address": master_address,
+            "master_port": master_port,
+            "rank_offset": rank_offset,
+            "world_size": world_size,
+            "group_name": group_name,
+            "backend": backend,
+        }
+        if stages := self._weight_update_stages():
+            payload["stages"] = stages
+        return self._make_request("init_weights_update_group", payload)
 
     def destroy_weights_update_group(self, group_name):
+        payload = {"group_name": group_name}
+        if stages := self._weight_update_stages():
+            payload["stages"] = stages
         try:
-            return self._make_request(
-                "destroy_weights_update_group",
-                {
-                    "group_name": group_name,
-                },
-            )
+            return self._make_request("destroy_weights_update_group", payload)
         except requests.exceptions.RequestException:
             # catch the case there the engine is just created and does not have the group.
             pass
@@ -591,8 +602,12 @@ def update_weights_from_distributed(
             "dtypes": [str(dtype).replace("torch.", "") for dtype in dtypes],
             "shapes": shapes,
             "group_name": group_name,
-            "flush_cache": flush_cache,
+            # engine-level flush is a no-op under the omni admin api, so the server-internal
+            # post-update flush must run instead
+            "flush_cache": True if self._omni_admin_api else flush_cache,
         }
+        if stages := self._weight_update_stages():
+            payload["stages"] = stages
         if weight_version is not None:
             payload["weight_version"] = weight_version
         return self._make_request(
@@ -615,10 +630,14 @@ def continue_generation(self):
 
     def begin_weight_update(self):
         """Open a weight-update session on the engine (restores packed weights for loading)."""
+        if self._omni_admin_api:
+            return  # no such route; the omni server runs its own update lifecycle
         return self._make_request("begin_weight_update", {})
 
     def end_weight_update(self):
         """Close the weight-update session (post-load + quant post-process on the full model)."""
+        if self._omni_admin_api:
+            return
         return self._make_request("end_weight_update", {})
 
     def update_weight_version(self, weight_version: str):
diff --git a/miles/rollout/generate_hub/sglang_omni.py b/miles/rollout/generate_hub/sglang_omni.py
index 4552985ce25..8f87791eb12 100644
--- a/miles/rollout/generate_hub/sglang_omni.py
+++ b/miles/rollout/generate_hub/sglang_omni.py
@@ -32,9 +32,9 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
     }, f"{sample.status=}"
     # the omni server declares return_routed_experts/return_indexer_topk in its protocol but
     # implements neither replay; fail loud instead of training on silently missing traces
-    assert not (args.use_rollout_routing_replay or args.use_rollout_indexer_replay), (
-        "sglang-omni rollout has no routing/indexer replay; unset --use-rollout-routing-replay / --use-rollout-indexer-replay"
-    )
+    assert not (
+        args.use_rollout_routing_replay or args.use_rollout_indexer_replay
+    ), "sglang-omni rollout has no routing/indexer replay; unset --use-rollout-routing-replay / --use-rollout-indexer-replay"
     url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate"
 
     prompt_ids = compute_prompt_ids_from_sample(input.state, sample)
diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py
index d4469f6de03..1f78520d78d 100644
--- a/miles/utils/arguments.py
+++ b/miles/utils/arguments.py
@@ -599,6 +599,27 @@ def add_rollout_arguments(parser):
                 nargs="+",
                 help="Address and ports of the external engines.",
             )
+            parser.add_argument(
+                "--rollout-external-admin-api",
+                choices=["sglang", "sglang-omni"],
+                default="sglang",
+                help=(
+                    "Admin/weight-update dialect of the external engines. 'sglang-omni' servers have no "
+                    "/flush_cache or /begin|end_weight_update routes and instead quiesce+flush internally "
+                    "around /update_weights_from_distributed."
+                ),
+            )
+            parser.add_argument(
+                "--rollout-weight-update-stages",
+                type=str,
+                default=None,
+                nargs="+",
+                help=(
+                    "Stage names to scope /init_weights_update_group, /update_weights_from_distributed and "
+                    "/destroy_weights_update_group to, for multi-stage (omni) servers; unset lets the server "
+                    "fan the op out to every registered stage, which breaks the NCCL world-size accounting."
+                ),
+            )
             parser.add_argument(
                 "--update-weight-transfer-mode",
                 choices=["broadcast", "p2p", "disk-delta"],
diff --git a/miles_plugins/models/qwen3_omni_thinker.py b/miles_plugins/models/qwen3_omni_thinker.py
index 9f2abe75630..84ce7470bba 100644
--- a/miles_plugins/models/qwen3_omni_thinker.py
+++ b/miles_plugins/models/qwen3_omni_thinker.py
@@ -80,7 +80,9 @@ def load_frozen_audio_encoder(omni_checkpoint: str | Path, device, dtype) -> tor
     encoder = encoder.to(device=device, dtype=dtype)
     encoder.requires_grad_(False)
     encoder.eval()
-    logger.info(f"loaded frozen Qwen3-Omni audio tower from {src} ({sum(p.numel() for p in encoder.parameters()):,} params)")
+    logger.info(
+        f"loaded frozen Qwen3-Omni audio tower from {src} ({sum(p.numel() for p in encoder.parameters()):,} params)"
+    )
     return encoder
 
 
@@ -108,7 +110,9 @@ def scatter_audio_embeddings(hidden, input_ids, audio_embeds, audio_token_id) ->
     interleave samples for b > 1.
     """
     assert hidden.dim() == 3 and input_ids.dim() == 2, f"{hidden.shape=} {input_ids.shape=}"
-    assert hidden.size(1) == 1 and input_ids.size(0) == 1, f"audio injection requires the packed [1, s] layout, got {input_ids.shape}"
+    assert (
+        hidden.size(1) == 1 and input_ids.size(0) == 1
+    ), f"audio injection requires the packed [1, s] layout, got {input_ids.shape}"
     mask = input_ids[0] == audio_token_id  # [s]
     num_positions = int(mask.sum())
     assert num_positions == audio_embeds.size(0), (
@@ -153,9 +157,9 @@ def forward(*fargs, **kwargs):
 
         assert not fargs, "audio injection expects keyword-only forward calls"
         assert "decoder_input" not in kwargs, "decoder_input already set upstream"
-        assert not getattr(args, "sequence_parallel", False), (
-            "audio injection scatters full-sequence embeddings; run with sequence-parallel off"
-        )
+        assert not getattr(
+            args, "sequence_parallel", False
+        ), "audio injection scatters full-sequence embeddings; run with sequence-parallel off"
         assert getattr(args, "context_parallel_size", 1) == 1, "audio injection requires context_parallel_size == 1"
 
         input_ids = kwargs["input_ids"]
diff --git a/tests/fast/backends/sglang_utils/test_omni_admin_api.py b/tests/fast/backends/sglang_utils/test_omni_admin_api.py
new file mode 100644
index 00000000000..53862fcc685
--- /dev/null
+++ b/tests/fast/backends/sglang_utils/test_omni_admin_api.py
@@ -0,0 +1,104 @@
+"""sglang-omni external-server admin contract: stages scoping + lifecycle no-ops.
+
+The omni server accepts /init_weights_update_group, /update_weights_from_distributed and
+/destroy_weights_update_group with an extra `stages` field (unset => the op fans out to
+EVERY registered stage and breaks the NCCL world-size accounting), quiesces/flushes
+internally during the update, and exposes no /flush_cache, /begin_weight_update or
+/end_weight_update routes.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="stage-a-cpu", labels=[])
+
+pytest.importorskip("sglang")
+from miles.backends.sglang_utils.sglang_engine import SGLangEngine
+
+
+def _engine(monkeypatch, calls, **args_overrides):
+    engine = SGLangEngine.__new__(SGLangEngine)
+    engine.node_rank = 0
+    engine.server_host = "fake-host"
+    engine.server_port = 1234
+    engine.args = SimpleNamespace(
+        rollout_external_admin_api=args_overrides.pop("rollout_external_admin_api", "sglang"),
+        rollout_weight_update_stages=args_overrides.pop("rollout_weight_update_stages", None),
+        **args_overrides,
+    )
+    monkeypatch.setattr(
+        SGLangEngine, "_make_request", lambda self, endpoint, payload=None: calls.append((endpoint, payload))
+    )
+    return engine
+
+
+def _init_group(engine):
+    return engine.init_weights_update_group("1.2.3.4", 29500, 1, 5, "miles-pp_0", "nccl")
+
+
+def test_stages_forwarded_on_all_weight_group_requests(monkeypatch):
+    calls = []
+    engine = _engine(monkeypatch, calls, rollout_weight_update_stages=["thinker"])
+
+    _init_group(engine)
+    engine.update_weights_from_distributed(["w"], ["torch.bfloat16"], [[2, 2]], "miles-pp_0")
+    engine.destroy_weights_update_group("miles-pp_0")
+
+    by_endpoint = dict(calls)
+    assert by_endpoint["init_weights_update_group"]["stages"] == ["thinker"]
+    assert by_endpoint["update_weights_from_distributed"]["stages"] == ["thinker"]
+    assert by_endpoint["destroy_weights_update_group"]["stages"] == ["thinker"]
+
+
+def test_stages_absent_by_default(monkeypatch):
+    calls = []
+    engine = _engine(monkeypatch, calls)
+
+    _init_group(engine)
+    engine.update_weights_from_distributed(["w"], ["torch.bfloat16"], [[2, 2]], "miles-pp_0")
+    engine.destroy_weights_update_group("miles-pp_0")
+
+    for _, payload in calls:
+        assert "stages" not in payload
+
+
+def test_omni_admin_api_noops_missing_lifecycle_routes(monkeypatch):
+    calls = []
+    engine = _engine(monkeypatch, calls, rollout_external_admin_api="sglang-omni")
+    monkeypatch.setattr(
+        "requests.get", lambda *a, **k: (_ for _ in ()).throw(AssertionError("flush_cache must not hit HTTP"))
+    )
+
+    engine.flush_cache()
+    engine.begin_weight_update()
+    engine.end_weight_update()
+
+    assert calls == [], f"omni admin api must not POST missing routes, got {calls}"
+
+
+def test_omni_admin_api_delegates_flush_to_update_request(monkeypatch):
+    calls = []
+    engine = _engine(monkeypatch, calls, rollout_external_admin_api="sglang-omni")
+
+    engine.update_weights_from_distributed(["w"], ["torch.bfloat16"], [[2, 2]], "miles-pp_0")
+
+    [(endpoint, payload)] = calls
+    assert endpoint == "update_weights_from_distributed"
+    # the engine-level flush is a no-op under omni, so the server-internal one must run
+    assert payload["flush_cache"] is True
+
+
+def test_sglang_admin_api_keeps_default_lifecycle(monkeypatch):
+    calls = []
+    engine = _engine(monkeypatch, calls)
+
+    engine.begin_weight_update()
+    engine.end_weight_update()
+    engine.update_weights_from_distributed(["w"], ["torch.bfloat16"], [[2, 2]], "miles-pp_0")
+
+    endpoints = [e for e, _ in calls]
+    assert endpoints == ["begin_weight_update", "end_weight_update", "update_weights_from_distributed"]
+    assert dict(calls)["update_weights_from_distributed"]["flush_cache"] is False
diff --git a/tests/fast/backends/training_utils/test_get_batch_multimodal_audio.py b/tests/fast/backends/training_utils/test_get_batch_multimodal_audio.py
index 1a3d588a6de..3398d1a564b 100644
--- a/tests/fast/backends/training_utils/test_get_batch_multimodal_audio.py
+++ b/tests/fast/backends/training_utils/test_get_batch_multimodal_audio.py
@@ -4,13 +4,12 @@
 
 import pytest
 import torch
+from tests.ci.ci_register import register_cpu_ci
 
 import miles.backends.training_utils.cp_utils as cp_utils_mod
 import miles.backends.training_utils.data as data_mod
 from miles.backends.training_utils.data import get_batch
 
-from tests.ci.ci_register import register_cpu_ci
-
 register_cpu_ci(est_time=10, suite="stage-a-cpu", labels=[])
 
 MEL = 8
diff --git a/tests/fast/test_prepare_avqa.py b/tests/fast/test_prepare_avqa.py
new file mode 100644
index 00000000000..4369e0a8f91
--- /dev/null
+++ b/tests/fast/test_prepare_avqa.py
@@ -0,0 +1,79 @@
+"""AVQA (r1aqa jsonl) -> miles jsonl conversion for the omni thinker audio-RL example."""
+
+import importlib.util
+import json
+from pathlib import Path
+
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[])
+
+_REPO = Path(__file__).resolve().parents[2]
+
+
+def _load_tool():
+    path = _REPO / "examples" / "omni_thinker" / "prepare_avqa.py"
+    spec = importlib.util.spec_from_file_location("prepare_avqa", path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+ROW = {
+    "id": 184,
+    "question_text": "What is the main source of sound in the video?",
+    "multi_choice": ["Car", "motorcycle", "siren", "gun fire"],
+    "answer": 2,
+    "dataset_name": "AVQA",
+    "audio_path": "./Joysw909/AVQA/VGG10000/-HIPq7T3eFI_11.wav",
+}
+
+
+def test_convert_row_maps_answer_index_to_letter_and_audio_path():
+    record = _load_tool().convert_row(ROW, audio_root="/data/avqa")
+
+    assert record["label"] == "C"
+    assert record["audios"] == ["/data/avqa/VGG10000/-HIPq7T3eFI_11.wav"]
+    assert record["metadata"]["choices"] == ["Car", "motorcycle", "siren", "gun fire"]
+    prompt = record["prompt"]
+    assert prompt.startswith("<audio>")
+    assert "A. Car" in prompt and "D. gun fire" in prompt
+    assert record["prompt"].count("<audio>") == 1
+
+
+def test_convert_row_reward_roundtrip_via_gpqa():
+    from miles.rollout.rm_hub.gpqa import compute_gpqa_reward
+
+    record = _load_tool().convert_row(ROW, audio_root="/data/avqa")
+
+    assert compute_gpqa_reward("The answer is C.", record["label"], metadata=record["metadata"]) == 1.0
+    assert compute_gpqa_reward("Answer: siren", record["label"], metadata=record["metadata"]) == 1.0
+    assert compute_gpqa_reward("Answer: B", record["label"], metadata=record["metadata"]) == 0.0
+
+
+def test_convert_file_writes_jsonl(tmp_path):
+    src = tmp_path / "train_r1aqa_line.json"
+    with open(src, "w") as f:
+        f.write(json.dumps(ROW) + "\n")
+        f.write(json.dumps({**ROW, "id": 185, "answer": 0}) + "\n")
+
+    out = tmp_path / "avqa.jsonl"
+    n = _load_tool().convert_file(src, out, audio_root="/data/avqa", max_samples=None)
+
+    assert n == 2
+    lines = [json.loads(line) for line in open(out)]
+    assert lines[0]["label"] == "C"
+    assert lines[1]["label"] == "A"
+
+
+def test_convert_file_respects_max_samples(tmp_path):
+    src = tmp_path / "train_r1aqa_line.json"
+    with open(src, "w") as f:
+        for i in range(5):
+            f.write(json.dumps({**ROW, "id": i}) + "\n")
+
+    out = tmp_path / "avqa.jsonl"
+    n = _load_tool().convert_file(src, out, audio_root="/data/avqa", max_samples=3)
+
+    assert n == 3
+    assert sum(1 for _ in open(out)) == 3
diff --git a/tools/extract_qwen3_omni_thinker.py b/tools/extract_qwen3_omni_thinker.py
index 575cc53d9e8..635898741f0 100644
--- a/tools/extract_qwen3_omni_thinker.py
+++ b/tools/extract_qwen3_omni_thinker.py
@@ -61,8 +61,14 @@ def extract(src, dst, shard_size_gb: float = 5.0) -> int:
         json.dump(synthesize_thinker_config(omni_config), f, indent=2)
 
     for fname in (
-        "tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt",
-        "special_tokens_map.json", "generation_config.json", "chat_template.json", "chat_template.jinja",
+        "tokenizer.json",
+        "tokenizer_config.json",
+        "vocab.json",
+        "merges.txt",
+        "special_tokens_map.json",
+        "generation_config.json",
+        "chat_template.json",
+        "chat_template.jinja",
     ):
         if (src / fname).exists():
             shutil.copy2(src / fname, dst / fname)

From e13d408568f9f143f9095d90bcff5a8f0538cf36 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Fri, 24 Jul 2026 03:03:06 +0800
Subject: [PATCH 13/32] [omni] external-init sanity skip for omni admin api;
 example actor_tp/server_tp accounting

- _init_external: sglang-omni exposes no /get_server_info; health check is the only preflight
- example: decouple actor GPUs (TP=EP=actor_tp) from ray total (actor_tp + omni_server_tp
  external attach bundles); disagg placement, colocate stays off (external servers have no
  sleep/offload routes)
---
 .../omni_thinker/run_qwen3_omni_thinker_grpo.py | 14 +++++++-------
 miles/backends/sglang_utils/sglang_engine.py    |  4 ++++
 .../sglang_utils/test_omni_admin_api.py         | 17 +++++++++++++++++
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
index e10beac52b8..8eabb4650ce 100644
--- a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
+++ b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
@@ -26,7 +26,7 @@ class ScriptArgs(U.ExecuteTrainConfig):
     mode: Literal["normal", "debug_minimal"] = "normal"
     sync_mode: Literal["skip", "distributed"] = "distributed"
     run_id: str = U.create_run_id()
-    num_gpus_per_node: int = 8
+    actor_tp: int = 4  # actor GPUs (TP=EP); ray also reserves omni_server_tp bundles for the external attach
     data_dir: str = "/root/datasets"
     model_dir: str = "/root/models"
     megatron_path: str = "/root/Megatron-LM"
@@ -48,7 +48,7 @@ def prepare(args: ScriptArgs):
     U.convert_checkpoint(
         model_name=THINKER_MODEL,
         megatron_model_type=MEGATRON_MODEL_TYPE,
-        num_gpus_per_node=args.num_gpus_per_node,
+        num_gpus_per_node=args.actor_tp,
         dir_dst=args.model_dir,
         hf_checkpoint=f"{args.model_dir}/{THINKER_MODEL}",
         megatron_path=args.megatron_path,
@@ -109,11 +109,11 @@ def execute(args: ScriptArgs):
         consistency_args += "--debug-skip-weight-update "
 
     perf_args = (
-        "--tensor-model-parallel-size 8 "
+        f"--tensor-model-parallel-size {args.actor_tp} "
         # no --sequence-parallel: audio injection scatters full-sequence embeddings
         "--pipeline-model-parallel-size 1 "
         "--context-parallel-size 1 "
-        "--expert-model-parallel-size 8 "
+        f"--expert-model-parallel-size {args.actor_tp} "
         "--expert-tensor-parallel-size 1 "
         "--recompute-granularity full "
         "--recompute-method uniform "
@@ -151,8 +151,8 @@ def execute(args: ScriptArgs):
         "--attention-softmax-in-fp32 "
         "--attention-backend flash "
         "--actor-num-nodes 1 "
-        f"--actor-num-gpus-per-node {args.num_gpus_per_node} "
-        f"--num-gpus-per-node {args.num_gpus_per_node} "
+        f"--actor-num-gpus-per-node {args.actor_tp} "
+        f"--num-gpus-per-node {args.actor_tp + args.omni_server_tp} "
     )
 
     train_args = (
@@ -170,7 +170,7 @@ def execute(args: ScriptArgs):
 
     U.execute_train(
         train_args=train_args,
-        num_gpus_per_node=args.num_gpus_per_node,
+        num_gpus_per_node=args.actor_tp + args.omni_server_tp,
         megatron_model_type=MEGATRON_MODEL_TYPE,
         train_script="train.py",
         megatron_path=args.megatron_path,
diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py
index a175794035e..f92ff2247db 100644
--- a/miles/backends/sglang_utils/sglang_engine.py
+++ b/miles/backends/sglang_utils/sglang_engine.py
@@ -240,6 +240,10 @@ def _sanity_check_server_args(actual_server_args, expect_server_args):
             api_key=None,
             is_process_alive=lambda: True,
         )
+        if self._omni_admin_api:
+            # omni servers expose no /get_server_info; health is the only preflight
+            logger.info("sglang-omni admin api: skipping /get_server_info sanity check")
+            return
         actual_server_args = _get_actual_server_args()
         _sanity_check_server_args(actual_server_args, expect_server_args)
 
diff --git a/tests/fast/backends/sglang_utils/test_omni_admin_api.py b/tests/fast/backends/sglang_utils/test_omni_admin_api.py
index 53862fcc685..069dfc53874 100644
--- a/tests/fast/backends/sglang_utils/test_omni_admin_api.py
+++ b/tests/fast/backends/sglang_utils/test_omni_admin_api.py
@@ -22,6 +22,7 @@
 def _engine(monkeypatch, calls, **args_overrides):
     engine = SGLangEngine.__new__(SGLangEngine)
     engine.node_rank = 0
+    engine.rank = 0
     engine.server_host = "fake-host"
     engine.server_port = 1234
     engine.args = SimpleNamespace(
@@ -91,6 +92,22 @@ def test_omni_admin_api_delegates_flush_to_update_request(monkeypatch):
     assert payload["flush_cache"] is True
 
 
+def test_omni_admin_api_external_init_skips_server_args_check(monkeypatch):
+    import miles.backends.sglang_utils.sglang_engine as engine_mod
+
+    calls = []
+    engine = _engine(monkeypatch, calls, rollout_external_admin_api="sglang-omni")
+    monkeypatch.setattr(engine_mod, "_wait_server_healthy", lambda **kwargs: calls.append("health"))
+    monkeypatch.setattr(
+        "requests.get",
+        lambda *a, **k: (_ for _ in ()).throw(AssertionError("omni server has no /get_server_info")),
+    )
+
+    engine._init_external({"host": "fake-host", "port": 1234}, external_engine_need_check_fields=["port"])
+
+    assert calls == ["health"]
+
+
 def test_sglang_admin_api_keeps_default_lifecycle(monkeypatch):
     calls = []
     engine = _engine(monkeypatch, calls)

From 28bb1c10ba4795e277c9801744c1537829f171cc Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Fri, 24 Jul 2026 03:05:24 +0800
Subject: [PATCH 14/32] [omni] thinker model args: ffn-hidden-size matches the
 omni config (768, inert for all-MoE)

---
 scripts/models/qwen3-omni-30B-A3B-thinker.sh | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/scripts/models/qwen3-omni-30B-A3B-thinker.sh b/scripts/models/qwen3-omni-30B-A3B-thinker.sh
index a071a2e6e32..de1700ae0e7 100644
--- a/scripts/models/qwen3-omni-30B-A3B-thinker.sh
+++ b/scripts/models/qwen3-omni-30B-A3B-thinker.sh
@@ -1,5 +1,8 @@
 # Megatron model args for the Qwen3-Omni-30B-A3B thinker (text MoE).
-# Architecturally Qwen3-30B-A3B; only vocab differs (152064 vs 151936, from the real config.json).
+# Architecturally Qwen3-30B-A3B; differs in vocab (152064 vs 151936) and the inert dense
+# ffn size: the omni thinker config stamps intermediate_size=768 (== moe_intermediate_size),
+# and every layer is MoE, so ffn-hidden-size never materializes weights - it must simply
+# match the HF config for hf_validate_args.
 
 NLAYERS="${MODEL_ARGS_NUM_LAYERS:-48}"
 FIRST_K_DENSE_REPLACE=0
@@ -25,7 +28,7 @@ MODEL_ARGS=(
    --kv-channels 128
    --num-layers $NLAYERS
    --hidden-size 2048
-   --ffn-hidden-size 6144
+   --ffn-hidden-size 768
 
    --normalization RMSNorm
    --position-embedding-type rope

From 491d5e02e575a171fd91cc8361772b91955b2ea2 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Fri, 24 Jul 2026 03:09:55 +0800
Subject: [PATCH 15/32] [omni] extract tool ships processor artifacts +
 unpacked chat_template.jinja

The trainer loads AutoProcessor/AutoTokenizer from the extracted dir: preprocessor_config
et al. make the omni processor resolvable there, and AutoTokenizer only auto-loads the
.jinja template variant (chat_template.json is processor-level), so unpack it.
---
 .../test_extract_qwen3_omni_thinker.py             | 14 ++++++++++++++
 tools/extract_qwen3_omni_thinker.py                | 11 +++++++++++
 2 files changed, 25 insertions(+)

diff --git a/tests/fast/backends/megatron_utils/test_extract_qwen3_omni_thinker.py b/tests/fast/backends/megatron_utils/test_extract_qwen3_omni_thinker.py
index 47d477c2678..14278824001 100644
--- a/tests/fast/backends/megatron_utils/test_extract_qwen3_omni_thinker.py
+++ b/tests/fast/backends/megatron_utils/test_extract_qwen3_omni_thinker.py
@@ -143,6 +143,20 @@ def test_extract_multi_shard_writes_streaming_index(tmp_path):
         assert torch.equal(out[f"model.layers.{i}.mlp.gate.weight"], torch.full((64, 64), float(i)))
 
 
+def test_extract_unpacks_chat_template_json_for_the_tokenizer(tmp_path):
+    tool = _load_extract_tool()
+    src, dst = tmp_path / "src", tmp_path / "dst"
+    _write_omni_src(src, {"thinker.model.norm.weight": torch.zeros(2)})
+    template = "{% for m in messages %}{{ m.content }}{% endfor %}"
+    with open(src / "chat_template.json", "w") as f:
+        json.dump({"chat_template": template}, f)
+
+    tool.extract(src, dst, shard_size_gb=1.0)
+
+    # AutoTokenizer auto-loads chat_template.jinja; the .json variant is processor-only
+    assert (dst / "chat_template.jinja").read_text() == template
+
+
 def test_extract_fails_loud_when_no_thinker_tensors(tmp_path):
     tool = _load_extract_tool()
     src, dst = tmp_path / "src", tmp_path / "dst"
diff --git a/tools/extract_qwen3_omni_thinker.py b/tools/extract_qwen3_omni_thinker.py
index 635898741f0..ebde57bbaf5 100644
--- a/tools/extract_qwen3_omni_thinker.py
+++ b/tools/extract_qwen3_omni_thinker.py
@@ -69,10 +69,21 @@ def extract(src, dst, shard_size_gb: float = 5.0) -> int:
         "generation_config.json",
         "chat_template.json",
         "chat_template.jinja",
+        # processor artifacts: the trainer loads AutoProcessor from this dir to expand
+        # audio/vision placeholders exactly like the rollout server
+        "preprocessor_config.json",
+        "processor_config.json",
+        "video_preprocessor_config.json",
     ):
         if (src / fname).exists():
             shutil.copy2(src / fname, dst / fname)
 
+    # chat_template.json is a processor-level file; AutoTokenizer only auto-loads the
+    # .jinja variant, and the trainer templates via the tokenizer
+    if not (src / "chat_template.jinja").exists() and (src / "chat_template.json").exists():
+        with open(src / "chat_template.json") as f:
+            (dst / "chat_template.jinja").write_text(json.load(f)["chat_template"])
+
     index_path = src / "model.safetensors.index.json"
     if index_path.exists():
         with open(index_path) as f:

From f1fe3cd68bae6f303c39dcb3dfeb6c78e63620ae Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Fri, 24 Jul 2026 03:12:30 +0800
Subject: [PATCH 16/32] [omni] example: TIS without --use-rollout-logprobs (now
 mutually exclusive upstream)

---
 examples/omni_thinker/run_qwen3_omni_thinker_grpo.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
index 8eabb4650ce..b64719410df 100644
--- a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
+++ b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
@@ -104,7 +104,9 @@ def execute(args: ScriptArgs):
     # frozen audio tower for placeholder-embedding injection (audio-input training)
     mm_args = f"--qwen3-omni-audio-encoder-path {args.model_dir}/{OMNI_MODEL} "
 
-    consistency_args = "--use-rollout-logprobs " "--get-mismatch-metrics " "--use-tis " "--tis-clip 2.0 "
+    # trainer recomputes old logprobs (needed by TIS); rollout logprobs feed the TIS weight
+    # and the mismatch metrics ("--use-rollout-logprobs" is mutually exclusive with TIS)
+    consistency_args = "--get-mismatch-metrics " "--use-tis " "--tis-clip 2.0 "
     if args.sync_mode == "skip":
         consistency_args += "--debug-skip-weight-update "
 

From 3d2e9e6a518c607db008bcac30273dac2f98e7fd Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Fri, 24 Jul 2026 03:14:25 +0800
Subject: [PATCH 17/32] [omni] example: wire the MIS helper for TIS + mismatch
 metrics (now required upstream)

---
 examples/omni_thinker/run_qwen3_omni_thinker_grpo.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
index b64719410df..d1945788c60 100644
--- a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
+++ b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
@@ -105,8 +105,14 @@ def execute(args: ScriptArgs):
     mm_args = f"--qwen3-omni-audio-encoder-path {args.model_dir}/{OMNI_MODEL} "
 
     # trainer recomputes old logprobs (needed by TIS); rollout logprobs feed the TIS weight
-    # and the mismatch metrics ("--use-rollout-logprobs" is mutually exclusive with TIS)
-    consistency_args = "--get-mismatch-metrics " "--use-tis " "--tis-clip 2.0 "
+    # and the mismatch metrics ("--use-rollout-logprobs" is mutually exclusive with TIS);
+    # the MIS helper is the upstream-maintained TIS impl + fine-grained mismatch metrics
+    consistency_args = (
+        "--get-mismatch-metrics "
+        "--use-tis "
+        "--custom-tis-function-path examples.train_infer_mismatch_helper.mis.compute_mis_weights_with_cp "
+        "--custom-config-path examples/train_infer_mismatch_helper/mis.yaml "
+    )
     if args.sync_mode == "skip":
         consistency_args += "--debug-skip-weight-update "
 

From f8c54b8130369f1894042b42957950b2c451b780 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Sun, 26 Jul 2026 03:09:10 +0800
Subject: [PATCH 18/32] [omni] multimodal payload benchmark (video transport
 cost is the scaling risk)

---
 examples/omni_thinker/bench_mm_payload.py | 127 ++++++++++++++++++++++
 tests/fast/test_bench_mm_payload.py       |  54 +++++++++
 2 files changed, 181 insertions(+)
 create mode 100644 examples/omni_thinker/bench_mm_payload.py
 create mode 100644 tests/fast/test_bench_mm_payload.py

diff --git a/examples/omni_thinker/bench_mm_payload.py b/examples/omni_thinker/bench_mm_payload.py
new file mode 100644
index 00000000000..3b092014d30
--- /dev/null
+++ b/examples/omni_thinker/bench_mm_payload.py
@@ -0,0 +1,127 @@
+"""Benchmark the processed-multimodal wire contract (serialize -> JSON -> deserialize).
+
+Quantifies the base64-tensor-in-JSON costs of `multimodal_train_inputs` per modality —
+the known heavy case is video (pixel_values_videos grows with frames x resolution) — so
+the transport decision (keep JSON vs move to a binary/shared-memory path) is made on
+numbers, not vibes.
+
+    python examples/omni_thinker/bench_mm_payload.py [--frames 8 16 32] [--seconds 10 30]
+    python examples/omni_thinker/bench_mm_payload.py --server http://<ip>:<port>  # + HTTP POST timing
+"""
+
+import argparse
+import json
+import time
+
+import torch
+
+from miles.rollout.generate_utils.generate_endpoint_utils import serialize_multimodal_train_inputs
+
+PATCH = 16
+TEMPORAL_PATCH = 2
+MERGE = 2  # spatial merge; grid counts are pre-merge patches per qwen vision geometry
+MEL_BINS = 128
+FRAMES_PER_SECOND_MEL = 100  # whisper-style feature extractor: 10ms hop
+
+
+def make_video_bundle(num_frames: int, height: int, width: int) -> dict[str, torch.Tensor]:
+    grid_t = num_frames // TEMPORAL_PATCH
+    grid_h, grid_w = height // PATCH, width // PATCH
+    patch_dim = 3 * TEMPORAL_PATCH * PATCH * PATCH
+    return {
+        "pixel_values_videos": torch.randn(grid_t * grid_h * grid_w, patch_dim, dtype=torch.float32),
+        "video_grid_thw": torch.tensor([[grid_t, grid_h, grid_w]], dtype=torch.long),
+        "video_second_per_grid": torch.tensor([0.5], dtype=torch.float32),
+    }
+
+
+def make_image_bundle(height: int, width: int) -> dict[str, torch.Tensor]:
+    grid_h, grid_w = height // PATCH, width // PATCH
+    patch_dim = 3 * TEMPORAL_PATCH * PATCH * PATCH
+    return {
+        "pixel_values": torch.randn(grid_h * grid_w, patch_dim, dtype=torch.float32),
+        "image_grid_thw": torch.tensor([[1, grid_h, grid_w]], dtype=torch.long),
+    }
+
+
+def make_audio_bundle(seconds: float) -> dict[str, torch.Tensor]:
+    frames = int(seconds * FRAMES_PER_SECOND_MEL)
+    return {
+        "input_features": torch.randn(1, MEL_BINS, frames, dtype=torch.float32),
+        "feature_attention_mask": torch.ones(1, frames, dtype=torch.long),
+    }
+
+
+def _deserialize(bundle: dict) -> dict[str, torch.Tensor]:
+    """Mirror the server decode (sglang-omni preprocessor.py)."""
+    import pybase64
+
+    out = {}
+    for name, spec in bundle["tensors"].items():
+        raw = bytearray(pybase64.b64decode(spec["data"]))
+        out[name] = torch.frombuffer(raw, dtype=getattr(torch, spec["dtype"])).reshape(spec["shape"])
+    return out
+
+
+def measure_roundtrip(tensors: dict[str, torch.Tensor], server: str | None = None) -> dict:
+    t0 = time.perf_counter()
+    bundle = serialize_multimodal_train_inputs(tensors)
+    t1 = time.perf_counter()
+    payload = json.dumps(bundle)
+    t2 = time.perf_counter()
+    decoded_bundle = json.loads(payload)
+    restored = _deserialize(decoded_bundle)
+    t3 = time.perf_counter()
+
+    report = {
+        "raw_bytes": sum(t.numel() * t.element_size() for t in tensors.values()),
+        "payload_bytes": len(payload),
+        "serialize_ms": (t1 - t0) * 1e3,
+        "json_dump_ms": (t2 - t1) * 1e3,
+        "deserialize_ms": (t3 - t2) * 1e3,
+        "roundtrip_equal": all(torch.equal(restored[k], tensors[k].cpu()) for k in tensors),
+    }
+
+    if server:
+        import requests
+
+        body = {
+            "input_ids": [1, 2, 3],
+            "sampling_params": {"max_new_tokens": 1, "temperature": 0.0},
+            "multimodal_train_inputs": bundle,
+        }
+        t4 = time.perf_counter()
+        resp = requests.post(f"{server}/generate", json=body, timeout=600)
+        report["http_ms"] = (time.perf_counter() - t4) * 1e3
+        report["http_status"] = resp.status_code
+    return report
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--frames", type=int, nargs="+", default=[8, 16, 32, 64])
+    parser.add_argument("--resolution", type=int, default=448)
+    parser.add_argument("--seconds", type=float, nargs="+", default=[10.0, 30.0])
+    parser.add_argument("--server", type=str, default=None, help="optional http://ip:port for POST timing")
+    args = parser.parse_args()
+
+    rows = [("image", f"{args.resolution}px", make_image_bundle(args.resolution, args.resolution))]
+    rows += [(f"video", f"{n}f@{args.resolution}px", make_video_bundle(n, args.resolution, args.resolution)) for n in args.frames]
+    rows += [("audio", f"{s:.0f}s", make_audio_bundle(s)) for s in args.seconds]
+
+    header = f"{'modality':8s} {'case':12s} {'raw_MB':>8s} {'json_MB':>8s} {'ser_ms':>8s} {'dump_ms':>8s} {'deser_ms':>9s}"
+    print(header)
+    for modality, case, tensors in rows:
+        r = measure_roundtrip(tensors, server=args.server)
+        line = (
+            f"{modality:8s} {case:12s} {r['raw_bytes'] / 2**20:8.1f} {r['payload_bytes'] / 2**20:8.1f} "
+            f"{r['serialize_ms']:8.1f} {r['json_dump_ms']:8.1f} {r['deserialize_ms']:9.1f}"
+        )
+        if args.server:
+            line += f"  http={r['http_ms']:.0f}ms({r['http_status']})"
+        assert r["roundtrip_equal"]
+        print(line)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/fast/test_bench_mm_payload.py b/tests/fast/test_bench_mm_payload.py
new file mode 100644
index 00000000000..a1837088100
--- /dev/null
+++ b/tests/fast/test_bench_mm_payload.py
@@ -0,0 +1,54 @@
+"""Multimodal payload benchmark: bundle synthesis + measurement plumbing."""
+
+import importlib.util
+from pathlib import Path
+
+import torch
+
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[])
+
+_REPO = Path(__file__).resolve().parents[2]
+
+
+def _load_tool():
+    path = _REPO / "examples" / "omni_thinker" / "bench_mm_payload.py"
+    spec = importlib.util.spec_from_file_location("bench_mm_payload", path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def test_video_bundle_shapes_follow_qwen_patch_geometry():
+    tool = _load_tool()
+    # 16 frames @ 448x448, patch 16, temporal_patch 2, merge 2 -> grid t=8, h=w=28
+    bundle = tool.make_video_bundle(num_frames=16, height=448, width=448)
+
+    grid = bundle["video_grid_thw"]
+    assert grid.tolist() == [[8, 28, 28]]
+    # flattened patches x (3 * temporal_patch * patch^2)
+    assert bundle["pixel_values_videos"].shape == (8 * 28 * 28, 3 * 2 * 16 * 16)
+    assert bundle["video_second_per_grid"].shape == (1,)
+
+
+def test_image_and_audio_bundles():
+    tool = _load_tool()
+    image = tool.make_image_bundle(height=448, width=448)
+    assert image["image_grid_thw"].tolist() == [[1, 28, 28]]
+    assert image["pixel_values"].shape == (1 * 28 * 28, 3 * 2 * 16 * 16)
+
+    audio = tool.make_audio_bundle(seconds=10.0)
+    assert audio["input_features"].shape[1] == 128  # mel bins
+    assert audio["feature_attention_mask"].shape == audio["input_features"].shape[:1] + audio["input_features"].shape[2:]
+
+
+def test_measure_roundtrip_reports_size_and_integrity():
+    tool = _load_tool()
+    bundle = {"pixel_values": torch.randn(64, 32)}
+
+    report = tool.measure_roundtrip(bundle)
+
+    assert report["payload_bytes"] > 64 * 32 * 4  # base64 expands raw fp32
+    assert report["serialize_ms"] >= 0 and report["deserialize_ms"] >= 0
+    assert report["roundtrip_equal"] is True

From 2dd77723d2f2ee2fba0fea441895181b3d7a3e48 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Sun, 26 Jul 2026 03:16:03 +0800
Subject: [PATCH 19/32] [omni] external attach: probe /health for the omni
 admin api

The default preflight polls /health_generate then /flush_cache; the omni server exposes
neither route, so the 404s retried forever and the external attach hung before rollout.
---
 miles/backends/sglang_utils/sglang_engine.py  | 45 +++++++---------
 .../sglang_utils/test_omni_admin_api.py       | 52 +++++++++++++++++++
 2 files changed, 71 insertions(+), 26 deletions(-)

diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py
index f92ff2247db..2924ff3049e 100644
--- a/miles/backends/sglang_utils/sglang_engine.py
+++ b/miles/backends/sglang_utils/sglang_engine.py
@@ -86,40 +86,32 @@ def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process:
     return p
 
 
-def _wait_server_healthy(base_url, api_key, is_process_alive):
+def _wait_server_healthy(base_url, api_key, is_process_alive, probe_paths=("health_generate", "flush_cache")):
+    """Poll `probe_paths` in order until each returns 200.
+
+    The sglang defaults probe /health_generate then /flush_cache (queue drained before
+    offload); servers without those routes (sglang-omni) must pass their own paths or the
+    404s retry forever.
+    """
     headers = {
         "Content-Type": "application/json; charset=utf-8",
         "Authorization": f"Bearer {api_key}",
     }
 
     with requests.Session() as session:
-        while True:
-            try:
-                response = session.get(f"{base_url}/health_generate", headers=headers)
-                if response.status_code == 200:
-                    break
-            except requests.RequestException:
-                pass
-
-            if not is_process_alive():
-                raise Exception("Server process terminated unexpectedly.")
-
-            time.sleep(2)
-
-        # use flush_cache to make sure the working queue is empty, so that we can do offload
-        while True:
-            try:
-                response = session.get(f"{base_url}/flush_cache", headers=headers)
-                if response.status_code == 200:
-                    break
-
-            except requests.RequestException:
-                pass
+        for path in probe_paths:
+            while True:
+                try:
+                    response = session.get(f"{base_url}/{path}", headers=headers)
+                    if response.status_code == 200:
+                        break
+                except requests.RequestException:
+                    pass
 
-            if not is_process_alive():
-                raise Exception("Server process terminated unexpectedly.")
+                if not is_process_alive():
+                    raise Exception("Server process terminated unexpectedly.")
 
-            time.sleep(2)
+                time.sleep(2)
 
 
 class SGLangEngine(RayActor):
@@ -239,6 +231,7 @@ def _sanity_check_server_args(actual_server_args, expect_server_args):
             base_url=f"http://{self.server_host}:{self.server_port}",
             api_key=None,
             is_process_alive=lambda: True,
+            probe_paths=("health",) if self._omni_admin_api else ("health_generate", "flush_cache"),
         )
         if self._omni_admin_api:
             # omni servers expose no /get_server_info; health is the only preflight
diff --git a/tests/fast/backends/sglang_utils/test_omni_admin_api.py b/tests/fast/backends/sglang_utils/test_omni_admin_api.py
index 069dfc53874..8093fe4cad6 100644
--- a/tests/fast/backends/sglang_utils/test_omni_admin_api.py
+++ b/tests/fast/backends/sglang_utils/test_omni_admin_api.py
@@ -108,6 +108,58 @@ def test_omni_admin_api_external_init_skips_server_args_check(monkeypatch):
     assert calls == ["health"]
 
 
+def test_omni_admin_api_external_init_probes_health_only(monkeypatch):
+    # the omni server exposes /health but neither /health_generate nor /flush_cache;
+    # probing the sglang defaults would retry 404s forever and hang the attach
+    import miles.backends.sglang_utils.sglang_engine as engine_mod
+
+    captured = {}
+
+    def fake_wait(**kwargs):
+        captured.update(kwargs)
+
+    monkeypatch.setattr(engine_mod, "_wait_server_healthy", fake_wait)
+
+    calls = []
+    omni = _engine(monkeypatch, calls, rollout_external_admin_api="sglang-omni")
+    omni._init_external({"host": "fake-host", "port": 1234}, external_engine_need_check_fields=[])
+    assert captured["probe_paths"] == ("health",)
+
+    monkeypatch.setattr(SGLangEngine, "_make_request", lambda self, endpoint, payload=None: None)
+    monkeypatch.setattr(
+        "requests.get",
+        lambda *a, **k: type("R", (), {"status_code": 200, "json": dict, "raise_for_status": lambda self: None})(),
+    )
+    default = _engine(monkeypatch, [], rollout_external_admin_api="sglang")
+    default._init_external({"host": "fake-host", "port": 1234}, external_engine_need_check_fields=[])
+    assert captured["probe_paths"] == ("health_generate", "flush_cache")
+
+
+def test_wait_server_healthy_hits_only_requested_probe_paths(monkeypatch):
+    from miles.backends.sglang_utils.sglang_engine import _wait_server_healthy
+
+    seen = []
+
+    class _FakeSession:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *a):
+            return False
+
+        def get(self, url, headers=None):
+            seen.append(url)
+            return type("R", (), {"status_code": 200})()
+
+    monkeypatch.setattr("requests.Session", _FakeSession)
+
+    _wait_server_healthy(
+        base_url="http://x:1", api_key=None, is_process_alive=lambda: True, probe_paths=("health",)
+    )
+
+    assert seen == ["http://x:1/health"]
+
+
 def test_sglang_admin_api_keeps_default_lifecycle(monkeypatch):
     calls = []
     engine = _engine(monkeypatch, calls)

From dec709d86a926e44a9c7b9e489e4df828c9a9f2d Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Sun, 26 Jul 2026 03:22:02 +0800
Subject: [PATCH 20/32] [omni] adapter: filter sampling params to the omni
 /generate schema

RolloutSamplingParams is extra=forbid; miles' detok-only keys (skip_special_tokens,
no_stop_trim, spaces_between_special_tokens) 422'd every request. Those flags never reach
the token/logprob streams, so drop them; unknown keys still fail loud.
---
 miles/rollout/generate_hub/sglang_omni.py     | 10 ++++
 .../rollout/generate_hub/test_sglang_omni.py  | 57 +++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/miles/rollout/generate_hub/sglang_omni.py b/miles/rollout/generate_hub/sglang_omni.py
index 8f87791eb12..2d89bcb0cce 100644
--- a/miles/rollout/generate_hub/sglang_omni.py
+++ b/miles/rollout/generate_hub/sglang_omni.py
@@ -12,6 +12,13 @@
 """
 
 from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput
+
+# RolloutSamplingParams on the omni server is extra="forbid"; this mirrors its schema
+_OMNI_SAMPLING_KEYS = frozenset(
+    ("temperature", "top_p", "top_k", "min_p", "repetition_penalty", "stop", "stop_token_ids", "seed", "max_new_tokens", "max_tokens")
+)
+# detok/text-shaping flags: they never reach the token or logprob streams the trainer consumes
+_DETOK_ONLY_KEYS = frozenset(("skip_special_tokens", "no_stop_trim", "spaces_between_special_tokens"))
 from miles.rollout.generate_utils.generate_endpoint_utils import (
     compute_prompt_ids_from_sample,
     compute_request_payload,
@@ -71,6 +78,9 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
     payload["return_omni_rollout"] = False
     # the trainer recompute cannot replay a repetition penalty (logprobs would diverge)
     payload["sampling_params"]["repetition_penalty"] = 1.0
+    unknown_keys = set(payload["sampling_params"]) - _OMNI_SAMPLING_KEYS - _DETOK_ONLY_KEYS
+    assert not unknown_keys, f"sampling params outside the omni /generate schema: {sorted(unknown_keys)}"
+    payload["sampling_params"] = {k: v for k, v in payload["sampling_params"].items() if k in _OMNI_SAMPLING_KEYS}
     if sample.metadata:
         payload["metadata"] = sample.metadata
 
diff --git a/tests/fast/rollout/generate_hub/test_sglang_omni.py b/tests/fast/rollout/generate_hub/test_sglang_omni.py
index 87cbf7fbc11..daee9954fff 100644
--- a/tests/fast/rollout/generate_hub/test_sglang_omni.py
+++ b/tests/fast/rollout/generate_hub/test_sglang_omni.py
@@ -273,6 +273,63 @@ async def fake_update(*args, **kwargs):
     assert "metadata" not in payload
 
 
+def test_sglang_omni_adapter_filters_sampling_params_to_omni_schema(monkeypatch):
+    # the omni RolloutSamplingParams is extra="forbid": miles' detok-only keys 422 the request
+    from miles.rollout.generate_hub import sglang_omni
+
+    captured = {}
+
+    async def fake_post(url, payload, headers=None):
+        captured["payload"] = payload
+        return {}
+
+    async def fake_update(*args, **kwargs):
+        return None
+
+    monkeypatch.setattr(sglang_omni, "post", fake_post)
+    monkeypatch.setattr(sglang_omni, "update_sample_from_response", fake_update)
+    monkeypatch.setattr(sglang_omni, "compute_prompt_ids_from_sample", lambda state, sample: [1, 2, 3])
+    generate_input = SimpleNamespace(
+        args=_adapter_args(),
+        sample=Sample(),
+        sampling_params={
+            "temperature": 1.0,
+            "top_p": 1.0,
+            "top_k": -1,
+            "max_new_tokens": 16,
+            "stop": None,
+            "stop_token_ids": None,
+            "skip_special_tokens": False,
+            "no_stop_trim": True,
+            "spaces_between_special_tokens": False,
+        },
+        state=None,
+    )
+
+    asyncio.run(sglang_omni.generate(generate_input))
+
+    sp = captured["payload"]["sampling_params"]
+    for detok_key in ("skip_special_tokens", "no_stop_trim", "spaces_between_special_tokens"):
+        assert detok_key not in sp
+    assert sp["temperature"] == 1.0 and sp["max_new_tokens"] == 16 and sp["top_k"] == -1
+    assert sp["repetition_penalty"] == 1.0
+
+
+def test_sglang_omni_adapter_rejects_unknown_sampling_keys(monkeypatch):
+    from miles.rollout.generate_hub import sglang_omni
+
+    monkeypatch.setattr(sglang_omni, "compute_prompt_ids_from_sample", lambda state, sample: [1, 2, 3])
+    generate_input = SimpleNamespace(
+        args=_adapter_args(),
+        sample=Sample(),
+        sampling_params={"temperature": 1.0, "max_new_tokens": 16, "min_new_tokens": 4},
+        state=None,
+    )
+
+    with pytest.raises(AssertionError, match="min_new_tokens"):
+        asyncio.run(sglang_omni.generate(generate_input))
+
+
 def test_sglang_omni_adapter_forwards_sample_metadata(monkeypatch):
     from miles.rollout.generate_hub import sglang_omni
 

From 2a20bf708aff684f4021ccbb094318dffb299d88 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Sun, 26 Jul 2026 03:38:44 +0800
Subject: [PATCH 21/32] [omni] post-rollout abort for external omni servers:
 pause(abort)+continue

Both abort paths listed workers via the router /workers API and POSTed /abort_request;
a bare omni server has neither route, so the cleanup after a fully successful rollout
404-crashed the job. pause_generation(mode=abort) + continue_generation is the
omni-native abort_all.
---
 .../generate_utils/generate_endpoint_utils.py | 18 ++++++++++
 .../inference_rollout_train.py                | 11 ++++--
 miles/rollout/sglang_rollout.py               | 29 ++++++++-------
 .../rollout/generate_hub/test_sglang_omni.py  | 35 +++++++++++++++++++
 4 files changed, 78 insertions(+), 15 deletions(-)

diff --git a/miles/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py
index 37472d553b5..a45a0631bb6 100644
--- a/miles/rollout/generate_utils/generate_endpoint_utils.py
+++ b/miles/rollout/generate_utils/generate_endpoint_utils.py
@@ -9,6 +9,7 @@
 import pybase64
 import torch
 
+from miles.utils.http_utils import post
 from miles.utils.lora import LORA_ADAPTER_NAME, is_lora_enabled
 from miles.utils.processing_utils import (
     call_processor,
@@ -62,6 +63,23 @@ def compute_prompt_ids_from_sample(state, sample, tools=None):
         return state.tokenizer.encode(prompt, add_special_tokens=False)
 
 
+def is_omni_external_admin(args) -> bool:
+    return getattr(args, "rollout_external_admin_api", "sglang") == "sglang-omni"
+
+
+async def abort_external_omni_requests(args) -> None:
+    """Abort in-flight generation on external sglang-omni servers.
+
+    A bare omni server exposes neither the router /workers listing nor /abort_request;
+    pause_generation(mode="abort") followed by continue_generation is the omni-native
+    equivalent of abort_all.
+    """
+    for addr in args.rollout_external_engine_addrs:
+        base_url = f"http://{addr}"
+        await post(f"{base_url}/pause_generation", {"mode": "abort"})
+        await post(f"{base_url}/continue_generation", {})
+
+
 def policy_uses_routing_key(args) -> bool:
     return args.sglang_router_policy in ("consistent_hashing", "manual")
 
diff --git a/miles/rollout/inference_rollout/inference_rollout_train.py b/miles/rollout/inference_rollout/inference_rollout_train.py
index ac7b9a742b9..cce05f34656 100644
--- a/miles/rollout/inference_rollout/inference_rollout_train.py
+++ b/miles/rollout/inference_rollout/inference_rollout_train.py
@@ -7,6 +7,7 @@
 from packaging.version import parse
 from tqdm import tqdm
 
+from miles.rollout.generate_utils.generate_endpoint_utils import abort_external_omni_requests, is_omni_external_admin
 from miles.rollout.base_types import RolloutFnTrainOutput
 from miles.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter
 from miles.rollout.generate_utils.prefill_logprobs import recompute_samples_rollout_logprobs_via_prefill
@@ -25,9 +26,13 @@ async def abort(state: GenerateState, pendings: set, rollout_id: int) -> list[li
     assert not state.aborted
     state.aborted = True
 
-    urls = await get_worker_urls(args)
-    logger.info(f"Abort request for {urls}")
-    await asyncio.gather(*[post(f"{url}/abort_request", {"abort_all": True}) for url in urls])
+    if is_omni_external_admin(args):
+        # bare omni servers have no router /workers nor /abort_request
+        await abort_external_omni_requests(args)
+    else:
+        urls = await get_worker_urls(args)
+        logger.info(f"Abort request for {urls}")
+        await asyncio.gather(*[post(f"{url}/abort_request", {"abort_all": True}) for url in urls])
 
     # Let the agent integration tear down its in-flight trials so they stop hitting
     # SGLang, instead of running on until their own max_seq_len / timeout.
diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py
index 7e6f60e1829..2b9103a480a 100644
--- a/miles/rollout/sglang_rollout.py
+++ b/miles/rollout/sglang_rollout.py
@@ -13,6 +13,7 @@
 from packaging.version import parse
 from tqdm import tqdm
 
+from miles.rollout.generate_utils.generate_endpoint_utils import abort_external_omni_requests, is_omni_external_admin
 from miles.rollout.base_types import GenerateFnInput, RolloutFnEvalOutput, RolloutFnTrainOutput
 from miles.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter
 from miles.rollout.inference_rollout.compatibility import load_generate_function
@@ -377,19 +378,23 @@ async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]:
     assert not state.aborted
     state.aborted = True
 
-    if parse(sglang_router.__version__) <= parse("0.2.1") or args.use_miles_router:
-        response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/list_workers")
-        urls = response["urls"]
+    if is_omni_external_admin(args):
+        # bare omni servers have no router /workers nor /abort_request
+        await abort_external_omni_requests(args)
     else:
-        response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers")
-        urls = [worker["url"] for worker in response["workers"]]
-
-    logger.info(f"Abort request for {urls}")
-    abort_tasks = [post(f"{url}/abort_request", {"abort_all": True}) for url in urls]
-    abort_results = await asyncio.gather(*abort_tasks, return_exceptions=True)
-    for url, result in zip(urls, abort_results, strict=False):
-        if isinstance(result, Exception):
-            logger.warning(f"Failed to abort worker at {url}: {result}")
+        if parse(sglang_router.__version__) <= parse("0.2.1") or args.use_miles_router:
+            response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/list_workers")
+            urls = response["urls"]
+        else:
+            response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers")
+            urls = [worker["url"] for worker in response["workers"]]
+
+        logger.info(f"Abort request for {urls}")
+        abort_tasks = [post(f"{url}/abort_request", {"abort_all": True}) for url in urls]
+        abort_results = await asyncio.gather(*abort_tasks, return_exceptions=True)
+        for url, result in zip(urls, abort_results, strict=False):
+            if isinstance(result, Exception):
+                logger.warning(f"Failed to abort worker at {url}: {result}")
 
     # Let the agent integration tear down its in-flight trials so they stop hitting
     # SGLang, instead of running on until their own max_seq_len / timeout.
diff --git a/tests/fast/rollout/generate_hub/test_sglang_omni.py b/tests/fast/rollout/generate_hub/test_sglang_omni.py
index daee9954fff..c4f426550e8 100644
--- a/tests/fast/rollout/generate_hub/test_sglang_omni.py
+++ b/tests/fast/rollout/generate_hub/test_sglang_omni.py
@@ -204,6 +204,41 @@ async def fail_post(url, payload, headers=None):
     assert output.samples is sample
 
 
+def test_abort_external_omni_requests_uses_pause_abort_continue(monkeypatch):
+    # a bare omni server has neither the router /workers listing nor /abort_request;
+    # pause(mode=abort) + continue is the omni-native equivalent
+    from miles.rollout.generate_utils import generate_endpoint_utils as geu
+
+    calls = []
+
+    async def fake_post(url, payload, **kwargs):
+        calls.append((url, payload))
+        return {}
+
+    monkeypatch.setattr(geu, "post", fake_post)
+    args = SimpleNamespace(
+        rollout_external_admin_api="sglang-omni",
+        rollout_external_engine_addrs=["10.0.0.1:30111", "10.0.0.2:30111"],
+    )
+
+    asyncio.run(geu.abort_external_omni_requests(args))
+
+    assert calls == [
+        ("http://10.0.0.1:30111/pause_generation", {"mode": "abort"}),
+        ("http://10.0.0.1:30111/continue_generation", {}),
+        ("http://10.0.0.2:30111/pause_generation", {"mode": "abort"}),
+        ("http://10.0.0.2:30111/continue_generation", {}),
+    ]
+
+
+def test_is_omni_external_admin(monkeypatch):
+    from miles.rollout.generate_utils import generate_endpoint_utils as geu
+
+    assert geu.is_omni_external_admin(SimpleNamespace(rollout_external_admin_api="sglang-omni")) is True
+    assert geu.is_omni_external_admin(SimpleNamespace(rollout_external_admin_api="sglang")) is False
+    assert geu.is_omni_external_admin(SimpleNamespace()) is False
+
+
 class TestHarness:
     """Through the real parse_args + mock sglang server harness (generation_fixtures)."""
 

From 9f0e4d40cf1ff91b3ad1423430ee966459af2173 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Sun, 26 Jul 2026 03:56:40 +0800
Subject: [PATCH 22/32] [omni] debug tier: bf16 grad accum (TP1 fits 2xH200;
 mcore requires SP for MoE+TP>1)

---
 .../run_qwen3_omni_thinker_grpo.py            |   5 +-
 .../megatron_utils/test_extract_thinker_vl.py | 211 ++++++++++++++++++
 tools/extract_qwen3_omni_thinker.py           | 199 +++++++++++++++--
 3 files changed, 395 insertions(+), 20 deletions(-)
 create mode 100644 tests/fast/backends/megatron_utils/test_extract_thinker_vl.py

diff --git a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
index d1945788c60..beef11a2d23 100644
--- a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
+++ b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
@@ -155,7 +155,10 @@ def execute(args: ScriptArgs):
     misc_args = (
         "--attention-dropout 0.0 "
         "--hidden-dropout 0.0 "
-        "--accumulate-allreduce-grads-in-fp32 "
+        # mcore requires SP for MoE+TP>1, and the audio injection needs SP off, so the
+        # debug tier runs TP1 (bf16 grad accum keeps 2xH200 within memory)
+        + ("" if debug_minimal else "--accumulate-allreduce-grads-in-fp32 ")
+        + 
         "--attention-softmax-in-fp32 "
         "--attention-backend flash "
         "--actor-num-nodes 1 "
diff --git a/tests/fast/backends/megatron_utils/test_extract_thinker_vl.py b/tests/fast/backends/megatron_utils/test_extract_thinker_vl.py
new file mode 100644
index 00000000000..22a48916d55
--- /dev/null
+++ b/tests/fast/backends/megatron_utils/test_extract_thinker_vl.py
@@ -0,0 +1,211 @@
+"""Pseudo-Qwen3-VL extraction: thinker (visual+text) masquerading as qwen3_vl_moe."""
+
+import importlib.util
+import json
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[])
+
+_REPO = Path(__file__).resolve().parents[4]
+
+
+def _load_tool():
+    path = _REPO / "tools" / "extract_qwen3_omni_thinker.py"
+    spec = importlib.util.spec_from_file_location("extract_qwen3_omni_thinker", path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def test_vl_name_map_text_and_visual():
+    m = _load_tool().map_thinker_param_name_vl
+    assert m("thinker.lm_head.weight") == "lm_head.weight"
+    assert m("thinker.model.embed_tokens.weight") == "model.language_model.embed_tokens.weight"
+    assert m("thinker.model.layers.3.self_attn.q_proj.weight") == "model.language_model.layers.3.self_attn.q_proj.weight"
+    assert m("thinker.model.layers.3.mlp.gate.weight") == "model.language_model.layers.3.mlp.gate.weight"
+    assert m("thinker.visual.patch_embed.proj.weight") == "model.visual.patch_embed.proj.weight"
+    assert m("thinker.visual.blocks.7.attn.qkv.weight") == "model.visual.blocks.7.attn.qkv.weight"
+    # merger renames: ln_q -> norm, mlp.0/2 -> linear_fc1/2, merger_list -> deepstack_merger_list
+    assert m("thinker.visual.merger.ln_q.weight") == "model.visual.merger.norm.weight"
+    assert m("thinker.visual.merger.mlp.0.bias") == "model.visual.merger.linear_fc1.bias"
+    assert m("thinker.visual.merger.mlp.2.weight") == "model.visual.merger.linear_fc2.weight"
+    assert m("thinker.visual.merger_list.1.ln_q.weight") == "model.visual.deepstack_merger_list.1.norm.weight"
+    assert m("thinker.visual.merger_list.2.mlp.2.bias") == "model.visual.deepstack_merger_list.2.linear_fc2.bias"
+    # drops
+    assert m("thinker.audio_tower.layers.0.fc1.weight") is None
+    assert m("talker.model.norm.weight") is None
+    assert m("code2wav.decoder.weight") is None
+
+
+def test_vl_expert_fusion_hub_layout():
+    tool = _load_tool()
+    # E=2 experts, I=3 (moe ffn), H=4 (hidden)
+    gate = [torch.arange(12, dtype=torch.float32).reshape(3, 4) + 100 * e for e in range(2)]
+    up = [torch.arange(12, dtype=torch.float32).reshape(3, 4) + 100 * e + 50 for e in range(2)]
+    down = [torch.arange(12, dtype=torch.float32).reshape(4, 3) + 100 * e for e in range(2)]
+
+    gate_up_fused, down_fused = tool.fuse_layer_experts_vl(gate, up, down)
+
+    # hub layout: gate_up [E, H, 2I], down [E, I, H]
+    assert gate_up_fused.shape == (2, 4, 6)
+    assert down_fused.shape == (2, 3, 4)
+    for e in range(2):
+        expected = torch.cat([gate[e], up[e]], dim=0).T  # [H, 2I]
+        assert torch.equal(gate_up_fused[e], expected)
+        assert torch.equal(down_fused[e], down[e].T)
+
+
+def test_vl_config_synthesis():
+    synth = _load_tool().synthesize_vl_config
+    omni = {
+        "model_type": "qwen3_omni_moe",
+        "thinker_config": {
+            "image_token_id": 151655,
+            "video_token_id": 151656,
+            "vision_start_token_id": 151652,
+            "vision_end_token_id": 151653,
+            "position_id_per_seconds": 13,
+            "audio_token_id": 151675,
+            "tie_word_embeddings": False,
+            "text_config": {
+                "model_type": "qwen3_omni_moe_text",
+                "vocab_size": 152064,
+                "hidden_size": 2048,
+                "num_hidden_layers": 48,
+                "rope_theta": 1000000,
+                "max_position_embeddings": 65536,
+                "rope_scaling": {"interleaved": True, "mrope_interleaved": True, "mrope_section": [24, 20, 20], "type": "default"},
+                "num_experts": 128,
+            },
+            "vision_config": {
+                "model_type": "qwen3_omni_moe_vision_encoder",
+                "depth": 27,
+                "hidden_size": 1152,
+                "deepstack_visual_indexes": [8, 16, 24],
+                "apply_vit_abs_pos_embed": True,
+                "image_size": 768,
+                "spatial_patch_size": 16,
+                "in_chans": 3,
+                "tokens_per_second": 2,
+                "out_hidden_size": 2048,
+                "spatial_merge_size": 2,
+                "patch_size": 16,
+                "temporal_patch_size": 2,
+            },
+            "audio_config": {"d_model": 1280},
+        },
+    }
+
+    cfg = synth(omni)
+
+    assert cfg["model_type"] == "qwen3_vl_moe"
+    assert cfg["architectures"] == ["Qwen3VLMoeForConditionalGeneration"]
+    assert cfg["tie_word_embeddings"] is False
+    for key in ("image_token_id", "video_token_id", "vision_start_token_id", "vision_end_token_id"):
+        assert cfg[key] == omni["thinker_config"][key]
+    text = cfg["text_config"]
+    assert text["model_type"] == "qwen3_vl_moe_text"
+    assert text["vocab_size"] == 152064 and text["rope_theta"] == 1000000
+    assert text["max_position_embeddings"] == 65536
+    assert text["rope_scaling"] == {"rope_type": "default", "mrope_section": [24, 20, 20], "mrope_interleaved": True}
+    assert text["bos_token_id"] == 151643 and text["eos_token_id"] == 151645
+    vision = cfg["vision_config"]
+    assert vision["model_type"] == "qwen3_vl_moe"
+    assert vision["num_position_embeddings"] == 2304  # (image_size // patch_size) ** 2
+    for gone in ("apply_vit_abs_pos_embed", "image_size", "spatial_patch_size", "in_chans", "tokens_per_second"):
+        assert gone not in vision
+    assert "audio_config" not in cfg
+    # sideband for the trainer-side omni video rope (no VL slot for it)
+    assert cfg["omni_sideband"] == {"position_id_per_seconds": 13, "audio_token_id": 151675}
+
+
+def _tiny_omni_src(src: Path):
+    src.mkdir(parents=True, exist_ok=True)
+    thinker = {
+        "image_token_id": 151655,
+        "video_token_id": 151656,
+        "vision_start_token_id": 151652,
+        "vision_end_token_id": 151653,
+        "position_id_per_seconds": 13,
+        "audio_token_id": 151675,
+        "tie_word_embeddings": False,
+        "text_config": {
+            "model_type": "qwen3_omni_moe_text",
+            "vocab_size": 32,
+            "num_experts": 2,
+            "rope_scaling": {"interleaved": True, "mrope_interleaved": True, "mrope_section": [24, 20, 20], "type": "default"},
+        },
+        "vision_config": {"model_type": "x", "image_size": 32, "patch_size": 16},
+    }
+    with open(src / "config.json", "w") as f:
+        json.dump({"thinker_config": thinker}, f)
+    tensors = {
+        "thinker.model.embed_tokens.weight": torch.randn(32, 4),
+        "thinker.lm_head.weight": torch.randn(32, 4),
+        "thinker.visual.merger.ln_q.weight": torch.randn(6),
+        "thinker.audio_tower.proj1.weight": torch.zeros(2, 2),
+    }
+    for e in range(2):
+        tensors[f"thinker.model.layers.0.mlp.experts.{e}.gate_proj.weight"] = torch.randn(3, 4)
+        tensors[f"thinker.model.layers.0.mlp.experts.{e}.up_proj.weight"] = torch.randn(3, 4)
+        tensors[f"thinker.model.layers.0.mlp.experts.{e}.down_proj.weight"] = torch.randn(4, 3)
+    save_file(tensors, src / "model.safetensors", metadata={"format": "pt"})
+    return tensors
+
+
+def test_extract_vl_roundtrip(tmp_path):
+    tool = _load_tool()
+    src, dst = tmp_path / "src", tmp_path / "dst"
+    tensors = _tiny_omni_src(src)
+
+    tool.extract(src, dst, shard_size_gb=1.0, variant="vl")
+
+    with open(dst / "config.json") as f:
+        cfg = json.load(f)
+    assert cfg["model_type"] == "qwen3_vl_moe"
+
+    out = {}
+    with safe_open(dst / "model.safetensors", framework="pt") as reader:
+        for key in reader.keys():
+            out[key] = reader.get_tensor(key)
+
+    assert "model.language_model.embed_tokens.weight" in out
+    assert "lm_head.weight" in out
+    assert "model.visual.merger.norm.weight" in out
+    assert not any("audio_tower" in k for k in out)
+    gate_up = out["model.language_model.layers.0.mlp.experts.gate_up_proj"]
+    down = out["model.language_model.layers.0.mlp.experts.down_proj"]
+    assert gate_up.shape == (2, 4, 6) and down.shape == (2, 3, 4)
+    expected_e1 = torch.cat(
+        [
+            tensors["thinker.model.layers.0.mlp.experts.1.gate_proj.weight"],
+            tensors["thinker.model.layers.0.mlp.experts.1.up_proj.weight"],
+        ],
+        dim=0,
+    ).T
+    assert torch.equal(gate_up[1], expected_e1)
+    # per-expert 2D names must NOT leak through
+    assert not any(".experts.0." in k or ".experts.1." in k for k in out)
+
+
+def test_extract_text_variant_unchanged(tmp_path):
+    tool = _load_tool()
+    src, dst = tmp_path / "src", tmp_path / "dst"
+    _tiny_omni_src(src)
+
+    tool.extract(src, dst, shard_size_gb=1.0, variant="text")
+
+    with open(dst / "config.json") as f:
+        assert json.load(f)["model_type"] == "qwen3_moe"
+    with safe_open(dst / "model.safetensors", framework="pt") as reader:
+        keys = list(reader.keys())
+    assert "model.embed_tokens.weight" in keys
+    assert any(".experts.0.gate_proj" in k for k in keys)  # text variant keeps per-expert 2D
+    assert not any("visual" in k for k in keys)
diff --git a/tools/extract_qwen3_omni_thinker.py b/tools/extract_qwen3_omni_thinker.py
index ebde57bbaf5..7d52d20cff6 100644
--- a/tools/extract_qwen3_omni_thinker.py
+++ b/tools/extract_qwen3_omni_thinker.py
@@ -1,19 +1,41 @@
-"""Extract a standalone Qwen3-MoE thinker checkpoint from a full Qwen3-Omni model.
+"""Extract a standalone thinker checkpoint from a full Qwen3-Omni model.
 
-miles loads HF via AutoBridge keyed on model_type; the composite omni checkpoint has
-no bridge, but the thinker text backbone is a plain Qwen3-MoE. This writes a
-self-contained HF dir (thinker text + lm_head, renamed, model_type=qwen3_moe) for the
-existing Qwen3MoEBridge. Audio/visual towers are dropped here on purpose: the trainer
-loads them frozen from the original omni checkpoint (see the audio-injection wrapper),
-and the rollout server keeps its own copies.
+Two variants, both keyed on model_type so miles' AutoBridge can load them:
 
-    python tools/extract_qwen3_omni_thinker.py --src <omni> --dst <thinker>
+- ``text`` (default): thinker text backbone + lm_head as a plain Qwen3-MoE
+  (model_type=qwen3_moe, per-expert 2D weights). Audio/visual towers are dropped on
+  purpose: the trainer loads the audio tower frozen from the original omni checkpoint
+  (see the audio-injection plugin), and the rollout server keeps its own copies.
+
+- ``vl``: thinker visual tower + text backbone masquerading as Qwen3-VL-MoE
+  (model_type=qwen3_vl_moe). Names are mechanically compatible except three merger
+  renames (ln_q->norm, mlp.0/2->linear_fc1/2, merger_list->deepstack_merger_list) and
+  the MoE experts, which the VL hub format fuses to 3D hub layout
+  (gate_up_proj [E,H,2I], down_proj [E,I,H]). This unlocks miles' existing
+  bridge-based Qwen3-VL training path (packed mrope + deepstack + CP) for image/video
+  input. The audio tower is dropped; omni-only rope fields ride along under
+  ``omni_sideband`` in config.json for the trainer-side video-rope override.
+
+    python tools/extract_qwen3_omni_thinker.py --src <omni> --dst <thinker> [--variant vl]
 """
 
 from __future__ import annotations
 
+import re
+
 _NON_TEXT_THINKER_SUBMODULES = ("audio_tower.", "visual.", "model.audio_tower.", "model.visual.")
 
+_VL_MERGER_RENAMES = (
+    (re.compile(r"^visual\.merger\.ln_q\."), "visual.merger.norm."),
+    (re.compile(r"^visual\.merger\.mlp\.0\."), "visual.merger.linear_fc1."),
+    (re.compile(r"^visual\.merger\.mlp\.2\."), "visual.merger.linear_fc2."),
+    (re.compile(r"^visual\.merger_list\.(\d+)\.ln_q\."), r"visual.deepstack_merger_list.\1.norm."),
+    (re.compile(r"^visual\.merger_list\.(\d+)\.mlp\.0\."), r"visual.deepstack_merger_list.\1.linear_fc1."),
+    (re.compile(r"^visual\.merger_list\.(\d+)\.mlp\.2\."), r"visual.deepstack_merger_list.\1.linear_fc2."),
+)
+
+_EXPERT_RE = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$")
+
 
 def map_thinker_param_name(name: str) -> str | None:
     """Full omni param name -> standalone Qwen3-MoE name, or None to drop."""
@@ -25,6 +47,44 @@ def map_thinker_param_name(name: str) -> str | None:
     return rest.replace("model.language_model.", "model.")
 
 
+def map_thinker_param_name_vl(name: str) -> str | None:
+    """Full omni param name -> pseudo Qwen3-VL-MoE name, or None to drop.
+
+    Expert weights keep their per-expert names here; `extract(variant="vl")` fuses them
+    into the 3D hub tensors afterwards.
+    """
+    if not name.startswith("thinker."):
+        return None
+    rest = name[len("thinker.") :].replace("model.language_model.", "model.")
+    if rest.startswith(("audio_tower.", "model.audio_tower.")):
+        return None
+    if rest == "lm_head.weight":
+        return rest
+    if rest.startswith("visual."):
+        for pattern, replacement in _VL_MERGER_RENAMES:
+            if pattern.search(rest):
+                rest = pattern.sub(replacement, rest)
+                break
+        return f"model.{rest}"
+    if rest.startswith("model."):
+        return rest.replace("model.", "model.language_model.", 1)
+    return None
+
+
+def fuse_layer_experts_vl(gates, ups, downs):
+    """Per-expert 2D weights -> fused hub-layout tensors.
+
+    gate/up: [I, H] each -> gate_up_proj [E, H, 2I]; down: [H, I] -> down_proj [E, I, H].
+    Hub layout is what transformers (Transpose check_dims), Megatron-Bridge import, and
+    sglang's fused loaders all expect.
+    """
+    import torch
+
+    gate_up = torch.stack([torch.cat([g, u], dim=0).T for g, u in zip(gates, ups, strict=True)], dim=0)
+    down = torch.stack([d.T for d in downs], dim=0)
+    return gate_up.contiguous(), down.contiguous()
+
+
 def synthesize_thinker_config(omni_config: dict) -> dict:
     """Plain Qwen3-MoE config dict from the full omni config."""
     thinker = omni_config.get("thinker_config", omni_config)
@@ -38,11 +98,92 @@ def synthesize_thinker_config(omni_config: dict) -> dict:
     return cfg
 
 
-def extract(src, dst, shard_size_gb: float = 5.0) -> int:
+# vision fields whose omni spelling has no qwen3_vl_moe slot (dropped after use)
+_VL_VISION_DROPPED = ("apply_vit_abs_pos_embed", "image_size", "spatial_patch_size", "in_chans", "tokens_per_second")
+
+
+def synthesize_vl_config(omni_config: dict) -> dict:
+    """Pseudo Qwen3-VL-MoE config from the full omni config.
+
+    Omni numerics are preserved (vocab 152064, rope_theta 1e6, max_position_embeddings
+    65536); rope_scaling is normalized to the VL whitelist; the vision config gains
+    num_position_embeddings (the pos_embed table size, (image_size/patch_size)^2) and
+    drops omni-only spellings. Fields with no VL slot that the trainer still needs
+    (omni video rope) ride under `omni_sideband`.
+    """
+    thinker = omni_config.get("thinker_config", omni_config)
+
+    text = dict(thinker["text_config"])
+    text["model_type"] = "qwen3_vl_moe_text"
+    rope_scaling = dict(text.get("rope_scaling") or {})
+    text["rope_scaling"] = {
+        "rope_type": rope_scaling.get("rope_type") or rope_scaling.get("type", "default"),
+        "mrope_section": rope_scaling["mrope_section"],
+        "mrope_interleaved": rope_scaling.get("mrope_interleaved", True),
+    }
+    text.setdefault("bos_token_id", 151643)
+    text.setdefault("eos_token_id", 151645)
+
+    vision = dict(thinker["vision_config"])
+    vision["model_type"] = "qwen3_vl_moe"
+    if "num_position_embeddings" not in vision and "image_size" in vision:
+        patch = vision.get("patch_size") or vision.get("spatial_patch_size")
+        vision["num_position_embeddings"] = (vision["image_size"] // patch) ** 2
+    for key in _VL_VISION_DROPPED:
+        vision.pop(key, None)
+
+    cfg = {
+        "model_type": "qwen3_vl_moe",
+        "architectures": ["Qwen3VLMoeForConditionalGeneration"],
+        "tie_word_embeddings": thinker.get("tie_word_embeddings", False),
+        "text_config": text,
+        "vision_config": vision,
+        "omni_sideband": {
+            "position_id_per_seconds": thinker.get("position_id_per_seconds"),
+            "audio_token_id": thinker.get("audio_token_id"),
+        },
+    }
+    for k in ("image_token_id", "video_token_id", "vision_start_token_id", "vision_end_token_id"):
+        cfg[k] = thinker[k]
+    return cfg
+
+
+class _VlExpertAccumulator:
+    """Collects per-expert 2D weights per layer and emits fused hub-layout tensors."""
+
+    def __init__(self, num_experts: int):
+        self.num_experts = num_experts
+        self._layers: dict[int, dict[str, dict[int, object]]] = {}
+
+    def offer(self, vl_name: str, tensor) -> list[tuple[str, object]]:
+        match = _EXPERT_RE.match(vl_name.replace("model.language_model.", "model.", 1))
+        if match is None:
+            return [(vl_name, tensor)]
+        layer, expert, kind = int(match.group(1)), int(match.group(2)), match.group(3)
+        buckets = self._layers.setdefault(layer, {"gate_proj": {}, "up_proj": {}, "down_proj": {}})
+        buckets[kind][expert] = tensor
+        if all(len(buckets[k]) == self.num_experts for k in ("gate_proj", "up_proj", "down_proj")):
+            del self._layers[layer]
+            order = range(self.num_experts)
+            gate_up, down = fuse_layer_experts_vl(
+                [buckets["gate_proj"][e] for e in order],
+                [buckets["up_proj"][e] for e in order],
+                [buckets["down_proj"][e] for e in order],
+            )
+            prefix = f"model.language_model.layers.{layer}.mlp.experts"
+            return [(f"{prefix}.gate_up_proj", gate_up), (f"{prefix}.down_proj", down)]
+        return []
+
+    def assert_drained(self):
+        assert not self._layers, f"incomplete expert layers left over: {sorted(self._layers)}"
+
+
+def extract(src, dst, shard_size_gb: float = 5.0, variant: str = "text") -> int:
     """Stream thinker tensors from src into sharded safetensors under dst.
 
     Each shard is written (and its buffer released) as soon as it fills, so peak RAM is
-    one shard, not the whole ~60 GB thinker. Returns the number of kept tensors.
+    one shard (plus, for the vl variant, at most a few layers of in-flight expert
+    weights). Returns the number of tensors written.
     """
     import json
     import os
@@ -52,13 +193,16 @@ def extract(src, dst, shard_size_gb: float = 5.0) -> int:
     from safetensors import safe_open
     from safetensors.torch import save_file
 
+    assert variant in ("text", "vl"), variant
     src, dst = Path(src), Path(dst)
     dst.mkdir(parents=True, exist_ok=True)
 
     with open(src / "config.json") as f:
         omni_config = json.load(f)
+    synthesize = synthesize_thinker_config if variant == "text" else synthesize_vl_config
+    out_config = synthesize(omni_config)
     with open(dst / "config.json", "w") as f:
-        json.dump(synthesize_thinker_config(omni_config), f, indent=2)
+        json.dump(out_config, f, indent=2)
 
     for fname in (
         "tokenizer.json",
@@ -91,6 +235,12 @@ def extract(src, dst, shard_size_gb: float = 5.0) -> int:
     else:
         shard_files = ["model.safetensors"]
 
+    map_name = map_thinker_param_name if variant == "text" else map_thinker_param_name_vl
+    accumulator = None
+    if variant == "vl":
+        thinker = omni_config.get("thinker_config", omni_config)
+        accumulator = _VlExpertAccumulator(num_experts=thinker["text_config"]["num_experts"])
+
     shard_size_bytes = int(shard_size_gb * (1024**3))
     out_index: dict[str, str] = {}
     shard_names: list[str] = []
@@ -111,18 +261,28 @@ def flush():
             total_size += t.numel() * t.element_size()
         buf, buf_bytes = {}, 0
 
+    def emit(out_name, tensor):
+        nonlocal buf_bytes, total_kept
+        buf[out_name] = tensor
+        buf_bytes += tensor.numel() * tensor.element_size()
+        total_kept += 1
+        if buf_bytes >= shard_size_bytes:
+            flush()
+
     for shard_file in shard_files:
         with safe_open(src / shard_file, framework="pt") as reader:
             for key in reader.keys():
-                new_key = map_thinker_param_name(key)
+                new_key = map_name(key)
                 if new_key is None:
                     continue
                 tensor = reader.get_tensor(key)
-                buf[new_key] = tensor
-                buf_bytes += tensor.numel() * tensor.element_size()
-                total_kept += 1
-                if buf_bytes >= shard_size_bytes:
-                    flush()
+                if accumulator is not None:
+                    for out_name, out_tensor in accumulator.offer(new_key, tensor):
+                        emit(out_name, out_tensor)
+                else:
+                    emit(new_key, tensor)
+    if accumulator is not None:
+        accumulator.assert_drained()
     flush()
 
     if total_kept == 0:
@@ -134,7 +294,7 @@ def flush():
         with open(dst / "model.safetensors.index.json", "w") as f:
             json.dump({"metadata": {"total_size": total_size}, "weight_map": out_index}, f, indent=2)
 
-    print(f"[done] {total_kept} thinker tensors -> {dst}")
+    print(f"[done] {total_kept} thinker tensors ({variant}) -> {dst}")
     return total_kept
 
 
@@ -145,8 +305,9 @@ def main() -> None:
     parser.add_argument("--src", required=True)
     parser.add_argument("--dst", required=True)
     parser.add_argument("--shard-size-gb", type=float, default=5.0)
+    parser.add_argument("--variant", choices=["text", "vl"], default="text")
     args = parser.parse_args()
-    extract(args.src, args.dst, shard_size_gb=args.shard_size_gb)
+    extract(args.src, args.dst, shard_size_gb=args.shard_size_gb, variant=args.variant)
 
 
 if __name__ == "__main__":

From fc5194bd3f0cfa3a1c58406016bc09578ba236ec Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Sun, 26 Jul 2026 10:47:57 +0800
Subject: [PATCH 23/32] [omni] sequence-parallel aware audio injection
 (contiguous-chunk scatter)

mcore hard-requires SP for MoE + TP>1 in training, so the no-SP constraint made every
2-GPU layout either invalid (TP2 no-SP) or OOM (TP1: full expert grad buffers). The
embedding output under SP is the rank's contiguous chunk; the scatter now slices the
placeholder mask and the encoder rows to that window. Debug tier drops the zero-coef
KL ref copy and trims the token budget.
---
 .../run_qwen3_omni_thinker_grpo.py            |  8 ++--
 miles_plugins/models/qwen3_omni_thinker.py    | 38 +++++++++++++-----
 tests/fast/test_qwen3_omni_audio_injection.py | 40 +++++++++++++------
 3 files changed, 61 insertions(+), 25 deletions(-)

diff --git a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
index beef11a2d23..9d0df4bd8dd 100644
--- a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
+++ b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
@@ -118,8 +118,10 @@ def execute(args: ScriptArgs):
 
     perf_args = (
         f"--tensor-model-parallel-size {args.actor_tp} "
-        # no --sequence-parallel: audio injection scatters full-sequence embeddings
-        "--pipeline-model-parallel-size 1 "
+        # the audio injection is sequence-parallel aware (contiguous-chunk scatter);
+        # mcore requires SP for MoE + TP>1
+        + ("--sequence-parallel " if args.actor_tp > 1 else "")
+        + "--pipeline-model-parallel-size 1 "
         "--context-parallel-size 1 "
         f"--expert-model-parallel-size {args.actor_tp} "
         "--expert-tensor-parallel-size 1 "
@@ -127,7 +129,7 @@ def execute(args: ScriptArgs):
         "--recompute-method uniform "
         "--recompute-num-layers 1 "
         "--use-dynamic-batch-size "
-        "--max-tokens-per-gpu 9216 "
+        f"--max-tokens-per-gpu {2048 if debug_minimal else 9216} "
         "--optimizer-cpu-offload "
         "--overlap-cpu-optimizer-d2h-h2d "
         "--use-precision-aware-optimizer "
diff --git a/miles_plugins/models/qwen3_omni_thinker.py b/miles_plugins/models/qwen3_omni_thinker.py
index 84ce7470bba..c0abb4278c9 100644
--- a/miles_plugins/models/qwen3_omni_thinker.py
+++ b/miles_plugins/models/qwen3_omni_thinker.py
@@ -102,23 +102,37 @@ def compute_audio_embeddings(encoder, input_features, feature_attention_mask, au
         return encoder(input_features, feature_lens=feature_lens).last_hidden_state
 
 
-def scatter_audio_embeddings(hidden, input_ids, audio_embeds, audio_token_id) -> torch.Tensor:
+def scatter_audio_embeddings(hidden, input_ids, audio_embeds, audio_token_id, sp_rank: int = 0, sp_size: int = 1) -> torch.Tensor:
     """Replace hidden rows at audio placeholder positions, out of place.
 
-    hidden: [s, b, h] (mcore embedding layout); input_ids: [b, s]; audio_embeds: [n, h].
-    Requires the packed b == 1 layout: masked_scatter fills s-major, which would
-    interleave samples for b > 1.
+    hidden: [s_local, b, h] (mcore embedding layout); input_ids: [b, s_global];
+    audio_embeds: [n_global, h]. Requires the packed b == 1 layout: masked_scatter fills
+    s-major, which would interleave samples for b > 1.
+
+    With sequence parallelism the embedding output is the rank's contiguous chunk
+    (s_local = s_global / sp_size, rows [sp_rank*s_local, (sp_rank+1)*s_local)); the
+    mask and the encoder outputs are sliced to that window, so every rank scatters
+    exactly its own audio rows.
     """
     assert hidden.dim() == 3 and input_ids.dim() == 2, f"{hidden.shape=} {input_ids.shape=}"
     assert (
         hidden.size(1) == 1 and input_ids.size(0) == 1
     ), f"audio injection requires the packed [1, s] layout, got {input_ids.shape}"
-    mask = input_ids[0] == audio_token_id  # [s]
+    mask = input_ids[0] == audio_token_id  # [s_global]
     num_positions = int(mask.sum())
     assert num_positions == audio_embeds.size(0), (
         f"audio placeholder/embedding mismatch: {num_positions} audio tokens in input_ids "
         f"vs {audio_embeds.size(0)} encoder outputs"
     )
+    if sp_size > 1:
+        s_local = hidden.size(0)
+        assert s_local * sp_size == mask.numel(), (
+            f"sequence-parallel chunking mismatch: local {s_local} x sp {sp_size} != global {mask.numel()}"
+        )
+        start = sp_rank * s_local
+        prior = int(mask[:start].sum())
+        mask = mask[start : start + s_local]
+        audio_embeds = audio_embeds[prior : prior + int(mask.sum())]
     audio_embeds = audio_embeds.to(device=hidden.device, dtype=hidden.dtype)
     return hidden.masked_scatter(mask.view(-1, 1, 1).to(hidden.device), audio_embeds)
 
@@ -157,12 +171,16 @@ def forward(*fargs, **kwargs):
 
         assert not fargs, "audio injection expects keyword-only forward calls"
         assert "decoder_input" not in kwargs, "decoder_input already set upstream"
-        assert not getattr(
-            args, "sequence_parallel", False
-        ), "audio injection scatters full-sequence embeddings; run with sequence-parallel off"
         assert getattr(args, "context_parallel_size", 1) == 1, "audio injection requires context_parallel_size == 1"
 
         input_ids = kwargs["input_ids"]
+        # with sequence parallelism the embedding output is this rank's contiguous chunk
+        sp_rank, sp_size = 0, 1
+        if getattr(args, "sequence_parallel", False) and getattr(args, "tensor_model_parallel_size", 1) > 1:
+            from megatron.core import parallel_state as mpu
+
+            sp_rank = mpu.get_tensor_model_parallel_rank()
+            sp_size = mpu.get_tensor_model_parallel_world_size()
         hidden = model.embedding(input_ids=input_ids, position_ids=kwargs.get("position_ids"))
         encoder = encoder_loader(device=hidden.device, dtype=hidden.dtype)
         audio_embeds = compute_audio_embeddings(
@@ -171,7 +189,9 @@ def forward(*fargs, **kwargs):
             audio_kwargs.get("feature_attention_mask"),
             audio_kwargs.get("audio_feature_lengths"),
         )
-        kwargs["decoder_input"] = scatter_audio_embeddings(hidden, input_ids, audio_embeds, audio_token_id)
+        kwargs["decoder_input"] = scatter_audio_embeddings(
+            hidden, input_ids, audio_embeds, audio_token_id, sp_rank=sp_rank, sp_size=sp_size
+        )
         return orig_forward(**kwargs)
 
     model.forward = forward
diff --git a/tests/fast/test_qwen3_omni_audio_injection.py b/tests/fast/test_qwen3_omni_audio_injection.py
index 17ea03d11e6..d36276d563b 100644
--- a/tests/fast/test_qwen3_omni_audio_injection.py
+++ b/tests/fast/test_qwen3_omni_audio_injection.py
@@ -229,19 +229,33 @@ def test_install_non_pre_process_swallows_audio_kwargs(tiny_encoder):
     assert "decoder_input" not in model.captured
 
 
-def test_install_rejects_sequence_parallel_with_audio(tiny_encoder):
-    model = _StubGPT()
-    _install(model, tiny_encoder, sequence_parallel=True)
-    input_features, feature_attention_mask, n_tokens = _audio_batch(tiny_encoder)
-    with pytest.raises(AssertionError, match="sequence.parallel"):
-        model.forward(
-            input_ids=_packed_input_ids([n_tokens]),
-            input_features=input_features,
-            feature_attention_mask=feature_attention_mask,
-        )
-    # text-only batches stay allowed under SP
-    model.forward(input_ids=_packed_input_ids([0]))
-    assert model.captured is not None
+def test_scatter_sequence_parallel_chunks_reconstruct_full_result():
+    input_ids = _packed_input_ids([2, 3])
+    s = input_ids.size(1)
+    if s % 2:  # pad to an even split like miles' tp-multiple padding does
+        input_ids = torch.cat([input_ids, torch.tensor([[7]])], dim=1)
+        s += 1
+    torch.manual_seed(4)
+    hidden_full = torch.randn(s, 1, HIDDEN)
+    audio_embeds = -torch.arange(1, 5 * HIDDEN + 1, dtype=torch.float32).reshape(5, HIDDEN)
+
+    full = scatter_audio_embeddings(hidden_full, input_ids, audio_embeds, AUDIO_TOKEN_ID)
+    chunk = s // 2
+    part0 = scatter_audio_embeddings(
+        hidden_full[:chunk], input_ids, audio_embeds, AUDIO_TOKEN_ID, sp_rank=0, sp_size=2
+    )
+    part1 = scatter_audio_embeddings(
+        hidden_full[chunk:], input_ids, audio_embeds, AUDIO_TOKEN_ID, sp_rank=1, sp_size=2
+    )
+
+    assert torch.equal(torch.cat([part0, part1], dim=0), full)
+
+
+def test_scatter_sequence_parallel_rejects_bad_chunking():
+    input_ids = _packed_input_ids([2])
+    hidden_local = torch.zeros(3, 1, HIDDEN)  # 3 * 2 != seq len
+    with pytest.raises(AssertionError, match="chunking"):
+        scatter_audio_embeddings(hidden_local, input_ids, torch.zeros(2, HIDDEN), AUDIO_TOKEN_ID, sp_rank=0, sp_size=2)
 
 
 # ------------------------------- checkpoint loader -------------------------------

From b7bf4eb54b021ed9872b72873122de66684d7d0a Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Sun, 26 Jul 2026 10:50:31 +0800
Subject: [PATCH 24/32] [omni] format

---
 examples/omni_thinker/bench_mm_payload.py     |  9 +++++++--
 .../run_qwen3_omni_thinker_grpo.py            |  7 ++-----
 miles/rollout/generate_hub/sglang_omni.py     | 13 ++++++++++++-
 .../inference_rollout_train.py                |  2 +-
 miles/rollout/sglang_rollout.py               |  2 +-
 miles_plugins/models/qwen3_omni_thinker.py    | 10 ++++++----
 .../megatron_utils/test_extract_thinker_vl.py | 19 +++++++++++++++----
 .../sglang_utils/test_omni_admin_api.py       |  4 +---
 tests/fast/test_bench_mm_payload.py           |  4 +++-
 9 files changed, 48 insertions(+), 22 deletions(-)

diff --git a/examples/omni_thinker/bench_mm_payload.py b/examples/omni_thinker/bench_mm_payload.py
index 3b092014d30..dc2c97aa068 100644
--- a/examples/omni_thinker/bench_mm_payload.py
+++ b/examples/omni_thinker/bench_mm_payload.py
@@ -106,10 +106,15 @@ def main() -> None:
     args = parser.parse_args()
 
     rows = [("image", f"{args.resolution}px", make_image_bundle(args.resolution, args.resolution))]
-    rows += [(f"video", f"{n}f@{args.resolution}px", make_video_bundle(n, args.resolution, args.resolution)) for n in args.frames]
+    rows += [
+        ("video", f"{n}f@{args.resolution}px", make_video_bundle(n, args.resolution, args.resolution))
+        for n in args.frames
+    ]
     rows += [("audio", f"{s:.0f}s", make_audio_bundle(s)) for s in args.seconds]
 
-    header = f"{'modality':8s} {'case':12s} {'raw_MB':>8s} {'json_MB':>8s} {'ser_ms':>8s} {'dump_ms':>8s} {'deser_ms':>9s}"
+    header = (
+        f"{'modality':8s} {'case':12s} {'raw_MB':>8s} {'json_MB':>8s} {'ser_ms':>8s} {'dump_ms':>8s} {'deser_ms':>9s}"
+    )
     print(header)
     for modality, case, tensors in rows:
         r = measure_roundtrip(tensors, server=args.server)
diff --git a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
index 9d0df4bd8dd..a5f650244d3 100644
--- a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
+++ b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
@@ -120,8 +120,7 @@ def execute(args: ScriptArgs):
         f"--tensor-model-parallel-size {args.actor_tp} "
         # the audio injection is sequence-parallel aware (contiguous-chunk scatter);
         # mcore requires SP for MoE + TP>1
-        + ("--sequence-parallel " if args.actor_tp > 1 else "")
-        + "--pipeline-model-parallel-size 1 "
+        + ("--sequence-parallel " if args.actor_tp > 1 else "") + "--pipeline-model-parallel-size 1 "
         "--context-parallel-size 1 "
         f"--expert-model-parallel-size {args.actor_tp} "
         "--expert-tensor-parallel-size 1 "
@@ -159,9 +158,7 @@ def execute(args: ScriptArgs):
         "--hidden-dropout 0.0 "
         # mcore requires SP for MoE+TP>1, and the audio injection needs SP off, so the
         # debug tier runs TP1 (bf16 grad accum keeps 2xH200 within memory)
-        + ("" if debug_minimal else "--accumulate-allreduce-grads-in-fp32 ")
-        + 
-        "--attention-softmax-in-fp32 "
+        + ("" if debug_minimal else "--accumulate-allreduce-grads-in-fp32 ") + "--attention-softmax-in-fp32 "
         "--attention-backend flash "
         "--actor-num-nodes 1 "
         f"--actor-num-gpus-per-node {args.actor_tp} "
diff --git a/miles/rollout/generate_hub/sglang_omni.py b/miles/rollout/generate_hub/sglang_omni.py
index 2d89bcb0cce..e3d6f61b414 100644
--- a/miles/rollout/generate_hub/sglang_omni.py
+++ b/miles/rollout/generate_hub/sglang_omni.py
@@ -15,7 +15,18 @@
 
 # RolloutSamplingParams on the omni server is extra="forbid"; this mirrors its schema
 _OMNI_SAMPLING_KEYS = frozenset(
-    ("temperature", "top_p", "top_k", "min_p", "repetition_penalty", "stop", "stop_token_ids", "seed", "max_new_tokens", "max_tokens")
+    (
+        "temperature",
+        "top_p",
+        "top_k",
+        "min_p",
+        "repetition_penalty",
+        "stop",
+        "stop_token_ids",
+        "seed",
+        "max_new_tokens",
+        "max_tokens",
+    )
 )
 # detok/text-shaping flags: they never reach the token or logprob streams the trainer consumes
 _DETOK_ONLY_KEYS = frozenset(("skip_special_tokens", "no_stop_trim", "spaces_between_special_tokens"))
diff --git a/miles/rollout/inference_rollout/inference_rollout_train.py b/miles/rollout/inference_rollout/inference_rollout_train.py
index cce05f34656..64626455a22 100644
--- a/miles/rollout/inference_rollout/inference_rollout_train.py
+++ b/miles/rollout/inference_rollout/inference_rollout_train.py
@@ -7,9 +7,9 @@
 from packaging.version import parse
 from tqdm import tqdm
 
-from miles.rollout.generate_utils.generate_endpoint_utils import abort_external_omni_requests, is_omni_external_admin
 from miles.rollout.base_types import RolloutFnTrainOutput
 from miles.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter
+from miles.rollout.generate_utils.generate_endpoint_utils import abort_external_omni_requests, is_omni_external_admin
 from miles.rollout.generate_utils.prefill_logprobs import recompute_samples_rollout_logprobs_via_prefill
 from miles.rollout.inference_rollout.inference_rollout_common import GenerateState, generate_and_rm_group
 from miles.utils import dumper_utils
diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py
index 2b9103a480a..bcc23e6f904 100644
--- a/miles/rollout/sglang_rollout.py
+++ b/miles/rollout/sglang_rollout.py
@@ -13,9 +13,9 @@
 from packaging.version import parse
 from tqdm import tqdm
 
-from miles.rollout.generate_utils.generate_endpoint_utils import abort_external_omni_requests, is_omni_external_admin
 from miles.rollout.base_types import GenerateFnInput, RolloutFnEvalOutput, RolloutFnTrainOutput
 from miles.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter
+from miles.rollout.generate_utils.generate_endpoint_utils import abort_external_omni_requests, is_omni_external_admin
 from miles.rollout.inference_rollout.compatibility import load_generate_function
 from miles.utils import dumper_utils
 from miles.utils.async_utils import run
diff --git a/miles_plugins/models/qwen3_omni_thinker.py b/miles_plugins/models/qwen3_omni_thinker.py
index c0abb4278c9..722b8b9c0f4 100644
--- a/miles_plugins/models/qwen3_omni_thinker.py
+++ b/miles_plugins/models/qwen3_omni_thinker.py
@@ -102,7 +102,9 @@ def compute_audio_embeddings(encoder, input_features, feature_attention_mask, au
         return encoder(input_features, feature_lens=feature_lens).last_hidden_state
 
 
-def scatter_audio_embeddings(hidden, input_ids, audio_embeds, audio_token_id, sp_rank: int = 0, sp_size: int = 1) -> torch.Tensor:
+def scatter_audio_embeddings(
+    hidden, input_ids, audio_embeds, audio_token_id, sp_rank: int = 0, sp_size: int = 1
+) -> torch.Tensor:
     """Replace hidden rows at audio placeholder positions, out of place.
 
     hidden: [s_local, b, h] (mcore embedding layout); input_ids: [b, s_global];
@@ -126,9 +128,9 @@ def scatter_audio_embeddings(hidden, input_ids, audio_embeds, audio_token_id, sp
     )
     if sp_size > 1:
         s_local = hidden.size(0)
-        assert s_local * sp_size == mask.numel(), (
-            f"sequence-parallel chunking mismatch: local {s_local} x sp {sp_size} != global {mask.numel()}"
-        )
+        assert (
+            s_local * sp_size == mask.numel()
+        ), f"sequence-parallel chunking mismatch: local {s_local} x sp {sp_size} != global {mask.numel()}"
         start = sp_rank * s_local
         prior = int(mask[:start].sum())
         mask = mask[start : start + s_local]
diff --git a/tests/fast/backends/megatron_utils/test_extract_thinker_vl.py b/tests/fast/backends/megatron_utils/test_extract_thinker_vl.py
index 22a48916d55..dd268f648b1 100644
--- a/tests/fast/backends/megatron_utils/test_extract_thinker_vl.py
+++ b/tests/fast/backends/megatron_utils/test_extract_thinker_vl.py
@@ -4,7 +4,6 @@
 import json
 from pathlib import Path
 
-import pytest
 import torch
 from safetensors import safe_open
 from safetensors.torch import save_file
@@ -28,7 +27,9 @@ def test_vl_name_map_text_and_visual():
     m = _load_tool().map_thinker_param_name_vl
     assert m("thinker.lm_head.weight") == "lm_head.weight"
     assert m("thinker.model.embed_tokens.weight") == "model.language_model.embed_tokens.weight"
-    assert m("thinker.model.layers.3.self_attn.q_proj.weight") == "model.language_model.layers.3.self_attn.q_proj.weight"
+    assert (
+        m("thinker.model.layers.3.self_attn.q_proj.weight") == "model.language_model.layers.3.self_attn.q_proj.weight"
+    )
     assert m("thinker.model.layers.3.mlp.gate.weight") == "model.language_model.layers.3.mlp.gate.weight"
     assert m("thinker.visual.patch_embed.proj.weight") == "model.visual.patch_embed.proj.weight"
     assert m("thinker.visual.blocks.7.attn.qkv.weight") == "model.visual.blocks.7.attn.qkv.weight"
@@ -81,7 +82,12 @@ def test_vl_config_synthesis():
                 "num_hidden_layers": 48,
                 "rope_theta": 1000000,
                 "max_position_embeddings": 65536,
-                "rope_scaling": {"interleaved": True, "mrope_interleaved": True, "mrope_section": [24, 20, 20], "type": "default"},
+                "rope_scaling": {
+                    "interleaved": True,
+                    "mrope_interleaved": True,
+                    "mrope_section": [24, 20, 20],
+                    "type": "default",
+                },
                 "num_experts": 128,
             },
             "vision_config": {
@@ -140,7 +146,12 @@ def _tiny_omni_src(src: Path):
             "model_type": "qwen3_omni_moe_text",
             "vocab_size": 32,
             "num_experts": 2,
-            "rope_scaling": {"interleaved": True, "mrope_interleaved": True, "mrope_section": [24, 20, 20], "type": "default"},
+            "rope_scaling": {
+                "interleaved": True,
+                "mrope_interleaved": True,
+                "mrope_section": [24, 20, 20],
+                "type": "default",
+            },
         },
         "vision_config": {"model_type": "x", "image_size": 32, "patch_size": 16},
     }
diff --git a/tests/fast/backends/sglang_utils/test_omni_admin_api.py b/tests/fast/backends/sglang_utils/test_omni_admin_api.py
index 8093fe4cad6..2356ed3d09b 100644
--- a/tests/fast/backends/sglang_utils/test_omni_admin_api.py
+++ b/tests/fast/backends/sglang_utils/test_omni_admin_api.py
@@ -153,9 +153,7 @@ def get(self, url, headers=None):
 
     monkeypatch.setattr("requests.Session", _FakeSession)
 
-    _wait_server_healthy(
-        base_url="http://x:1", api_key=None, is_process_alive=lambda: True, probe_paths=("health",)
-    )
+    _wait_server_healthy(base_url="http://x:1", api_key=None, is_process_alive=lambda: True, probe_paths=("health",))
 
     assert seen == ["http://x:1/health"]
 
diff --git a/tests/fast/test_bench_mm_payload.py b/tests/fast/test_bench_mm_payload.py
index a1837088100..7c9b0368343 100644
--- a/tests/fast/test_bench_mm_payload.py
+++ b/tests/fast/test_bench_mm_payload.py
@@ -40,7 +40,9 @@ def test_image_and_audio_bundles():
 
     audio = tool.make_audio_bundle(seconds=10.0)
     assert audio["input_features"].shape[1] == 128  # mel bins
-    assert audio["feature_attention_mask"].shape == audio["input_features"].shape[:1] + audio["input_features"].shape[2:]
+    assert (
+        audio["feature_attention_mask"].shape == audio["input_features"].shape[:1] + audio["input_features"].shape[2:]
+    )
 
 
 def test_measure_roundtrip_reports_size_and_integrity():

From 97a3fc5601b3061b174ff272e964529b2231e911 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Mon, 27 Jul 2026 12:09:01 +0800
Subject: [PATCH 25/32] [omni] split out PR#6 content + tidy imports

- extract tool: drop the --variant vl path (belongs to the video/image PR #6); this
  branch keeps only the text-variant + processor artifacts
- remove test_extract_thinker_vl.py (lives on #6)
- sglang_omni.py: move the sampling-key frozensets below the import block
- audio injection: note the TP-redundant tower forward tradeoff under SP
---
 miles/rollout/generate_hub/sglang_omni.py     |  16 +-
 miles_plugins/models/qwen3_omni_thinker.py    |   3 +
 .../megatron_utils/test_extract_thinker_vl.py | 222 ------------------
 tools/extract_qwen3_omni_thinker.py           | 199 ++--------------
 4 files changed, 30 insertions(+), 410 deletions(-)
 delete mode 100644 tests/fast/backends/megatron_utils/test_extract_thinker_vl.py

diff --git a/miles/rollout/generate_hub/sglang_omni.py b/miles/rollout/generate_hub/sglang_omni.py
index e3d6f61b414..37da9ee8a67 100644
--- a/miles/rollout/generate_hub/sglang_omni.py
+++ b/miles/rollout/generate_hub/sglang_omni.py
@@ -12,6 +12,14 @@
 """
 
 from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput
+from miles.rollout.generate_utils.generate_endpoint_utils import (
+    compute_prompt_ids_from_sample,
+    compute_request_payload,
+    compute_routing_headers,
+    update_sample_from_response,
+)
+from miles.utils.http_utils import post
+from miles.utils.types import Sample
 
 # RolloutSamplingParams on the omni server is extra="forbid"; this mirrors its schema
 _OMNI_SAMPLING_KEYS = frozenset(
@@ -30,14 +38,6 @@
 )
 # detok/text-shaping flags: they never reach the token or logprob streams the trainer consumes
 _DETOK_ONLY_KEYS = frozenset(("skip_special_tokens", "no_stop_trim", "spaces_between_special_tokens"))
-from miles.rollout.generate_utils.generate_endpoint_utils import (
-    compute_prompt_ids_from_sample,
-    compute_request_payload,
-    compute_routing_headers,
-    update_sample_from_response,
-)
-from miles.utils.http_utils import post
-from miles.utils.types import Sample
 
 
 async def generate(input: GenerateFnInput) -> GenerateFnOutput:
diff --git a/miles_plugins/models/qwen3_omni_thinker.py b/miles_plugins/models/qwen3_omni_thinker.py
index 722b8b9c0f4..fe326a47cba 100644
--- a/miles_plugins/models/qwen3_omni_thinker.py
+++ b/miles_plugins/models/qwen3_omni_thinker.py
@@ -185,6 +185,9 @@ def forward(*fargs, **kwargs):
             sp_size = mpu.get_tensor_model_parallel_world_size()
         hidden = model.embedding(input_ids=input_ids, position_ids=kwargs.get("position_ids"))
         encoder = encoder_loader(device=hidden.device, dtype=hidden.dtype)
+        # NOTE: every TP rank re-encodes the whole audio and then scatters only its own SP
+        # chunk -> TP_size x redundant tower forward. Fine here (frozen ~0.6B tower, small
+        # audio); if audio grows or TP scales, encode once on rank 0 and broadcast.
         audio_embeds = compute_audio_embeddings(
             encoder,
             input_features,
diff --git a/tests/fast/backends/megatron_utils/test_extract_thinker_vl.py b/tests/fast/backends/megatron_utils/test_extract_thinker_vl.py
deleted file mode 100644
index dd268f648b1..00000000000
--- a/tests/fast/backends/megatron_utils/test_extract_thinker_vl.py
+++ /dev/null
@@ -1,222 +0,0 @@
-"""Pseudo-Qwen3-VL extraction: thinker (visual+text) masquerading as qwen3_vl_moe."""
-
-import importlib.util
-import json
-from pathlib import Path
-
-import torch
-from safetensors import safe_open
-from safetensors.torch import save_file
-
-from tests.ci.ci_register import register_cpu_ci
-
-register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[])
-
-_REPO = Path(__file__).resolve().parents[4]
-
-
-def _load_tool():
-    path = _REPO / "tools" / "extract_qwen3_omni_thinker.py"
-    spec = importlib.util.spec_from_file_location("extract_qwen3_omni_thinker", path)
-    mod = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(mod)
-    return mod
-
-
-def test_vl_name_map_text_and_visual():
-    m = _load_tool().map_thinker_param_name_vl
-    assert m("thinker.lm_head.weight") == "lm_head.weight"
-    assert m("thinker.model.embed_tokens.weight") == "model.language_model.embed_tokens.weight"
-    assert (
-        m("thinker.model.layers.3.self_attn.q_proj.weight") == "model.language_model.layers.3.self_attn.q_proj.weight"
-    )
-    assert m("thinker.model.layers.3.mlp.gate.weight") == "model.language_model.layers.3.mlp.gate.weight"
-    assert m("thinker.visual.patch_embed.proj.weight") == "model.visual.patch_embed.proj.weight"
-    assert m("thinker.visual.blocks.7.attn.qkv.weight") == "model.visual.blocks.7.attn.qkv.weight"
-    # merger renames: ln_q -> norm, mlp.0/2 -> linear_fc1/2, merger_list -> deepstack_merger_list
-    assert m("thinker.visual.merger.ln_q.weight") == "model.visual.merger.norm.weight"
-    assert m("thinker.visual.merger.mlp.0.bias") == "model.visual.merger.linear_fc1.bias"
-    assert m("thinker.visual.merger.mlp.2.weight") == "model.visual.merger.linear_fc2.weight"
-    assert m("thinker.visual.merger_list.1.ln_q.weight") == "model.visual.deepstack_merger_list.1.norm.weight"
-    assert m("thinker.visual.merger_list.2.mlp.2.bias") == "model.visual.deepstack_merger_list.2.linear_fc2.bias"
-    # drops
-    assert m("thinker.audio_tower.layers.0.fc1.weight") is None
-    assert m("talker.model.norm.weight") is None
-    assert m("code2wav.decoder.weight") is None
-
-
-def test_vl_expert_fusion_hub_layout():
-    tool = _load_tool()
-    # E=2 experts, I=3 (moe ffn), H=4 (hidden)
-    gate = [torch.arange(12, dtype=torch.float32).reshape(3, 4) + 100 * e for e in range(2)]
-    up = [torch.arange(12, dtype=torch.float32).reshape(3, 4) + 100 * e + 50 for e in range(2)]
-    down = [torch.arange(12, dtype=torch.float32).reshape(4, 3) + 100 * e for e in range(2)]
-
-    gate_up_fused, down_fused = tool.fuse_layer_experts_vl(gate, up, down)
-
-    # hub layout: gate_up [E, H, 2I], down [E, I, H]
-    assert gate_up_fused.shape == (2, 4, 6)
-    assert down_fused.shape == (2, 3, 4)
-    for e in range(2):
-        expected = torch.cat([gate[e], up[e]], dim=0).T  # [H, 2I]
-        assert torch.equal(gate_up_fused[e], expected)
-        assert torch.equal(down_fused[e], down[e].T)
-
-
-def test_vl_config_synthesis():
-    synth = _load_tool().synthesize_vl_config
-    omni = {
-        "model_type": "qwen3_omni_moe",
-        "thinker_config": {
-            "image_token_id": 151655,
-            "video_token_id": 151656,
-            "vision_start_token_id": 151652,
-            "vision_end_token_id": 151653,
-            "position_id_per_seconds": 13,
-            "audio_token_id": 151675,
-            "tie_word_embeddings": False,
-            "text_config": {
-                "model_type": "qwen3_omni_moe_text",
-                "vocab_size": 152064,
-                "hidden_size": 2048,
-                "num_hidden_layers": 48,
-                "rope_theta": 1000000,
-                "max_position_embeddings": 65536,
-                "rope_scaling": {
-                    "interleaved": True,
-                    "mrope_interleaved": True,
-                    "mrope_section": [24, 20, 20],
-                    "type": "default",
-                },
-                "num_experts": 128,
-            },
-            "vision_config": {
-                "model_type": "qwen3_omni_moe_vision_encoder",
-                "depth": 27,
-                "hidden_size": 1152,
-                "deepstack_visual_indexes": [8, 16, 24],
-                "apply_vit_abs_pos_embed": True,
-                "image_size": 768,
-                "spatial_patch_size": 16,
-                "in_chans": 3,
-                "tokens_per_second": 2,
-                "out_hidden_size": 2048,
-                "spatial_merge_size": 2,
-                "patch_size": 16,
-                "temporal_patch_size": 2,
-            },
-            "audio_config": {"d_model": 1280},
-        },
-    }
-
-    cfg = synth(omni)
-
-    assert cfg["model_type"] == "qwen3_vl_moe"
-    assert cfg["architectures"] == ["Qwen3VLMoeForConditionalGeneration"]
-    assert cfg["tie_word_embeddings"] is False
-    for key in ("image_token_id", "video_token_id", "vision_start_token_id", "vision_end_token_id"):
-        assert cfg[key] == omni["thinker_config"][key]
-    text = cfg["text_config"]
-    assert text["model_type"] == "qwen3_vl_moe_text"
-    assert text["vocab_size"] == 152064 and text["rope_theta"] == 1000000
-    assert text["max_position_embeddings"] == 65536
-    assert text["rope_scaling"] == {"rope_type": "default", "mrope_section": [24, 20, 20], "mrope_interleaved": True}
-    assert text["bos_token_id"] == 151643 and text["eos_token_id"] == 151645
-    vision = cfg["vision_config"]
-    assert vision["model_type"] == "qwen3_vl_moe"
-    assert vision["num_position_embeddings"] == 2304  # (image_size // patch_size) ** 2
-    for gone in ("apply_vit_abs_pos_embed", "image_size", "spatial_patch_size", "in_chans", "tokens_per_second"):
-        assert gone not in vision
-    assert "audio_config" not in cfg
-    # sideband for the trainer-side omni video rope (no VL slot for it)
-    assert cfg["omni_sideband"] == {"position_id_per_seconds": 13, "audio_token_id": 151675}
-
-
-def _tiny_omni_src(src: Path):
-    src.mkdir(parents=True, exist_ok=True)
-    thinker = {
-        "image_token_id": 151655,
-        "video_token_id": 151656,
-        "vision_start_token_id": 151652,
-        "vision_end_token_id": 151653,
-        "position_id_per_seconds": 13,
-        "audio_token_id": 151675,
-        "tie_word_embeddings": False,
-        "text_config": {
-            "model_type": "qwen3_omni_moe_text",
-            "vocab_size": 32,
-            "num_experts": 2,
-            "rope_scaling": {
-                "interleaved": True,
-                "mrope_interleaved": True,
-                "mrope_section": [24, 20, 20],
-                "type": "default",
-            },
-        },
-        "vision_config": {"model_type": "x", "image_size": 32, "patch_size": 16},
-    }
-    with open(src / "config.json", "w") as f:
-        json.dump({"thinker_config": thinker}, f)
-    tensors = {
-        "thinker.model.embed_tokens.weight": torch.randn(32, 4),
-        "thinker.lm_head.weight": torch.randn(32, 4),
-        "thinker.visual.merger.ln_q.weight": torch.randn(6),
-        "thinker.audio_tower.proj1.weight": torch.zeros(2, 2),
-    }
-    for e in range(2):
-        tensors[f"thinker.model.layers.0.mlp.experts.{e}.gate_proj.weight"] = torch.randn(3, 4)
-        tensors[f"thinker.model.layers.0.mlp.experts.{e}.up_proj.weight"] = torch.randn(3, 4)
-        tensors[f"thinker.model.layers.0.mlp.experts.{e}.down_proj.weight"] = torch.randn(4, 3)
-    save_file(tensors, src / "model.safetensors", metadata={"format": "pt"})
-    return tensors
-
-
-def test_extract_vl_roundtrip(tmp_path):
-    tool = _load_tool()
-    src, dst = tmp_path / "src", tmp_path / "dst"
-    tensors = _tiny_omni_src(src)
-
-    tool.extract(src, dst, shard_size_gb=1.0, variant="vl")
-
-    with open(dst / "config.json") as f:
-        cfg = json.load(f)
-    assert cfg["model_type"] == "qwen3_vl_moe"
-
-    out = {}
-    with safe_open(dst / "model.safetensors", framework="pt") as reader:
-        for key in reader.keys():
-            out[key] = reader.get_tensor(key)
-
-    assert "model.language_model.embed_tokens.weight" in out
-    assert "lm_head.weight" in out
-    assert "model.visual.merger.norm.weight" in out
-    assert not any("audio_tower" in k for k in out)
-    gate_up = out["model.language_model.layers.0.mlp.experts.gate_up_proj"]
-    down = out["model.language_model.layers.0.mlp.experts.down_proj"]
-    assert gate_up.shape == (2, 4, 6) and down.shape == (2, 3, 4)
-    expected_e1 = torch.cat(
-        [
-            tensors["thinker.model.layers.0.mlp.experts.1.gate_proj.weight"],
-            tensors["thinker.model.layers.0.mlp.experts.1.up_proj.weight"],
-        ],
-        dim=0,
-    ).T
-    assert torch.equal(gate_up[1], expected_e1)
-    # per-expert 2D names must NOT leak through
-    assert not any(".experts.0." in k or ".experts.1." in k for k in out)
-
-
-def test_extract_text_variant_unchanged(tmp_path):
-    tool = _load_tool()
-    src, dst = tmp_path / "src", tmp_path / "dst"
-    _tiny_omni_src(src)
-
-    tool.extract(src, dst, shard_size_gb=1.0, variant="text")
-
-    with open(dst / "config.json") as f:
-        assert json.load(f)["model_type"] == "qwen3_moe"
-    with safe_open(dst / "model.safetensors", framework="pt") as reader:
-        keys = list(reader.keys())
-    assert "model.embed_tokens.weight" in keys
-    assert any(".experts.0.gate_proj" in k for k in keys)  # text variant keeps per-expert 2D
-    assert not any("visual" in k for k in keys)
diff --git a/tools/extract_qwen3_omni_thinker.py b/tools/extract_qwen3_omni_thinker.py
index 7d52d20cff6..ebde57bbaf5 100644
--- a/tools/extract_qwen3_omni_thinker.py
+++ b/tools/extract_qwen3_omni_thinker.py
@@ -1,41 +1,19 @@
-"""Extract a standalone thinker checkpoint from a full Qwen3-Omni model.
+"""Extract a standalone Qwen3-MoE thinker checkpoint from a full Qwen3-Omni model.
 
-Two variants, both keyed on model_type so miles' AutoBridge can load them:
+miles loads HF via AutoBridge keyed on model_type; the composite omni checkpoint has
+no bridge, but the thinker text backbone is a plain Qwen3-MoE. This writes a
+self-contained HF dir (thinker text + lm_head, renamed, model_type=qwen3_moe) for the
+existing Qwen3MoEBridge. Audio/visual towers are dropped here on purpose: the trainer
+loads them frozen from the original omni checkpoint (see the audio-injection wrapper),
+and the rollout server keeps its own copies.
 
-- ``text`` (default): thinker text backbone + lm_head as a plain Qwen3-MoE
-  (model_type=qwen3_moe, per-expert 2D weights). Audio/visual towers are dropped on
-  purpose: the trainer loads the audio tower frozen from the original omni checkpoint
-  (see the audio-injection plugin), and the rollout server keeps its own copies.
-
-- ``vl``: thinker visual tower + text backbone masquerading as Qwen3-VL-MoE
-  (model_type=qwen3_vl_moe). Names are mechanically compatible except three merger
-  renames (ln_q->norm, mlp.0/2->linear_fc1/2, merger_list->deepstack_merger_list) and
-  the MoE experts, which the VL hub format fuses to 3D hub layout
-  (gate_up_proj [E,H,2I], down_proj [E,I,H]). This unlocks miles' existing
-  bridge-based Qwen3-VL training path (packed mrope + deepstack + CP) for image/video
-  input. The audio tower is dropped; omni-only rope fields ride along under
-  ``omni_sideband`` in config.json for the trainer-side video-rope override.
-
-    python tools/extract_qwen3_omni_thinker.py --src <omni> --dst <thinker> [--variant vl]
+    python tools/extract_qwen3_omni_thinker.py --src <omni> --dst <thinker>
 """
 
 from __future__ import annotations
 
-import re
-
 _NON_TEXT_THINKER_SUBMODULES = ("audio_tower.", "visual.", "model.audio_tower.", "model.visual.")
 
-_VL_MERGER_RENAMES = (
-    (re.compile(r"^visual\.merger\.ln_q\."), "visual.merger.norm."),
-    (re.compile(r"^visual\.merger\.mlp\.0\."), "visual.merger.linear_fc1."),
-    (re.compile(r"^visual\.merger\.mlp\.2\."), "visual.merger.linear_fc2."),
-    (re.compile(r"^visual\.merger_list\.(\d+)\.ln_q\."), r"visual.deepstack_merger_list.\1.norm."),
-    (re.compile(r"^visual\.merger_list\.(\d+)\.mlp\.0\."), r"visual.deepstack_merger_list.\1.linear_fc1."),
-    (re.compile(r"^visual\.merger_list\.(\d+)\.mlp\.2\."), r"visual.deepstack_merger_list.\1.linear_fc2."),
-)
-
-_EXPERT_RE = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$")
-
 
 def map_thinker_param_name(name: str) -> str | None:
     """Full omni param name -> standalone Qwen3-MoE name, or None to drop."""
@@ -47,44 +25,6 @@ def map_thinker_param_name(name: str) -> str | None:
     return rest.replace("model.language_model.", "model.")
 
 
-def map_thinker_param_name_vl(name: str) -> str | None:
-    """Full omni param name -> pseudo Qwen3-VL-MoE name, or None to drop.
-
-    Expert weights keep their per-expert names here; `extract(variant="vl")` fuses them
-    into the 3D hub tensors afterwards.
-    """
-    if not name.startswith("thinker."):
-        return None
-    rest = name[len("thinker.") :].replace("model.language_model.", "model.")
-    if rest.startswith(("audio_tower.", "model.audio_tower.")):
-        return None
-    if rest == "lm_head.weight":
-        return rest
-    if rest.startswith("visual."):
-        for pattern, replacement in _VL_MERGER_RENAMES:
-            if pattern.search(rest):
-                rest = pattern.sub(replacement, rest)
-                break
-        return f"model.{rest}"
-    if rest.startswith("model."):
-        return rest.replace("model.", "model.language_model.", 1)
-    return None
-
-
-def fuse_layer_experts_vl(gates, ups, downs):
-    """Per-expert 2D weights -> fused hub-layout tensors.
-
-    gate/up: [I, H] each -> gate_up_proj [E, H, 2I]; down: [H, I] -> down_proj [E, I, H].
-    Hub layout is what transformers (Transpose check_dims), Megatron-Bridge import, and
-    sglang's fused loaders all expect.
-    """
-    import torch
-
-    gate_up = torch.stack([torch.cat([g, u], dim=0).T for g, u in zip(gates, ups, strict=True)], dim=0)
-    down = torch.stack([d.T for d in downs], dim=0)
-    return gate_up.contiguous(), down.contiguous()
-
-
 def synthesize_thinker_config(omni_config: dict) -> dict:
     """Plain Qwen3-MoE config dict from the full omni config."""
     thinker = omni_config.get("thinker_config", omni_config)
@@ -98,92 +38,11 @@ def synthesize_thinker_config(omni_config: dict) -> dict:
     return cfg
 
 
-# vision fields whose omni spelling has no qwen3_vl_moe slot (dropped after use)
-_VL_VISION_DROPPED = ("apply_vit_abs_pos_embed", "image_size", "spatial_patch_size", "in_chans", "tokens_per_second")
-
-
-def synthesize_vl_config(omni_config: dict) -> dict:
-    """Pseudo Qwen3-VL-MoE config from the full omni config.
-
-    Omni numerics are preserved (vocab 152064, rope_theta 1e6, max_position_embeddings
-    65536); rope_scaling is normalized to the VL whitelist; the vision config gains
-    num_position_embeddings (the pos_embed table size, (image_size/patch_size)^2) and
-    drops omni-only spellings. Fields with no VL slot that the trainer still needs
-    (omni video rope) ride under `omni_sideband`.
-    """
-    thinker = omni_config.get("thinker_config", omni_config)
-
-    text = dict(thinker["text_config"])
-    text["model_type"] = "qwen3_vl_moe_text"
-    rope_scaling = dict(text.get("rope_scaling") or {})
-    text["rope_scaling"] = {
-        "rope_type": rope_scaling.get("rope_type") or rope_scaling.get("type", "default"),
-        "mrope_section": rope_scaling["mrope_section"],
-        "mrope_interleaved": rope_scaling.get("mrope_interleaved", True),
-    }
-    text.setdefault("bos_token_id", 151643)
-    text.setdefault("eos_token_id", 151645)
-
-    vision = dict(thinker["vision_config"])
-    vision["model_type"] = "qwen3_vl_moe"
-    if "num_position_embeddings" not in vision and "image_size" in vision:
-        patch = vision.get("patch_size") or vision.get("spatial_patch_size")
-        vision["num_position_embeddings"] = (vision["image_size"] // patch) ** 2
-    for key in _VL_VISION_DROPPED:
-        vision.pop(key, None)
-
-    cfg = {
-        "model_type": "qwen3_vl_moe",
-        "architectures": ["Qwen3VLMoeForConditionalGeneration"],
-        "tie_word_embeddings": thinker.get("tie_word_embeddings", False),
-        "text_config": text,
-        "vision_config": vision,
-        "omni_sideband": {
-            "position_id_per_seconds": thinker.get("position_id_per_seconds"),
-            "audio_token_id": thinker.get("audio_token_id"),
-        },
-    }
-    for k in ("image_token_id", "video_token_id", "vision_start_token_id", "vision_end_token_id"):
-        cfg[k] = thinker[k]
-    return cfg
-
-
-class _VlExpertAccumulator:
-    """Collects per-expert 2D weights per layer and emits fused hub-layout tensors."""
-
-    def __init__(self, num_experts: int):
-        self.num_experts = num_experts
-        self._layers: dict[int, dict[str, dict[int, object]]] = {}
-
-    def offer(self, vl_name: str, tensor) -> list[tuple[str, object]]:
-        match = _EXPERT_RE.match(vl_name.replace("model.language_model.", "model.", 1))
-        if match is None:
-            return [(vl_name, tensor)]
-        layer, expert, kind = int(match.group(1)), int(match.group(2)), match.group(3)
-        buckets = self._layers.setdefault(layer, {"gate_proj": {}, "up_proj": {}, "down_proj": {}})
-        buckets[kind][expert] = tensor
-        if all(len(buckets[k]) == self.num_experts for k in ("gate_proj", "up_proj", "down_proj")):
-            del self._layers[layer]
-            order = range(self.num_experts)
-            gate_up, down = fuse_layer_experts_vl(
-                [buckets["gate_proj"][e] for e in order],
-                [buckets["up_proj"][e] for e in order],
-                [buckets["down_proj"][e] for e in order],
-            )
-            prefix = f"model.language_model.layers.{layer}.mlp.experts"
-            return [(f"{prefix}.gate_up_proj", gate_up), (f"{prefix}.down_proj", down)]
-        return []
-
-    def assert_drained(self):
-        assert not self._layers, f"incomplete expert layers left over: {sorted(self._layers)}"
-
-
-def extract(src, dst, shard_size_gb: float = 5.0, variant: str = "text") -> int:
+def extract(src, dst, shard_size_gb: float = 5.0) -> int:
     """Stream thinker tensors from src into sharded safetensors under dst.
 
     Each shard is written (and its buffer released) as soon as it fills, so peak RAM is
-    one shard (plus, for the vl variant, at most a few layers of in-flight expert
-    weights). Returns the number of tensors written.
+    one shard, not the whole ~60 GB thinker. Returns the number of kept tensors.
     """
     import json
     import os
@@ -193,16 +52,13 @@ def extract(src, dst, shard_size_gb: float = 5.0, variant: str = "text") -> int:
     from safetensors import safe_open
     from safetensors.torch import save_file
 
-    assert variant in ("text", "vl"), variant
     src, dst = Path(src), Path(dst)
     dst.mkdir(parents=True, exist_ok=True)
 
     with open(src / "config.json") as f:
         omni_config = json.load(f)
-    synthesize = synthesize_thinker_config if variant == "text" else synthesize_vl_config
-    out_config = synthesize(omni_config)
     with open(dst / "config.json", "w") as f:
-        json.dump(out_config, f, indent=2)
+        json.dump(synthesize_thinker_config(omni_config), f, indent=2)
 
     for fname in (
         "tokenizer.json",
@@ -235,12 +91,6 @@ def extract(src, dst, shard_size_gb: float = 5.0, variant: str = "text") -> int:
     else:
         shard_files = ["model.safetensors"]
 
-    map_name = map_thinker_param_name if variant == "text" else map_thinker_param_name_vl
-    accumulator = None
-    if variant == "vl":
-        thinker = omni_config.get("thinker_config", omni_config)
-        accumulator = _VlExpertAccumulator(num_experts=thinker["text_config"]["num_experts"])
-
     shard_size_bytes = int(shard_size_gb * (1024**3))
     out_index: dict[str, str] = {}
     shard_names: list[str] = []
@@ -261,28 +111,18 @@ def flush():
             total_size += t.numel() * t.element_size()
         buf, buf_bytes = {}, 0
 
-    def emit(out_name, tensor):
-        nonlocal buf_bytes, total_kept
-        buf[out_name] = tensor
-        buf_bytes += tensor.numel() * tensor.element_size()
-        total_kept += 1
-        if buf_bytes >= shard_size_bytes:
-            flush()
-
     for shard_file in shard_files:
         with safe_open(src / shard_file, framework="pt") as reader:
             for key in reader.keys():
-                new_key = map_name(key)
+                new_key = map_thinker_param_name(key)
                 if new_key is None:
                     continue
                 tensor = reader.get_tensor(key)
-                if accumulator is not None:
-                    for out_name, out_tensor in accumulator.offer(new_key, tensor):
-                        emit(out_name, out_tensor)
-                else:
-                    emit(new_key, tensor)
-    if accumulator is not None:
-        accumulator.assert_drained()
+                buf[new_key] = tensor
+                buf_bytes += tensor.numel() * tensor.element_size()
+                total_kept += 1
+                if buf_bytes >= shard_size_bytes:
+                    flush()
     flush()
 
     if total_kept == 0:
@@ -294,7 +134,7 @@ def emit(out_name, tensor):
         with open(dst / "model.safetensors.index.json", "w") as f:
             json.dump({"metadata": {"total_size": total_size}, "weight_map": out_index}, f, indent=2)
 
-    print(f"[done] {total_kept} thinker tensors ({variant}) -> {dst}")
+    print(f"[done] {total_kept} thinker tensors -> {dst}")
     return total_kept
 
 
@@ -305,9 +145,8 @@ def main() -> None:
     parser.add_argument("--src", required=True)
     parser.add_argument("--dst", required=True)
     parser.add_argument("--shard-size-gb", type=float, default=5.0)
-    parser.add_argument("--variant", choices=["text", "vl"], default="text")
     args = parser.parse_args()
-    extract(args.src, args.dst, shard_size_gb=args.shard_size_gb, variant=args.variant)
+    extract(args.src, args.dst, shard_size_gb=args.shard_size_gb)
 
 
 if __name__ == "__main__":

From 33ae6614d10e71c65d2cf23ec4f4f4b8933de672 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Sun, 26 Jul 2026 03:40:46 +0800
Subject: [PATCH 26/32] [omni-vl] Video-R1-260k prep: mixed image/video MCQ ->
 miles jsonl (gpqa reward)

---
 examples/omni_thinker/prepare_video_r1.py | 99 +++++++++++++++++++++++
 tests/fast/test_prepare_video_r1.py       | 93 +++++++++++++++++++++
 2 files changed, 192 insertions(+)
 create mode 100644 examples/omni_thinker/prepare_video_r1.py
 create mode 100644 tests/fast/test_prepare_video_r1.py

diff --git a/examples/omni_thinker/prepare_video_r1.py b/examples/omni_thinker/prepare_video_r1.py
new file mode 100644
index 00000000000..81d514781fe
--- /dev/null
+++ b/examples/omni_thinker/prepare_video_r1.py
@@ -0,0 +1,99 @@
+"""Convert Video-R1-260k (mixed image/video MCQ with verifiable answers) into miles jsonl.
+
+Rows carry lettered options and `<answer>X</answer>` solutions, so `--rm-type gpqa`
+scores them directly. Both media columns are always present (one empty), pairing with
+`--multimodal-keys '{"video": "videos", "image": "images"}'`.
+
+    python examples/omni_thinker/prepare_video_r1.py \
+        --src <dir>/Video-R1-260k.json --dst <data_dir>/video_r1.jsonl \
+        --media-root <dir> [--data-types video image] [--path-prefixes CLEVRER] [--max-samples 5000]
+"""
+
+import argparse
+import json
+import re
+from pathlib import Path
+
+_PLACEHOLDER = {"video": "<video>", "image": "<image>"}
+_PROMPT_TEMPLATE = (
+    "{placeholder}{problem}\n{options}\n"
+    "Watch carefully and choose the best option. Respond with the option letter in the form 'Answer: <letter>'."
+)
+_ANSWER_RE = re.compile(r"<answer>\s*([A-Za-z])\s*</answer>")
+
+
+def convert_row(row: dict, media_root) -> dict:
+    match = _ANSWER_RE.search(row["solution"])
+    assert match, f"no letter answer in solution: {row['solution']!r} (problem_id={row.get('problem_id')})"
+    letter = match.group(1).upper()
+    media_path = str(Path(media_root) / row["path"].removeprefix("./"))
+    data_type = row["data_type"]
+    assert data_type in _PLACEHOLDER, f"unsupported data_type: {data_type}"
+    return {
+        "prompt": _PROMPT_TEMPLATE.format(
+            placeholder=_PLACEHOLDER[data_type],
+            problem=row["problem"].strip(),
+            options="\n".join(row["options"]),
+        ),
+        "videos": [media_path] if data_type == "video" else [],
+        "images": [media_path] if data_type == "image" else [],
+        "label": letter,
+        "metadata": {"choices": list(row["options"])},
+    }
+
+
+def convert_file(
+    src,
+    dst,
+    media_root,
+    data_types: tuple[str, ...] = ("video", "image"),
+    path_prefixes: tuple[str, ...] | None = None,
+    problem_types: tuple[str, ...] = ("multiple choice",),
+    require_media: bool = False,
+    max_samples: int | None = None,
+) -> int:
+    with open(src) as f:
+        rows = json.load(f)
+
+    n = 0
+    with open(dst, "w") as fout:
+        for row in rows:
+            if max_samples is not None and n >= max_samples:
+                break
+            if row["data_type"] not in data_types or row["problem_type"] not in problem_types:
+                continue
+            rel = row["path"].removeprefix("./")
+            if path_prefixes and not rel.startswith(tuple(path_prefixes)):
+                continue
+            if require_media and not (Path(media_root) / rel).exists():
+                continue
+            fout.write(json.dumps(convert_row(row, media_root)) + "\n")
+            n += 1
+    assert n > 0, f"no rows converted from {src}"
+    return n
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--src", required=True, help="Video-R1-260k.json")
+    parser.add_argument("--dst", required=True)
+    parser.add_argument("--media-root", required=True, help="dir the media zips were extracted into")
+    parser.add_argument("--data-types", nargs="+", default=["video", "image"])
+    parser.add_argument("--path-prefixes", nargs="+", default=None, help="e.g. CLEVRER to match downloaded zips")
+    parser.add_argument("--require-media", action="store_true", help="skip rows whose media file is absent")
+    parser.add_argument("--max-samples", type=int, default=None)
+    args = parser.parse_args()
+    n = convert_file(
+        args.src,
+        args.dst,
+        media_root=args.media_root,
+        data_types=tuple(args.data_types),
+        path_prefixes=tuple(args.path_prefixes) if args.path_prefixes else None,
+        require_media=args.require_media,
+        max_samples=args.max_samples,
+    )
+    print(f"[done] {n} samples -> {args.dst}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/fast/test_prepare_video_r1.py b/tests/fast/test_prepare_video_r1.py
new file mode 100644
index 00000000000..78e51f294ed
--- /dev/null
+++ b/tests/fast/test_prepare_video_r1.py
@@ -0,0 +1,93 @@
+"""Video-R1-260k (image+video MCQ) -> miles jsonl conversion for the VL thinker example."""
+
+import importlib.util
+import json
+from pathlib import Path
+
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[])
+
+_REPO = Path(__file__).resolve().parents[2]
+
+
+def _load_tool():
+    path = _REPO / "examples" / "omni_thinker" / "prepare_video_r1.py"
+    spec = importlib.util.spec_from_file_location("prepare_video_r1", path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+VIDEO_ROW = {
+    "problem_id": 2,
+    "problem": "What appears on the screen?",
+    "data_type": "video",
+    "problem_type": "multiple choice",
+    "options": ["A. A notification", "B. A command", "C. A warning", "D. An update"],
+    "solution": "<answer>A</answer>",
+    "path": "./CLEVRER/video_validation/video_10000.mp4",
+    "data_source": "",
+}
+IMAGE_ROW = {
+    "problem_id": 1,
+    "problem": "What role does the cylinder play?",
+    "data_type": "image",
+    "problem_type": "multiple choice",
+    "options": ["A) focus", "B) weight", "C) casing", "D) medium"],
+    "solution": "<answer>D</answer>",
+    "path": "./Knowledge/ArxivQA/images/1702.jpg",
+    "data_source": "",
+}
+
+
+def test_convert_video_row_builds_video_prompt_and_label():
+    record = _load_tool().convert_row(VIDEO_ROW, media_root="/data/vr1")
+
+    assert record["label"] == "A"
+    assert record["videos"] == ["/data/vr1/CLEVRER/video_validation/video_10000.mp4"]
+    assert record["images"] == []
+    prompt = record["prompt"]
+    assert prompt.startswith("<video>") and "<image>" not in prompt
+    assert "A. A notification" in prompt and "D. An update" in prompt
+    assert record["metadata"]["choices"] == VIDEO_ROW["options"]
+
+
+def test_convert_image_row_builds_image_prompt():
+    record = _load_tool().convert_row(IMAGE_ROW, media_root="/data/vr1")
+
+    assert record["label"] == "D"
+    assert record["images"] == ["/data/vr1/Knowledge/ArxivQA/images/1702.jpg"]
+    assert record["videos"] == []
+    assert record["prompt"].startswith("<image>")
+
+
+def test_reward_roundtrip_via_gpqa():
+    from miles.rollout.rm_hub.gpqa import compute_gpqa_reward
+
+    record = _load_tool().convert_row(VIDEO_ROW, media_root="/data/vr1")
+
+    assert compute_gpqa_reward("The answer is A.", record["label"], metadata=record["metadata"]) == 1.0
+    assert compute_gpqa_reward("Answer: B", record["label"], metadata=record["metadata"]) == 0.0
+
+
+def test_convert_file_filters_types_sources_and_existing_media(tmp_path):
+    tool = _load_tool()
+    src = tmp_path / "Video-R1-260k.json"
+    other = {**VIDEO_ROW, "problem_id": 3, "problem_type": "numerical", "solution": "<answer>7</answer>"}
+    missing = {**VIDEO_ROW, "problem_id": 4, "path": "./CLEVRER/missing.mp4"}
+    off_source = {**VIDEO_ROW, "problem_id": 5, "path": "./LLaVA-Video-178K/x.mp4"}
+    with open(src, "w") as f:
+        json.dump([VIDEO_ROW, IMAGE_ROW, other, missing, off_source], f)
+    media_root = tmp_path / "media"
+    (media_root / "CLEVRER" / "video_validation").mkdir(parents=True)
+    (media_root / "CLEVRER" / "video_validation" / "video_10000.mp4").write_bytes(b"x")
+
+    out = tmp_path / "vr1.jsonl"
+    n = tool.convert_file(
+        src, out, media_root=media_root, data_types=("video",), path_prefixes=("CLEVRER",), require_media=True
+    )
+
+    assert n == 1
+    [record] = [json.loads(line) for line in open(out)]
+    assert record["label"] == "A" and record["videos"][0].endswith("video_10000.mp4")

From a0bcdd89b118b511f77d8c8a724aaca0a5df686e Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Sun, 26 Jul 2026 04:12:19 +0800
Subject: [PATCH 27/32] [omni-vl] omni TM-RoPE override for the bridge trainer
 + debug alloc headroom

- miles_plugins/models/qwen3_omni_thinker_vl.py: video segments reroute to the ported
  omni get_rope_index (parity-tested against the HF reference on mixed layouts);
  image/text delegate to the bridge formula (proven identical); video_second_per_grid
  popped pre-forward and consumed via an ordered cursor; provider freeze knobs cleared
- example: expandable_segments for MoE dispatch allgather headroom
---
 .../run_qwen3_omni_thinker_grpo.py            |   2 +
 miles_plugins/models/qwen3_omni_thinker_vl.py | 192 ++++++++++++++++++
 tests/fast/test_qwen3_omni_vl_rope.py         | 153 ++++++++++++++
 3 files changed, 347 insertions(+)
 create mode 100644 miles_plugins/models/qwen3_omni_thinker_vl.py
 create mode 100644 tests/fast/test_qwen3_omni_vl_rope.py

diff --git a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
index a5f650244d3..5382a1e2d52 100644
--- a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
+++ b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
@@ -186,6 +186,8 @@ def execute(args: ScriptArgs):
         megatron_path=args.megatron_path,
         extra_env_vars={
             "FLASHINFER_DISABLE_VERSION_CHECK": "1",
+            # MoE token-dispatch allgathers die on fragmentation headroom otherwise
+            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
             "PYTHONPATH": f"{args.megatron_path}",
         },
     )
diff --git a/miles_plugins/models/qwen3_omni_thinker_vl.py b/miles_plugins/models/qwen3_omni_thinker_vl.py
new file mode 100644
index 00000000000..5ef805f2a8b
--- /dev/null
+++ b/miles_plugins/models/qwen3_omni_thinker_vl.py
@@ -0,0 +1,192 @@
+"""Omni video TM-RoPE + kwarg hygiene for the pseudo-Qwen3-VL thinker.
+
+The pseudo-VL checkpoint (tools/extract_qwen3_omni_thinker.py --variant vl) trains on
+miles' bridge Qwen3-VL path. Image/text positions are identical between the omni and VL
+formulas, but video differs: the omni server rolls out with TM-RoPE
+(t_index = k * second_per_grid * position_id_per_seconds, one contiguous block per
+video), while the VL/bridge formula splits per frame around timestamp text tokens the
+omni processor never emits — silently wrong positions after frame 1. This plugin swaps
+the video branch to the omni formula and keeps delegating image-only segments to the
+bridge implementation.
+
+It also pops `video_second_per_grid` (omni processor output with no slot in the bridge
+forward) and un-freezes the language model on providers whose defaults freeze it.
+
+Install via `install_omni_vl(args)` — wired behind --qwen3-omni-vl in model_provider.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import threading
+from pathlib import Path
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+_tls = threading.local()
+_PATCHED = "_miles_omni_vl_patched"
+
+
+def omni_video_rope_index(
+    spatial_merge_size: int,
+    image_token_id: int,
+    video_token_id: int,
+    vision_start_token_id: int,
+    input_ids: torch.Tensor,
+    image_grid_thw: torch.Tensor | None,
+    video_grid_thw: torch.Tensor | None,
+    video_second_per_grid: torch.Tensor | None,
+    position_id_per_seconds: float,
+) -> torch.Tensor:
+    """Qwen3-Omni thinker TM-RoPE positions for one [1, s] row (text/image/video, no audio).
+
+    Port of HF Qwen3OmniMoe get_rope_index minus the audio branches: video temporal
+    index strides by second_per_grid * position_id_per_seconds within one contiguous
+    block; images use t = const; text runs sequentially from the running max + 1.
+    """
+    ids = input_ids[0]
+    tokens = ids.tolist()
+    device = input_ids.device
+    image_idx = video_idx = 0
+    st = 0
+    pos_chunks: list[torch.Tensor] = []
+
+    def _vision_block(st_idx: float, grid, t_index: torch.Tensor) -> torch.Tensor:
+        t, h, w = int(grid[0]), int(grid[1]) // spatial_merge_size, int(grid[2]) // spatial_merge_size
+        t_pos = t_index.view(t, 1).expand(t, h * w).flatten()
+        h_pos = torch.arange(h).view(1, h, 1).expand(t, h, w).flatten()
+        w_pos = torch.arange(w).view(1, 1, w).expand(t, h, w).flatten()
+        return torch.stack([t_pos, h_pos.float(), w_pos.float()]) + st_idx
+
+    while st < len(tokens):
+        st_idx = pos_chunks[-1].max().item() + 1 if pos_chunks else 0
+        try:
+            next_vision = tokens.index(vision_start_token_id, st)
+        except ValueError:
+            next_vision = len(tokens)
+
+        text_len = next_vision - st
+        if text_len > 0:
+            pos_chunks.append(torch.arange(text_len).view(1, -1).expand(3, -1).float() + st_idx)
+            st_idx += text_len
+        if next_vision >= len(tokens):
+            break
+
+        # vision_start token itself is ordinary text
+        pos_chunks.append(torch.arange(1).view(1, -1).expand(3, -1).float() + st_idx)
+        st_idx += 1
+        marker = tokens[next_vision + 1]
+        if marker == image_token_id:
+            grid = image_grid_thw[image_idx]
+            t_index = (torch.arange(int(grid[0])) * 1 * position_id_per_seconds).float()
+            block = _vision_block(st_idx, grid, t_index)
+            image_idx += 1
+        elif marker == video_token_id:
+            grid = video_grid_thw[video_idx]
+            second = float(video_second_per_grid[video_idx])
+            t_index = (torch.arange(int(grid[0])) * second * position_id_per_seconds).float()
+            block = _vision_block(st_idx, grid, t_index)
+            video_idx += 1
+        else:
+            raise AssertionError(f"vision_start not followed by an image/video token: {marker}")
+        pos_chunks.append(block)
+        st = next_vision + 1 + block.shape[1]
+
+    positions = torch.cat(pos_chunks, dim=1)
+    assert positions.shape[1] == len(tokens), f"{positions.shape=} vs {len(tokens)=}"
+    return positions.view(3, 1, -1).to(device)
+
+
+def _make_omni_aware_rope_index(orig_get_rope_index, position_id_per_seconds: float):
+    """Bridge-signature get_rope_index that reroutes video segments to the omni formula.
+
+    Consumes `video_second_per_grid` sequentially from _tls (segments are processed in
+    order by the qwen3_vl packed-mrope patch, so a running cursor stays aligned with the
+    grid slices it hands us).
+    """
+
+    def omni_aware(merge, img_id, vid_id, vstart, input_ids, image_grid_thw=None, video_grid_thw=None, attention_mask=None):
+        if video_grid_thw is None or video_grid_thw.numel() == 0:
+            return orig_get_rope_index(
+                merge, img_id, vid_id, vstart, input_ids,
+                image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask,
+            )
+        seconds_all = getattr(_tls, "video_second_per_grid", None)
+        assert seconds_all is not None, (
+            "video tokens reached the trainer without video_second_per_grid; the omni processor "
+            "always emits it - check the multimodal_train_inputs plumbing"
+        )
+        cursor = getattr(_tls, "video_cursor", 0)
+        n_videos = video_grid_thw.shape[0]
+        seconds = seconds_all[cursor : cursor + n_videos]
+        assert seconds.numel() == n_videos, f"video_second_per_grid exhausted: {cursor=} {n_videos=} {seconds_all.numel()=}"
+        _tls.video_cursor = cursor + n_videos
+        positions = omni_video_rope_index(
+            merge, img_id, vid_id, vstart, input_ids,
+            image_grid_thw, video_grid_thw, seconds, position_id_per_seconds,
+        )
+        return positions, None
+
+    return omni_aware
+
+
+def resolve_position_id_per_seconds(hf_checkpoint: str | Path) -> float:
+    """The pseudo-VL config has no VL slot for this; the extractor stashes it in omni_sideband."""
+    with open(Path(hf_checkpoint) / "config.json") as f:
+        config = json.load(f)
+    sideband = config.get("omni_sideband") or {}
+    value = sideband.get("position_id_per_seconds")
+    assert value is not None, (
+        f"omni_sideband.position_id_per_seconds missing from {hf_checkpoint}/config.json - "
+        "re-extract with tools/extract_qwen3_omni_thinker.py --variant vl"
+    )
+    return float(value)
+
+
+def install_omni_vl(args) -> None:
+    """Install the omni video rope override + video_second_per_grid pop on the bridge model."""
+    import importlib
+
+    import miles_plugins.models.qwen3_vl as qwen3_vl_patch
+
+    model_mod = importlib.import_module("megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model")
+    if getattr(model_mod, _PATCHED, False):
+        return
+
+    position_id_per_seconds = resolve_position_id_per_seconds(args.hf_checkpoint)
+
+    orig_build = qwen3_vl_patch._build_packed_positions
+
+    def build_with_omni_video(model, parsed, kwargs, orig_get_rope_index):
+        return orig_build(model, parsed, kwargs, _make_omni_aware_rope_index(orig_get_rope_index, position_id_per_seconds))
+
+    qwen3_vl_patch._build_packed_positions = build_with_omni_video
+
+    Qwen3VLModel = model_mod.Qwen3VLModel
+    inner_forward = Qwen3VLModel.forward
+
+    def forward(self, *fargs, **kwargs):
+        seconds = kwargs.pop("video_second_per_grid", None)
+        _tls.video_second_per_grid = seconds
+        _tls.video_cursor = 0
+        try:
+            return inner_forward(self, *fargs, **kwargs)
+        finally:
+            _tls.video_second_per_grid = None
+            _tls.video_cursor = 0
+
+    Qwen3VLModel.forward = forward
+    setattr(model_mod, _PATCHED, True)
+    logger.info(f"omni VL patch installed (position_id_per_seconds={position_id_per_seconds})")
+
+
+def unfreeze_provider(provider) -> None:
+    """Some megatron-bridge versions default Qwen3VLMoEModelProvider to a frozen LM;
+    RL must train the backbone, so force-clear every freeze knob that exists."""
+    for field in ("freeze_language_model", "freeze_vision_model", "freeze_vision_projection"):
+        if getattr(provider, field, False):
+            logger.warning(f"pseudo-VL provider had {field}=True; clearing it for RL training")
+            setattr(provider, field, False)
diff --git a/tests/fast/test_qwen3_omni_vl_rope.py b/tests/fast/test_qwen3_omni_vl_rope.py
new file mode 100644
index 00000000000..5f98384fbb7
--- /dev/null
+++ b/tests/fast/test_qwen3_omni_vl_rope.py
@@ -0,0 +1,153 @@
+"""Omni TM-RoPE port for the pseudo-VL trainer: parity with the HF reference."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=30, suite="stage-a-cpu", labels=[])
+
+from miles_plugins.models.qwen3_omni_thinker_vl import (
+    _make_omni_aware_rope_index,
+    _tls,
+    omni_video_rope_index,
+)
+
+MERGE = 2
+IMG, VID, VSTART, VEND = 151655, 151656, 151652, 151653
+POS_PER_SEC = 13.0
+
+
+def _hf_reference():
+    from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+        Qwen3OmniMoePreTrainedModelForConditionalGeneration,
+    )
+
+    ref = Qwen3OmniMoePreTrainedModelForConditionalGeneration.__new__(
+        Qwen3OmniMoePreTrainedModelForConditionalGeneration
+    )
+    config = SimpleNamespace(
+        image_token_id=IMG,
+        video_token_id=VID,
+        audio_token_id=151675,
+        vision_start_token_id=VSTART,
+        audio_start_token_id=151669,
+        position_id_per_seconds=POS_PER_SEC,
+        vision_config=SimpleNamespace(spatial_merge_size=MERGE),
+    )
+    object.__setattr__(ref, "config", config)
+    object.__setattr__(ref, "spatial_merge_size", MERGE)
+    return ref
+
+
+def _build_ids(segments) -> tuple[torch.Tensor, list, list, list]:
+    """segments: list of ('text', n) | ('image', (t,h,w)) | ('video', (t,h,w), seconds)."""
+    ids, image_grids, video_grids, seconds = [], [], [], []
+    text_token = 7
+    for seg in segments:
+        if seg[0] == "text":
+            ids += [text_token] * seg[1]
+        elif seg[0] == "image":
+            t, h, w = seg[1]
+            image_grids.append([t, h, w])
+            ids += [VSTART] + [IMG] * (t * h * w // MERGE**2) + [VEND]
+        else:
+            t, h, w = seg[1]
+            video_grids.append([t, h, w])
+            seconds.append(seg[2])
+            ids += [VSTART] + [VID] * (t * h * w // MERGE**2) + [VEND]
+    return torch.tensor([ids]), image_grids, video_grids, seconds
+
+
+@pytest.mark.parametrize(
+    "segments",
+    [
+        [("text", 5), ("video", (4, 4, 6), 0.5), ("text", 3)],
+        [("text", 2), ("image", (1, 4, 4)), ("text", 1), ("video", (2, 6, 4), 2.0), ("text", 4)],
+        [("video", (2, 4, 4), 1.0), ("video", (3, 4, 4), 0.25), ("text", 2)],
+        [("text", 6)],
+        [("image", (1, 6, 6)), ("text", 2)],
+    ],
+)
+def test_omni_video_rope_matches_hf_reference(segments):
+    input_ids, image_grids, video_grids, seconds = _build_ids(segments)
+    image_grid_thw = torch.tensor(image_grids) if image_grids else None
+    video_grid_thw = torch.tensor(video_grids) if video_grids else None
+    second_per_grid = torch.tensor(seconds) if seconds else None
+
+    ours = omni_video_rope_index(
+        MERGE, IMG, VID, VSTART, input_ids, image_grid_thw, video_grid_thw, second_per_grid, POS_PER_SEC
+    )
+
+    ref = _hf_reference()
+    expected, _ = ref.get_rope_index(
+        input_ids=input_ids,
+        image_grid_thw=image_grid_thw,
+        video_grid_thw=video_grid_thw,
+        attention_mask=torch.ones_like(input_ids),
+        use_audio_in_video=False,
+        audio_seqlens=None,
+        second_per_grids=second_per_grid,
+    )
+
+    assert ours.shape == expected.shape, f"{ours.shape=} {expected.shape=}"
+    assert torch.allclose(ours.float(), expected.float()), (
+        f"positions diverge from the HF omni reference for {segments}"
+    )
+
+
+def test_omni_aware_rope_delegates_images_and_consumes_seconds_in_order():
+    calls = []
+
+    def fake_orig(merge, img_id, vid_id, vstart, ids, image_grid_thw=None, video_grid_thw=None, attention_mask=None):
+        calls.append("orig")
+        return torch.zeros(3, 1, ids.shape[1]), None
+
+    omni_aware = _make_omni_aware_rope_index(fake_orig, POS_PER_SEC)
+    _tls.video_second_per_grid = torch.tensor([0.5, 2.0])
+    _tls.video_cursor = 0
+    try:
+        # image-only segment -> delegate to the bridge implementation
+        ids_img, image_grids, _, _ = _build_ids([("image", (1, 4, 4)), ("text", 1)])
+        omni_aware(MERGE, IMG, VID, VSTART, ids_img, torch.tensor(image_grids), None, None)
+        assert calls == ["orig"]
+
+        # two video segments consume seconds sequentially
+        ids_v1, _, grids_v1, _ = _build_ids([("video", (2, 4, 4), 0.5)])
+        pos1, _ = omni_aware(MERGE, IMG, VID, VSTART, ids_v1, None, torch.tensor(grids_v1), None)
+        assert _tls.video_cursor == 1
+        ids_v2, _, grids_v2, _ = _build_ids([("video", (2, 4, 4), 2.0)])
+        pos2, _ = omni_aware(MERGE, IMG, VID, VSTART, ids_v2, None, torch.tensor(grids_v2), None)
+        assert _tls.video_cursor == 2
+        # different seconds -> different temporal strides
+        assert not torch.equal(pos1, pos2)
+        # stride check: frame 2 of video 1 starts at t = 0.5 * 13 relative to block start
+        t_axis = pos1[0, 0]
+        block = t_axis[1:9]  # exactly the 8 video tokens (skip vision_start, stop before vision_end)
+        assert block.max() - block.min() == pytest.approx(0.5 * POS_PER_SEC)
+    finally:
+        _tls.video_second_per_grid = None
+        _tls.video_cursor = 0
+
+
+def test_omni_aware_rope_fails_loud_without_seconds():
+    omni_aware = _make_omni_aware_rope_index(lambda *a, **k: (None, None), POS_PER_SEC)
+    _tls.video_second_per_grid = None
+    ids, _, grids, _ = _build_ids([("video", (2, 4, 4), 0.5)])
+    with pytest.raises(AssertionError, match="video_second_per_grid"):
+        omni_aware(MERGE, IMG, VID, VSTART, ids, None, torch.tensor(grids), None)
+
+
+def test_resolve_position_id_per_seconds(tmp_path):
+    from miles_plugins.models.qwen3_omni_thinker_vl import resolve_position_id_per_seconds
+
+    import json
+
+    (tmp_path / "config.json").write_text(json.dumps({"omni_sideband": {"position_id_per_seconds": 13}}))
+    assert resolve_position_id_per_seconds(tmp_path) == 13.0
+
+    (tmp_path / "config.json").write_text(json.dumps({}))
+    with pytest.raises(AssertionError, match="omni_sideband"):
+        resolve_position_id_per_seconds(tmp_path)

From cdad2e77335b38c03d7dfb5c030ed818b61818bb Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Sun, 26 Jul 2026 04:22:09 +0800
Subject: [PATCH 28/32] [omni-vl] debug tier: drop the zero-coef KL ref copy;
 revert expandable_segments (NCCL-incompatible)

---
 examples/omni_thinker/run_qwen3_omni_thinker_grpo.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
index 5382a1e2d52..95a0477b6b9 100644
--- a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
+++ b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py
@@ -136,10 +136,10 @@ def execute(args: ScriptArgs):
 
     grpo_args = (
         "--advantage-estimator grpo "
-        "--use-kl-loss "
-        "--kl-loss-coef 0.00 "
-        "--kl-loss-type low_var_kl "
-        "--entropy-coef 0.00 "
+        # kl coef is 0, but --use-kl-loss still loads a full ref-model copy (~33 GB/GPU
+        # at TP1); the debug tier drops it to fit 2 GPUs
+        + ("" if debug_minimal else "--use-kl-loss --kl-loss-coef 0.00 --kl-loss-type low_var_kl ")
+        + "--entropy-coef 0.00 "
         "--eps-clip 0.2 "
         "--eps-clip-high 0.28 "
     )
@@ -186,8 +186,6 @@ def execute(args: ScriptArgs):
         megatron_path=args.megatron_path,
         extra_env_vars={
             "FLASHINFER_DISABLE_VERSION_CHECK": "1",
-            # MoE token-dispatch allgathers die on fragmentation headroom otherwise
-            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
             "PYTHONPATH": f"{args.megatron_path}",
         },
     )

From 2e12b12049d7706102448692e2733b9db5c709e2 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Sun, 26 Jul 2026 10:47:57 +0800
Subject: [PATCH 29/32] [omni-vl] trainer wiring + example: --qwen3-omni-vl
 installs the TM-RoPE video override, clears provider freeze knobs; Video-R1
 example (CLEVRER video / Chart image tiers, external omni server, sync skip
 until the rename shim is wired); pseudo-VL -> omni server name map with
 extraction roundtrip tests

---
 .../run_qwen3_omni_thinker_vl_grpo.py         | 190 ++++++++++++++++++
 .../backends/megatron_utils/model_provider.py |   8 +
 miles/utils/arguments.py                      |  10 +
 miles_plugins/models/qwen3_omni_thinker_vl.py |  34 ++++
 tests/fast/test_qwen3_omni_vl_rope.py         |  35 ++++
 5 files changed, 277 insertions(+)
 create mode 100644 examples/omni_thinker/run_qwen3_omni_thinker_vl_grpo.py

diff --git a/examples/omni_thinker/run_qwen3_omni_thinker_vl_grpo.py b/examples/omni_thinker/run_qwen3_omni_thinker_vl_grpo.py
new file mode 100644
index 00000000000..5aa8066c179
--- /dev/null
+++ b/examples/omni_thinker/run_qwen3_omni_thinker_vl_grpo.py
@@ -0,0 +1,190 @@
+"""GRPO on the Qwen3-Omni thinker with image/video input (pseudo-Qwen3-VL bridge route).
+
+The thinker (visual tower + text backbone) is extracted as a qwen3_vl_moe checkpoint and
+trained on miles' bridge Qwen3-VL path; --qwen3-omni-vl installs the omni TM-RoPE video
+override (image positions already match; see miles_plugins/models/qwen3_omni_thinker_vl.py).
+Rollout runs on a standalone sglang-omni server through the processed-multimodal contract.
+
+Data: Video-R1-260k MCQ rows scored by --rm-type gpqa. --modality video uses the CLEVRER
+zip; --modality image uses the Chart zips (single-source subsets keep the download small).
+
+v1 is off-policy (--debug-skip-weight-update): the bridge exports pseudo-VL names, and the
+rename shim into the omni server's thinker.* namespace is not wired into the update path
+yet (pseudo_vl_to_omni_server_name is ready + tested).
+"""
+
+import os
+from dataclasses import dataclass
+from typing import Literal
+
+import miles.utils.external_utils.command_utils as U
+
+OMNI_MODEL = "Qwen3-Omni-30B-A3B-Instruct"
+THINKER_VL_MODEL = "Qwen3-Omni-30B-A3B-Thinker-VL"
+MEGATRON_MODEL_TYPE = "qwen3-omni-30B-A3B-thinker"
+VIDEO_R1 = "Video-R1/Video-R1-data"
+_MODALITY_SOURCES = {"video": ["CLEVRER"], "image": ["Chart"]}
+
+
+@dataclass
+class ScriptArgs(U.ExecuteTrainConfig):
+    mode: Literal["normal", "debug_minimal"] = "normal"
+    modality: Literal["video", "image"] = "video"
+    run_id: str = U.create_run_id()
+    actor_tp: int = 4
+    data_dir: str = "/root/datasets"
+    model_dir: str = "/root/models"
+    megatron_path: str = "/root/Megatron-LM"
+    omni_server_ip: str = "127.0.0.1"
+    omni_server_port: int = 30000
+    omni_server_tp: int = 4
+    max_samples: int = 5120
+    extra_args: str = ""
+
+
+def prepare(args: ScriptArgs):
+    U.exec_command(f"mkdir -p {args.model_dir} {args.data_dir}")
+    U.exec_command(f"hf download Qwen/{OMNI_MODEL} --local-dir {args.model_dir}/{OMNI_MODEL}")
+    repo = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    U.exec_command(
+        f"python {repo}/tools/extract_qwen3_omni_thinker.py --variant vl "
+        f"--src {args.model_dir}/{OMNI_MODEL} --dst {args.model_dir}/{THINKER_VL_MODEL}"
+    )
+    vr1_dir = f"{args.data_dir}/Video-R1-data"
+    sources = _MODALITY_SOURCES[args.modality]
+    includes = " ".join(f'--include "{src}/*" ' for src in sources)
+    U.exec_command(f"hf download {VIDEO_R1} --repo-type dataset Video-R1-260k.json --local-dir {vr1_dir}")
+    U.exec_command(f"hf download {VIDEO_R1} --repo-type dataset {includes} --local-dir {vr1_dir}")
+    for src in sources:
+        U.exec_command(f"bash -c 'cd {vr1_dir}/{src} && for z in *.zip; do unzip -n -q $z; done'")
+    U.exec_command(
+        f"python {repo}/examples/omni_thinker/prepare_video_r1.py "
+        f"--src {vr1_dir}/Video-R1-260k.json --dst {args.data_dir}/video_r1_{args.modality}.jsonl "
+        f"--media-root {vr1_dir} --data-types {args.modality} "
+        f"--path-prefixes {' '.join(sources)} --require-media --max-samples {args.max_samples}"
+    )
+
+
+def execute(args: ScriptArgs):
+    ckpt_args = (
+        f"--hf-checkpoint {args.model_dir}/{THINKER_VL_MODEL}/ "
+        f"--load {args.model_dir}/{THINKER_VL_MODEL}/ "
+        "--megatron-to-hf-mode bridge "
+        "--qwen3-omni-vl "
+    )
+
+    debug_minimal = args.mode == "debug_minimal"
+    multimodal_keys = '{"video": "videos", "image": "images"}'
+    rollout_args = (
+        "--custom-generate-function-path miles.rollout.generate_hub.sglang_omni.generate "
+        f"--prompt-data {args.data_dir}/video_r1_{args.modality}.jsonl "
+        "--input-key prompt "
+        "--label-key label "
+        f"--multimodal-keys '{multimodal_keys}' "
+        "--apply-chat-template "
+        "--rollout-shuffle "
+        "--rm-type gpqa "
+        f"--num-rollout {6 if debug_minimal else args.max_samples // 32} "
+        f"--rollout-batch-size {8 if debug_minimal else 32} "
+        f"--n-samples-per-prompt {4 if debug_minimal else 8} "
+        f"--rollout-max-response-len {100 if debug_minimal else 2048} "
+        "--rollout-temperature 1 "
+        f"--global-batch-size {32 if debug_minimal else 256} "
+        "--balance-data "
+        f"--sglang-router-ip {args.omni_server_ip} "
+        f"--sglang-router-port {args.omni_server_port} "
+        "--rollout-external "
+        f"--rollout-external-engine-addrs {args.omni_server_ip}:{args.omni_server_port} "
+        "--rollout-external-admin-api sglang-omni "
+        "--rollout-weight-update-stages thinker "
+        f"--rollout-num-gpus {args.omni_server_tp} "
+        f"--rollout-num-gpus-per-engine {args.omni_server_tp} "
+    )
+
+    consistency_args = (
+        "--get-mismatch-metrics "
+        "--use-tis "
+        "--custom-tis-function-path examples.train_infer_mismatch_helper.mis.compute_mis_weights_with_cp "
+        "--custom-config-path examples/train_infer_mismatch_helper/mis.yaml "
+        "--debug-skip-weight-update "  # v1: pseudo-VL -> omni-server rename shim not wired yet
+    )
+
+    # mirrors the geo3k megatron/bridge tier (SP stays on: no audio injection in VL runs)
+    perf_args = (
+        f"--tensor-model-parallel-size {args.actor_tp} "
+        "--sequence-parallel "
+        "--pipeline-model-parallel-size 1 "
+        "--context-parallel-size 1 "
+        "--expert-model-parallel-size 1 "
+        "--expert-tensor-parallel-size 1 "
+        "--recompute-granularity full "
+        "--recompute-method uniform "
+        "--recompute-num-layers 1 "
+        "--use-dynamic-batch-size "
+        f"--max-tokens-per-gpu {2048 if debug_minimal else 4096} "
+        "--optimizer-cpu-offload "
+        "--overlap-cpu-optimizer-d2h-h2d "
+        "--use-precision-aware-optimizer "
+    )
+
+    grpo_args = (
+        "--advantage-estimator grpo "
+        + ("" if debug_minimal else "--use-kl-loss --kl-loss-coef 0.00 --kl-loss-type low_var_kl ")
+        + "--entropy-coef 0.00 "
+        "--eps-clip 0.2 "
+        "--eps-clip-high 0.28 "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+    )
+
+    misc_args = (
+        "--attention-dropout 0.0 "
+        "--hidden-dropout 0.0 "
+        + ("" if debug_minimal else "--accumulate-allreduce-grads-in-fp32 ")
+        + "--attention-softmax-in-fp32 "
+        "--attention-backend flash "
+        "--actor-num-nodes 1 "
+        f"--actor-num-gpus-per-node {args.actor_tp} "
+        f"--num-gpus-per-node {args.actor_tp + args.omni_server_tp} "
+    )
+
+    train_args = (
+        f"{ckpt_args} "
+        f"{rollout_args} "
+        f"{consistency_args} "
+        f"{optimizer_args} "
+        f"{grpo_args} "
+        f"{U.get_default_wandb_args(__file__, run_id=args.run_id)} "
+        f"{perf_args} "
+        f"{misc_args} "
+        f"{args.extra_args} "
+    )
+
+    U.execute_train(
+        train_args=train_args,
+        num_gpus_per_node=args.actor_tp + args.omni_server_tp,
+        megatron_model_type=MEGATRON_MODEL_TYPE,
+        train_script="train.py",
+        megatron_path=args.megatron_path,
+        extra_env_vars={
+            "FLASHINFER_DISABLE_VERSION_CHECK": "1",
+            "PYTHONPATH": f"{args.megatron_path}",
+        },
+    )
+
+
+@U.dataclass_cli
+def main(args: ScriptArgs):
+    prepare(args)
+    execute(args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/miles/backends/megatron_utils/model_provider.py b/miles/backends/megatron_utils/model_provider.py
index 5f06fd60932..7991dae48b1 100644
--- a/miles/backends/megatron_utils/model_provider.py
+++ b/miles/backends/megatron_utils/model_provider.py
@@ -166,6 +166,10 @@ def wrapped_model_provider(
         bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)
         provider = bridge.to_megatron_provider(load_weights=False)
         _apply_bridge_runtime_config(provider, args)
+        if getattr(args, "qwen3_omni_vl", False):
+            from miles_plugins.models.qwen3_omni_thinker_vl import unfreeze_provider
+
+            unfreeze_provider(provider)
         provider.finalize()
 
         def wrapped_bridge_provider(
@@ -335,6 +339,10 @@ def _maybe_install_omni_audio_injection(args: argparse.Namespace, model) -> None
         from miles_plugins.models.qwen3_omni_thinker import install_audio_injection
 
         install_audio_injection(model, args)
+    if getattr(args, "qwen3_omni_vl", False):
+        from miles_plugins.models.qwen3_omni_thinker_vl import install_omni_vl
+
+        install_omni_vl(args)
 
 
 def _maybe_install_witness(
diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py
index 1f78520d78d..df40baaa24e 100644
--- a/miles/utils/arguments.py
+++ b/miles/utils/arguments.py
@@ -844,6 +844,16 @@ def add_data_arguments(parser):
                     'JSON string for multimodal data mapping media types to data keys. Example: \'{"image": "image_file"}\''
                 ),
             )
+            parser.add_argument(
+                "--qwen3-omni-vl",
+                action="store_true",
+                default=False,
+                help=(
+                    "Train the pseudo-Qwen3-VL omni thinker (tools/extract_qwen3_omni_thinker.py "
+                    "--variant vl): installs the omni TM-RoPE video override on the bridge model "
+                    "and clears provider freeze knobs. See miles_plugins/models/qwen3_omni_thinker_vl.py."
+                ),
+            )
             parser.add_argument(
                 "--qwen3-omni-audio-encoder-path",
                 type=str,
diff --git a/miles_plugins/models/qwen3_omni_thinker_vl.py b/miles_plugins/models/qwen3_omni_thinker_vl.py
index 5ef805f2a8b..70e37910a0e 100644
--- a/miles_plugins/models/qwen3_omni_thinker_vl.py
+++ b/miles_plugins/models/qwen3_omni_thinker_vl.py
@@ -150,6 +150,7 @@ def install_omni_vl(args) -> None:
     """Install the omni video rope override + video_second_per_grid pop on the bridge model."""
     import importlib
 
+    import miles_plugins.megatron_bridge  # noqa: F401  installs the packed-mrope patch first
     import miles_plugins.models.qwen3_vl as qwen3_vl_patch
 
     model_mod = importlib.import_module("megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model")
@@ -183,6 +184,39 @@ def forward(self, *fargs, **kwargs):
     logger.info(f"omni VL patch installed (position_id_per_seconds={position_id_per_seconds})")
 
 
+_VL_TO_OMNI_MERGER = (
+    ("model.visual.merger.norm.", "thinker.visual.merger.ln_q."),
+    ("model.visual.merger.linear_fc1.", "thinker.visual.merger.mlp.0."),
+    ("model.visual.merger.linear_fc2.", "thinker.visual.merger.mlp.2."),
+)
+
+
+def pseudo_vl_to_omni_server_name(name: str) -> str | None:
+    """Bridge-exported pseudo-VL HF name -> the omni server's thinker param name.
+
+    The reverse of the extraction rename map, for pushing trainer weights into a server
+    that runs the omni-named model. Returns None for tensors the server has no slot for.
+    Expert tensors stay fused (the sglang omni loader consumes hub-layout gate_up_proj /
+    down_proj directly).
+    """
+    import re as _re
+
+    if name == "lm_head.weight":
+        return "thinker.lm_head.weight"
+    for vl_prefix, omni_prefix in _VL_TO_OMNI_MERGER:
+        if name.startswith(vl_prefix):
+            return omni_prefix + name[len(vl_prefix) :]
+    match = _re.match(r"^model\.visual\.deepstack_merger_list\.(\d+)\.(norm|linear_fc1|linear_fc2)\.(.+)$", name)
+    if match:
+        part = {"norm": "ln_q", "linear_fc1": "mlp.0", "linear_fc2": "mlp.2"}[match.group(2)]
+        return f"thinker.visual.merger_list.{match.group(1)}.{part}.{match.group(3)}"
+    if name.startswith("model.visual."):
+        return "thinker." + name[len("model.") :]
+    if name.startswith("model.language_model."):
+        return "thinker.model." + name[len("model.language_model.") :]
+    return None
+
+
 def unfreeze_provider(provider) -> None:
     """Some megatron-bridge versions default Qwen3VLMoEModelProvider to a frozen LM;
     RL must train the backbone, so force-clear every freeze knob that exists."""
diff --git a/tests/fast/test_qwen3_omni_vl_rope.py b/tests/fast/test_qwen3_omni_vl_rope.py
index 5f98384fbb7..440f07d45b8 100644
--- a/tests/fast/test_qwen3_omni_vl_rope.py
+++ b/tests/fast/test_qwen3_omni_vl_rope.py
@@ -140,6 +140,41 @@ def test_omni_aware_rope_fails_loud_without_seconds():
         omni_aware(MERGE, IMG, VID, VSTART, ids, None, torch.tensor(grids), None)
 
 
+def test_pseudo_vl_to_omni_server_names_roundtrip_the_extraction_map():
+    import importlib.util
+    from pathlib import Path
+
+    from miles_plugins.models.qwen3_omni_thinker_vl import pseudo_vl_to_omni_server_name
+
+    repo = Path(__file__).resolve().parents[2]
+    spec = importlib.util.spec_from_file_location("extract_tool", repo / "tools" / "extract_qwen3_omni_thinker.py")
+    tool = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(tool)
+
+    omni_names = [
+        "thinker.lm_head.weight",
+        "thinker.model.embed_tokens.weight",
+        "thinker.model.layers.3.self_attn.q_proj.weight",
+        "thinker.visual.patch_embed.proj.weight",
+        "thinker.visual.blocks.7.attn.qkv.weight",
+        "thinker.visual.merger.ln_q.weight",
+        "thinker.visual.merger.mlp.0.bias",
+        "thinker.visual.merger.mlp.2.weight",
+        "thinker.visual.merger_list.1.ln_q.weight",
+        "thinker.visual.merger_list.2.mlp.2.bias",
+    ]
+    for omni in omni_names:
+        vl = tool.map_thinker_param_name_vl(omni)
+        assert pseudo_vl_to_omni_server_name(vl) == omni, f"{omni} -> {vl} did not roundtrip"
+
+    # fused experts map onto the sglang fused loader names (thinker. + model.)
+    assert (
+        pseudo_vl_to_omni_server_name("model.language_model.layers.0.mlp.experts.gate_up_proj")
+        == "thinker.model.layers.0.mlp.experts.gate_up_proj"
+    )
+    assert pseudo_vl_to_omni_server_name("not_a_model_tensor") is None
+
+
 def test_resolve_position_id_per_seconds(tmp_path):
     from miles_plugins.models.qwen3_omni_thinker_vl import resolve_position_id_per_seconds
 

From 289b6b5416e2c231d78e0f6c34e352d7cda3d749 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Sun, 26 Jul 2026 10:50:29 +0800
Subject: [PATCH 30/32] [omni-vl] format

---
 miles_plugins/models/qwen3_omni_thinker_vl.py | 33 +++++++++++++++----
 tests/fast/test_qwen3_omni_vl_rope.py         | 17 ++++------
 2 files changed, 32 insertions(+), 18 deletions(-)

diff --git a/miles_plugins/models/qwen3_omni_thinker_vl.py b/miles_plugins/models/qwen3_omni_thinker_vl.py
index 70e37910a0e..0f9e689b5ca 100644
--- a/miles_plugins/models/qwen3_omni_thinker_vl.py
+++ b/miles_plugins/models/qwen3_omni_thinker_vl.py
@@ -108,11 +108,19 @@ def _make_omni_aware_rope_index(orig_get_rope_index, position_id_per_seconds: fl
     grid slices it hands us).
     """
 
-    def omni_aware(merge, img_id, vid_id, vstart, input_ids, image_grid_thw=None, video_grid_thw=None, attention_mask=None):
+    def omni_aware(
+        merge, img_id, vid_id, vstart, input_ids, image_grid_thw=None, video_grid_thw=None, attention_mask=None
+    ):
         if video_grid_thw is None or video_grid_thw.numel() == 0:
             return orig_get_rope_index(
-                merge, img_id, vid_id, vstart, input_ids,
-                image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask,
+                merge,
+                img_id,
+                vid_id,
+                vstart,
+                input_ids,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                attention_mask=attention_mask,
             )
         seconds_all = getattr(_tls, "video_second_per_grid", None)
         assert seconds_all is not None, (
@@ -122,11 +130,20 @@ def omni_aware(merge, img_id, vid_id, vstart, input_ids, image_grid_thw=None, vi
         cursor = getattr(_tls, "video_cursor", 0)
         n_videos = video_grid_thw.shape[0]
         seconds = seconds_all[cursor : cursor + n_videos]
-        assert seconds.numel() == n_videos, f"video_second_per_grid exhausted: {cursor=} {n_videos=} {seconds_all.numel()=}"
+        assert (
+            seconds.numel() == n_videos
+        ), f"video_second_per_grid exhausted: {cursor=} {n_videos=} {seconds_all.numel()=}"
         _tls.video_cursor = cursor + n_videos
         positions = omni_video_rope_index(
-            merge, img_id, vid_id, vstart, input_ids,
-            image_grid_thw, video_grid_thw, seconds, position_id_per_seconds,
+            merge,
+            img_id,
+            vid_id,
+            vstart,
+            input_ids,
+            image_grid_thw,
+            video_grid_thw,
+            seconds,
+            position_id_per_seconds,
         )
         return positions, None
 
@@ -162,7 +179,9 @@ def install_omni_vl(args) -> None:
     orig_build = qwen3_vl_patch._build_packed_positions
 
     def build_with_omni_video(model, parsed, kwargs, orig_get_rope_index):
-        return orig_build(model, parsed, kwargs, _make_omni_aware_rope_index(orig_get_rope_index, position_id_per_seconds))
+        return orig_build(
+            model, parsed, kwargs, _make_omni_aware_rope_index(orig_get_rope_index, position_id_per_seconds)
+        )
 
     qwen3_vl_patch._build_packed_positions = build_with_omni_video
 
diff --git a/tests/fast/test_qwen3_omni_vl_rope.py b/tests/fast/test_qwen3_omni_vl_rope.py
index 440f07d45b8..4517e67cca6 100644
--- a/tests/fast/test_qwen3_omni_vl_rope.py
+++ b/tests/fast/test_qwen3_omni_vl_rope.py
@@ -4,16 +4,11 @@
 
 import pytest
 import torch
-
 from tests.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=30, suite="stage-a-cpu", labels=[])
 
-from miles_plugins.models.qwen3_omni_thinker_vl import (
-    _make_omni_aware_rope_index,
-    _tls,
-    omni_video_rope_index,
-)
+from miles_plugins.models.qwen3_omni_thinker_vl import _make_omni_aware_rope_index, _tls, omni_video_rope_index
 
 MERGE = 2
 IMG, VID, VSTART, VEND = 151655, 151656, 151652, 151653
@@ -93,9 +88,9 @@ def test_omni_video_rope_matches_hf_reference(segments):
     )
 
     assert ours.shape == expected.shape, f"{ours.shape=} {expected.shape=}"
-    assert torch.allclose(ours.float(), expected.float()), (
-        f"positions diverge from the HF omni reference for {segments}"
-    )
+    assert torch.allclose(
+        ours.float(), expected.float()
+    ), f"positions diverge from the HF omni reference for {segments}"
 
 
 def test_omni_aware_rope_delegates_images_and_consumes_seconds_in_order():
@@ -176,10 +171,10 @@ def test_pseudo_vl_to_omni_server_names_roundtrip_the_extraction_map():
 
 
 def test_resolve_position_id_per_seconds(tmp_path):
-    from miles_plugins.models.qwen3_omni_thinker_vl import resolve_position_id_per_seconds
-
     import json
 
+    from miles_plugins.models.qwen3_omni_thinker_vl import resolve_position_id_per_seconds
+
     (tmp_path / "config.json").write_text(json.dumps({"omni_sideband": {"position_id_per_seconds": 13}}))
     assert resolve_position_id_per_seconds(tmp_path) == 13.0
 

From 5bbc781ecf12811273aac85bf75ac8d74be96178 Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Sun, 26 Jul 2026 10:52:37 +0800
Subject: [PATCH 31/32] [omni-vl] on-policy weight-sync shim: pseudo-VL export
 names -> thinker.* (frozen visual skipped)

v1 trains the text backbone only (visual tower + projection frozen), so rollout and
training vision features stay identical and the thinker stage is the only sync target;
the bridge iterator renames exports behind --qwen3-omni-vl and fails loud on unmapped
tensors.
---
 .../hf_weight_iterator_bridge.py              | 19 +++++++++++++++++
 miles_plugins/models/qwen3_omni_thinker_vl.py | 17 +++++++++------
 tests/fast/test_qwen3_omni_vl_rope.py         | 21 +++++++++++++++++++
 3 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py
index cc15ecee73e..76ffba88b74 100644
--- a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py
+++ b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py
@@ -12,6 +12,22 @@
 from .hf_weight_iterator_base import HfWeightIteratorBase
 
 
+def rename_named_weights_for_omni_server(named_weights):
+    """Pseudo-VL export names -> the omni server's thinker.* namespace.
+
+    The frozen visual tower is skipped (its weights never change and the thinker stage
+    has no slot for them); every text/lm_head tensor must map, or we fail loud.
+    """
+    from miles_plugins.models.qwen3_omni_thinker_vl import pseudo_vl_to_omni_server_name
+
+    for hf_name, weight, megatron_name in named_weights:
+        if hf_name.startswith("model.visual."):
+            continue
+        omni_name = pseudo_vl_to_omni_server_name(hf_name)
+        assert omni_name is not None, f"no omni-server slot for exported tensor {hf_name}"
+        yield omni_name, weight, megatron_name
+
+
 class HfWeightIteratorBridge(HfWeightIteratorBase):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -65,6 +81,9 @@ def get_hf_weight_chunks(self, megatron_local_weights, weight_type: str = "base"
             elif weight_type == "lora":
                 named_weights = ((h, w, m) for h, w, m in named_weights if is_lora_weight_name(h))
 
+            if getattr(self.args, "qwen3_omni_vl", False):
+                named_weights = rename_named_weights_for_omni_server(named_weights)
+
             groups = get_atomic_update_groups(self.args, self.model_name)
             units = _stream_atomic_units(named_weights, groups)
             yield from _chunk_atomic_units_by_size(units, chunk_size=self.args.update_weight_buffer_size)
diff --git a/miles_plugins/models/qwen3_omni_thinker_vl.py b/miles_plugins/models/qwen3_omni_thinker_vl.py
index 0f9e689b5ca..d05c3548e76 100644
--- a/miles_plugins/models/qwen3_omni_thinker_vl.py
+++ b/miles_plugins/models/qwen3_omni_thinker_vl.py
@@ -237,9 +237,14 @@ def pseudo_vl_to_omni_server_name(name: str) -> str | None:
 
 
 def unfreeze_provider(provider) -> None:
-    """Some megatron-bridge versions default Qwen3VLMoEModelProvider to a frozen LM;
-    RL must train the backbone, so force-clear every freeze knob that exists."""
-    for field in ("freeze_language_model", "freeze_vision_model", "freeze_vision_projection"):
-        if getattr(provider, field, False):
-            logger.warning(f"pseudo-VL provider had {field}=True; clearing it for RL training")
-            setattr(provider, field, False)
+    """RL trains the text backbone; the visual tower + projection stay frozen (v1), which
+    keeps rollout/training vision features identical and makes text-only weight sync
+    sufficient. Some megatron-bridge versions default freeze_language_model=True, which
+    would silently train nothing but mergers."""
+    if getattr(provider, "freeze_language_model", False):
+        logger.warning("pseudo-VL provider had freeze_language_model=True; clearing it for RL training")
+        provider.freeze_language_model = False
+    for field in ("freeze_vision_model", "freeze_vision_projection"):
+        if hasattr(provider, field) and not getattr(provider, field):
+            logger.info(f"freezing {field} (v1 trains the text backbone only)")
+            setattr(provider, field, True)
diff --git a/tests/fast/test_qwen3_omni_vl_rope.py b/tests/fast/test_qwen3_omni_vl_rope.py
index 4517e67cca6..5372b6af438 100644
--- a/tests/fast/test_qwen3_omni_vl_rope.py
+++ b/tests/fast/test_qwen3_omni_vl_rope.py
@@ -181,3 +181,24 @@ def test_resolve_position_id_per_seconds(tmp_path):
     (tmp_path / "config.json").write_text(json.dumps({}))
     with pytest.raises(AssertionError, match="omni_sideband"):
         resolve_position_id_per_seconds(tmp_path)
+
+
+def test_rename_named_weights_for_omni_server_skips_frozen_visual():
+    from miles.backends.megatron_utils.update_weight.hf_weight_iterator_bridge import (
+        rename_named_weights_for_omni_server,
+    )
+
+    named = [
+        ("lm_head.weight", "W0", "m0"),
+        ("model.language_model.layers.0.self_attn.q_proj.weight", "W1", "m1"),
+        ("model.visual.blocks.0.attn.qkv.weight", "W2", "m2"),
+        ("model.language_model.layers.0.mlp.experts.gate_up_proj", "W3", "m3"),
+    ]
+
+    out = list(rename_named_weights_for_omni_server(iter(named)))
+
+    assert out == [
+        ("thinker.lm_head.weight", "W0", "m0"),
+        ("thinker.model.layers.0.self_attn.q_proj.weight", "W1", "m1"),
+        ("thinker.model.layers.0.mlp.experts.gate_up_proj", "W3", "m3"),
+    ]

From 039d1c1e08ae4f2f2837d409afb74faf55ddc35c Mon Sep 17 00:00:00 2001
From: yxs <xuesongyey@gmail.com>
Date: Mon, 27 Jul 2026 12:10:39 +0800
Subject: [PATCH 32/32] [omni-vl] extract --variant vl + config synthesis +
 test (moved here from #1)

The pseudo-Qwen3-VL extraction path (map_thinker_param_name_vl, synthesize_vl_config,
fuse_layer_experts_vl, _VlExpertAccumulator) belongs to this PR, not the audio PR #1;
carry it here explicitly so #1 stays the clean audio milestone and #6 owns the VL delta.
---
 .../megatron_utils/test_extract_thinker_vl.py | 222 ++++++++++++++++++
 tools/extract_qwen3_omni_thinker.py           | 199 ++++++++++++++--
 2 files changed, 402 insertions(+), 19 deletions(-)
 create mode 100644 tests/fast/backends/megatron_utils/test_extract_thinker_vl.py

diff --git a/tests/fast/backends/megatron_utils/test_extract_thinker_vl.py b/tests/fast/backends/megatron_utils/test_extract_thinker_vl.py
new file mode 100644
index 00000000000..dd268f648b1
--- /dev/null
+++ b/tests/fast/backends/megatron_utils/test_extract_thinker_vl.py
@@ -0,0 +1,222 @@
+"""Pseudo-Qwen3-VL extraction: thinker (visual+text) masquerading as qwen3_vl_moe."""
+
+import importlib.util
+import json
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+from tests.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[])
+
+_REPO = Path(__file__).resolve().parents[4]
+
+
+def _load_tool():
+    path = _REPO / "tools" / "extract_qwen3_omni_thinker.py"
+    spec = importlib.util.spec_from_file_location("extract_qwen3_omni_thinker", path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def test_vl_name_map_text_and_visual():
+    m = _load_tool().map_thinker_param_name_vl
+    assert m("thinker.lm_head.weight") == "lm_head.weight"
+    assert m("thinker.model.embed_tokens.weight") == "model.language_model.embed_tokens.weight"
+    assert (
+        m("thinker.model.layers.3.self_attn.q_proj.weight") == "model.language_model.layers.3.self_attn.q_proj.weight"
+    )
+    assert m("thinker.model.layers.3.mlp.gate.weight") == "model.language_model.layers.3.mlp.gate.weight"
+    assert m("thinker.visual.patch_embed.proj.weight") == "model.visual.patch_embed.proj.weight"
+    assert m("thinker.visual.blocks.7.attn.qkv.weight") == "model.visual.blocks.7.attn.qkv.weight"
+    # merger renames: ln_q -> norm, mlp.0/2 -> linear_fc1/2, merger_list -> deepstack_merger_list
+    assert m("thinker.visual.merger.ln_q.weight") == "model.visual.merger.norm.weight"
+    assert m("thinker.visual.merger.mlp.0.bias") == "model.visual.merger.linear_fc1.bias"
+    assert m("thinker.visual.merger.mlp.2.weight") == "model.visual.merger.linear_fc2.weight"
+    assert m("thinker.visual.merger_list.1.ln_q.weight") == "model.visual.deepstack_merger_list.1.norm.weight"
+    assert m("thinker.visual.merger_list.2.mlp.2.bias") == "model.visual.deepstack_merger_list.2.linear_fc2.bias"
+    # drops
+    assert m("thinker.audio_tower.layers.0.fc1.weight") is None
+    assert m("talker.model.norm.weight") is None
+    assert m("code2wav.decoder.weight") is None
+
+
+def test_vl_expert_fusion_hub_layout():
+    tool = _load_tool()
+    # E=2 experts, I=3 (moe ffn), H=4 (hidden)
+    gate = [torch.arange(12, dtype=torch.float32).reshape(3, 4) + 100 * e for e in range(2)]
+    up = [torch.arange(12, dtype=torch.float32).reshape(3, 4) + 100 * e + 50 for e in range(2)]
+    down = [torch.arange(12, dtype=torch.float32).reshape(4, 3) + 100 * e for e in range(2)]
+
+    gate_up_fused, down_fused = tool.fuse_layer_experts_vl(gate, up, down)
+
+    # hub layout: gate_up [E, H, 2I], down [E, I, H]
+    assert gate_up_fused.shape == (2, 4, 6)
+    assert down_fused.shape == (2, 3, 4)
+    for e in range(2):
+        expected = torch.cat([gate[e], up[e]], dim=0).T  # [H, 2I]
+        assert torch.equal(gate_up_fused[e], expected)
+        assert torch.equal(down_fused[e], down[e].T)
+
+
+def test_vl_config_synthesis():
+    synth = _load_tool().synthesize_vl_config
+    omni = {
+        "model_type": "qwen3_omni_moe",
+        "thinker_config": {
+            "image_token_id": 151655,
+            "video_token_id": 151656,
+            "vision_start_token_id": 151652,
+            "vision_end_token_id": 151653,
+            "position_id_per_seconds": 13,
+            "audio_token_id": 151675,
+            "tie_word_embeddings": False,
+            "text_config": {
+                "model_type": "qwen3_omni_moe_text",
+                "vocab_size": 152064,
+                "hidden_size": 2048,
+                "num_hidden_layers": 48,
+                "rope_theta": 1000000,
+                "max_position_embeddings": 65536,
+                "rope_scaling": {
+                    "interleaved": True,
+                    "mrope_interleaved": True,
+                    "mrope_section": [24, 20, 20],
+                    "type": "default",
+                },
+                "num_experts": 128,
+            },
+            "vision_config": {
+                "model_type": "qwen3_omni_moe_vision_encoder",
+                "depth": 27,
+                "hidden_size": 1152,
+                "deepstack_visual_indexes": [8, 16, 24],
+                "apply_vit_abs_pos_embed": True,
+                "image_size": 768,
+                "spatial_patch_size": 16,
+                "in_chans": 3,
+                "tokens_per_second": 2,
+                "out_hidden_size": 2048,
+                "spatial_merge_size": 2,
+                "patch_size": 16,
+                "temporal_patch_size": 2,
+            },
+            "audio_config": {"d_model": 1280},
+        },
+    }
+
+    cfg = synth(omni)
+
+    assert cfg["model_type"] == "qwen3_vl_moe"
+    assert cfg["architectures"] == ["Qwen3VLMoeForConditionalGeneration"]
+    assert cfg["tie_word_embeddings"] is False
+    for key in ("image_token_id", "video_token_id", "vision_start_token_id", "vision_end_token_id"):
+        assert cfg[key] == omni["thinker_config"][key]
+    text = cfg["text_config"]
+    assert text["model_type"] == "qwen3_vl_moe_text"
+    assert text["vocab_size"] == 152064 and text["rope_theta"] == 1000000
+    assert text["max_position_embeddings"] == 65536
+    assert text["rope_scaling"] == {"rope_type": "default", "mrope_section": [24, 20, 20], "mrope_interleaved": True}
+    assert text["bos_token_id"] == 151643 and text["eos_token_id"] == 151645
+    vision = cfg["vision_config"]
+    assert vision["model_type"] == "qwen3_vl_moe"
+    assert vision["num_position_embeddings"] == 2304  # (image_size // patch_size) ** 2
+    for gone in ("apply_vit_abs_pos_embed", "image_size", "spatial_patch_size", "in_chans", "tokens_per_second"):
+        assert gone not in vision
+    assert "audio_config" not in cfg
+    # sideband for the trainer-side omni video rope (no VL slot for it)
+    assert cfg["omni_sideband"] == {"position_id_per_seconds": 13, "audio_token_id": 151675}
+
+
+def _tiny_omni_src(src: Path):
+    src.mkdir(parents=True, exist_ok=True)
+    thinker = {
+        "image_token_id": 151655,
+        "video_token_id": 151656,
+        "vision_start_token_id": 151652,
+        "vision_end_token_id": 151653,
+        "position_id_per_seconds": 13,
+        "audio_token_id": 151675,
+        "tie_word_embeddings": False,
+        "text_config": {
+            "model_type": "qwen3_omni_moe_text",
+            "vocab_size": 32,
+            "num_experts": 2,
+            "rope_scaling": {
+                "interleaved": True,
+                "mrope_interleaved": True,
+                "mrope_section": [24, 20, 20],
+                "type": "default",
+            },
+        },
+        "vision_config": {"model_type": "x", "image_size": 32, "patch_size": 16},
+    }
+    with open(src / "config.json", "w") as f:
+        json.dump({"thinker_config": thinker}, f)
+    tensors = {
+        "thinker.model.embed_tokens.weight": torch.randn(32, 4),
+        "thinker.lm_head.weight": torch.randn(32, 4),
+        "thinker.visual.merger.ln_q.weight": torch.randn(6),
+        "thinker.audio_tower.proj1.weight": torch.zeros(2, 2),
+    }
+    for e in range(2):
+        tensors[f"thinker.model.layers.0.mlp.experts.{e}.gate_proj.weight"] = torch.randn(3, 4)
+        tensors[f"thinker.model.layers.0.mlp.experts.{e}.up_proj.weight"] = torch.randn(3, 4)
+        tensors[f"thinker.model.layers.0.mlp.experts.{e}.down_proj.weight"] = torch.randn(4, 3)
+    save_file(tensors, src / "model.safetensors", metadata={"format": "pt"})
+    return tensors
+
+
+def test_extract_vl_roundtrip(tmp_path):
+    tool = _load_tool()
+    src, dst = tmp_path / "src", tmp_path / "dst"
+    tensors = _tiny_omni_src(src)
+
+    tool.extract(src, dst, shard_size_gb=1.0, variant="vl")
+
+    with open(dst / "config.json") as f:
+        cfg = json.load(f)
+    assert cfg["model_type"] == "qwen3_vl_moe"
+
+    out = {}
+    with safe_open(dst / "model.safetensors", framework="pt") as reader:
+        for key in reader.keys():
+            out[key] = reader.get_tensor(key)
+
+    assert "model.language_model.embed_tokens.weight" in out
+    assert "lm_head.weight" in out
+    assert "model.visual.merger.norm.weight" in out
+    assert not any("audio_tower" in k for k in out)
+    gate_up = out["model.language_model.layers.0.mlp.experts.gate_up_proj"]
+    down = out["model.language_model.layers.0.mlp.experts.down_proj"]
+    assert gate_up.shape == (2, 4, 6) and down.shape == (2, 3, 4)
+    expected_e1 = torch.cat(
+        [
+            tensors["thinker.model.layers.0.mlp.experts.1.gate_proj.weight"],
+            tensors["thinker.model.layers.0.mlp.experts.1.up_proj.weight"],
+        ],
+        dim=0,
+    ).T
+    assert torch.equal(gate_up[1], expected_e1)
+    # per-expert 2D names must NOT leak through
+    assert not any(".experts.0." in k or ".experts.1." in k for k in out)
+
+
+def test_extract_text_variant_unchanged(tmp_path):
+    tool = _load_tool()
+    src, dst = tmp_path / "src", tmp_path / "dst"
+    _tiny_omni_src(src)
+
+    tool.extract(src, dst, shard_size_gb=1.0, variant="text")
+
+    with open(dst / "config.json") as f:
+        assert json.load(f)["model_type"] == "qwen3_moe"
+    with safe_open(dst / "model.safetensors", framework="pt") as reader:
+        keys = list(reader.keys())
+    assert "model.embed_tokens.weight" in keys
+    assert any(".experts.0.gate_proj" in k for k in keys)  # text variant keeps per-expert 2D
+    assert not any("visual" in k for k in keys)
diff --git a/tools/extract_qwen3_omni_thinker.py b/tools/extract_qwen3_omni_thinker.py
index ebde57bbaf5..7d52d20cff6 100644
--- a/tools/extract_qwen3_omni_thinker.py
+++ b/tools/extract_qwen3_omni_thinker.py
@@ -1,19 +1,41 @@
-"""Extract a standalone Qwen3-MoE thinker checkpoint from a full Qwen3-Omni model.
+"""Extract a standalone thinker checkpoint from a full Qwen3-Omni model.
 
-miles loads HF via AutoBridge keyed on model_type; the composite omni checkpoint has
-no bridge, but the thinker text backbone is a plain Qwen3-MoE. This writes a
-self-contained HF dir (thinker text + lm_head, renamed, model_type=qwen3_moe) for the
-existing Qwen3MoEBridge. Audio/visual towers are dropped here on purpose: the trainer
-loads them frozen from the original omni checkpoint (see the audio-injection wrapper),
-and the rollout server keeps its own copies.
+Two variants, both keyed on model_type so miles' AutoBridge can load them:
 
-    python tools/extract_qwen3_omni_thinker.py --src <omni> --dst <thinker>
+- ``text`` (default): thinker text backbone + lm_head as a plain Qwen3-MoE
+  (model_type=qwen3_moe, per-expert 2D weights). Audio/visual towers are dropped on
+  purpose: the trainer loads the audio tower frozen from the original omni checkpoint
+  (see the audio-injection plugin), and the rollout server keeps its own copies.
+
+- ``vl``: thinker visual tower + text backbone masquerading as Qwen3-VL-MoE
+  (model_type=qwen3_vl_moe). Names are mechanically compatible except three merger
+  renames (ln_q->norm, mlp.0/2->linear_fc1/2, merger_list->deepstack_merger_list) and
+  the MoE experts, which the VL hub format fuses to 3D hub layout
+  (gate_up_proj [E,H,2I], down_proj [E,I,H]). This unlocks miles' existing
+  bridge-based Qwen3-VL training path (packed mrope + deepstack + CP) for image/video
+  input. The audio tower is dropped; omni-only rope fields ride along under
+  ``omni_sideband`` in config.json for the trainer-side video-rope override.
+
+    python tools/extract_qwen3_omni_thinker.py --src <omni> --dst <thinker> [--variant vl]
 """
 
 from __future__ import annotations
 
+import re
+
 _NON_TEXT_THINKER_SUBMODULES = ("audio_tower.", "visual.", "model.audio_tower.", "model.visual.")
 
+_VL_MERGER_RENAMES = (
+    (re.compile(r"^visual\.merger\.ln_q\."), "visual.merger.norm."),
+    (re.compile(r"^visual\.merger\.mlp\.0\."), "visual.merger.linear_fc1."),
+    (re.compile(r"^visual\.merger\.mlp\.2\."), "visual.merger.linear_fc2."),
+    (re.compile(r"^visual\.merger_list\.(\d+)\.ln_q\."), r"visual.deepstack_merger_list.\1.norm."),
+    (re.compile(r"^visual\.merger_list\.(\d+)\.mlp\.0\."), r"visual.deepstack_merger_list.\1.linear_fc1."),
+    (re.compile(r"^visual\.merger_list\.(\d+)\.mlp\.2\."), r"visual.deepstack_merger_list.\1.linear_fc2."),
+)
+
+_EXPERT_RE = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$")
+
 
 def map_thinker_param_name(name: str) -> str | None:
     """Full omni param name -> standalone Qwen3-MoE name, or None to drop."""
@@ -25,6 +47,44 @@ def map_thinker_param_name(name: str) -> str | None:
     return rest.replace("model.language_model.", "model.")
 
 
+def map_thinker_param_name_vl(name: str) -> str | None:
+    """Full omni param name -> pseudo Qwen3-VL-MoE name, or None to drop.
+
+    Expert weights keep their per-expert names here; `extract(variant="vl")` fuses them
+    into the 3D hub tensors afterwards.
+    """
+    if not name.startswith("thinker."):
+        return None
+    rest = name[len("thinker.") :].replace("model.language_model.", "model.")
+    if rest.startswith(("audio_tower.", "model.audio_tower.")):
+        return None
+    if rest == "lm_head.weight":
+        return rest
+    if rest.startswith("visual."):
+        for pattern, replacement in _VL_MERGER_RENAMES:
+            if pattern.search(rest):
+                rest = pattern.sub(replacement, rest)
+                break
+        return f"model.{rest}"
+    if rest.startswith("model."):
+        return rest.replace("model.", "model.language_model.", 1)
+    return None
+
+
+def fuse_layer_experts_vl(gates, ups, downs):
+    """Per-expert 2D weights -> fused hub-layout tensors.
+
+    gate/up: [I, H] each -> gate_up_proj [E, H, 2I]; down: [H, I] -> down_proj [E, I, H].
+    Hub layout is what transformers (Transpose check_dims), Megatron-Bridge import, and
+    sglang's fused loaders all expect.
+    """
+    import torch
+
+    gate_up = torch.stack([torch.cat([g, u], dim=0).T for g, u in zip(gates, ups, strict=True)], dim=0)
+    down = torch.stack([d.T for d in downs], dim=0)
+    return gate_up.contiguous(), down.contiguous()
+
+
 def synthesize_thinker_config(omni_config: dict) -> dict:
     """Plain Qwen3-MoE config dict from the full omni config."""
     thinker = omni_config.get("thinker_config", omni_config)
@@ -38,11 +98,92 @@ def synthesize_thinker_config(omni_config: dict) -> dict:
     return cfg
 
 
-def extract(src, dst, shard_size_gb: float = 5.0) -> int:
+# vision fields whose omni spelling has no qwen3_vl_moe slot (dropped after use)
+_VL_VISION_DROPPED = ("apply_vit_abs_pos_embed", "image_size", "spatial_patch_size", "in_chans", "tokens_per_second")
+
+
+def synthesize_vl_config(omni_config: dict) -> dict:
+    """Pseudo Qwen3-VL-MoE config from the full omni config.
+
+    Omni numerics are preserved (vocab 152064, rope_theta 1e6, max_position_embeddings
+    65536); rope_scaling is normalized to the VL whitelist; the vision config gains
+    num_position_embeddings (the pos_embed table size, (image_size/patch_size)^2) and
+    drops omni-only spellings. Fields with no VL slot that the trainer still needs
+    (omni video rope) ride under `omni_sideband`.
+    """
+    thinker = omni_config.get("thinker_config", omni_config)
+
+    text = dict(thinker["text_config"])
+    text["model_type"] = "qwen3_vl_moe_text"
+    rope_scaling = dict(text.get("rope_scaling") or {})
+    text["rope_scaling"] = {
+        "rope_type": rope_scaling.get("rope_type") or rope_scaling.get("type", "default"),
+        "mrope_section": rope_scaling["mrope_section"],
+        "mrope_interleaved": rope_scaling.get("mrope_interleaved", True),
+    }
+    text.setdefault("bos_token_id", 151643)
+    text.setdefault("eos_token_id", 151645)
+
+    vision = dict(thinker["vision_config"])
+    vision["model_type"] = "qwen3_vl_moe"
+    if "num_position_embeddings" not in vision and "image_size" in vision:
+        patch = vision.get("patch_size") or vision.get("spatial_patch_size")
+        vision["num_position_embeddings"] = (vision["image_size"] // patch) ** 2
+    for key in _VL_VISION_DROPPED:
+        vision.pop(key, None)
+
+    cfg = {
+        "model_type": "qwen3_vl_moe",
+        "architectures": ["Qwen3VLMoeForConditionalGeneration"],
+        "tie_word_embeddings": thinker.get("tie_word_embeddings", False),
+        "text_config": text,
+        "vision_config": vision,
+        "omni_sideband": {
+            "position_id_per_seconds": thinker.get("position_id_per_seconds"),
+            "audio_token_id": thinker.get("audio_token_id"),
+        },
+    }
+    for k in ("image_token_id", "video_token_id", "vision_start_token_id", "vision_end_token_id"):
+        cfg[k] = thinker[k]
+    return cfg
+
+
+class _VlExpertAccumulator:
+    """Collects per-expert 2D weights per layer and emits fused hub-layout tensors."""
+
+    def __init__(self, num_experts: int):
+        self.num_experts = num_experts
+        self._layers: dict[int, dict[str, dict[int, object]]] = {}
+
+    def offer(self, vl_name: str, tensor) -> list[tuple[str, object]]:
+        match = _EXPERT_RE.match(vl_name.replace("model.language_model.", "model.", 1))
+        if match is None:
+            return [(vl_name, tensor)]
+        layer, expert, kind = int(match.group(1)), int(match.group(2)), match.group(3)
+        buckets = self._layers.setdefault(layer, {"gate_proj": {}, "up_proj": {}, "down_proj": {}})
+        buckets[kind][expert] = tensor
+        if all(len(buckets[k]) == self.num_experts for k in ("gate_proj", "up_proj", "down_proj")):
+            del self._layers[layer]
+            order = range(self.num_experts)
+            gate_up, down = fuse_layer_experts_vl(
+                [buckets["gate_proj"][e] for e in order],
+                [buckets["up_proj"][e] for e in order],
+                [buckets["down_proj"][e] for e in order],
+            )
+            prefix = f"model.language_model.layers.{layer}.mlp.experts"
+            return [(f"{prefix}.gate_up_proj", gate_up), (f"{prefix}.down_proj", down)]
+        return []
+
+    def assert_drained(self):
+        assert not self._layers, f"incomplete expert layers left over: {sorted(self._layers)}"
+
+
+def extract(src, dst, shard_size_gb: float = 5.0, variant: str = "text") -> int:
     """Stream thinker tensors from src into sharded safetensors under dst.
 
     Each shard is written (and its buffer released) as soon as it fills, so peak RAM is
-    one shard, not the whole ~60 GB thinker. Returns the number of kept tensors.
+    one shard (plus, for the vl variant, at most a few layers of in-flight expert
+    weights). Returns the number of tensors written.
     """
     import json
     import os
@@ -52,13 +193,16 @@ def extract(src, dst, shard_size_gb: float = 5.0) -> int:
     from safetensors import safe_open
     from safetensors.torch import save_file
 
+    assert variant in ("text", "vl"), variant
     src, dst = Path(src), Path(dst)
     dst.mkdir(parents=True, exist_ok=True)
 
     with open(src / "config.json") as f:
         omni_config = json.load(f)
+    synthesize = synthesize_thinker_config if variant == "text" else synthesize_vl_config
+    out_config = synthesize(omni_config)
     with open(dst / "config.json", "w") as f:
-        json.dump(synthesize_thinker_config(omni_config), f, indent=2)
+        json.dump(out_config, f, indent=2)
 
     for fname in (
         "tokenizer.json",
@@ -91,6 +235,12 @@ def extract(src, dst, shard_size_gb: float = 5.0) -> int:
     else:
         shard_files = ["model.safetensors"]
 
+    map_name = map_thinker_param_name if variant == "text" else map_thinker_param_name_vl
+    accumulator = None
+    if variant == "vl":
+        thinker = omni_config.get("thinker_config", omni_config)
+        accumulator = _VlExpertAccumulator(num_experts=thinker["text_config"]["num_experts"])
+
     shard_size_bytes = int(shard_size_gb * (1024**3))
     out_index: dict[str, str] = {}
     shard_names: list[str] = []
@@ -111,18 +261,28 @@ def flush():
             total_size += t.numel() * t.element_size()
         buf, buf_bytes = {}, 0
 
+    def emit(out_name, tensor):
+        nonlocal buf_bytes, total_kept
+        buf[out_name] = tensor
+        buf_bytes += tensor.numel() * tensor.element_size()
+        total_kept += 1
+        if buf_bytes >= shard_size_bytes:
+            flush()
+
     for shard_file in shard_files:
         with safe_open(src / shard_file, framework="pt") as reader:
             for key in reader.keys():
-                new_key = map_thinker_param_name(key)
+                new_key = map_name(key)
                 if new_key is None:
                     continue
                 tensor = reader.get_tensor(key)
-                buf[new_key] = tensor
-                buf_bytes += tensor.numel() * tensor.element_size()
-                total_kept += 1
-                if buf_bytes >= shard_size_bytes:
-                    flush()
+                if accumulator is not None:
+                    for out_name, out_tensor in accumulator.offer(new_key, tensor):
+                        emit(out_name, out_tensor)
+                else:
+                    emit(new_key, tensor)
+    if accumulator is not None:
+        accumulator.assert_drained()
     flush()
 
     if total_kept == 0:
@@ -134,7 +294,7 @@ def flush():
         with open(dst / "model.safetensors.index.json", "w") as f:
             json.dump({"metadata": {"total_size": total_size}, "weight_map": out_index}, f, indent=2)
 
-    print(f"[done] {total_kept} thinker tensors -> {dst}")
+    print(f"[done] {total_kept} thinker tensors ({variant}) -> {dst}")
     return total_kept
 
 
@@ -145,8 +305,9 @@ def main() -> None:
     parser.add_argument("--src", required=True)
     parser.add_argument("--dst", required=True)
     parser.add_argument("--shard-size-gb", type=float, default=5.0)
+    parser.add_argument("--variant", choices=["text", "vl"], default="text")
     args = parser.parse_args()
-    extract(args.src, args.dst, shard_size_gb=args.shard_size_gb)
+    extract(args.src, args.dst, shard_size_gb=args.shard_size_gb, variant=args.variant)
 
 
 if __name__ == "__main__":