[Feature] Add tensor parallelism support for ming-omni (#270)

* [Core] Add TP driver-worker infrastructure for omni thinker

- Add TP follower module with register_omni_models(), follower_worker_loop(),
  and spawn_followers() for TP ranks > 0
- Expose tp_cpu_group and tp_size properties on ModelWorker
- Spawn TP follower processes in create_sglang_ar_engine before rank 0 init
- Manage follower lifecycle (stop signal + terminate) in OmniEngine
- Broadcast serialization-safe ModelWorkerBatch to followers each step

* [Bugfix] Fix TP runtime issues (daemon, pickle, broadcast, model config)

- Add model_override_args for Ming TP>1 (flattened llm_config fields)
- Fix broadcast_pyobj list wrapping and ServerArgs field name
- Set daemon=False for TP stages to allow follower subprocess spawning
- Remove mp.Event from follower spawn (unpicklable in nested spawn)

* [TP] Serialization-safe follower batch with device relocation

- Add serialization.py with make_follower_batch() and first-call pickle
  verification that diagnoses unpicklable fields
- Strip reqs, sampling_info, launch_done before broadcast
- Add relocate_batch_tensors() to move all tensors (including nested
  multimodal inputs like mrope_position_delta) to follower device
- Extract patch_batch_for_follower() to module scope for testability
- Pass real vocab_size to SamplingBatchInfo stub

* [Test] TP validation, serialization tests, and review hardening

- Add test_tp_follower.py: registration, picklability, batch patching
- Add test_tp_batch_serialization.py: pickle round-trip regression tests
- Add scripts/test_ming_tp.py: TP=1 vs TP=2 output consistency harness

* [Config] Expose TP in Qwen3 pipeline config and example entrypoints

- Inject server_args_overrides into thinker stage in Qwen3 pipeline configs
- Add --tp-size, --cpu-offload-gb, --mem-fraction-static CLI flags to
  run_qwen3_omni_server.py and run_qwen3_omni_text_first.py

* [Style] Pre-commit formatting fixes

* [TP] Sync page tables and harden follower runtime

* [TP] Mirror multimodal thinker forward paths on followers

* [Bugfix] Tear down torch.distributed in MingImageEncoder cleanup

* [TP] Enable tensor parallelism for Ming-flash-omni vision encoder

* [TP] fix follower shutdown ordering

* [Style] fix lint

* [Style] restore note attribution

* [TP] Fix Qwen3 speech TP validation and factory test mocks

* [TP] Fix _DummyEngine stub to accept follower_processes

* [chore] trim comments

* fix

* [TP] address review on PR #270

* [Test] Fix black formatting

Author: edwingao28

examples/run_ming_omni_server.py

@@ -23,6 +23,7 @@
 
 import argparse
 import logging
+import multiprocessing as mp
 import os
 
 from sglang_omni.models.ming_omni.config import MingOmniPipelineConfig
@@ -100,6 +101,7 @@ def main() -> None:
     overrides = {}
     if args.tp_size and args.tp_size > 1:
         overrides["tp_size"] = args.tp_size
+        overrides["disable_custom_all_reduce"] = True
     if args.quantization:
         overrides["quantization"] = args.quantization
     if args.cpu_offload_gb:
@@ -122,4 +124,5 @@ def main() -> None:
 
 
 if __name__ == "__main__":
+    mp.set_start_method("spawn", force=True)
     main()

examples/run_ming_omni_speech.py

@@ -75,6 +75,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--timeout", type=float, default=300.0)
     parser.add_argument("--cpu-offload-gb", type=float, default=0)
     parser.add_argument("--mem-fraction-static", type=float, default=None)
+    parser.add_argument(
+        "--tp-size", type=int, default=1, help="Tensor parallel size for thinker"
+    )
     return parser.parse_args()
 
 
@@ -89,6 +92,8 @@ async def main_async(args: argparse.Namespace) -> None:
     }
 
     overrides = {}
+    if args.tp_size > 1:
+        overrides["tp_size"] = args.tp_size
     if args.cpu_offload_gb:
         overrides["cpu_offload_gb"] = args.cpu_offload_gb
     if args.mem_fraction_static is not None:

examples/run_qwen3_omni_server.py

@@ -44,7 +44,19 @@ def parse_args() -> argparse.Namespace:
         default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
         help="Hugging Face model id or local path",
     )
-    parser.add_argument("--thinker-max-seq-len", type=int, default=8192)
+    parser.add_argument("--thinker-max-seq-len", type=int, default=None)
+    parser.add_argument(
+        "--cpu-offload-gb",
+        type=int,
+        default=0,
+        help="GB of model weights to offload to CPU",
+    )
+    parser.add_argument(
+        "--mem-fraction-static",
+        type=float,
+        default=None,
+        help="Fraction of GPU memory for KV cache",
+    )
 
     # Pipeline options
     parser.add_argument(
@@ -71,11 +83,26 @@ def parse_args() -> argparse.Namespace:
 def main() -> None:
     args = parse_args()
 
+    overrides = {}
+    if args.cpu_offload_gb:
+        overrides["cpu_offload_gb"] = args.cpu_offload_gb
+    if args.mem_fraction_static is not None:
+        overrides["mem_fraction_static"] = args.mem_fraction_static
+
     config = Qwen3OmniPipelineConfig(
         model_path=args.model_path,
         relay_backend=args.relay_backend,
+        server_args_overrides=overrides or None,
     )
 
+    # Override thinker_max_seq_len in stage executor args if provided
+    if args.thinker_max_seq_len is not None:
+        for stage in config.stages:
+            if stage.name == "thinker":
+                if stage.executor.args is None:
+                    stage.executor.args = {}
+                stage.executor.args["thinker_max_seq_len"] = args.thinker_max_seq_len
+
     launch_server(
         config,
         host=args.host,

examples/run_qwen3_omni_text_first.py

@@ -27,8 +27,7 @@ def parse_args() -> argparse.Namespace:
         help="Hugging Face model id",
     )
     parser.add_argument("--prompt", type=str, default="Describe this input.")
-    parser.add_argument("--dtype", type=str, default="bfloat16")
-    parser.add_argument("--thinker-max-seq-len", type=int, default=8192)
+    parser.add_argument("--thinker-max-seq-len", type=int, default=None)
     parser.add_argument("--max-new-tokens", type=int, default=1024)
     parser.add_argument("--temperature", type=float, default=0.8)
     parser.add_argument("--image-path", type=str, default=None)
@@ -40,14 +39,42 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--relay-backend", type=str, default="nixl", choices=["nixl", "shm"]
     )
+    parser.add_argument(
+        "--cpu-offload-gb",
+        type=int,
+        default=0,
+        help="GB of model weights to offload to CPU",
+    )
+    parser.add_argument(
+        "--mem-fraction-static",
+        type=float,
+        default=None,
+        help="Fraction of GPU memory for KV cache",
+    )
     return parser.parse_args()
 
 
 async def main_async(args: argparse.Namespace) -> None:
+    overrides = {}
+    if args.cpu_offload_gb:
+        overrides["cpu_offload_gb"] = args.cpu_offload_gb
+    if args.mem_fraction_static is not None:
+        overrides["mem_fraction_static"] = args.mem_fraction_static
+
     config = Qwen3OmniPipelineConfig(
         model_path=args.model_path,
         relay_backend=args.relay_backend,
+        server_args_overrides=overrides or None,
     )
+
+    # Override thinker_max_seq_len in stage executor args if provided
+    if args.thinker_max_seq_len is not None:
+        for stage in config.stages:
+            if stage.name == "thinker":
+                if stage.executor.args is None:
+                    stage.executor.args = {}
+                stage.executor.args["thinker_max_seq_len"] = args.thinker_max_seq_len
+
     runner = build_pipeline_runner(config)
 
     await runner.start()

scripts/test_ming_tp.py

@@ -0,0 +1,181 @@
+#!/usr/bin/env python3
+"""Validate Ming Omni thinker output consistency across TP configurations.
+
+Usage:
+    python scripts/test_ming_tp.py run --tp 1 --cpu-offload-gb 150
+    python scripts/test_ming_tp.py run --tp 2 --cpu-offload-gb 40
+    python scripts/test_ming_tp.py compare tp1_results.json tp2_results.json
+"""
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import logging
+import multiprocessing as mp
+import os
+import sys
+
+logging.basicConfig(
+    level=os.environ.get("LOGLEVEL", "INFO").upper(),
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+TEST_PROMPTS = [
+    "What is 1+1?",
+    "What is the capital of France?",
+    "What is the capital of Japan?",
+    "Explain quantum computing in one sentence.",
+]
+
+
+async def run_thinker(
+    tp_size: int,
+    cpu_offload_gb: int,
+    mem_fraction: float,
+    output_file: str,
+    attention_backend: str | None = None,
+):
+    from sglang_omni.models.ming_omni.config import MingOmniPipelineConfig
+    from sglang_omni.pipeline.mp_runner import MultiProcessPipelineRunner
+    from sglang_omni.proto import OmniRequest
+
+    overrides = {
+        "tp_size": tp_size,
+        "cpu_offload_gb": cpu_offload_gb,
+        "mem_fraction_static": mem_fraction,
+    }
+    if attention_backend is not None:
+        overrides["attention_backend"] = attention_backend
+
+    config = MingOmniPipelineConfig(
+        model_path="inclusionAI/Ming-flash-omni-2.0",
+        relay_backend="shm",
+        server_args_overrides=overrides,
+    )
+
+    runner = MultiProcessPipelineRunner(config)
+    logger.info(
+        "Starting pipeline with TP=%d, cpu_offload_gb=%d, attention_backend=%s ...",
+        tp_size,
+        cpu_offload_gb,
+        attention_backend,
+    )
+    await runner.start(timeout=600)
+
+    results = []
+    try:
+        for i, prompt in enumerate(TEST_PROMPTS):
+            logger.info("[%d/%d] Prompt: %s", i + 1, len(TEST_PROMPTS), prompt)
+            request = {
+                "messages": [
+                    {
+                        "role": "system",
+                        "content": "You are a friendly AI assistant. Please answer concisely.",
+                    },
+                    {"role": "user", "content": prompt},
+                ],
+                "audios": [],
+            }
+            result = await asyncio.wait_for(
+                runner.coordinator.submit(
+                    f"tp-test-{i}",
+                    OmniRequest(
+                        inputs=request,
+                        params={"max_new_tokens": 64, "temperature": 0.0},
+                    ),
+                ),
+                timeout=120,
+            )
+            text = ""
+            if isinstance(result, dict):
+                for stage_name, payload in result.items():
+                    data = (
+                        payload
+                        if isinstance(payload, dict)
+                        else getattr(payload, "data", {})
+                    )
+                    if isinstance(data, dict) and "text" in data:
+                        text = data["text"]
+                        break
+            assert text, f"Empty output for prompt: {prompt}"
+            results.append({"prompt": prompt, "output": text})
+            logger.info("  Output: %s", text[:200])
+    finally:
+        await runner.stop()
+
+    with open(output_file, "w") as f:
+        json.dump(
+            {"tp_size": tp_size, "results": results}, f, indent=2, ensure_ascii=False
+        )
+    logger.info("Results saved to %s", output_file)
+
+
+def compare_outputs(file1: str, file2: str):
+    with open(file1) as f:
+        data1 = json.load(f)
+    with open(file2) as f:
+        data2 = json.load(f)
+
+    print(f"\n{'='*60}")
+    print(f"Comparing TP={data1['tp_size']} vs TP={data2['tp_size']}")
+    print(f"{'='*60}")
+
+    all_match = True
+    for r1, r2 in zip(data1["results"], data2["results"]):
+        match = r1["output"].strip() == r2["output"].strip()
+        status = "MATCH" if match else "MISMATCH"
+        if not match:
+            all_match = False
+        print(f"\n[{status}] Prompt: {r1['prompt']}")
+        print(f"  TP={data1['tp_size']}: {r1['output'][:120]}")
+        print(f"  TP={data2['tp_size']}: {r2['output'][:120]}")
+
+    print(f"\n{'='*60}")
+    if all_match:
+        print("ALL OUTPUTS MATCH - TP validation PASSED")
+    else:
+        print("OUTPUTS DIFFER - TP validation FAILED, needs investigation")
+    print(f"{'='*60}")
+    return all_match
+
+
+def main():
+    mp.set_start_method("spawn", force=True)
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    sub = parser.add_subparsers(dest="cmd")
+
+    run_p = sub.add_parser("run")
+    run_p.add_argument("--tp", type=int, required=True)
+    run_p.add_argument("--cpu-offload-gb", type=int, default=80)
+    run_p.add_argument("--mem-fraction", type=float, default=0.80)
+    run_p.add_argument("--attention-backend", type=str, default=None)
+    run_p.add_argument("--output", type=str, default=None)
+
+    cmp_p = sub.add_parser("compare")
+    cmp_p.add_argument("file1")
+    cmp_p.add_argument("file2")
+
+    args = parser.parse_args()
+
+    if args.cmd == "run":
+        output = args.output or f"tp{args.tp}_results.json"
+        asyncio.run(
+            run_thinker(
+                args.tp,
+                args.cpu_offload_gb,
+                args.mem_fraction,
+                output,
+                args.attention_backend,
+            )
+        )
+    elif args.cmd == "compare":
+        sys.exit(0 if compare_outputs(args.file1, args.file2) else 1)
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()

sglang_omni/engines/ar/sglang_backend/model_runner.py

@@ -39,7 +39,9 @@ def __init__(
         weight_prefix: str | None = None,
     ) -> None:
         self._weight_prefix = weight_prefix
-        self._register_omni_model()
+        from sglang_omni.models.sglang_registry import register_omni_models_in_sglang
+
+        register_omni_models_in_sglang()
 
         port_args = PortArgs.init_new(server_args)
         tp_size = server_args.tp_size
@@ -61,26 +63,3 @@ def __init__(
             nccl_port=nccl_port,
             server_args=server_args,
         )
-
-    def _register_omni_model(self):
-        # Register sglang_omni model classes directly in SGLang's model registry.
-        from sglang.srt.models.registry import ModelRegistry
-
-        from sglang_omni.models.fishaudio_s2_pro.sglang_model import (
-            S2ProSGLangTextModel,
-        )
-        from sglang_omni.models.ming_omni.thinker import (
-            BailingMM2Config,
-            BailingMoeV2ForCausalLM,
-        )
-        from sglang_omni.models.qwen3_omni.talker import Qwen3OmniTalker
-
-        ModelRegistry.models["S2ProSGLangTextModel"] = S2ProSGLangTextModel
-        ModelRegistry.models["Qwen3OmniTalker"] = Qwen3OmniTalker
-        ModelRegistry.models["BailingMoeV2ForCausalLM"] = BailingMoeV2ForCausalLM
-
-        # Register BailingMM2Config with AutoConfig so SGLang can load
-        # config.json from HF repos missing configuration_bailingmm2.py.
-        from transformers import AutoConfig
-
-        AutoConfig.register("bailingmm_moe_v2_lite", BailingMM2Config)

sglang_omni/engines/ar/sglang_backend/model_worker.py

@@ -53,6 +53,15 @@ def __init__(
         )[0]
         set_random_seed(self.random_seed)
 
+    @property
+    def tp_cpu_group(self):
+        """NCCL CPU process group for TP broadcast operations."""
+        return self.model_runner.tp_group.cpu_group
+
+    @property
+    def tp_size(self) -> int:
+        return self.server_args.tp_size
+
     def _init_model_config(self):
         from sglang.srt.configs.model_config import ModelConfig