From f7cdc45977958cd8901edccb11d8f481bca97f0c Mon Sep 17 00:00:00 2001 From: Haco <923390377@qq.com> Date: Sat, 4 Jul 2026 11:17:25 +0800 Subject: [PATCH] [BugFix] deploy schema stage runtime overrides for CLI diffusion flags no longer works (#4101) Signed-off-by: xiaohajiayou <923390377@qq.com> --- tests/test_config_factory.py | 191 +++++++++++++++++++++++++++++++ vllm_omni/config/stage_config.py | 87 +++++++++++--- 2 files changed, 262 insertions(+), 16 deletions(-) diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 63a67d652..23d36c338 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -1588,6 +1588,197 @@ class TestAuraOmniDeploy: assert stages[3].yaml_engine_args["model_arch"] == "Qwen3TTSCode2Wav" +class TestDeployCliOverrideFlow: + """Test deploy-YAML baselines overridden by CLI runtime overrides.""" + + def test_diffusion_deploy_fields_can_be_overridden_by_cli(self, tmp_path): + deploy_path = tmp_path / "diffusion_stage.yaml" + deploy_path.write_text( + """ +async_chunk: false +stages: + - stage_id: 0 + devices: "4,5" + parallel_config: + pipeline_parallel_size: 1 + data_parallel_size: 1 + tensor_parallel_size: 1 + enable_expert_parallel: false + sequence_parallel_size: 1 + ulysses_degree: 1 + ring_degree: 1 + cfg_parallel_size: 1 + vae_patch_parallel_size: 1 + use_hsdp: false + hsdp_shard_size: -1 + hsdp_replicate_size: 1 + engine_args: + cache_backend: cache_dit + diffusion_attention_backend: FLASH_ATTN + diffusion_kv_cache_dtype: auto + step_execution: false + vae_use_tiling: false + enable_cpu_offload: false + max_generated_image_size: 1048576 + tts_max_instructions_length: 1000 +""", + encoding="utf-8", + ) + + pipeline = PipelineConfig( + model_type="test_diffusion", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + final_output=True, + final_output_type="image", + ), + ), + ) + deploy = load_deploy_config(deploy_path) + stages = merge_pipeline_deploy(pipeline, deploy) + stage = stages[0] + + assert stage.yaml_engine_args["parallel_config"]["pipeline_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["data_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["enable_expert_parallel"] is False + assert stage.yaml_engine_args["parallel_config"]["sequence_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["ulysses_degree"] == 1 + assert stage.yaml_engine_args["parallel_config"]["ring_degree"] == 1 + assert stage.yaml_engine_args["parallel_config"]["cfg_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["vae_patch_parallel_size"] == 1 + assert stage.yaml_engine_args["cache_backend"] == "cache_dit" + assert stage.yaml_engine_args["parallel_config"]["use_hsdp"] is False + assert stage.yaml_engine_args["parallel_config"]["hsdp_shard_size"] == -1 + assert stage.yaml_engine_args["parallel_config"]["hsdp_replicate_size"] == 1 + assert stage.yaml_engine_args["diffusion_attention_backend"] == "FLASH_ATTN" + assert stage.yaml_engine_args["diffusion_kv_cache_dtype"] == "auto" + assert stage.yaml_engine_args["step_execution"] is False + assert stage.yaml_engine_args["vae_use_tiling"] is False + assert stage.yaml_engine_args["enable_cpu_offload"] is False + assert stage.yaml_engine_args["max_generated_image_size"] == 1048576 + assert stage.yaml_engine_args["tts_max_instructions_length"] == 1000 + + stage.runtime_overrides = StageConfigFactory._merge_cli_overrides( + stage, + { + "pipeline_parallel_size": 2, + "data_parallel_size": 3, + "tensor_parallel_size": 4, + "enable_expert_parallel": True, + "sequence_parallel_size": 24, + "ulysses_degree": 2, + "ring_degree": 4, + "cfg_parallel_size": 2, + "vae_patch_parallel_size": 2, + "cache_backend": "tea_cache", + "use_hsdp": True, + "hsdp_shard_size": 8, + "hsdp_replicate_size": 2, + "diffusion_attention_backend": "SAGE_ATTN", + "diffusion_kv_cache_dtype": "fp8", + "step_execution": True, + "vae_use_tiling": True, + "enable_cpu_offload": True, + "max_generated_image_size": 2097152, + "tts_max_instructions_length": 2000, + }, + ) + + omega_config = stage.to_omegaconf() + + assert omega_config.engine_args.cache_backend == "tea_cache" + assert omega_config.engine_args.diffusion_attention_backend == "SAGE_ATTN" + assert omega_config.engine_args.diffusion_kv_cache_dtype == "fp8" + assert omega_config.engine_args.step_execution is True + assert omega_config.engine_args.vae_use_tiling is True + assert omega_config.engine_args.enable_cpu_offload is True + assert omega_config.engine_args.max_generated_image_size == 2097152 + assert omega_config.engine_args.tts_max_instructions_length == 2000 + assert omega_config.engine_args.parallel_config.pipeline_parallel_size == 2 + assert omega_config.engine_args.parallel_config.data_parallel_size == 3 + assert omega_config.engine_args.parallel_config.tensor_parallel_size == 4 + assert omega_config.engine_args.parallel_config.enable_expert_parallel is True + assert omega_config.engine_args.parallel_config.sequence_parallel_size == 24 + assert omega_config.engine_args.parallel_config.ulysses_degree == 2 + assert omega_config.engine_args.parallel_config.ring_degree == 4 + assert omega_config.engine_args.parallel_config.cfg_parallel_size == 2 + assert omega_config.engine_args.parallel_config.vae_patch_parallel_size == 2 + assert omega_config.engine_args.parallel_config.use_hsdp is True + assert omega_config.engine_args.parallel_config.hsdp_shard_size == 8 + assert omega_config.engine_args.parallel_config.hsdp_replicate_size == 2 + + def test_llm_deploy_fields_can_be_overridden_by_cli(self, tmp_path): + deploy_path = tmp_path / "llm_stage.yaml" + deploy_path.write_text( + """ +async_chunk: false +stages: + - stage_id: 0 + devices: "0,1" + engine_args: + tensor_parallel_size: 1 + enable_expert_parallel: false + gpu_memory_utilization: 0.5 + max_num_seqs: 16 + max_num_batched_tokens: 1024 + max_model_len: 4096 + enforce_eager: false +""", + encoding="utf-8", + ) + + pipeline = PipelineConfig( + model_type="test_llm", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="thinker", + execution_type=StageExecutionType.LLM_AR, + final_output=True, + final_output_type="text", + ), + ), + ) + deploy = load_deploy_config(deploy_path) + stages = merge_pipeline_deploy(pipeline, deploy) + stage = stages[0] + + assert stage.yaml_engine_args["tensor_parallel_size"] == 1 + assert stage.yaml_engine_args["enable_expert_parallel"] is False + assert stage.yaml_engine_args["gpu_memory_utilization"] == 0.5 + assert stage.yaml_engine_args["max_num_seqs"] == 16 + assert stage.yaml_engine_args["max_num_batched_tokens"] == 1024 + assert stage.yaml_engine_args["max_model_len"] == 4096 + assert stage.yaml_engine_args["enforce_eager"] is False + + stage.runtime_overrides = StageConfigFactory._merge_cli_overrides( + stage, + { + "tensor_parallel_size": 2, + "enable_expert_parallel": True, + "gpu_memory_utilization": 0.9, + "max_num_seqs": 32, + "max_num_batched_tokens": 2048, + "max_model_len": 8192, + "enforce_eager": True, + }, + ) + + omega_config = stage.to_omegaconf() + + assert omega_config.engine_args.tensor_parallel_size == 2 + assert omega_config.engine_args.enable_expert_parallel is True + assert omega_config.engine_args.gpu_memory_utilization == 0.9 + assert omega_config.engine_args.max_num_seqs == 32 + assert omega_config.engine_args.max_num_batched_tokens == 2048 + assert omega_config.engine_args.max_model_len == 8192 + assert omega_config.engine_args.enforce_eager is True + + class TestSentinelDefaultPrecedence: """Caller-typed (non-None) values win over YAML; None values fall through to YAML / dataclass defaults (#3035).""" diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index a269216da..f66023bc4 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -63,7 +63,12 @@ def build_stage_runtime_overrides( if internal_keys is None: from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys - internal_keys = internal_blacklist_keys() | SHARED_FIELDS + # Some fields are modeled as orchestrator-owned for top-level CLI + # parsing, but are also legitimate deploy-time stage overrides. Keep + # the default blacklist for true orchestrator/shared fields while + # allowing any field explicitly represented by the deploy schema to + # continue flowing into per-stage overrides. + internal_keys = (internal_blacklist_keys() | SHARED_FIELDS) - deploy_runtime_override_keys() result: dict[str, Any] = {} @@ -303,7 +308,7 @@ class StageDeployConfig: the top level of ``DeployConfig`` and propagated to every stage. """ - # === Omni fields === + # === Omni stage wrapper fields === # Stage identity and Omni runtime placement. stage_id: int devices: str | None = None @@ -316,22 +321,33 @@ class StageDeployConfig: default_sampling_params: dict[str, Any] | None = None subtalker_sampling_params: dict[str, Any] | None = None - # === vLLM EngineArgs fields === - # Parallelism and scheduler/memory capacity. + # === Generic stage engine fields === + # Parallelism, scheduler, and memory-capacity controls. tensor_parallel_size: int | None = None + enable_expert_parallel: bool | None = None gpu_memory_utilization: float | None = None max_num_seqs: int | None = None max_num_batched_tokens: int | None = None max_model_len: int | None = None - # Execution, scheduling, and KV/cache behavior. + # Generic execution, scheduling, and KV/cache behavior. enforce_eager: bool | None = None async_scheduling: bool | None = None disable_hybrid_kv_cache_manager: bool | None = None mm_processor_cache_gb: float | None = None - # Diffusion parallel_config deploy override fields. - enable_expert_parallel: bool | None = None + # Generic compilation, profiling, tokenizer/config parsing, and model + # loading controls. + compilation_config: dict[str, Any] | None = None + profiler_config: dict[str, Any] | None = None + skip_mm_profiling: bool | None = None + enable_flashinfer_autotune: bool | None = None + config_format: str | None = None + load_format: str | None = None + tokenizer_mode: str | None = None + + # === Diffusion stage runtime fields === + # Diffusion parallel_config deploy/runtime override fields. ulysses_degree: int | None = None ulysses_mode: str | None = None ring_degree: int | None = None @@ -343,16 +359,44 @@ class StageDeployConfig: hsdp_shard_size: int | None = None hsdp_replicate_size: int | None = None - # Compilation, profiling, tokenizer/config parsing, and model loading. - compilation_config: dict[str, Any] | None = None - profiler_config: dict[str, Any] | None = None - skip_mm_profiling: bool | None = None - enable_flashinfer_autotune: bool | None = None - config_format: str | None = None - load_format: str | None = None - tokenizer_mode: str | None = None + # Diffusion model loading and adapter construction. + model_class_name: str | None = None + diffusion_load_format: str | None = None + diffusers_load_kwargs: dict[str, Any] | None = None + diffusers_call_kwargs: dict[str, Any] | None = None + diffusion_quantization_config: str | None = None + diffusion_attention_backend: str | None = None + diffusion_attention_config: dict[str, Any] | None = None - # Pass-through vLLM EngineArgs fields that are not represented above. + # Diffusion execution, cache, and VAE behavior. + cache_backend: str | None = None + cache_config: dict[str, Any] | None = None + enable_cache_dit_summary: bool | None = None + step_execution: bool | None = None + vae_use_slicing: bool | None = None + vae_use_tiling: bool | None = None + boundary_ratio: float | None = None + flow_shift: float | None = None + diffusion_kv_cache_dtype: str | None = None + diffusion_kv_cache_skip_steps: str | None = None + diffusion_kv_cache_skip_layers: str | None = None + auxiliary_text_encoder: str | None = None + + # Runtime optimizations used by diffusion loading/execution. + enable_multithread_weight_load: bool | None = None + num_weight_load_threads: int | None = None + enable_cpu_offload: bool | None = None + enable_layerwise_offload: bool | None = None + + # Diffusion-specific debug and observability knobs. + enable_diffusion_pipeline_profiler: bool | None = None + + # Modality/service constraints consumed outside the core engine config. + max_generated_image_size: int | None = None + tts_max_instructions_length: int | None = None + + # === Pass-through stage engine fields === + # Pass-through stage engine args that are not represented above. engine_extras: dict[str, Any] = field(default_factory=dict) @@ -409,6 +453,17 @@ _STAGE_RESERVED_KEYS = frozenset( _STAGE_DEPLOY_FIELDS = {f.name: f for f in fields(StageDeployConfig) if f.name not in _STAGE_RESERVED_KEYS} +def deploy_runtime_override_keys() -> frozenset[str]: + """Return deploy-schema fields that are valid CLI/runtime overrides. + + These keys form the positive contract for stage override propagation: + stage-scoped deploy knobs plus top-level pipeline-wide engine settings. + They must remain overridable even if they are also modeled on + ``OrchestratorArgs`` for top-level CLI parsing. + """ + return frozenset(_STAGE_DEPLOY_FIELDS) | frozenset(_PIPELINE_WIDE_ENGINE_FIELDS) + + def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: """Parse a single stage entry from deploy YAML into StageDeployConfig.""" # Get the non-reserved keys for this stage