Skip to content

[Config] Remove invalid LLM-only engine_args from diffusion stage configs #2622

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hsliuustc0106 merged 1 commit into from
Apr 10, 2026
Merged

[Config] Remove invalid LLM-only engine_args from diffusion stage configs #2622

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 0 additions & 6 deletions tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,15 +47,9 @@ stage_args:
engine_args:
model_stage: dit
max_num_seqs: 1
gpu_memory_utilization: 0.45
enforce_eager: true
trust_remote_code: true
engine_output_type: image
distributed_executor_backend: mp
enable_prefix_caching: false
max_num_batched_tokens: 32768
tensor_parallel_size: 1
Copy link
Copy Markdown
Collaborator

@princepride princepride Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why we remove tensor_parallel_size? Seems someone just edit tensor_parallel_size place, right?

Copy link
Copy Markdown
Contributor Author

@ianliuy ianliuy Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good observation! Yes tensor_parallel_size was effectively "moved" to parallel_config.tensor_parallel_size when DiffusionParallelConfig was introduced in PR #189. Since then, OmniDiffusionConfig no longer has a top-level tensor_parallel_size field, so from_kwargs() silently drops it (data.py L692-694). The value here is also 1, which matches the DiffusionParallelConfig default. Issue #2635 also tracks this inconsistency.

load_format: dummy
omni_kv_config:
need_recv_cache: true
engine_input_source: [0]
Expand Down
6 changes: 0 additions & 6 deletions tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,15 +46,9 @@ stage_args:
engine_args:
model_stage: dit
max_num_seqs: 1
gpu_memory_utilization: 0.45
enforce_eager: true
trust_remote_code: true
engine_output_type: image
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
tensor_parallel_size: 1
Copy link
Copy Markdown
Collaborator

@princepride princepride Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here

Copy link
Copy Markdown
Contributor Author

@ianliuy ianliuy Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same reasoning as above.

load_format: dummy
omni_kv_config:
need_recv_cache: true
engine_input_source: [0]
Expand Down
68 changes: 68 additions & 0 deletions tests/test_diffusion_config_fields.py
View file
Open in desktop
Copy link
Copy Markdown
Collaborator

@princepride princepride Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If future users add new special fields, how should we maintain this UT, considering these new fields might only apply to a specific model?

Copy link
Copy Markdown
Contributor Author

@ianliuy ianliuy Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good question. Two cases:

  • If the new field is added to the OmniDiffusionConfig dataclass, the test picks it up automatically via fields(OmniDiffusionConfig) no changes needed.
  • If it's consumed outside the dataclass (like model_stage by the stage init layer, or quantization via from_kwargs() backwards-compat), add it to the allowlist in this test.

I'll add a comment in the test to document this but before I do, I'd love to hear your thoughts on whether this approach works for you.

Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Ensure diffusion stage YAML configs only use valid OmniDiffusionConfig fields.

Regression test for https://github.com/vllm-project/vllm-omni/issues/2563
"""

from dataclasses import fields
from pathlib import Path

import pytest
import yaml

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

try:
from vllm_omni.diffusion.data import OmniDiffusionConfig
except Exception:
OmniDiffusionConfig = None


@pytest.mark.skipif(
OmniDiffusionConfig is None,
reason="OmniDiffusionConfig could not be imported (missing torch?)",
)
def test_diffusion_stage_configs_only_contain_valid_fields():
"""Diffusion stage engine_args must only contain OmniDiffusionConfig fields.

Regression test for https://github.com/vllm-project/vllm-omni/issues/2563
"""
# Scan both main configs and test configs
repo_root = Path(__file__).parent.parent
config_dirs = [
repo_root / "vllm_omni" / "model_executor" / "stage_configs",
]
# Also scan test directories recursively
test_dir = repo_root / "tests"

yaml_paths: list[Path] = []
for config_dir in config_dirs:
yaml_paths.extend(sorted(config_dir.glob("*.yaml")))
yaml_paths.extend(sorted(test_dir.rglob("*.yaml")))

valid_fields = {f.name for f in fields(OmniDiffusionConfig)}
# model_stage is consumed by the stage init layer, not OmniDiffusionConfig
valid_fields.add("model_stage")
# model_arch is consumed by the stage init layer for diffusion model class resolution
valid_fields.add("model_arch")
# "quantization" is mapped to "quantization_config" by from_kwargs() backwards-compat
valid_fields.add("quantization")
Copy link
Copy Markdown
Collaborator

@lishunyang12 lishunyang12 Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: worth a short block comment above the allowlist explaining the maintenance policy — entries here are fields consumed outside OmniDiffusionConfig (e.g. by extract_stage_metadata or the backwards-compat path in from_kwargs). Saves the next maintainer a git blame.


invalid_entries: list[tuple[str, set[str]]] = []
for yaml_path in yaml_paths:
with open(yaml_path) as fh:
config = yaml.safe_load(fh)

stages = config.get("stage_args", config.get("stages", []))
for stage in stages:
if stage.get("stage_type") != "diffusion":
continue
engine_args = stage.get("engine_args", {})
invalid = set(engine_args.keys()) - valid_fields
if invalid:
invalid_entries.append((yaml_path.relative_to(repo_root), invalid))

assert not invalid_entries, "Diffusion stage configs contain fields not in OmniDiffusionConfig:\n" + "\n".join(
f" {name}: {sorted(bad)}" for name, bad in invalid_entries
)
5 changes: 0 additions & 5 deletions vllm_omni/model_executor/stage_configs/bagel.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,14 +52,9 @@ stage_args:
engine_args:
model_stage: dit
max_num_seqs: 1
gpu_memory_utilization: 0.45
enforce_eager: true
trust_remote_code: true
engine_output_type: image
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
tensor_parallel_size: 1
omni_kv_config:
need_recv_cache: true
engine_input_source: [0]
Expand Down
5 changes: 0 additions & 5 deletions vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,14 +45,9 @@ stage_args:
engine_args:
model_stage: dit
max_num_seqs: 1
gpu_memory_utilization: 0.45
enforce_eager: true
trust_remote_code: true
engine_output_type: image
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
tensor_parallel_size: 1
omni_kv_config:
need_recv_cache: true
engine_input_source: [0]
Expand Down
5 changes: 0 additions & 5 deletions vllm_omni/model_executor/stage_configs/bagel_single_stage.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,14 +9,9 @@ stage_args:
engine_args:
model_stage: dit
max_num_seqs: 1
gpu_memory_utilization: 0.45
enforce_eager: true
trust_remote_code: true
engine_output_type: image
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
tensor_parallel_size: 1

final_output: true
final_output_type: image
Expand Down
5 changes: 0 additions & 5 deletions vllm_omni/model_executor/stage_configs/bagel_think.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,14 +49,9 @@ stage_args:
engine_args:
model_stage: dit
max_num_seqs: 1
gpu_memory_utilization: 0.45
enforce_eager: true
trust_remote_code: true
engine_output_type: image
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
tensor_parallel_size: 1
omni_kv_config:
need_recv_cache: true
engine_input_source: [0]
Expand Down
5 changes: 0 additions & 5 deletions vllm_omni/model_executor/stage_configs/bagel_usp2.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,14 +45,9 @@ stage_args:
max_batch_size: 1
engine_args:
model_stage: dit
gpu_memory_utilization: 0.45
enforce_eager: true
trust_remote_code: true
engine_output_type: image
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
tensor_parallel_size: 1
parallel_config:
ulysses_degree: 2
# ring_degree: 2
Expand Down
4 changes: 0 additions & 4 deletions vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,9 @@ stage_args:
engine_args:
max_num_seqs: 1
model_stage: dit
gpu_memory_utilization: 0.65
enforce_eager: true
trust_remote_code: true
engine_output_type: image
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
parallel_config:
tensor_parallel_size: 4
enable_expert_parallel: true
Expand Down
4 changes: 0 additions & 4 deletions vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit_2gpu_fp8.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,9 @@ stage_args:
max_batch_size: 1
engine_args:
model_stage: dit
gpu_memory_utilization: 0.9
enforce_eager: true
trust_remote_code: true
engine_output_type: image
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
quantization: "fp8"
parallel_config:
tensor_parallel_size: 2
Expand Down
4 changes: 0 additions & 4 deletions vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,12 +50,8 @@ stage_args:
max_batch_size: 1
engine_args:
model_stage: diffusion
gpu_memory_utilization: 0.9
enforce_eager: true
engine_output_type: image
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
vae_use_slicing: false
vae_use_tiling: false
cache_backend: null
Expand Down
2 changes: 0 additions & 2 deletions vllm_omni/model_executor/stage_configs/omnivoice.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,8 @@ stage_args:
engine_args:
model_stage: dit
model_class_name: "OmniVoicePipeline"
gpu_memory_utilization: 0.5
enforce_eager: true
trust_remote_code: true
engine_output_type: audio
Copy link
Copy Markdown
Collaborator

@princepride princepride Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should maintain it, correct me if I am wrong

Copy link
Copy Markdown
Contributor Author

@ianliuy ianliuy Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These are safe to remove this is a stage_type: diffusion stage, and OmniDiffusionConfig has neither field. from_kwargs() silently drops them (data.py L692-694). For engine_output_type specifically, extract_stage_metadata() also hardcodes it to None for all diffusion stages (stage_init_utils.py L171). Audio routing is handled by the stage-level final_output_type: audio (which is preserved) and the SupportAudioOutput interface on OmniVoicePipeline.

distributed_executor_backend: "mp"
dtype: "float32"
final_output: true
Expand Down
Loading