From 765002061878d6c260009f0b13e042cfbb491cc6 Mon Sep 17 00:00:00 2001 From: NumberWan Date: Thu, 24 Sep 2026 18:08:46 +0800 Subject: [PATCH] [BugFix][Cosmos3] Register opt-in T2I pipeline so deploy-config applies (#7971) Signed-off-by: NumberWan --- recipes/cosmos3/Cosmos3-Nano.md | 11 ++- recipes/cosmos3/Cosmos3-Super.md | 8 +- tests/config/test_config_factory.py | 99 +++++++++++++++++++ vllm_omni/config/pipeline_registry.py | 14 ++- vllm_omni/deploy/cosmos3_omni.yaml | 30 ++++++ .../model_executor/models/cosmos3/pipeline.py | 55 ++++++++--- 6 files changed, 193 insertions(+), 24 deletions(-) create mode 100644 vllm_omni/deploy/cosmos3_omni.yaml diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index 856d05b5a..162e358d4 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -126,7 +126,10 @@ speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to `0` removes the streak cap. To run **without** guardrails (you are responsible for license compliance), -add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use +add `--no-guardrails` (no token/`cosmos-guardrail` needed), or pass +[`vllm_omni/deploy/cosmos3_omni.yaml`](../../vllm_omni/deploy/cosmos3_omni.yaml) +via `--deploy-config` (same opt-in overlay as Super; `pipeline: cosmos3_omni_deploy`). +For extra GPUs use `--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`; `--enable-layerwise-offload` reduces VRAM on smaller GPUs; `--vae-fast-path channels_last` speeds up the Wan VAE video decode by switching the @@ -728,8 +731,10 @@ vllm serve nvidia/Cosmos3-Nano \ ``` To run **without** guardrails (you are responsible for license compliance), -add `--no-guardrails` (no token/`cosmos-guardrail` needed). For tensor parallel -add `--tensor-parallel-size 8`. `--quantization fp8` and +add `--no-guardrails` (no token/`cosmos-guardrail` needed), or pass +[`vllm_omni/deploy/cosmos3_omni.yaml`](../../vllm_omni/deploy/cosmos3_omni.yaml) +via `--deploy-config` (same opt-in overlay as Super; `pipeline: cosmos3_omni_deploy`). +For tensor parallel add `--tensor-parallel-size 8`. `--quantization fp8` and `--enable-layerwise-offload` are not supported on NPU. The pipeline auto-resolves from `model_index.json`; pass `--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly. diff --git a/recipes/cosmos3/Cosmos3-Super.md b/recipes/cosmos3/Cosmos3-Super.md index 1e5dbf7f0..01f64e398 100644 --- a/recipes/cosmos3/Cosmos3-Super.md +++ b/recipes/cosmos3/Cosmos3-Super.md @@ -57,7 +57,13 @@ vllm serve nvidia/Cosmos3-Super \ Guardrails are on by default (gated `nvidia/Cosmos-1.0-Guardrail` — `pip install cosmos-guardrail`, accept the license, set `HF_TOKEN`); add `--no-guardrails` to -disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs; +disable, or pass +[`vllm_omni/deploy/cosmos3_omni.yaml`](../../vllm_omni/deploy/cosmos3_omni.yaml) +via `--deploy-config` on either multi-GPU command above (opt-in omni deploy +overlay, `pipeline: cosmos3_omni_deploy`; without that `pipeline:` key +`--deploy-config` cannot apply for Cosmos3). Same YAML works for Super / Nano +(T2I and video); topology `final_output_type` matches the CLI default (`video`). +`--enable-layerwise-offload` reduces VRAM on smaller GPUs; `--vae-fast-path channels_last` speeds up the Wan VAE video decode by switching the decoder convolutions to channels-last kernels (output no longer bit-identical to diffusers; the default `lossless` fast path is bit-exact, see diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index f9a5b4ea8..7f9f51eaa 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -694,6 +694,105 @@ class TestCosmos3PolicyPipeline: assert stage.engine_args.model_config.policy_server_config.action_space == "joint_position" +class TestCosmos3OmniDeployPipeline: + """Omni deploy yaml selects an opt-in topology so --deploy-config applies (#6874).""" + + def test_registered_without_auto_capturing_cosmos3_omni(self): + assert "cosmos3_omni_deploy" in OMNI_PIPELINES + # Same shared HF metadata as policy: must not register cosmos3_omni for + # auto-detect, or T2I/video/policy would collide. + assert "cosmos3_omni" not in OMNI_PIPELINES + pipeline = OMNI_PIPELINES["cosmos3_omni_deploy"] + assert pipeline.hf_architectures == () + assert pipeline.diffusers_class_name is None + # Align with CLI Cosmos3OmniDiffusersPipeline default (video), not T2I-only. + assert pipeline.stages[0].final_output_type == "video" + + def test_deploy_yaml_applies_guardrails_false(self): + deploy = load_deploy_config(get_deploy_config_path("cosmos3_omni.yaml")) + assert deploy.pipeline == "cosmos3_omni_deploy" + assert deploy.stages[0].devices is None + + stages = merge_pipeline_deploy(OMNI_PIPELINES["cosmos3_omni_deploy"], deploy) + assert len(stages) == 1 + stage = stages[0].to_omegaconf() + + assert stage.stage_type == "diffusion" + assert stage.final_output_type == "video" + assert stage.engine_args.model_class_name == "Cosmos3OmniDiffusersPipeline" + assert stage.engine_args.model_config.guardrails is False + assert "devices" not in stage.runtime + + def test_deploy_yaml_accepts_documented_multi_gpu_cli(self): + from vllm_omni.engine.stage_init_utils import _check_stage_device_layout + + deploy = load_deploy_config(get_deploy_config_path("cosmos3_omni.yaml")) + stages, _ = StageConfigFactory._create_legacy_from_registry( + OMNI_PIPELINES["cosmos3_omni_deploy"], + { + "cfg_parallel_size": 2, + "use_hsdp": True, + "hsdp_shard_size": 2, + }, + user_deploy_config=deploy, + ) + assert len(stages) == 1 + omega = stages[0].to_omegaconf() + assert omega.engine_args.model_config.guardrails is False + assert omega.engine_args.parallel_config.cfg_parallel_size == 2 + assert omega.engine_args.parallel_config.use_hsdp is True + assert omega.engine_args.parallel_config.hsdp_shard_size == 2 + assert "devices" not in omega.runtime + _check_stage_device_layout(omega, dict(omega.engine_args)) + + def test_get_pipeline_config_honours_deploy_yaml_pipeline_key(self): + deploy_path = get_deploy_config_path("cosmos3_omni.yaml") + # No HF download: explicit pipeline: key is highest priority. + pipeline = StageConfigFactory.get_pipeline_config( + model="nvidia/Cosmos3-Super-Text2Image", + trust_remote_code=True, + deploy_config_path=deploy_path, + ) + assert pipeline is not None + assert pipeline.model_type == "cosmos3_omni_deploy" + + def test_without_deploy_yaml_stays_on_unregistered_fallback(self): + # Shared HF metadata must not auto-select a pipeline; only the deploy + # yaml ``pipeline:`` key selects cosmos3_omni_deploy and merges guardrails. + assert "cosmos3_omni" not in OMNI_PIPELINES + + class FakeCosmos3Config(PretrainedConfig): + model_type = "cosmos3_omni" + + fake_config = FakeCosmos3Config() + fake_config.architectures = ["Cosmos3OmniDiffusersPipeline"] + model = "nvidia/Cosmos3-Super-Text2Image" + deploy_path = get_deploy_config_path("cosmos3_omni.yaml") + + StageConfigFactory.get_hf_config.cache_clear() + StageConfigFactory.try_infer_model_type.cache_clear() + with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config): + assert ( + StageConfigFactory.get_pipeline_config( + model=model, + trust_remote_code=True, + deploy_config_path=None, + ) + is None + ) + + pipeline = StageConfigFactory.get_pipeline_config( + model=model, + trust_remote_code=True, + deploy_config_path=deploy_path, + ) + assert pipeline is not None + assert pipeline.model_type == "cosmos3_omni_deploy" + + stages = merge_pipeline_deploy(pipeline, load_deploy_config(deploy_path)) + assert stages[0].to_omegaconf().engine_args.model_config.guardrails is False + + class TestStagePipelineConfig: def test_frozen(self): s = StagePipelineConfig(stage_id=0, model_stage="a") diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 997a6991f..5d7a7f4bd 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -49,7 +49,10 @@ from vllm_omni.model_executor.models.bagel.pipeline import ( BAGEL_THINK_PIPELINE, ) from vllm_omni.model_executor.models.breeze_tts_2.pipeline import BREEZE_TTS_2_PIPELINE -from vllm_omni.model_executor.models.cosmos3.pipeline import COSMOS3_POLICY_PIPELINE +from vllm_omni.model_executor.models.cosmos3.pipeline import ( + COSMOS3_OMNI_DEPLOY_PIPELINE, + COSMOS3_POLICY_PIPELINE, +) from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE @@ -154,11 +157,12 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "Gr00tN1d7": GR00T_N1D7_PIPELINE, "pi0": PI0_PIPELINE, "pi05": PI05_PIPELINE, - # Cosmos3 policy checkpoints share HF metadata with the T2I/video Cosmos3 - # checkpoints (which stay on the single-stage diffusion fallback), so this - # entry is only reachable through a deploy yaml's ``pipeline:`` key - # (see deploy/cosmos3_policy_droid.yaml). + # Cosmos3 policy / omni-deploy topologies share HF metadata with video + # Cosmos3 checkpoints (which stay on the single-stage diffusion fallback), + # so these entries are only reachable through a deploy yaml's ``pipeline:`` + # key (see deploy/cosmos3_policy_droid.yaml and deploy/cosmos3_omni.yaml). "cosmos3_policy": COSMOS3_POLICY_PIPELINE, + "cosmos3_omni_deploy": COSMOS3_OMNI_DEPLOY_PIPELINE, "gepard": GEPARD_PIPELINE, "glm_image": GLM_IMAGE_PIPELINE, "hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE, diff --git a/vllm_omni/deploy/cosmos3_omni.yaml b/vllm_omni/deploy/cosmos3_omni.yaml new file mode 100644 index 000000000..4995c52f5 --- /dev/null +++ b/vllm_omni/deploy/cosmos3_omni.yaml @@ -0,0 +1,30 @@ +# Cosmos3 omni deploy overlay: single co-located diffusion stage (both MoT +# towers in one process). Topology (COSMOS3_OMNI_DEPLOY_PIPELINE) is declared in +# vllm_omni/model_executor/models/cosmos3/pipeline.py. +# +# T2I / video / policy Cosmos3 checkpoints share HF metadata +# (model_type=cosmos3_omni, model_index.json _class_name= +# Cosmos3OmniDiffusersPipeline). This pipeline cannot be auto-detected and must +# be selected explicitly so --deploy-config is applied (#6874): +# +# vllm serve nvidia/Cosmos3-Super --omni \ +# --deploy-config /absolute/path/to/vllm_omni/deploy/cosmos3_omni.yaml +# +# Same overlay works for Nano (including video). Without --deploy-config, +# behaviour is unchanged (CLI single-stage fallback). --no-guardrails remains +# the CLI-only equivalent of model_config.guardrails below. Per-checkpoint +# knobs can still be layered with --stage-overrides. + +pipeline: cosmos3_omni_deploy +# async_chunk defaults to True, which the deploy validator rejects for a +# single-stage pipeline (no next-stage processor). Must be false here. +async_chunk: false +distributed_executor_backend: mp + +stages: + - stage_id: 0 + model_class_name: Cosmos3OmniDiffusersPipeline + model_config: + # Gated nvidia/Cosmos-1.0-Guardrail is often unavailable; default this + # recommended layout off. Equivalent to launching with --no-guardrails. + guardrails: false diff --git a/vllm_omni/model_executor/models/cosmos3/pipeline.py b/vllm_omni/model_executor/models/cosmos3/pipeline.py index d26b4f46c..6ed61b585 100644 --- a/vllm_omni/model_executor/models/cosmos3/pipeline.py +++ b/vllm_omni/model_executor/models/cosmos3/pipeline.py @@ -1,21 +1,30 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Cosmos3 policy single-stage topology for online OpenPI serving. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Cosmos3 opt-in topologies for policy (OpenPI) and omni deploy YAML. -Cosmos3 policy checkpoints (e.g. ``nvidia/Cosmos3-Nano-Policy-DROID``) are one -diffusion stage: robot observation -> action chunk. They are -registered in ``OMNI_PIPELINES`` because online OpenPI serving -(``/v1/realtime/robot/openpi``) requires a registered pipeline so the deploy -yaml's ``model_config.policy_server_config`` reaches the websocket handshake. +Both topologies declare neither ``hf_architectures`` nor +``diffusers_class_name``. Every Cosmos3 checkpoint (T2I, video, policy) +shares HF metadata (``model_type=cosmos3_omni``, ``model_index.json`` +``_class_name=Cosmos3OmniDiffusersPipeline``). Auto-registering any of them +under that metadata would capture the others. Select explicitly via a deploy +yaml ``pipeline:`` key. -This pipeline declares neither ``hf_architectures`` nor -``diffusers_class_name``: policy checkpoints share their HF metadata -(``model_type=cosmos3_omni``, ``model_index.json`` ``_class_name= -Cosmos3OmniDiffusersPipeline``) with the T2I/video Cosmos3 checkpoints, which -must keep resolving through the default single-stage diffusion fallback. -Select this pipeline explicitly via a deploy yaml ``pipeline:`` key, e.g. -``vllm serve nvidia/Cosmos3-Nano-Policy-DROID --omni --deploy-config -/absolute/path/to/vllm-omni/vllm_omni/deploy/cosmos3_policy_droid.yaml``. +Policy (``cosmos3_policy``):: + + vllm serve nvidia/Cosmos3-Nano-Policy-DROID --omni \\ + --deploy-config .../vllm_omni/deploy/cosmos3_policy_droid.yaml + +Omni deploy overlay (``cosmos3_omni_deploy``) — makes ``--deploy-config`` +apply stage / ``model_config`` (e.g. ``guardrails: false``) instead of being +silently dropped (#6874). Same overlay works for Super / Nano (T2I and video). +``final_output_type`` matches the CLI Cosmos3OmniDiffusersPipeline default +(``video``):: + + vllm serve nvidia/Cosmos3-Super --omni \\ + --deploy-config .../vllm_omni/deploy/cosmos3_omni.yaml + +Without ``--deploy-config``, T2I/video keep the default single-stage +diffusion CLI fallback. ``--no-guardrails`` remains the CLI-only path. """ from vllm_omni.config.stage_config import ( @@ -39,3 +48,19 @@ COSMOS3_POLICY_PIPELINE = PipelineConfig( ), ), ) + +COSMOS3_OMNI_DEPLOY_PIPELINE = PipelineConfig( + model_type="cosmos3_omni_deploy", + model_arch="Cosmos3OmniDiffusersPipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="video", + model_arch="Cosmos3OmniDiffusersPipeline", + ), + ), +)