Repository navigation
[Config] Add structured VllmOmniConfig classes - #4425
Conversation
Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Signed-off-by: zwhzzz0821 <2831474076@qq.com>
Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Signed-off-by: Acerak01-fy <wfy2003324@163.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
Posted a few inline compatibility comments for the new structured config path.
| } | ||
| ) | ||
|
|
||
| _PIPELINE_DEPLOY_CLI_FIELDS = ( |
There was a problem hiding this comment.
This duplicates the legacy pipeline-wide propagation list, but it is already missing active fields such as active_stream_window and custom_voice_dir from stage_config._PIPELINE_WIDE_ENGINE_FIELDS. Because from_registry() uses this smaller list as part of its compatibility baseline, existing pipeline-wide knobs can disappear from the structured config and later cutover would regress features like custom voice loading. Can we reuse/export the legacy list or add parity coverage for all pipeline-wide engine fields?
| for topology in pipeline.stages | ||
| ) | ||
|
|
||
| worker_backend = ( |
There was a problem hiding this comment.
worker_backend and distributed_executor_backend are different settings. The former is the orchestrator backend (multi_process/ray), while the latter is an engine/diffusion executor backend (mp, ray, etc.). Several deploy YAMLs set distributed_executor_backend: mp, so this can produce orchestrator_config.worker_backend == "mp", which is not accepted by the orchestrator CLI path. This should read cli_overrides.get("worker_backend") and otherwise keep the orchestrator default.
| or deploy_for_registry.distributed_executor_backend | ||
| or "multi_process" | ||
| ) | ||
| orchestrator_config = OrchestratorConfig( |
There was a problem hiding this comment.
Most orchestrator-scoped CLI overrides are not projected into OrchestratorConfig here. Fields declared above, such as stage_init_timeout, init_timeout, ray_address, omni_lb_policy, omni_heartbeat_timeout, shm_threshold_bytes, and batch_timeout, will silently stay at defaults even when present in cli_overrides. Please map these explicitly or derive them from the orchestrator config fields so this structured layer preserves CLI behavior.
Signed-off-by: zwhzzz0821 <2831474076@qq.com>
|
Thanks for the review! I’ve fixed the issues you mentioned. @hsliuustc0106 |
alex-jw-brooks
left a comment
alex-jw-brooks
left a comment
There was a problem hiding this comment.
Thank you for working on this @Acerak01-fy @zwhzzz0821! Some thoughts.
This will also definitely overlap with @kushanam's RFC and PR for Composable Parallelism. I will try to take another look at that RFC and review that PR in the next day or so so that we can coordinate 🙏
|
|
||
| _ASYNC_AR_SCHEDULER = "vllm_omni.core.sched.omni_ar_scheduler.OmniARAsyncScheduler" | ||
| _SYNC_AR_SCHEDULER = "vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler" | ||
| _GENERATION_SCHEDULER = "vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler" |
There was a problem hiding this comment.
We should avoid referring to modules / paths as strings like this since it can make it difficult to catch bugs if things are refactored in the future.
I think it would be better to do something like what we have in _resolve_scheduler which maps to a str enum value that can be mapped to the corresponding class to avoid potential lazy import errors
| "tool_call_parser", | ||
| "tool_parser_plugin", | ||
| "uvicorn_log_level", | ||
| } |
There was a problem hiding this comment.
List like this are really difficult to maintain, it would be best to avoid adding another one.
Where did this list come from? E.g., for some of them like enable_tool_template, I don't see references to them anywhere in vLLM or vLLM Omni, unless I'm missing something?
There was a problem hiding this comment.
The original list was trying to prevent server/orchestrator-level CLI arguments from being treated as per-stage engine config.
cli_overrides is a mixed bag: it can contain stage/runtime overrides, but it can also contain server-level arguments from the API/CLI layer, such as host, port, api_key, SSL options, chat template options, etc. Those fields are not stage config. If we pass them through as stage overrides, they can accidentally end up inside a stage config, especially a diffusion stage where fields like host and
port also exist for the diffusion worker/runtime. So the old _ORCHESTRATOR_ONLY_CLI_KEYS blacklist was mainly there to filter out these non-stage CLI/server fields.
Updated. I removed the hard-coded _ORCHESTRATOR_ONLY_CLI_KEYS list. _stage_cli_overrides() now reuses build_stage_runtime_overrides() for the existing stage/global CLI override routing, then filters
global stage overrides through fields derived from OmniEngineArgs and deploy typed fields instead of maintaining another blacklist.
I also added a regression test to make sure server-level keys such as host, port, and api_key do not leak into diffusion stage config, while explicit per-stage overrides like stage_0_host /
stage_0_port still work.
|
|
||
|
|
||
| def _resolve_execution_mode(execution_type: StageExecutionType) -> tuple[StageType, str | None]: | ||
| return _EXECUTION_TYPE_TO_STAGE_WORKER.get(execution_type, (StageType.LLM, None)) |
There was a problem hiding this comment.
I think we should not have a default fallback value here and raise an error if the execution type isn't in the _EXECUTION_TYPE_TO_STAGE_WORKER map, since that shouldn't be an expected behavior 🙂
| next_stage_proc = topology.async_chunk_process_next_stage_input_func | ||
| elif not async_chunk and topology.sync_process_input_func: | ||
| input_proc = topology.sync_process_input_func | ||
| return input_proc, next_stage_proc |
There was a problem hiding this comment.
For utils like this, where it's copied but functionally unmodified, can you import the original instead of copying it? Otherwise we may have regressions if the behavior is modified prior to migrating to unified config format
| return merged | ||
|
|
||
|
|
||
| _PLATFORM_DEEP_MERGE_KEYS = frozenset( |
There was a problem hiding this comment.
Same for things like this that are copied over and renamed, as it makes it harder to see what the actual diff in the code actually is
| "tokenizer_mode", | ||
| ) | ||
|
|
||
| _DIFFUSION_STAGE_ENGINE_FIELDS = _DIFFUSION_CONFIG_FIELDS - {"model", "stage_id"} |
There was a problem hiding this comment.
Similar comment on the lists here - keeping lists of field names like this gets messy quickly since it can get out of sync with the underlying objects
| deploy: DeployConfig, | ||
| engine: Mapping[str, Any], | ||
| ) -> CacheConfig: | ||
| return CacheConfig( |
There was a problem hiding this comment.
Can you add a comment that this is currently not the diffusion cache config? This is a bit of a confusing point because the DiffusionConfig has its own cache_config attr for the diffusion cache type (i.e., teacache / cache dit etc) but they are separate 😅
I can help clean up the typing for this part in a separate PR so it will be more clear here later on too
| parallel_config.get("data_parallel_size"), | ||
| engine.get("data_parallel_size"), | ||
| deploy.data_parallel_size, | ||
| default=1, |
There was a problem hiding this comment.
I think this initialization pattern needs to be revisited. I understand the intention, but passing things like default=1 when the corresponding config already has a default defined in the dataclass is an antipattern, and it also makes it so that if the priority changed, we would need to go through each of the things using _first_defined and change the calls one at a time
| engine: Mapping[str, Any], | ||
| ) -> ModelConfig: | ||
| return ModelConfig( | ||
| enable_sleep_mode=_first_defined(engine.get("enable_sleep_mode"), default=False), |
There was a problem hiding this comment.
Similar to the other comment, this initialization pattern looks pretty opaque to me, especially when there are default fallback values in .gets. Part of the reason the current config stack has been hard to work with is that defaults are defined in different places in the code, and there's not a very clear place when they get merged together, I think this is the main reason the unification is needed to begin with 🙂
| def _build_model_config( | ||
| topology: StagePipelineConfig, | ||
| stage_deploy: StageDeployConfig | None, | ||
| engine: Mapping[str, Any], |
There was a problem hiding this comment.
Can you make the typing of the engine arg being passed here stricter? Mapping is pretty nondescriptive, it would be easy to misconfigure it by mistake
Signed-off-by: Acerak01-fy <wfy2003324@163.com>
c2e8f27 to
4062f9d
Compare
|
Thanks for the review! I’ve fixed the issues you mentioned. @alex-jw-brooks |
| __all__ = [ | ||
| "OmniModelConfig", | ||
| "LoRAConfig", | ||
| "VllmOmniConfig", |
There was a problem hiding this comment.
I think we should rearranage the order of the all
There was a problem hiding this comment.
I reordered all into logical export groups and added short comments to make the grouping explicit. Please check if this looks reasonable.
| master_port: int | None = None | ||
| host: str | None = None | ||
| port: int | None = None | ||
| scheduler_port: int = 5555 |
There was a problem hiding this comment.
what is scheduler port?
There was a problem hiding this comment.
I checked the runtime path and couldn’t find any current consumer for scheduler_port; diffusion scheduling is in-process today. I removed it from the structured DiffusionConfig surface instead of inventing new semantics for it.
| """Top-level structured Omni config built once from registry inputs.""" | ||
|
|
||
| pipeline_config: PipelineConfig | ||
| stage_configs: tuple[VllmOmniStageConfig, ...] |
There was a problem hiding this comment.
does diffusionconfig belong one type of the stage_configs?
There was a problem hiding this comment.
DiffusionConfig is not a separate stage config type; it is an optional subconfig on VllmOmniStageConfig, populated only for StageExecutionType.DIFFUSION. Added a field comment to make that explicit.
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
Added a few concise parity comments for the structured config layer.
| if self.ulysses_mode not in {"strict", "advanced_uaa"}: | ||
| raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'") | ||
|
|
||
| base_world_size = ( |
There was a problem hiding this comment.
Can we align this with DiffusionParallelConfig before treating it as the parity target? The current diffusion config does not include vae_patch_parallel_size in the world-size product, supports hsdp_shard_size=-1 auto-sharding, and rejects HSDP with TP/DP. This version would change device sizing/validation on cutover.
| quantization_config = _build_quantization_config(deploy, engine.quantization) | ||
| parallel_config = _build_parallel_config(deploy, engine.parallel) | ||
|
|
||
| stage_config = VllmOmniStageConfig( |
There was a problem hiding this comment.
active_stream_window and custom_voice_dir are copied onto deploy_for_registry, but I don't see them retained in VllmOmniStageConfig or any subconfig here. Legacy _build_engine_args propagates them per stage, and downstream consumes custom_voice_dir; could we add parity coverage or a clear home for these fields?
| stage_deploy: StageDeployConfig | None, | ||
| engine: _ModelEngineOverrides, | ||
| ) -> ModelConfig: | ||
| return ModelConfig( |
There was a problem hiding this comment.
Legacy build_engine_args_dict() derives has_sampling_extra_args from default_sampling_params.extra_args. This structured path leaves ModelConfig.has_sampling_extra_args at the dataclass default, so deploys with sampling extra_args won't be represented accurately.
| if quantization_config is not None: | ||
| diffusion_kwargs["quantization_config"] = _copy_value(quantization_config) | ||
|
|
||
| return DiffusionConfig(**{k: v for k, v in diffusion_kwargs.items() if v is not None}) |
There was a problem hiding this comment.
Building DiffusionConfig directly bypasses OmniDiffusionConfig.from_kwargs() normalization. That drops existing aliases/backcompat behavior such as diffusion_attention_backend, legacy kv_cache_*, static_lora_scale, and the cache-backend env fallback. Could we reuse that normalization or add equivalent parity tests?
|
@hsliuustc0106 @alex-jw-brooks we have addressed the review comments and updated the PR accordingly. Could you please take another look when you have a chance? |
|
@hsliuustc0106 @alex-jw-brooks I have already revised the code based on the latest feedback. Please take another look. |
Signed-off-by: Wang fuyin <wfy2003324@163.com>
|
|
||
| from vllm_omni.config.stage_config import ( | ||
| _DEPLOY_DIR, | ||
| _PIPELINE_REGISTRY, |
There was a problem hiding this comment.
This breaks current CI: _PIPELINE_REGISTRY is no longer exported from stage_config on main. Please use pipeline_registry.OMNI_PIPELINES here and resolve callable entries the same way StageConfigFactory.resolve_pipeline_config() does.
| VllmOmniDiffusionStageConfig, | ||
| VllmOmniGenerationStageConfig, | ||
| ) | ||
| from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES |
There was a problem hiding this comment.
This should use OMNI_PIPELINES; _OMNI_PIPELINES is not exported anymore. Please also replace the _PIPELINE_REGISTRY references below with the new registry API so these tests match current main.
| "DeployConfig", | ||
| "PipelineConfig", | ||
| "StageConfig", | ||
| "StageConfigFactory", |
There was a problem hiding this comment.
StageConfigFactory is listed in __all__ but is never imported into this module, so from vllm_omni.config import StageConfigFactory will fail. Please import it from config_factory or remove this export.
Signed-off-by: Acerak01-fy <wfy2003324@163.com>
hsliuustc0106
added
ready
There was a problem hiding this comment.
@Acerak01-fy Generally LGTM, left some comments
| return copy.deepcopy(value) | ||
|
|
||
|
|
||
| def _copy_if_not_none(value: Any) -> Any: |
| - _DIFFUSION_RUNTIME_CONFIG_FIELDS | ||
| - _DIFFUSION_PARALLEL_CONFIG_FIELDS | ||
| ) | ||
| _DIFFUSION_MOVED_SHARED_FIELDS = frozenset( |
There was a problem hiding this comment.
Is this used anywhere?
| "PipelineConfig", | ||
| "StageDeployConfig", | ||
| "DeployConfig", | ||
| "StageType", |
There was a problem hiding this comment.
nit: StageType appears twice
Signed-off-by: Acerak01-fy <wfy2003324@163.com>
|
@yuanheng-zhao Thank you for your carefulness. I have already corrected these minor issues. |
Gaohan123
added
ready
Signed-off-by: Acerak01-fy <wfy2003324@163.com>
|
LGTM. Since this PR is additive and does not switch runtime consumers yet, I do not think it conflicts with the HunyuanImage3 NPU EP/SP fix . #4344 |
yes! |
hsliuustc0106
added
ready
…election
Push all AR-Diffusion/DreamZero specifics out of the public diffusion base into
the deploy config + experimental package, so the public diff is a tiny generic seam:
- diffusion_engine.py: resolve_engine_class/make_engine are now a generic factory
("default" / DiffusionEngine subclass / import-path string). Removed the hardcoded
"ar_diffusion" branch and its experimental import — no coupling to experimental/.
- data.py: removed default_engine_backend_for_model, resolve_default_engine_backend,
_AR_DIFFUSION_ENGINE_MODEL_CLASSES, the two engine_backend-rewrite sites in
enrich_config + from_kwargs, the ar_diffusion_kv_config field, and the structured-
config mirror note. Public engine-routing addition is now just the generic
engine_backend / diffusion_model_runner_cls fields. from_kwargs returns to its
clean cls(**filtered_kwargs) tail (no vllm-project#4425 conflict).
- dreamzero.yaml: select the engine via engine_backend=<ARDiffusionEngine qualname>.
- runner: KV overrides read from model_config["ar_diffusion_kv_config"] (defaults work).
- tests: generic dispatch tests (qualname -> ARDiffusionEngine); dropped per-model
default tests. 51 passed; E2E DreamZero generates via the yaml with no env.
Trade-off: a bare load without the deploy yaml no longer auto-selects the engine;
DreamZero is selected via dreamzero.yaml (its single supported path).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…commit_chunk docs) Close inline review findings 2-5 from hsliuustc0106: - Inline 2: delete the two unused project_kv methods in causal_wan_model.py (_kv_populate_cross inlines the identical projection); fix the stale comment. - Inline 3: gate DreamZero->BDE engine/runner routing behind BDE_KV_ENABLE via resolve_default_engine_backend(), so non-BDE deploys stay on base DiffusionEngine and base is selectable again. +4 unit tests. - Inline 4: docstrings clarifying the two commit_chunk paths (test-only vs the real advance inside BDEKVState.update_kv_cache). - Inline 5: vllm-project#4425 structured-config coordination is blocked (not in this tree); deferral marker left at the field defs + a note in the remaining-work doc. Validated on H200 (DreamZero-DROID): tests/bde 77 passing; routing gate confirmed end-to-end (KV-off -> base engine, BDE_KV_ENABLE=1 -> BDE); bit-exact parity (max_abs_diff=0, PSNR=inf) preserved. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
…election
Push all AR-Diffusion/DreamZero specifics out of the public diffusion base into
the deploy config + experimental package, so the public diff is a tiny generic seam:
- diffusion_engine.py: resolve_engine_class/make_engine are now a generic factory
("default" / DiffusionEngine subclass / import-path string). Removed the hardcoded
"ar_diffusion" branch and its experimental import — no coupling to experimental/.
- data.py: removed default_engine_backend_for_model, resolve_default_engine_backend,
_AR_DIFFUSION_ENGINE_MODEL_CLASSES, the two engine_backend-rewrite sites in
enrich_config + from_kwargs, the ar_diffusion_kv_config field, and the structured-
config mirror note. Public engine-routing addition is now just the generic
engine_backend / diffusion_model_runner_cls fields. from_kwargs returns to its
clean cls(**filtered_kwargs) tail (no vllm-project#4425 conflict).
- dreamzero.yaml: select the engine via engine_backend=<ARDiffusionEngine qualname>.
- runner: KV overrides read from model_config["ar_diffusion_kv_config"] (defaults work).
- tests: generic dispatch tests (qualname -> ARDiffusionEngine); dropped per-model
default tests. 51 passed; E2E DreamZero generates via the yaml with no env.
Trade-off: a bare load without the deploy yaml no longer auto-selects the engine;
DreamZero is selected via dreamzero.yaml (its single supported path).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Signed-off-by: Acerak01-fy <wfy2003324@163.com> Signed-off-by: zwhzzz0821 <2831474076@qq.com> Signed-off-by: Wang fuyin <wfy2003324@163.com> Co-authored-by: zwhzzz0821 <2831474076@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
7 participants
RFC: #4021
Co-authored-by: zwhzzz0821 2831474076@qq.com
Motivation
This PR implements RFC #4021 Phase 2 / PR 4 by adding the new structured
VllmOmniConfiglayer in an additive way.The goal of this phase is to introduce the new config object model without switching existing runtime consumers yet. The new layer is built from the current pipeline registry and deploy YAML path, then tested against the legacy
merge_pipeline_deploy()behavior so future cutover PRs have a stable compatibility baseline.Design
The new top-level config separates pipeline topology, per-stage settings, and orchestrator-only settings:
pipeline_configstage_configsorchestrator_configStage configs are represented as concrete stage realizations instead of one monolithic stage config type:
StageExecutionType.LLM_ARVllmOmniARStageConfigStageExecutionType.LLM_GENERATIONVllmOmniGenerationStageConfigStageExecutionType.DIFFUSIONVllmOmniDiffusionStageConfigStageConfigTypeis the public union type for these realizations.VllmOmniConfig.from_registry()routes each stage through_STAGE_CONFIG_BUILDERSusingStagePipelineConfig.execution_type.Summary of Changes
vllm_omni/config/omni_config.pyVllmOmniConfigwithpipeline_config,stage_configs, andorchestrator_configBaseVllmOmniStageConfigplus concrete AR, generation, and diffusion stage configsStageConfigTypefor the concrete stage config unionOmniStageParallelConfigfrom diffusion-specificOmniStageDiffusionParallelConfig_DiffusionConfigProjectionattached only to diffusion stagesVllmOmniConfig.from_registry()vllm_omni.configtests/config/test_omni_config.pyRFC Mapping
Phase 2 / PR 4 requires:
VllmOmniConfig.from_registry().merge_pipeline_deploy()for registered models.This PR follows that scope. Existing consumers are not changed to use
VllmOmniConfigyet.Structured Config Classes
Added the following public structured config classes in
vllm_omni/config/omni_config.py:OmniStageModelConfigOmniStageLoadConfigOmniStageCacheConfigOmniStageSchedulerConfigOmniStageConnectorConfigOmniStageRuntimeConfigOmniStageParallelConfigOmniStageDiffusionParallelConfigVllmOmniOrchestratorConfigBaseVllmOmniStageConfigVllmOmniARStageConfigVllmOmniGenerationStageConfigVllmOmniDiffusionStageConfigStageConfigTypeVllmOmniConfigDiffusion-only config fields are kept in the internal
_DiffusionConfigProjection. This avoids exporting a publicDiffusionConfigname while still preserving the current diffusion config surface forVllmOmniDiffusionStageConfig.Stage Realization Details
Common stage fields live on
BaseVllmOmniStageConfig:stage_pipeline_configmodel_configload_configcache_configscheduler_configconnector_configruntime_configparallel_configquantization_configVllmOmniDiffusionStageConfigadds:diffusion_configparallel_config: OmniStageDiffusionParallelConfigAR and generation stages use the common
OmniStageParallelConfig. Diffusion stages useOmniStageDiffusionParallelConfig.The selection is driven by pipeline topology:
Parallel Config Details
OmniStageParallelConfigcontains only common parallel settings:pipeline_parallel_sizedata_parallel_sizetensor_parallel_sizeenable_expert_parallelworld_sizeOmniStageDiffusionParallelConfigextends the common config with diffusion-specific parallel settings:sequence_parallel_sizeulysses_degreering_degreeulysses_modecfg_parallel_sizevae_patch_parallel_sizeuse_hsdpmask_sp_paddinghsdp_shard_sizehsdp_replicate_sizeQuantization Config Details
quantization_configis typed through the internal_QuantizationConfigTypealias:This matches the existing runtime behavior while avoiding an unbounded
Anyinterface.Compatibility Details
from_registry()constructs the new structured config from the existing registry/deploy path:_PIPELINE_REGISTRYvllm_omni/deployThe implementation preserves the current
PipelineConfigobject directly, including fields added after the RFC draft such ashf_config_predicate.The structured config remains additive. Existing runtime consumers are not switched over in this PR.
Tests
Added coverage in
tests/config/test_omni_config.pyfor:VllmOmniConfig.from_registry()parity withmerge_pipeline_deploy()._OMNI_PIPELINES.stage_by_id()lookup.StageExecutionType.OmniStageParallelConfig.OmniStageDiffusionParallelConfig.world_sizeandsequence_parallel_sizebehavior.sequence_parallel_sizecompatibility.mask_sp_paddingpreservation.StageConfig.to_omegaconf().build_engine_args_dict().Validation
Ran:
Result:
Follow-ups
This PR intentionally does not switch existing consumers to
VllmOmniConfig.Consumer cutover, removal of legacy compatibility paths, and runtime adoption are left for later RFC #4021 phases.