Zhou TaichangandClaude Opus 5
aff7d64948
Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
1602 lines
62 KiB
Python
1602 lines
62 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Tests for the additive structured Omni config."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import fields
|
|
from inspect import Parameter, signature
|
|
from multiprocessing.reduction import ForkingPickler
|
|
from pathlib import Path
|
|
|
|
import msgspec
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
from pydantic.fields import FieldInfo
|
|
from transformers import Qwen3OmniMoeConfig
|
|
from vllm.config import CacheConfig as VllmCacheConfig
|
|
from vllm.config import CompilationConfig as VllmCompilationConfig
|
|
from vllm.config import LoadConfig as VllmLoadConfig
|
|
from vllm.config import ParallelConfig as VllmParallelConfig
|
|
from vllm.config import ProfilerConfig as VllmProfilerConfig
|
|
from vllm.config import SchedulerConfig as VllmSchedulerConfig
|
|
|
|
from tests.helpers.stage_config import get_deploy_config_path
|
|
from vllm_omni.config import omni_config as omni_config_module
|
|
from vllm_omni.config.omni_config import (
|
|
BaseVllmOmniStageConfig,
|
|
OmniStageCacheConfig,
|
|
OmniStageConnectorConfig,
|
|
OmniStageDiffusionParallelConfig,
|
|
OmniStageLoadConfig,
|
|
OmniStageModelConfig,
|
|
OmniStageParallelConfig,
|
|
OmniStageRuntimeConfig,
|
|
OmniStageSchedulerConfig,
|
|
VllmOmniARStageConfig,
|
|
VllmOmniConfig,
|
|
VllmOmniDiffusionStageConfig,
|
|
VllmOmniGenerationStageConfig,
|
|
)
|
|
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config
|
|
from vllm_omni.config.stage_config import (
|
|
_STAGE_DEPLOY_FIELDS,
|
|
PIPELINE_WIDE_ENGINE_FIELDS,
|
|
DeployConfig,
|
|
PipelineConfig,
|
|
StageDeployConfig,
|
|
StageExecutionType,
|
|
StagePipelineConfig,
|
|
load_deploy_config,
|
|
merge_pipeline_deploy,
|
|
)
|
|
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
|
from vllm_omni.engine.stage_engine_startup import _serialize_stage_config
|
|
from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
_DEPLOY_DIR = Path(__file__).parents[2] / "vllm_omni" / "deploy"
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _stable_test_platform(monkeypatch):
|
|
from vllm_omni import platforms
|
|
|
|
platform = platforms.current_omni_platform
|
|
monkeypatch.setattr(platform, "device_name", "cpu", raising=False)
|
|
monkeypatch.setattr(platform, "device_type", "cpu", raising=False)
|
|
|
|
|
|
def _load_default_deploy(pipeline: PipelineConfig) -> DeployConfig:
|
|
if pipeline.default_deploy_config_name is not None:
|
|
return load_deploy_config(_DEPLOY_DIR / pipeline.default_deploy_config_name)
|
|
return DeployConfig()
|
|
|
|
|
|
def _resolve_pipeline_or_skip(model_type: str, hf_config=None) -> PipelineConfig:
|
|
pipeline = resolve_pipeline_config(model_type, hf_config)
|
|
if pipeline is None:
|
|
pytest.skip(f"Pipeline {model_type!r} requires an HF config to resolve")
|
|
return pipeline
|
|
|
|
|
|
def _from_pipeline_key(
|
|
model_type: str,
|
|
hf_config=None,
|
|
deploy_config_path: str | None = None,
|
|
cli_overrides: dict | None = None,
|
|
) -> VllmOmniConfig:
|
|
return VllmOmniConfig.from_pipeline_config(
|
|
_resolve_pipeline_or_skip(model_type, hf_config),
|
|
deploy_config_path=deploy_config_path,
|
|
cli_overrides=cli_overrides,
|
|
)
|
|
|
|
|
|
def test_minimax_h3_text_encoder_tp_targets_only_structured_stage_zero() -> None:
|
|
config = _from_pipeline_key(
|
|
"minimax_h3_disaggregated",
|
|
cli_overrides={"text_encoder_tp_size": 4},
|
|
)
|
|
|
|
assert config.stage_by_id(0).parallel_config.tensor_parallel_size == 4
|
|
diffusion = config.stage_by_id(1)
|
|
assert diffusion.parallel_config.tensor_parallel_size == 1
|
|
assert diffusion.parallel_config.text_encoder_tp_size == 1
|
|
|
|
|
|
def test_duplex_session_capacity_propagates_to_every_structured_stage() -> None:
|
|
omni_config = _from_pipeline_key("personaplex")
|
|
|
|
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [2, 2]
|
|
assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["duplex", "duplex"]
|
|
|
|
|
|
def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -> None:
|
|
deploy_path = tmp_path / "personaplex-turn.yaml"
|
|
deploy_path.write_text("session_mode: turn\n", encoding="utf-8")
|
|
|
|
omni_config = _from_pipeline_key("personaplex", deploy_config_path=str(deploy_path))
|
|
|
|
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [1, 1]
|
|
assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["turn", "turn"]
|
|
|
|
|
|
@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
|
|
def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str):
|
|
pipeline = _resolve_pipeline_or_skip(model_type)
|
|
legacy_deploy = _load_default_deploy(pipeline)
|
|
|
|
legacy_stages = merge_pipeline_deploy(pipeline, legacy_deploy)
|
|
omni_config = VllmOmniConfig.from_pipeline_config(pipeline)
|
|
|
|
assert omni_config.pipeline_config is pipeline
|
|
assert len(omni_config.stage_configs) == len(legacy_stages)
|
|
|
|
for legacy_stage, omni_stage in zip(legacy_stages, omni_config.stage_configs, strict=True):
|
|
assert omni_config.stage_by_id(legacy_stage.stage_id) is omni_stage
|
|
|
|
assert omni_stage.stage_pipeline_config is pipeline.get_stage(legacy_stage.stage_id)
|
|
assert omni_stage.model_config.default_sampling_params == legacy_stage.yaml_extras.get(
|
|
"default_sampling_params"
|
|
)
|
|
assert omni_stage.connector_config.output_connectors == legacy_stage.yaml_extras.get("output_connectors")
|
|
assert omni_stage.connector_config.input_connectors == legacy_stage.yaml_extras.get("input_connectors")
|
|
assert omni_stage.runtime_config.devices == legacy_stage.yaml_runtime.get("devices")
|
|
assert omni_stage.runtime_config.num_replicas == legacy_stage.yaml_runtime.get("num_replicas", 1)
|
|
|
|
engine_args = legacy_stage.yaml_engine_args
|
|
assert omni_stage.model_config.duplex_max_sessions == engine_args.get("duplex_max_sessions", 1)
|
|
assert omni_stage.model_config.session_mode == engine_args.get("session_mode", "turn")
|
|
assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
|
|
assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
|
|
assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
|
|
assert omni_stage.cache_config.gpu_memory_utilization == engine_args.get("gpu_memory_utilization")
|
|
assert omni_stage.cache_config.enable_prefix_caching == engine_args.get("enable_prefix_caching")
|
|
expected_disable_hybrid = engine_args.get("disable_hybrid_kv_cache_manager")
|
|
if omni_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_GENERATION:
|
|
expected_disable_hybrid = True if expected_disable_hybrid is None else expected_disable_hybrid
|
|
assert omni_stage.cache_config.disable_hybrid_kv_cache_manager == expected_disable_hybrid
|
|
assert omni_stage.scheduler_config.max_num_seqs == engine_args.get("max_num_seqs")
|
|
assert omni_stage.scheduler_config.max_num_batched_tokens == engine_args.get("max_num_batched_tokens")
|
|
assert omni_stage.scheduler_config.enable_chunked_prefill == engine_args.get("enable_chunked_prefill")
|
|
assert omni_stage.scheduler_config.async_scheduling == engine_args.get("async_scheduling")
|
|
legacy_parallel_config = engine_args.get("parallel_config") or {}
|
|
assert omni_stage.parallel_config.tensor_parallel_size == legacy_parallel_config.get(
|
|
"tensor_parallel_size",
|
|
engine_args.get("tensor_parallel_size", 1),
|
|
)
|
|
|
|
if omni_stage.stage_pipeline_config.execution_type == StageExecutionType.DIFFUSION:
|
|
assert isinstance(omni_stage, VllmOmniDiffusionStageConfig)
|
|
assert omni_stage.diffusion_config is not None
|
|
assert omni_stage.diffusion_config.stage_id == legacy_stage.stage_id
|
|
assert omni_stage.diffusion_config.model_arch == engine_args.get("model_arch")
|
|
elif omni_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_AR:
|
|
assert isinstance(omni_stage, VllmOmniARStageConfig)
|
|
assert not hasattr(omni_stage, "diffusion_config")
|
|
else:
|
|
assert isinstance(omni_stage, VllmOmniGenerationStageConfig)
|
|
assert not hasattr(omni_stage, "diffusion_config")
|
|
|
|
|
|
def test_stage_by_id_raises_for_unknown_stage():
|
|
omni_config = _from_pipeline_key("qwen3_tts")
|
|
|
|
with pytest.raises(KeyError, match="no stage 99"):
|
|
omni_config.stage_by_id(99)
|
|
|
|
|
|
def test_resolve_execution_mode_rejects_unknown_execution_type():
|
|
with pytest.raises(ValueError, match="Unsupported stage execution type"):
|
|
omni_config_module._resolve_execution_mode("unknown_execution_type")
|
|
|
|
|
|
def test_from_pipeline_config_preserves_current_pipeline_config_object():
|
|
omni_config = _from_pipeline_key("minicpmo_4_5")
|
|
pipeline = _resolve_pipeline_or_skip("minicpmo_4_5")
|
|
|
|
assert omni_config.pipeline_config is pipeline
|
|
assert not hasattr(omni_config, "pipeline")
|
|
assert "hf_config_predicate" in {f.name for f in fields(PipelineConfig)}
|
|
assert omni_config.pipeline_config.hf_config_predicate is pipeline.hf_config_predicate
|
|
|
|
|
|
def test_from_pipeline_config_normalizes_stage_engine_extras_without_expanding_stage_deploy_config():
|
|
assert not hasattr(StageDeployConfig, "model_config")
|
|
assert not hasattr(StageDeployConfig, "parallel_config")
|
|
|
|
stage = _from_pipeline_key("dreamzero", deploy_config_path="dreamzero_tp1_cfg2").stage_by_id(0)
|
|
|
|
assert isinstance(stage, VllmOmniDiffusionStageConfig)
|
|
assert stage.parallel_config.tensor_parallel_size == 1
|
|
assert stage.parallel_config.cfg_parallel_size == 2
|
|
assert stage.diffusion_config.model_config["default_robot_embodiment"] == "roboarena"
|
|
|
|
|
|
@pytest.mark.parametrize("disabled", [True, False])
|
|
def test_frontend_log_stats_flag_is_not_an_unowned_stage_argument(disabled):
|
|
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
|
|
|
config = _from_pipeline_key("dots_tts", cli_overrides={"disable_log_stats": disabled})
|
|
assert config.stage_configs
|
|
engine_args = build_engine_args_dict_from_omni_stage_config(config.stage_by_id(0), model="test-model")
|
|
assert "disable_log_stats" not in engine_args
|
|
|
|
|
|
def test_from_pipeline_config_applies_cli_overrides_without_stage_config_runtime_bridge():
|
|
omni_config = _from_pipeline_key(
|
|
"qwen3_tts",
|
|
cli_overrides={
|
|
"stage_0_max_num_seqs": 7,
|
|
"stage_1_tensor_parallel_size": 2,
|
|
},
|
|
)
|
|
|
|
stage0 = omni_config.stage_by_id(0)
|
|
stage1 = omni_config.stage_by_id(1)
|
|
|
|
assert stage0.scheduler_config.max_num_seqs == 7
|
|
assert stage1.parallel_config.tensor_parallel_size == 2
|
|
assert stage1.runtime_config.num_gpus == stage1.parallel_config.world_size
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"cli_overrides",
|
|
[
|
|
{"enable_lora": True},
|
|
{"stage_0_enable_lora": True},
|
|
],
|
|
ids=["global", "stage-scoped"],
|
|
)
|
|
def test_from_pipeline_config_rejects_explicit_unowned_engine_cli_fields(cli_overrides):
|
|
with pytest.raises(ValueError, match=r"no structured config owner: enable_lora"):
|
|
_from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("engine_extras", "unowned_field"),
|
|
[
|
|
({"enable_lora": True}, "enable_lora"),
|
|
({"parallel_config": {"cfg_parallel_size": 2}}, "parallel_config.cfg_parallel_size"),
|
|
],
|
|
ids=["top-level", "nested-parallel-config"],
|
|
)
|
|
def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras, unowned_field):
|
|
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
|
|
deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0, engine_extras=engine_extras)])
|
|
|
|
with pytest.raises(ValueError, match=rf"no structured config owner: {unowned_field}"):
|
|
VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
|
|
|
|
|
|
@pytest.mark.parametrize("stage_id", [0, 2], ids=["ar", "generation"])
|
|
def test_llm_additional_config_roundtrip_and_isolation(stage_id, monkeypatch):
|
|
from vllm_omni.engine import stage_init_utils
|
|
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
|
|
|
# Worker discovery requires hardware support; this test covers config transport.
|
|
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None)
|
|
additional_config = {"backend_options": {"enabled": True}}
|
|
pipeline = _resolve_pipeline_or_skip("minicpmo_4_5")
|
|
deploy = DeployConfig(
|
|
stages=[StageDeployConfig(stage_id=i, engine_extras={"additional_config": additional_config}) for i in (0, 2)]
|
|
)
|
|
config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
|
|
stage = config.stage_by_id(stage_id)
|
|
assert stage.runtime_config.additional_config == additional_config
|
|
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
|
|
assert engine_args["additional_config"] == additional_config
|
|
|
|
engine_args["additional_config"]["backend_options"]["enabled"] = False
|
|
assert stage.runtime_config.additional_config["backend_options"]["enabled"] is True
|
|
stage.runtime_config.additional_config["backend_options"]["enabled"] = False
|
|
assert config.stage_by_id(2 if stage_id == 0 else 0).runtime_config.additional_config == additional_config
|
|
assert additional_config["backend_options"]["enabled"] is True
|
|
|
|
|
|
@pytest.mark.parametrize("deploy_name", ["minicpmo_4_5", "minicpmo_4_5_2gpu", "minicpmo_4_5_3gpu"])
|
|
def test_minicpmo_npu_additional_config_reaches_engine_args(monkeypatch, deploy_name):
|
|
from vllm_omni.engine import stage_init_utils
|
|
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
monkeypatch.setattr(current_omni_platform, "device_name", "npu")
|
|
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None)
|
|
stage = _from_pipeline_key("minicpmo_4_5", deploy_config_path=deploy_name).stage_by_id(2)
|
|
expected = {"code2wav_enable_npu_graph": True, "code2wav_max_npu_graphs": 32}
|
|
assert stage.runtime_config.additional_config == expected
|
|
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
|
|
assert engine_args["additional_config"] == expected
|
|
|
|
|
|
def test_diffusion_additional_config_keeps_diffusion_owner():
|
|
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
|
|
|
additional_config = {"torchair_graph_config": {"enabled": True}}
|
|
stage = _from_pipeline_key("dreamzero", cli_overrides={"additional_config": additional_config}).stage_by_id(0)
|
|
assert stage.runtime_config.additional_config is None
|
|
assert stage.diffusion_config.additional_config == additional_config
|
|
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
|
|
assert engine_args["additional_config"] == additional_config
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("engine_extras", "cli_overrides"),
|
|
[
|
|
({"requires_full_payload_input": False}, {}),
|
|
({}, {"stage_1_requires_full_payload_input": False}),
|
|
],
|
|
ids=["deploy", "stage-cli"],
|
|
)
|
|
def test_from_pipeline_config_rejects_full_payload_input_capability_overrides(
|
|
engine_extras,
|
|
cli_overrides,
|
|
):
|
|
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
|
|
deploy = DeployConfig(
|
|
stages=[
|
|
StageDeployConfig(
|
|
stage_id=1,
|
|
engine_extras=engine_extras,
|
|
)
|
|
]
|
|
)
|
|
|
|
with pytest.raises(ValueError, match=r"no structured config owner: requires_full_payload_input"):
|
|
VllmOmniConfig.from_pipeline_config(
|
|
pipeline,
|
|
user_deploy_config=deploy,
|
|
cli_overrides=cli_overrides,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("cli_overrides", "stage_id"),
|
|
[
|
|
({"kv_cache_dtype": "fp8"}, 0),
|
|
({"stage_0_kv_cache_dtype": "fp8"}, 0),
|
|
({"stage_1_kv_cache_dtype": "fp8"}, 1),
|
|
],
|
|
ids=["global", "ar-stage", "generation-stage"],
|
|
)
|
|
def test_from_pipeline_config_accepts_upstream_cache_dtype_for_llm_stages(cli_overrides, stage_id):
|
|
stage = _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides).stage_by_id(stage_id)
|
|
|
|
assert stage.cache_config.cache_dtype == "fp8"
|
|
assert "cache_dtype" in stage.cache_config._omni_explicit_fields
|
|
|
|
|
|
@pytest.mark.parametrize("field_name", ["data_parallel_master_ip", "data_parallel_address"])
|
|
def test_from_pipeline_config_normalizes_nested_parallel_config_aliases(field_name):
|
|
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
|
|
deploy = DeployConfig(
|
|
stages=[
|
|
StageDeployConfig(
|
|
stage_id=0,
|
|
engine_extras={"parallel_config": {field_name: "10.0.0.1"}},
|
|
)
|
|
]
|
|
)
|
|
|
|
stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
|
|
|
|
assert stage.parallel_config.data_parallel_master_ip == "10.0.0.1"
|
|
assert "data_parallel_master_ip" in stage.parallel_config._omni_explicit_fields
|
|
|
|
|
|
def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_stage():
|
|
stage = _from_pipeline_key(
|
|
"dreamzero",
|
|
deploy_config_path="dreamzero_tp1_cfg2",
|
|
cli_overrides={"kv_cache_dtype": "fp8"},
|
|
).stage_by_id(0)
|
|
|
|
assert isinstance(stage, VllmOmniDiffusionStageConfig)
|
|
assert stage.diffusion_config.diffusion_kv_cache_dtype == "fp8"
|
|
|
|
|
|
def test_stage_cli_field_selection_defers_ownership_validation_until_sources_are_merged():
|
|
assert omni_config_module._stage_cli_overrides(0, {"enable_lora": True}) == {"enable_lora": True}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("model_type", "diffusion_stage_ids"),
|
|
[
|
|
("bagel", {1}),
|
|
("bagel_think", {1}),
|
|
("bagel_single_stage", {0}),
|
|
],
|
|
)
|
|
def test_step_execution_cli_is_scoped_to_diffusion_stages(model_type, diffusion_stage_ids):
|
|
omni_config = _from_pipeline_key(model_type, cli_overrides={"step_execution": True})
|
|
|
|
for stage in omni_config.stage_configs:
|
|
if stage.stage_id in diffusion_stage_ids:
|
|
assert isinstance(stage, VllmOmniDiffusionStageConfig)
|
|
assert stage.diffusion_config.step_execution is True
|
|
else:
|
|
assert not hasattr(stage, "diffusion_config")
|
|
|
|
|
|
def test_stage_scoped_step_execution_ignores_bagel_ar_stage():
|
|
omni_config = _from_pipeline_key(
|
|
"bagel",
|
|
cli_overrides={
|
|
"stage_0_step_execution": True,
|
|
"stage_1_step_execution": True,
|
|
},
|
|
)
|
|
|
|
assert omni_config.stage_by_id(1).diffusion_config.step_execution is True
|
|
|
|
|
|
def test_runtime_num_gpus_is_derived_from_parallel_world_size():
|
|
omni_config = _from_pipeline_key("hunyuan_image3_dit")
|
|
stage = omni_config.stage_by_id(0)
|
|
|
|
assert stage.parallel_config.tensor_parallel_size == 4
|
|
assert stage.parallel_config.world_size == 4
|
|
assert stage.runtime_config.num_gpus == 4
|
|
|
|
|
|
def test_runtime_num_gpus_ignores_stale_runtime_override():
|
|
omni_config = _from_pipeline_key(
|
|
"hunyuan_image3_dit",
|
|
cli_overrides={
|
|
"stage_0_num_gpus": 1,
|
|
},
|
|
)
|
|
stage = omni_config.stage_by_id(0)
|
|
|
|
assert stage.parallel_config.world_size == 4
|
|
assert stage.runtime_config.num_gpus == 4
|
|
|
|
|
|
def test_from_pipeline_config_does_not_route_server_cli_keys_to_diffusion_stage():
|
|
omni_config = _from_pipeline_key(
|
|
"dreamzero",
|
|
deploy_config_path="dreamzero_tp1_cfg2",
|
|
cli_overrides={
|
|
"host": "0.0.0.0",
|
|
"port": 8000,
|
|
"api_key": "secret",
|
|
"stage_0_host": "127.0.0.1",
|
|
"stage_0_port": 23456,
|
|
},
|
|
)
|
|
|
|
stage = omni_config.stage_by_id(0)
|
|
|
|
assert isinstance(stage, VllmOmniDiffusionStageConfig)
|
|
assert stage.diffusion_config.host == "127.0.0.1"
|
|
assert stage.diffusion_config.port == 23456
|
|
assert not hasattr(stage.diffusion_config, "api_key")
|
|
|
|
|
|
def test_pipeline_deploy_cli_fields_reuse_legacy_pipeline_wide_engine_fields():
|
|
assert omni_config_module._PIPELINE_DEPLOY_CLI_FIELDS is PIPELINE_WIDE_ENGINE_FIELDS
|
|
assert "active_stream_window" in omni_config_module._PIPELINE_DEPLOY_CLI_FIELDS
|
|
assert "custom_voice_dir" in omni_config_module._PIPELINE_DEPLOY_CLI_FIELDS
|
|
|
|
|
|
def test_pipeline_wide_model_fields_are_retained_on_structured_stage_configs(tmp_path):
|
|
custom_voice_dir = tmp_path / "voices"
|
|
omni_config = _from_pipeline_key(
|
|
"qwen3_tts",
|
|
cli_overrides={
|
|
"active_stream_window": 2,
|
|
"custom_voice_dir": str(custom_voice_dir),
|
|
},
|
|
)
|
|
|
|
assert {stage.model_config.active_stream_window for stage in omni_config.stage_configs} == {2}
|
|
assert {stage.model_config.custom_voice_dir for stage in omni_config.stage_configs} == {str(custom_voice_dir)}
|
|
|
|
|
|
def test_stage_deploy_engine_fields_reuse_legacy_stage_deploy_fields():
|
|
assert omni_config_module._STAGE_DEPLOY_ENGINE_FIELDS == tuple(_STAGE_DEPLOY_FIELDS)
|
|
assert "tensor_parallel_size" in omni_config_module._STAGE_DEPLOY_ENGINE_FIELDS
|
|
assert "stage_id" not in omni_config_module._STAGE_DEPLOY_ENGINE_FIELDS
|
|
|
|
|
|
def test_public_config_exports_use_stage_specific_sub_config_names():
|
|
import vllm_omni.config as config_pkg
|
|
|
|
generic_names = {
|
|
"CacheConfig",
|
|
"ConnectorConfig",
|
|
"LoadConfig",
|
|
"ModelConfig",
|
|
"OrchestratorConfig",
|
|
"ParallelConfig",
|
|
"RuntimeConfig",
|
|
"SchedulerConfig",
|
|
}
|
|
|
|
assert generic_names.isdisjoint(config_pkg.__all__)
|
|
assert {
|
|
"OmniStageCacheConfig",
|
|
"OmniStageConnectorConfig",
|
|
"OmniStageDiffusionParallelConfig",
|
|
"OmniStageLoadConfig",
|
|
"OmniStageModelConfig",
|
|
"VllmOmniOrchestratorConfig",
|
|
"OmniStageParallelConfig",
|
|
"OmniStageRuntimeConfig",
|
|
"OmniStageSchedulerConfig",
|
|
"StageConfigType",
|
|
}.issubset(config_pkg.__all__)
|
|
|
|
|
|
def test_from_pipeline_config_keeps_worker_backend_separate_from_distributed_executor_backend():
|
|
omni_config = _from_pipeline_key("dreamzero", deploy_config_path="dreamzero_tp1_cfg2")
|
|
|
|
stage = omni_config.stage_by_id(0)
|
|
assert isinstance(stage, VllmOmniDiffusionStageConfig)
|
|
assert stage.diffusion_config.distributed_executor_backend == "mp"
|
|
assert omni_config.orchestrator_config.worker_backend == "multi_process"
|
|
|
|
|
|
def test_from_pipeline_config_maps_orchestrator_cli_overrides():
|
|
omni_config = _from_pipeline_key(
|
|
"qwen3_tts",
|
|
cli_overrides={
|
|
"stage_init_timeout": 1200,
|
|
"init_timeout": 1800,
|
|
"worker_backend": "ray",
|
|
"ray_address": "ray://127.0.0.1:10001",
|
|
"omni_master_address": "127.0.0.1",
|
|
"omni_master_port": 12345,
|
|
"omni_dp_size_local": 2,
|
|
"omni_lb_policy": "round_robin",
|
|
"omni_heartbeat_timeout": 9.5,
|
|
"batch_timeout": 3,
|
|
},
|
|
)
|
|
|
|
orchestrator_config = omni_config.orchestrator_config
|
|
assert orchestrator_config.stage_init_timeout == 1200
|
|
assert orchestrator_config.init_timeout == 1800
|
|
assert orchestrator_config.worker_backend == "ray"
|
|
assert orchestrator_config.ray_address == "ray://127.0.0.1:10001"
|
|
assert orchestrator_config.omni_master_address == "127.0.0.1"
|
|
assert orchestrator_config.omni_master_port == 12345
|
|
assert orchestrator_config.omni_dp_size_local == 2
|
|
assert orchestrator_config.omni_lb_policy == "round_robin"
|
|
assert orchestrator_config.omni_heartbeat_timeout == 9.5
|
|
assert orchestrator_config.batch_timeout == 3
|
|
|
|
|
|
def test_from_pipeline_config_records_loaded_deploy_path_on_orchestrator_config():
|
|
omni_config = _from_pipeline_key("dreamzero", deploy_config_path="dreamzero_tp1_cfg2")
|
|
|
|
assert omni_config.pipeline_config.model_type == "dreamzero"
|
|
assert omni_config.orchestrator_config.deploy_config_path == str(_DEPLOY_DIR / "dreamzero_tp1_cfg2.yaml")
|
|
|
|
|
|
def test_from_pipeline_config_dispatches_async_chunk_processors_without_mutating_topology():
|
|
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
|
|
|
|
async_config = _from_pipeline_key("qwen3_tts")
|
|
assert async_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_async_chunk")
|
|
assert async_config.stage_by_id(1).custom_process_input_func is None
|
|
assert async_config.stage_by_id(1).model_config.requires_full_payload_input is True
|
|
|
|
sync_config = _from_pipeline_key("qwen3_tts", cli_overrides={"async_chunk": False})
|
|
assert sync_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
|
|
assert sync_config.stage_by_id(1).custom_process_input_func.endswith("talker2code2wav_token_only")
|
|
assert sync_config.stage_by_id(1).model_config.requires_full_payload_input is True
|
|
|
|
assert pipeline.get_stage(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
|
|
assert pipeline.get_stage(1).custom_process_input_func is None
|
|
|
|
|
|
def test_joyai_code2wav_waits_for_full_payload():
|
|
config = _from_pipeline_key("joyai_vl_interaction")
|
|
talker = config.stage_by_id(1)
|
|
code2wav = config.stage_by_id(2)
|
|
|
|
assert talker.custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
|
|
assert code2wav.custom_process_input_func.endswith("talker2code2wav_token_only")
|
|
assert code2wav.connector_config.async_chunk is False
|
|
assert code2wav.model_config.requires_full_payload_input is True
|
|
|
|
|
|
def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
|
|
assert not hasattr(BaseVllmOmniStageConfig, "from_stage_config")
|
|
assert not hasattr(BaseVllmOmniStageConfig, "to_legacy_stage_config")
|
|
|
|
public_fields = {f.name for f in fields(BaseVllmOmniStageConfig)}
|
|
|
|
assert public_fields == {
|
|
"stage_pipeline_config",
|
|
"model_config",
|
|
"load_config",
|
|
"cache_config",
|
|
"scheduler_config",
|
|
"connector_config",
|
|
"runtime_config",
|
|
"parallel_config",
|
|
"compilation_config",
|
|
"profiler_config",
|
|
"quantization_config",
|
|
}
|
|
assert "diffusion_config" not in public_fields
|
|
assert {f.name for f in fields(VllmOmniDiffusionStageConfig)} == public_fields | {"diffusion_config"}
|
|
assert {f.name for f in fields(VllmOmniARStageConfig)} == public_fields
|
|
assert {f.name for f in fields(VllmOmniGenerationStageConfig)} == public_fields
|
|
|
|
|
|
def test_runtime_config_fields_match_structured_runtime_scope():
|
|
assert {f.name for f in fields(OmniStageRuntimeConfig)} == {
|
|
"additional_config",
|
|
"distributed_executor_backend",
|
|
"worker_cls",
|
|
"devices",
|
|
"num_replicas",
|
|
"env",
|
|
"num_gpus",
|
|
"log_level",
|
|
"log_stats",
|
|
}
|
|
|
|
|
|
def test_upstream_config_compatible_fields_materialize_mapping_inputs():
|
|
compilation_config = {"backend": "eager"}
|
|
profiler_config = {"profiler": "cuda"}
|
|
|
|
stage_config = VllmOmniARStageConfig(
|
|
stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
|
|
compilation_config=compilation_config,
|
|
profiler_config=profiler_config,
|
|
)
|
|
|
|
assert isinstance(stage_config.compilation_config, VllmCompilationConfig)
|
|
assert stage_config.compilation_config.backend == "eager"
|
|
assert isinstance(stage_config.profiler_config, VllmProfilerConfig)
|
|
assert stage_config.profiler_config.profiler == "cuda"
|
|
|
|
|
|
def test_upstream_config_compatible_fields_validate_mapping_inputs():
|
|
with pytest.raises(ValidationError):
|
|
VllmOmniARStageConfig(
|
|
stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
|
|
profiler_config={"delay_iterations": -1},
|
|
)
|
|
|
|
|
|
def test_upstream_config_compatible_fields_keep_prebuilt_config_types():
|
|
compilation_config = VllmCompilationConfig(backend="eager")
|
|
profiler_config = VllmProfilerConfig(profiler="cuda")
|
|
|
|
stage_config = VllmOmniARStageConfig(
|
|
stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
|
|
compilation_config=compilation_config,
|
|
profiler_config=profiler_config,
|
|
)
|
|
|
|
assert stage_config.compilation_config is compilation_config
|
|
assert stage_config.profiler_config is profiler_config
|
|
|
|
|
|
def test_sub_config_fields_match_structured_scopes():
|
|
assert {f.name for f in fields(OmniStageModelConfig)} == {
|
|
"model",
|
|
"model_arch",
|
|
"revision",
|
|
"tokenizer_revision",
|
|
"code_revision",
|
|
"seed",
|
|
"logits_processors",
|
|
"trust_remote_code",
|
|
"dtype",
|
|
"attention_backend",
|
|
"attention_config",
|
|
"moe_backend",
|
|
"hf_overrides",
|
|
"limit_mm_per_prompt",
|
|
"interleave_mm_strings",
|
|
"media_io_kwargs",
|
|
"active_stream_window",
|
|
"session_mode",
|
|
"duplex_max_sessions",
|
|
"enable_sleep_mode",
|
|
"default_sampling_params",
|
|
"subtalker_sampling_params",
|
|
"silence_ban_frames",
|
|
"has_sampling_extra_args",
|
|
"custom_voice_dir",
|
|
"task_type",
|
|
"codec_frame_rate_hz",
|
|
"enforce_eager",
|
|
"max_cudagraph_capture_size",
|
|
"enable_flashinfer_autotune",
|
|
"enable_multithread_weight_load",
|
|
"num_weight_load_threads",
|
|
"disable_autocast",
|
|
# Per-stage checkpoint resolution for repos whose stages live in
|
|
# subfolders (e.g. Audex): mirrors StagePipelineConfig on the
|
|
# legacy engine-args path.
|
|
"model_subdir",
|
|
"tokenizer_subdir",
|
|
"requires_full_payload_input",
|
|
}
|
|
vllm_load_fields = {f.name for f in fields(VllmLoadConfig)}
|
|
assert issubclass(OmniStageLoadConfig, VllmLoadConfig)
|
|
assert {f.name for f in fields(OmniStageLoadConfig)} == vllm_load_fields | {
|
|
"tokenizer",
|
|
"skip_tokenizer_init",
|
|
"tokenizer_mode",
|
|
"config_format",
|
|
"skip_mm_profiling",
|
|
}
|
|
assert OmniStageLoadConfig(load_format="PT").load_format == "pt"
|
|
assert issubclass(OmniStageCacheConfig, VllmCacheConfig)
|
|
assert {f.name for f in fields(OmniStageCacheConfig)} == {f.name for f in fields(VllmCacheConfig)} | {
|
|
"disable_hybrid_kv_cache_manager",
|
|
"mm_processor_cache_gb",
|
|
"mamba_ssm_cache_dtype",
|
|
}
|
|
assert issubclass(OmniStageSchedulerConfig, VllmSchedulerConfig)
|
|
assert {f.name for f in fields(OmniStageSchedulerConfig)} == {f.name for f in fields(VllmSchedulerConfig)} | {
|
|
"max_model_len",
|
|
}
|
|
assert {f.name for f in fields(OmniStageConnectorConfig)} == {
|
|
"async_chunk",
|
|
"omni_kv_config",
|
|
"stage_connector",
|
|
"output_connectors",
|
|
"input_connectors",
|
|
}
|
|
vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
|
|
assert issubclass(OmniStageParallelConfig, VllmParallelConfig)
|
|
assert {f.name for f in fields(OmniStageParallelConfig)} == vllm_parallel_fields
|
|
assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == vllm_parallel_fields | {
|
|
"sequence_parallel_size",
|
|
"ulysses_degree",
|
|
"ring_degree",
|
|
"allgather_degree",
|
|
"ulysses_mode",
|
|
"ulysses_a2a_permute",
|
|
"cfg_parallel_size",
|
|
"vae_patch_parallel_size",
|
|
"vae_parallel_mode",
|
|
"text_encoder_tp_size",
|
|
"use_hsdp",
|
|
"mask_sp_padding",
|
|
"hsdp_shard_size",
|
|
"hsdp_replicate_size",
|
|
}
|
|
|
|
|
|
def test_inherited_sub_configs_initialize_transport_safe_derived_fields():
|
|
scheduler_config = OmniStageSchedulerConfig(max_num_batched_tokens=4096)
|
|
assert scheduler_config.max_num_encoder_input_tokens == 4096
|
|
assert scheduler_config.encoder_cache_size == 4096
|
|
|
|
unresolved_scheduler_config = OmniStageSchedulerConfig()
|
|
assert unresolved_scheduler_config.max_num_encoder_input_tokens is None
|
|
assert unresolved_scheduler_config.encoder_cache_size is None
|
|
|
|
parallel_config = OmniStageParallelConfig(data_parallel_size=3, data_parallel_rank=2)
|
|
assert parallel_config.data_parallel_index == 2
|
|
|
|
diffusion_parallel_config = OmniStageDiffusionParallelConfig(data_parallel_size=3, data_parallel_rank=2)
|
|
assert diffusion_parallel_config.data_parallel_index == 2
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"config_cls",
|
|
[
|
|
OmniStageLoadConfig,
|
|
OmniStageCacheConfig,
|
|
OmniStageSchedulerConfig,
|
|
OmniStageParallelConfig,
|
|
OmniStageDiffusionParallelConfig,
|
|
],
|
|
)
|
|
def test_inherited_sub_configs_are_keyword_only(config_cls):
|
|
assert all(parameter.kind is Parameter.KEYWORD_ONLY for parameter in signature(config_cls).parameters.values())
|
|
|
|
with pytest.raises(TypeError):
|
|
config_cls(1)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"config_cls",
|
|
[
|
|
OmniStageLoadConfig,
|
|
OmniStageCacheConfig,
|
|
OmniStageSchedulerConfig,
|
|
OmniStageParallelConfig,
|
|
OmniStageDiffusionParallelConfig,
|
|
],
|
|
)
|
|
def test_inherited_sub_configs_remain_msgpack_transport_safe(config_cls):
|
|
config = config_cls()
|
|
|
|
uninitialized_fields = [
|
|
config_field.name
|
|
for config_field in fields(config)
|
|
if isinstance(getattr(config, config_field.name), FieldInfo)
|
|
]
|
|
assert uninitialized_fields == []
|
|
|
|
msgspec.msgpack.encode(_serialize_stage_config(config))
|
|
|
|
|
|
def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_safe():
|
|
omni_config = _from_pipeline_key("qwen3_tts")
|
|
|
|
for stage_config in omni_config.stage_configs:
|
|
payload = {
|
|
"stage_id": stage_config.stage_id,
|
|
"stage_config": _serialize_stage_config(stage_config),
|
|
}
|
|
msgspec.msgpack.encode(payload)
|
|
|
|
|
|
def test_diffusion_parallel_config_fields_cover_legacy_surface():
|
|
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
|
|
|
legacy_fields = {f.name for f in fields(DiffusionParallelConfig)}
|
|
structured_fields = {f.name for f in fields(OmniStageDiffusionParallelConfig)}
|
|
expected_upstream_fields = {"mask_sp_padding"}
|
|
vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
|
|
|
|
assert legacy_fields | expected_upstream_fields <= structured_fields
|
|
assert structured_fields - legacy_fields - expected_upstream_fields - vllm_parallel_fields == set()
|
|
|
|
|
|
def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
|
|
cfg = OmniStageDiffusionParallelConfig(
|
|
pipeline_parallel_size=2,
|
|
data_parallel_size=3,
|
|
tensor_parallel_size=4,
|
|
cfg_parallel_size=3,
|
|
mask_sp_padding=True,
|
|
)
|
|
|
|
assert cfg.pipeline_parallel_size == 2
|
|
assert cfg.data_parallel_size == 3
|
|
assert cfg.cfg_parallel_size == 3
|
|
assert cfg.mask_sp_padding is True
|
|
assert cfg.world_size == 72
|
|
|
|
|
|
def test_parallel_config_derived_fields_are_not_init_inputs():
|
|
with pytest.raises(ValidationError):
|
|
OmniStageParallelConfig(world_size=4)
|
|
|
|
with pytest.raises(ValidationError):
|
|
OmniStageDiffusionParallelConfig(world_size=4)
|
|
|
|
with pytest.raises(ValidationError):
|
|
OmniStageDiffusionParallelConfig(sequence_parallel_size=2)
|
|
|
|
|
|
def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae_patch_parallel():
|
|
cfg = OmniStageDiffusionParallelConfig(
|
|
tensor_parallel_size=2,
|
|
cfg_parallel_size=2,
|
|
vae_patch_parallel_size=4,
|
|
)
|
|
|
|
assert cfg.vae_patch_parallel_size == 4
|
|
assert cfg.world_size == 4
|
|
|
|
|
|
def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
|
|
cfg = OmniStageDiffusionParallelConfig(
|
|
ulysses_degree=4,
|
|
use_hsdp=True,
|
|
hsdp_shard_size=-1,
|
|
hsdp_replicate_size=2,
|
|
)
|
|
|
|
assert cfg.hsdp_shard_size == 2
|
|
assert cfg.world_size == 4
|
|
|
|
|
|
def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp():
|
|
with pytest.raises(ValueError, match="not compatible with TP"):
|
|
OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
|
|
|
|
with pytest.raises(ValueError, match="not compatible with DP"):
|
|
OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
|
|
|
|
|
|
def test_from_pipeline_config_preserves_legacy_pp_dp_for_world_size():
|
|
cfg = _from_pipeline_key("hunyuan_image3_dit").stage_by_id(0).parallel_config
|
|
|
|
assert cfg.pipeline_parallel_size == 1
|
|
assert cfg.data_parallel_size == 1
|
|
assert cfg.tensor_parallel_size == 4
|
|
assert cfg.world_size == 4
|
|
|
|
|
|
def test_from_pipeline_config_derives_sequence_parallel_size_from_degrees(tmp_path):
|
|
deploy_path = tmp_path / "dreamzero_derived_parallel.yaml"
|
|
deploy_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"pipeline: dreamzero",
|
|
"async_chunk: false",
|
|
"stages:",
|
|
" - stage_id: 0",
|
|
" parallel_config:",
|
|
" sequence_parallel_size: 99",
|
|
" ulysses_degree: 2",
|
|
" ring_degree: 3",
|
|
]
|
|
)
|
|
)
|
|
|
|
stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
|
|
|
|
assert isinstance(stage, VllmOmniDiffusionStageConfig)
|
|
assert stage.parallel_config.sequence_parallel_size == 6
|
|
assert stage.parallel_config.world_size == 6
|
|
|
|
|
|
def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degree(tmp_path):
|
|
deploy_path = tmp_path / "dreamzero_allgather_parallel.yaml"
|
|
deploy_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"pipeline: dreamzero",
|
|
"async_chunk: false",
|
|
"stages:",
|
|
" - stage_id: 0",
|
|
" parallel_config:",
|
|
" sequence_parallel_size: 99",
|
|
" allgather_degree: 2",
|
|
]
|
|
)
|
|
)
|
|
|
|
stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
|
|
|
|
assert isinstance(stage, VllmOmniDiffusionStageConfig)
|
|
assert stage.parallel_config.allgather_degree == 2
|
|
assert stage.parallel_config.sequence_parallel_size == 2
|
|
assert stage.parallel_config.world_size == 2
|
|
|
|
|
|
def test_from_pipeline_config_preserves_deployed_ulysses_a2a_permute() -> None:
|
|
pipeline = _resolve_pipeline_or_skip("dreamzero")
|
|
deploy = DeployConfig(
|
|
stages=[
|
|
StageDeployConfig(
|
|
stage_id=0,
|
|
ulysses_a2a_permute=True,
|
|
)
|
|
]
|
|
)
|
|
|
|
stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
|
|
|
|
assert isinstance(stage, VllmOmniDiffusionStageConfig)
|
|
assert stage.parallel_config.ulysses_a2a_permute is True
|
|
|
|
|
|
def test_diffusion_parallel_config_accepts_four_way_guidance_parallelism():
|
|
cfg = OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
|
|
|
|
assert cfg.cfg_parallel_size == 4
|
|
assert cfg.world_size == 4
|
|
|
|
|
|
def test_diffusion_parallel_config_rejects_allgather_with_ulysses_or_ring():
|
|
with pytest.raises(ValidationError):
|
|
OmniStageDiffusionParallelConfig(allgather_degree=2, ulysses_degree=2)
|
|
|
|
|
|
def test_stage_realizations_use_stage_specific_parallel_config_types():
|
|
qwen_config = _from_pipeline_key("qwen3_tts")
|
|
ar_stage = qwen_config.stage_by_id(0)
|
|
generation_stage = qwen_config.stage_by_id(1)
|
|
diffusion_stage = _from_pipeline_key("hunyuan_image3_dit").stage_by_id(0)
|
|
|
|
assert isinstance(ar_stage, VllmOmniARStageConfig)
|
|
assert type(ar_stage.parallel_config) is OmniStageParallelConfig
|
|
assert not hasattr(ar_stage.parallel_config, "cfg_parallel_size")
|
|
assert not hasattr(ar_stage.parallel_config, "sequence_parallel_size")
|
|
assert not hasattr(ar_stage.parallel_config, "ulysses_degree")
|
|
|
|
assert isinstance(generation_stage, VllmOmniGenerationStageConfig)
|
|
assert type(generation_stage.parallel_config) is OmniStageParallelConfig
|
|
assert not hasattr(generation_stage.parallel_config, "cfg_parallel_size")
|
|
assert not hasattr(generation_stage.parallel_config, "sequence_parallel_size")
|
|
assert not hasattr(generation_stage.parallel_config, "ulysses_degree")
|
|
|
|
assert isinstance(diffusion_stage, VllmOmniDiffusionStageConfig)
|
|
assert isinstance(diffusion_stage.parallel_config, OmniStageDiffusionParallelConfig)
|
|
assert diffusion_stage.parallel_config.cfg_parallel_size == 1
|
|
assert diffusion_stage.parallel_config.sequence_parallel_size == 1
|
|
assert diffusion_stage.parallel_config.ulysses_degree == 1
|
|
|
|
|
|
def test_from_pipeline_config_preserves_diffusion_parallel_mask_sp_padding(tmp_path):
|
|
deploy_path = tmp_path / "dreamzero_mask_sp_padding.yaml"
|
|
deploy_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"pipeline: dreamzero",
|
|
"async_chunk: false",
|
|
"stages:",
|
|
" - stage_id: 0",
|
|
" parallel_config:",
|
|
" mask_sp_padding: true",
|
|
]
|
|
)
|
|
)
|
|
|
|
stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
|
|
|
|
assert isinstance(stage, VllmOmniDiffusionStageConfig)
|
|
assert stage.parallel_config.mask_sp_padding is True
|
|
|
|
|
|
def test_from_pipeline_config_routes_regional_compile_dynamic(tmp_path):
|
|
deploy_path = tmp_path / "dreamzero_compile.yaml"
|
|
deploy_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"pipeline: dreamzero",
|
|
"async_chunk: false",
|
|
"stages:",
|
|
" - stage_id: 0",
|
|
" diffusion_compile_granularity: regional",
|
|
" diffusion_compile_dynamic: false",
|
|
]
|
|
)
|
|
)
|
|
|
|
configured_stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
|
|
overridden_stage = _from_pipeline_key(
|
|
"dreamzero",
|
|
deploy_config_path=str(deploy_path),
|
|
cli_overrides={
|
|
"diffusion_compile_granularity": "full",
|
|
"diffusion_compile_dynamic": True,
|
|
},
|
|
).stage_by_id(0)
|
|
|
|
assert configured_stage.diffusion_config.diffusion_compile_granularity == "regional"
|
|
assert configured_stage.diffusion_config.diffusion_compile_dynamic is False
|
|
assert overridden_stage.diffusion_config.diffusion_compile_granularity == "full"
|
|
assert overridden_stage.diffusion_config.diffusion_compile_dynamic is True
|
|
|
|
|
|
def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
|
|
with pytest.raises(ValidationError, match="diffusion_compile_dynamic"):
|
|
omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
|
|
|
|
|
|
def test_from_pipeline_config_routes_ltx2_conv_vae_extra(tmp_path):
|
|
deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml"
|
|
deploy_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"pipeline: dreamzero",
|
|
"async_chunk: false",
|
|
"stages:",
|
|
" - stage_id: 0",
|
|
" extras:",
|
|
" ltx2_use_conv_vae: true",
|
|
]
|
|
)
|
|
)
|
|
|
|
stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
|
|
|
|
assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
|
|
|
|
|
|
def test_stage_override_routes_ltx2_conv_vae_extra():
|
|
stage = _from_pipeline_key(
|
|
"dreamzero",
|
|
cli_overrides={"stage_0_extras": {"ltx2_use_conv_vae": True}},
|
|
).stage_by_id(0)
|
|
|
|
assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
|
|
|
|
|
|
def test_structured_diffusion_config_rejects_invalid_compile_granularity():
|
|
with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
|
|
omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
|
|
|
|
|
|
def test_from_pipeline_config_matches_stage_config_to_omegaconf_behavior_for_representative_stage():
|
|
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
|
|
legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0]
|
|
omega_stage = legacy_stage.to_omegaconf()
|
|
omni_stage = _from_pipeline_key("qwen3_tts").stage_by_id(legacy_stage.stage_id)
|
|
|
|
assert omega_stage.stage_id == omni_stage.stage_id
|
|
assert omega_stage.stage_type == omni_stage.stage_type.value
|
|
assert omega_stage.engine_input_source == omni_stage.input_sources
|
|
assert omega_stage.final_output == omni_stage.final_output
|
|
assert omega_stage.final_output_type == omni_stage.final_output_type
|
|
assert omega_stage.is_comprehension == omni_stage.is_comprehension
|
|
assert omega_stage.engine_args.model_stage == omni_stage.model_stage
|
|
assert omega_stage.engine_args.worker_type == omni_stage.worker_type
|
|
assert omega_stage.engine_args.scheduler_cls == omni_stage.scheduler_cls
|
|
assert omega_stage.runtime.process is True
|
|
assert omega_stage.runtime.requires_multimodal_data == omni_stage.requires_multimodal_data
|
|
|
|
|
|
def test_from_pipeline_config_uses_hf_config_for_callable_resolver():
|
|
hf_config = Qwen3OmniMoeConfig()
|
|
hf_config.enable_audio_output = False
|
|
|
|
omni_config = _from_pipeline_key("qwen3_omni_moe", hf_config=hf_config)
|
|
|
|
assert omni_config.pipeline_config.model_type == "qwen3_omni_moe_thinker_only"
|
|
assert len(omni_config.stage_configs) == 1
|
|
assert omni_config.orchestrator_config.deploy_config_path is None
|
|
|
|
thinker = omni_config.stage_configs[0]
|
|
assert thinker.model_stage == "thinker"
|
|
assert thinker.model_config.default_sampling_params == {"detokenize": True}
|
|
|
|
|
|
def test_from_pipeline_config_accepts_pre_resolved_pipeline():
|
|
resolved_pipeline = PipelineConfig(
|
|
model_type="callable_resolved_variant",
|
|
stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),),
|
|
)
|
|
|
|
omni_config = VllmOmniConfig.from_pipeline_config(resolved_pipeline)
|
|
|
|
assert omni_config.pipeline_config is resolved_pipeline
|
|
|
|
|
|
def test_from_pipeline_config_prefers_loaded_user_deploy_config(monkeypatch):
|
|
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
|
|
user_deploy_config = DeployConfig(
|
|
stages=[StageDeployConfig(stage_id=0, max_num_seqs=7)],
|
|
)
|
|
monkeypatch.setattr(
|
|
omni_config_module,
|
|
"load_deploy_config",
|
|
lambda _path: pytest.fail("default deploy config should not be loaded"),
|
|
)
|
|
|
|
omni_config = VllmOmniConfig.from_pipeline_config(
|
|
pipeline,
|
|
user_deploy_config=user_deploy_config,
|
|
)
|
|
|
|
assert omni_config.stage_by_id(0).scheduler_config.max_num_seqs == 7
|
|
|
|
|
|
def test_from_pipeline_config_default_deploy_name_ignores_cwd(monkeypatch, tmp_path):
|
|
default_name = "pipeline_default.yaml"
|
|
(tmp_path / default_name).write_text("stages: []\n", encoding="utf-8")
|
|
monkeypatch.chdir(tmp_path)
|
|
pipeline = PipelineConfig(
|
|
model_type="pipeline_with_default",
|
|
default_deploy_config_name=default_name,
|
|
stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),),
|
|
)
|
|
loaded_paths = []
|
|
|
|
def _load_deploy_config(path):
|
|
loaded_paths.append(Path(path))
|
|
return DeployConfig()
|
|
|
|
monkeypatch.setattr(omni_config_module, "load_deploy_config", _load_deploy_config)
|
|
|
|
omni_config = VllmOmniConfig.from_pipeline_config(pipeline)
|
|
|
|
assert omni_config.orchestrator_config.deploy_config_path == str(_DEPLOY_DIR / default_name)
|
|
assert loaded_paths == [_DEPLOY_DIR / default_name]
|
|
|
|
|
|
def test_from_pipeline_config_uses_resolved_deploy_pipeline():
|
|
deploy_path = get_deploy_config_path("aura_omni.yaml")
|
|
pipeline = _resolve_pipeline_or_skip("aura_omni")
|
|
|
|
omni_config = VllmOmniConfig.from_pipeline_config(
|
|
pipeline,
|
|
deploy_config_path=str(deploy_path),
|
|
)
|
|
|
|
assert omni_config.pipeline_config.model_type == "aura_omni"
|
|
assert [stage.model_stage for stage in omni_config.stage_configs] == [
|
|
"asr",
|
|
"aura",
|
|
"qwen3_tts",
|
|
"code2wav",
|
|
]
|
|
|
|
|
|
def test_from_pipeline_config_matches_to_omegaconf_diffusion_parallel_config():
|
|
pipeline = _resolve_pipeline_or_skip("hunyuan_image3_dit")
|
|
legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0]
|
|
omega_stage = legacy_stage.to_omegaconf()
|
|
omni_stage = _from_pipeline_key("hunyuan_image3_dit").stage_by_id(legacy_stage.stage_id)
|
|
|
|
assert (
|
|
omega_stage.engine_args.parallel_config.pipeline_parallel_size
|
|
== omni_stage.parallel_config.pipeline_parallel_size
|
|
)
|
|
assert omega_stage.engine_args.parallel_config.data_parallel_size == omni_stage.parallel_config.data_parallel_size
|
|
assert (
|
|
omega_stage.engine_args.parallel_config.tensor_parallel_size == omni_stage.parallel_config.tensor_parallel_size
|
|
)
|
|
assert (
|
|
omega_stage.engine_args.parallel_config.sequence_parallel_size
|
|
== omni_stage.parallel_config.sequence_parallel_size
|
|
)
|
|
assert omega_stage.engine_args.parallel_config.cfg_parallel_size == omni_stage.parallel_config.cfg_parallel_size
|
|
assert (
|
|
omega_stage.engine_args.parallel_config.vae_patch_parallel_size
|
|
== omni_stage.parallel_config.vae_patch_parallel_size
|
|
)
|
|
|
|
|
|
def test_from_pipeline_config_matches_build_engine_args_dict_behavior_for_representative_stage(monkeypatch):
|
|
from vllm_omni.engine import stage_init_utils
|
|
|
|
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None)
|
|
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
|
|
legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0]
|
|
omega_stage = legacy_stage.to_omegaconf()
|
|
legacy_engine_args = build_legacy_engine_args_dict(
|
|
omega_stage,
|
|
model="/tmp/qwen3-tts",
|
|
stage_connector_spec={"name": "SharedMemoryConnector", "extra": {}},
|
|
)
|
|
omni_stage = _from_pipeline_key("qwen3_tts").stage_by_id(legacy_stage.stage_id)
|
|
|
|
assert legacy_engine_args["model"] == "/tmp/qwen3-tts"
|
|
assert legacy_engine_args["stage_id"] == omni_stage.stage_id
|
|
assert legacy_engine_args["model_stage"] == omni_stage.model_stage
|
|
assert legacy_engine_args["worker_type"] == omni_stage.worker_type
|
|
assert legacy_engine_args["scheduler_cls"] == omni_stage.scheduler_cls
|
|
assert legacy_engine_args["stage_connector_spec"] == {"name": "SharedMemoryConnector", "extra": {}}
|
|
assert legacy_engine_args["has_sampling_extra_args"] == bool(
|
|
(omni_stage.model_config.default_sampling_params or {}).get("extra_args")
|
|
)
|
|
assert omni_stage.model_config.has_sampling_extra_args == legacy_engine_args["has_sampling_extra_args"]
|
|
|
|
|
|
def test_from_pipeline_config_derives_has_sampling_extra_args_from_stage_defaults():
|
|
stage = _from_pipeline_key("voxtral_tts").stage_by_id(0)
|
|
|
|
assert (stage.model_config.default_sampling_params or {}).get("extra_args")
|
|
assert stage.model_config.has_sampling_extra_args is True
|
|
|
|
|
|
def test_diffusion_config_preserves_existing_coercion_hooks():
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.data import AttentionConfig, DiffusionCacheConfig
|
|
|
|
cfg = omni_config_module._DiffusionConfigProjection(
|
|
dtype="float32",
|
|
cache_config={"rel_l1_thresh": 0.3},
|
|
diffusion_attention_config={"default": "flash_attn"},
|
|
diffusion_kv_cache_skip_steps="0-2,4",
|
|
diffusion_kv_cache_skip_layers=[1, 3],
|
|
)
|
|
|
|
assert cfg.dtype is torch.float32
|
|
assert isinstance(cfg.cache_config, DiffusionCacheConfig)
|
|
assert isinstance(cfg.diffusion_attention_config, AttentionConfig)
|
|
assert cfg.diffusion_attention_config.default.backend == "flash_attn"
|
|
assert cfg.diffusion_kv_cache_skip_step_indices == {0, 1, 2, 4}
|
|
assert cfg.diffusion_kv_cache_skip_layer_indices == {1, 3}
|
|
assert cfg.max_cpu_loras == 1
|
|
|
|
|
|
def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "TEA_CACHE")
|
|
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
|
|
|
|
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
|
|
diffusion_attention_backend="flash_attn",
|
|
fa_deterministic=True,
|
|
kv_cache_dtype="fp8",
|
|
kv_cache_skip_steps="0-1",
|
|
kv_cache_skip_layers=[2],
|
|
static_lora_scale=0.25,
|
|
diffusion_kv_mode="paged_scheduler",
|
|
diffusion_kv_max_rows_per_request=2,
|
|
diffusers_load_kwargs=None,
|
|
diffusers_call_kwargs=None,
|
|
)
|
|
|
|
assert cfg.diffusion_attention_config.default.backend == "flash_attn"
|
|
assert cfg.fa_deterministic is True
|
|
assert cfg.diffusion_kv_cache_dtype == "fp8"
|
|
assert cfg.diffusion_kv_cache_skip_step_indices == {0, 1}
|
|
assert cfg.diffusion_kv_cache_skip_layer_indices == {2}
|
|
assert cfg.lora_scale == 0.25
|
|
assert cfg.cache_backend == "tea_cache"
|
|
assert cfg.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
|
|
assert cfg.diffusers_load_kwargs == {}
|
|
assert cfg.diffusers_call_kwargs == {}
|
|
|
|
|
|
def test_diffusion_config_none_values_preserve_dataclass_defaults():
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
normalized = OmniDiffusionConfig.normalize_init_kwargs(
|
|
{
|
|
"lora_scale": None,
|
|
"enable_sleep_mode": None,
|
|
"diffusers_load_kwargs": None,
|
|
}
|
|
)
|
|
|
|
assert "lora_scale" not in normalized
|
|
assert "enable_sleep_mode" not in normalized
|
|
assert normalized["diffusers_load_kwargs"] == {}
|
|
|
|
config = OmniDiffusionConfig.from_kwargs(
|
|
lora_scale=None,
|
|
enable_sleep_mode=None,
|
|
)
|
|
assert config.lora_scale == 1.0
|
|
assert config.enable_sleep_mode is False
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("canonical_key", "alias_key", "canonical_value", "alias_value"),
|
|
[
|
|
("lora_scale", "static_lora_scale", 0.75, 0.25),
|
|
(
|
|
"quantization_config",
|
|
"diffusion_quantization_config",
|
|
{"method": "canonical"},
|
|
{"method": "diffusion-alias"},
|
|
),
|
|
(
|
|
"quantization_config",
|
|
"quantization",
|
|
{"method": "canonical"},
|
|
"legacy-alias",
|
|
),
|
|
("diffusion_kv_cache_dtype", "kv_cache_dtype", "fp8", "fp16"),
|
|
("diffusion_kv_cache_skip_steps", "kv_cache_skip_steps", "0-1", "2-3"),
|
|
("diffusion_kv_cache_skip_layers", "kv_cache_skip_layers", "1-2", "3-4"),
|
|
("streaming_output", "diffusion_streaming_output", False, True),
|
|
],
|
|
)
|
|
def test_diffusion_alias_conflicts_prefer_canonical_key(
|
|
canonical_key,
|
|
alias_key,
|
|
canonical_value,
|
|
alias_value,
|
|
):
|
|
from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs
|
|
|
|
normalized = normalize_omni_diffusion_kwargs(
|
|
{
|
|
canonical_key: canonical_value,
|
|
alias_key: alias_value,
|
|
}
|
|
)
|
|
|
|
assert normalized[canonical_key] == canonical_value
|
|
assert alias_key not in normalized
|
|
|
|
|
|
def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch):
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
|
|
deploy_path = tmp_path / "dreamzero_diffusion_aliases.yaml"
|
|
deploy_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"pipeline: dreamzero",
|
|
"async_chunk: false",
|
|
"stages:",
|
|
" - stage_id: 0",
|
|
" diffusion_attention_backend: flash_attn",
|
|
" fa_deterministic: true",
|
|
" diffusion_kv_mode: paged_scheduler",
|
|
" diffusion_kv_max_rows_per_request: 2",
|
|
" kv_transfer_config:",
|
|
" kv_connector: MooncakeConnector",
|
|
" kv_role: kv_consumer",
|
|
" engine_id: dit-engine-1",
|
|
]
|
|
)
|
|
)
|
|
|
|
stage = _from_pipeline_key(
|
|
"dreamzero",
|
|
deploy_config_path=str(deploy_path),
|
|
).stage_by_id(0)
|
|
|
|
assert isinstance(stage, VllmOmniDiffusionStageConfig)
|
|
assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn"
|
|
assert stage.diffusion_config.fa_deterministic is True
|
|
assert stage.diffusion_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
|
|
assert stage.diffusion_config.diffusion_kv_max_rows_per_request == 2
|
|
assert stage.diffusion_config.kv_transfer_config.engine_id == "dit-engine-1"
|
|
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
|
|
|
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
|
|
od_config = OmniDiffusionConfig.from_kwargs(**engine_args)
|
|
assert od_config.kv_transfer_config.engine_id == "dit-engine-1"
|
|
|
|
|
|
def test_from_pipeline_config_forwards_fastvideo_vsa_topk(tmp_path, monkeypatch):
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
|
|
deploy_path = tmp_path / "dreamzero_fastvideo_vsa.yaml"
|
|
deploy_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"pipeline: dreamzero",
|
|
"async_chunk: false",
|
|
"stages:",
|
|
" - stage_id: 0",
|
|
" diffusion_attention_backend: FASTVIDEO_VSA",
|
|
" fastvideo_vsa_topk: 96",
|
|
]
|
|
)
|
|
)
|
|
|
|
stage = _from_pipeline_key(
|
|
"dreamzero",
|
|
deploy_config_path=str(deploy_path),
|
|
).stage_by_id(0)
|
|
|
|
attention_config = stage.diffusion_config.diffusion_attention_config
|
|
assert attention_config.default is not None
|
|
assert attention_config.default.backend == "FASTVIDEO_VSA"
|
|
assert attention_config.default.fastvideo_vsa_topk == 96
|
|
|
|
|
|
def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
|
|
deploy_path = tmp_path / "dreamzero_reserved_diffusion_kv.yaml"
|
|
deploy_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"pipeline: dreamzero",
|
|
"async_chunk: false",
|
|
"stages:",
|
|
" - stage_id: 0",
|
|
" diffusion_kv_mode: paged_worker_local",
|
|
]
|
|
)
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="reserved but not implemented"):
|
|
_from_pipeline_key(
|
|
"dreamzero",
|
|
deploy_config_path=str(deploy_path),
|
|
)
|
|
|
|
|
|
def test_diffusion_config_field_classification_covers_current_fields():
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
classified_fields = (
|
|
omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS
|
|
| omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS
|
|
| omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
|
|
)
|
|
|
|
assert classified_fields == {f.name for f in fields(omni_config_module._DiffusionConfigProjection)}
|
|
assert {f.name for f in fields(OmniDiffusionConfig)} <= (
|
|
classified_fields | omni_config_module._DIFFUSION_MOVED_SHARED_FIELDS
|
|
)
|
|
assert {
|
|
"enable_prompt_embed_cache",
|
|
"enable_session_state_manager",
|
|
"prompt_embed_cache_size",
|
|
"diffusion_kv_cache_dtype",
|
|
"diffusion_kv_mode",
|
|
"diffusion_kv_max_rows_per_request",
|
|
} <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
|
|
assert {
|
|
"revision",
|
|
"trust_remote_code",
|
|
"distributed_executor_backend",
|
|
} <= omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS
|
|
assert "prompt_file_path" in omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS
|
|
|
|
|
|
def test_diffusion_config_projection_keeps_mapping_quantization_config_serializable():
|
|
quantization_config = {
|
|
"method": "example_quant",
|
|
"weights": "weights.bin",
|
|
}
|
|
|
|
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config)
|
|
|
|
assert cfg.quantization_config == quantization_config
|
|
|
|
|
|
def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch):
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
quantization_config = {"method": "int8", "activation_scheme": "dynamic"}
|
|
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
|
|
quantization_config=quantization_config,
|
|
)
|
|
|
|
# Exercise the terminal construction performed by the future typed startup
|
|
# path without probing ports or loading remote model metadata.
|
|
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
|
|
monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None)
|
|
cfg.enrich_config()
|
|
|
|
assert cfg.quantization_config is not None
|
|
assert cfg.quantization_config.get_name() == "int8"
|
|
|
|
|
|
def test_compact_offload_config_reaches_terminal_config(monkeypatch):
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload_strategy
|
|
|
|
compact_config = {
|
|
"mode": "layer",
|
|
"components": ["dit", "text_encoder"],
|
|
"layer_options": {
|
|
"dit": {"weight_transfer": "rank-local", "resident_layers": 12},
|
|
"text_encoder": {"weight_transfer": "allgather"},
|
|
},
|
|
}
|
|
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
|
|
diffusion_offload_config=compact_config,
|
|
)
|
|
|
|
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
|
|
monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None)
|
|
cfg.enrich_config()
|
|
|
|
assert cfg.diffusion_offload_config == compact_config
|
|
assert cfg.extras == {}
|
|
assert resolve_offload_strategy(cfg) is OffloadStrategy.DISTRIBUTED_LAYER_WISE
|
|
restored = ForkingPickler.loads(ForkingPickler.dumps(cfg))
|
|
assert resolve_offload_strategy(restored) is OffloadStrategy.DISTRIBUTED_LAYER_WISE
|
|
|
|
|
|
def test_global_diffusion_offload_config_targets_only_diffusion_stages():
|
|
compact_config = {"mode": "layer", "components": ["dit"]}
|
|
|
|
config = _from_pipeline_key(
|
|
"minimax_h3_disaggregated",
|
|
cli_overrides={"diffusion_offload_config": compact_config},
|
|
)
|
|
|
|
assert not hasattr(config.stage_by_id(0), "diffusion_config")
|
|
assert config.stage_by_id(1).diffusion_config.diffusion_offload_config == compact_config
|
|
|
|
|
|
def test_explicit_llm_stage_diffusion_offload_override_is_rejected():
|
|
with pytest.raises(ValueError, match="no structured config owner: diffusion_offload_config"):
|
|
_from_pipeline_key(
|
|
"minimax_h3_disaggregated",
|
|
cli_overrides={
|
|
"stage_0_diffusion_offload_config": {"mode": "layer", "components": ["dit"]},
|
|
},
|
|
)
|
|
|
|
|
|
def test_compact_offload_config_is_validated_during_projection():
|
|
with pytest.raises(ValueError, match="Unknown diffusion offload mode"):
|
|
omni_config_module._DiffusionConfigProjection.from_kwargs(
|
|
diffusion_offload_config={
|
|
"mode": "layerwise",
|
|
"components": ["dit"],
|
|
}
|
|
)
|