[Config] Read engine args from VllmOmniConfig (#5678)
Signed-off-by: Acerak01-fy <wfy2003324@163.com> Signed-off-by: Wang fuyin <wfy2003324@163.com>
This commit is contained in:
Wang fuyin authored and GitHub committed
1 parent
6ec9033065
commit
3d7fc3b9ba
6 files changed
+1090
-72
No files matched your search
@@ -0,0 +1,510 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import sys
|
||||
import types
|
||||
from dataclasses import fields, replace
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from vllm_omni.config.omni_config import VllmOmniConfig
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config
|
||||
from vllm_omni.config.stage_config import (
|
||||
DeployConfig,
|
||||
PipelineConfig,
|
||||
StageDeployConfig,
|
||||
StageExecutionType,
|
||||
StagePipelineConfig,
|
||||
StageType,
|
||||
build_stage_runtime_overrides,
|
||||
load_deploy_config,
|
||||
merge_pipeline_deploy,
|
||||
)
|
||||
from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
|
||||
from vllm_omni.engine import stage_init_utils
|
||||
from vllm_omni.engine.arg_utils import OmniEngineArgs
|
||||
from vllm_omni.engine.stage_init_utils import (
|
||||
build_engine_args_dict,
|
||||
build_engine_args_dict_from_omni_stage_config,
|
||||
build_legacy_engine_args_dict,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_DEPLOY_DIR = Path(__file__).parents[2] / "vllm_omni" / "deploy"
|
||||
|
||||
|
||||
def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str, object]:
|
||||
backend_fields = fields(config_cls)
|
||||
backend_field_names = {backend_field.name for backend_field in backend_fields}
|
||||
backend_config = config_cls(
|
||||
**{name: copy.deepcopy(value) for name, value in engine_args.items() if name in backend_field_names}
|
||||
)
|
||||
effective_values: dict[str, object] = {}
|
||||
for backend_field in backend_fields:
|
||||
value = getattr(backend_config, backend_field.name)
|
||||
if isinstance(value, FieldInfo):
|
||||
value = value.get_default(call_default_factory=True)
|
||||
effective_values[backend_field.name] = value
|
||||
return effective_values
|
||||
|
||||
|
||||
_LLM_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniEngineArgs))
|
||||
_DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusionConfig))
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _stable_engine_arg_environment(monkeypatch, tmp_path):
|
||||
from vllm_omni import platforms
|
||||
|
||||
monkeypatch.delenv("VLLM_USE_FLASHINFER_MOE_FP16", raising=False)
|
||||
xcodec_model = tmp_path / "xcodec-model"
|
||||
xcodec_model.mkdir()
|
||||
monkeypatch.setenv("XCODEC1_PATH", str(xcodec_model))
|
||||
platform = platforms.current_omni_platform
|
||||
monkeypatch.setattr(platform, "device_name", "cpu", raising=False)
|
||||
monkeypatch.setattr(platform, "device_type", "cpu", raising=False)
|
||||
monkeypatch.setattr(platform, "is_rocm", lambda: False)
|
||||
monkeypatch.setattr(
|
||||
OmniDiffusionConfig,
|
||||
"_resolve_master_port",
|
||||
lambda config: config.master_port or 29500,
|
||||
)
|
||||
|
||||
def _resolve_test_worker(engine_args):
|
||||
worker_type = engine_args.get("worker_type")
|
||||
if worker_type is not None and engine_args.get("worker_cls") in (None, "auto"):
|
||||
engine_args["worker_cls"] = f"test.{worker_type}.Worker"
|
||||
|
||||
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", _resolve_test_worker)
|
||||
|
||||
|
||||
def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, str]:
|
||||
parent_model = tmp_path / "parent-model"
|
||||
stage_model = tmp_path / "stage-model"
|
||||
diffusion_model = tmp_path / "diffusion-model"
|
||||
for root in (parent_model, stage_model, diffusion_model):
|
||||
root.mkdir()
|
||||
(stage_model / "ar-model").mkdir()
|
||||
(stage_model / "ar-tokenizer").mkdir()
|
||||
|
||||
pipeline = PipelineConfig(
|
||||
model_type="typed-engine-args",
|
||||
model_arch="PipelineModel",
|
||||
stages=(
|
||||
StagePipelineConfig(
|
||||
stage_id=0,
|
||||
model_stage="thinker",
|
||||
execution_type=StageExecutionType.LLM_AR,
|
||||
requires_multimodal_data=True,
|
||||
hf_config_name="thinker_config",
|
||||
engine_output_type="hidden_states",
|
||||
model_arch="Qwen3OmniMoeForConditionalGeneration",
|
||||
model_subdir="ar-model",
|
||||
tokenizer_subdir="ar-tokenizer",
|
||||
retains_state_across_chunks=True,
|
||||
async_chunk_process_next_stage_input_func="operator.add",
|
||||
),
|
||||
StagePipelineConfig(
|
||||
stage_id=1,
|
||||
model_stage="talker",
|
||||
execution_type=StageExecutionType.LLM_GENERATION,
|
||||
input_sources=(0,),
|
||||
engine_output_type="audio_tokens",
|
||||
),
|
||||
StagePipelineConfig(
|
||||
stage_id=2,
|
||||
model_stage="dit",
|
||||
execution_type=StageExecutionType.DIFFUSION,
|
||||
input_sources=(1,),
|
||||
final_output=True,
|
||||
final_output_type="image",
|
||||
model_arch="TypedDiffusionPipeline",
|
||||
),
|
||||
),
|
||||
)
|
||||
deploy = DeployConfig(
|
||||
async_chunk=True,
|
||||
trust_remote_code=True,
|
||||
dtype="float16",
|
||||
distributed_executor_backend="mp",
|
||||
stages=[
|
||||
StageDeployConfig(
|
||||
stage_id=0,
|
||||
tensor_parallel_size=2,
|
||||
max_num_seqs=8,
|
||||
default_sampling_params={"extra_args": {"voice": "test"}},
|
||||
engine_extras={
|
||||
"model": str(stage_model),
|
||||
"attention_backend": "FLASHINFER",
|
||||
"hf_overrides": {"rope_scaling": {"factor": 2}},
|
||||
"limit_mm_per_prompt": {"audio": 1},
|
||||
"logits_processors": ["test.custom.LogitsProcessor"],
|
||||
"kv_cache_memory_bytes": 1024,
|
||||
"omni_kv_config": {"need_send_cache": True},
|
||||
"max_cudagraph_capture_size": 0,
|
||||
"worker_cls": "test.custom.Worker",
|
||||
},
|
||||
),
|
||||
StageDeployConfig(
|
||||
stage_id=1,
|
||||
max_num_seqs=4,
|
||||
),
|
||||
StageDeployConfig(
|
||||
stage_id=2,
|
||||
tensor_parallel_size=2,
|
||||
vae_parallel_mode="spatial_shard_height",
|
||||
diffusion_attention_config={
|
||||
"default": {"backend": "FLASH_ATTN"},
|
||||
"per_role": {"cross": {"backend": "TORCH_SDPA"}},
|
||||
},
|
||||
engine_extras={
|
||||
"model": str(diffusion_model),
|
||||
"engine_backend": "test.diffusion.Engine",
|
||||
"enable_session_state_manager": True,
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
return pipeline, deploy, str(parent_model)
|
||||
|
||||
|
||||
def _legacy_and_typed_stages(
|
||||
pipeline: PipelineConfig,
|
||||
deploy: DeployConfig,
|
||||
model: str,
|
||||
cli_overrides: dict[str, object] | None = None,
|
||||
):
|
||||
resolved_cli_overrides = {"model": model, **(cli_overrides or {})}
|
||||
legacy_stage_configs = merge_pipeline_deploy(
|
||||
pipeline,
|
||||
copy.deepcopy(deploy),
|
||||
)
|
||||
for stage in legacy_stage_configs:
|
||||
stage.runtime_overrides = build_stage_runtime_overrides(
|
||||
stage.stage_id,
|
||||
resolved_cli_overrides,
|
||||
)
|
||||
legacy_stages = [stage.to_omegaconf() for stage in legacy_stage_configs]
|
||||
omni_config = VllmOmniConfig.from_pipeline_config(
|
||||
pipeline,
|
||||
user_deploy_config=copy.deepcopy(deploy),
|
||||
cli_overrides=resolved_cli_overrides,
|
||||
)
|
||||
return legacy_stages, omni_config
|
||||
|
||||
|
||||
def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
|
||||
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
|
||||
legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
|
||||
connector_spec = {"name": "SharedMemoryConnector", "extra": {"mode": "test"}}
|
||||
cli_tokenizer = "/external/tokenizer"
|
||||
|
||||
typed_args_by_stage = {}
|
||||
for stage_id in (0, 1):
|
||||
legacy_args = build_legacy_engine_args_dict(
|
||||
legacy_stages[stage_id],
|
||||
model,
|
||||
stage_connector_spec=connector_spec,
|
||||
cli_tokenizer=cli_tokenizer,
|
||||
)
|
||||
typed_args = build_engine_args_dict_from_omni_stage_config(
|
||||
omni_config.stage_by_id(stage_id),
|
||||
model,
|
||||
stage_connector_spec=connector_spec,
|
||||
cli_tokenizer=cli_tokenizer,
|
||||
)
|
||||
typed_args_by_stage[stage_id] = typed_args
|
||||
|
||||
assert {name: typed_args[name] for name in legacy_args} == legacy_args
|
||||
|
||||
thinker_args = typed_args_by_stage[0]
|
||||
assert thinker_args["model"] == str(tmp_path / "stage-model" / "ar-model")
|
||||
assert thinker_args["tokenizer"] == str(tmp_path / "stage-model" / "ar-tokenizer")
|
||||
assert thinker_args["stage_id"] == 0
|
||||
assert thinker_args["stage_connector_spec"] == connector_spec
|
||||
assert thinker_args["has_sampling_extra_args"] is True
|
||||
assert thinker_args["omni_kv_config"] == {"need_send_cache": True}
|
||||
assert thinker_args["worker_cls"] == "test.custom.Worker"
|
||||
assert thinker_args["attention_backend"] == "FLASHINFER"
|
||||
assert thinker_args["trust_remote_code"] is True
|
||||
assert thinker_args["dtype"] == "float16"
|
||||
assert thinker_args["distributed_executor_backend"] == "mp"
|
||||
assert thinker_args["tensor_parallel_size"] == 2
|
||||
assert thinker_args["hf_overrides"] == {"rope_scaling": {"factor": 2}}
|
||||
assert thinker_args["limit_mm_per_prompt"] == {"audio": 1}
|
||||
assert thinker_args["logits_processors"] == ["test.custom.LogitsProcessor"]
|
||||
assert thinker_args["retains_state_across_chunks"] is True
|
||||
assert thinker_args["kv_cache_memory_bytes"] == 1024
|
||||
assert thinker_args["max_cudagraph_capture_size"] == 0
|
||||
|
||||
talker_args = typed_args_by_stage[1]
|
||||
assert talker_args["model"] == model
|
||||
assert talker_args["tokenizer"] == cli_tokenizer
|
||||
assert talker_args["worker_cls"] == "test.generation.Worker"
|
||||
assert talker_args["disable_hybrid_kv_cache_manager"] is True
|
||||
assert talker_args["enable_prefix_caching"] is False
|
||||
|
||||
|
||||
def test_typed_diffusion_engine_args_use_structured_diffusion_config(tmp_path):
|
||||
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
|
||||
legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
|
||||
|
||||
legacy_args = build_legacy_engine_args_dict(legacy_stages[2], model)
|
||||
typed_args = build_engine_args_dict_from_omni_stage_config(
|
||||
omni_config.stage_by_id(2),
|
||||
model,
|
||||
)
|
||||
|
||||
assert typed_args["stage_id"] == legacy_args["stage_id"] == 2
|
||||
assert typed_args["model"] == legacy_args["model"] == str(tmp_path / "diffusion-model")
|
||||
assert omni_config.stage_by_id(2).diffusion_config.model == str(tmp_path / "diffusion-model")
|
||||
assert typed_args["model_stage"] == legacy_args["model_stage"] == "dit"
|
||||
assert typed_args["model_arch"] == legacy_args["model_arch"] == "TypedDiffusionPipeline"
|
||||
assert typed_args["model_class_name"] == legacy_args["model_class_name"] == "TypedDiffusionPipeline"
|
||||
assert typed_args["engine_backend"] == legacy_args["engine_backend"] == "test.diffusion.Engine"
|
||||
assert typed_args["enable_session_state_manager"] is legacy_args["enable_session_state_manager"] is True
|
||||
assert typed_args["parallel_config"]["tensor_parallel_size"] == 2
|
||||
assert typed_args["parallel_config"]["vae_parallel_mode"] == "spatial_shard_height"
|
||||
assert isinstance(typed_args["diffusion_attention_config"], AttentionConfig)
|
||||
assert typed_args["diffusion_attention_config"].default.backend == "FLASH_ATTN"
|
||||
assert typed_args["diffusion_attention_config"].per_role["cross"].backend == "TORCH_SDPA"
|
||||
|
||||
|
||||
def test_typed_engine_args_preserve_explicit_backend_default_overrides(tmp_path):
|
||||
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
|
||||
deploy.enable_prefix_caching = False
|
||||
deploy.enable_chunked_prefill = False
|
||||
thinker_deploy = deploy.stages[0]
|
||||
thinker_deploy.gpu_memory_utilization = 0.75
|
||||
thinker_deploy.disable_hybrid_kv_cache_manager = False
|
||||
thinker_deploy.async_scheduling = False
|
||||
legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
|
||||
|
||||
legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model)
|
||||
typed_args = build_engine_args_dict_from_omni_stage_config(
|
||||
omni_config.stage_by_id(0),
|
||||
model,
|
||||
)
|
||||
expected = {
|
||||
"gpu_memory_utilization": 0.75,
|
||||
"enable_prefix_caching": False,
|
||||
"disable_hybrid_kv_cache_manager": False,
|
||||
"max_num_seqs": 8,
|
||||
"enable_chunked_prefill": False,
|
||||
"async_scheduling": False,
|
||||
}
|
||||
|
||||
assert {name: legacy_args[name] for name in expected} == expected
|
||||
assert {name: typed_args[name] for name in expected} == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stage_scoped", [False, True], ids=["global", "stage-scoped"])
|
||||
def test_typed_engine_args_preserve_inherited_model_and_load_cli_fields(tmp_path, stage_scoped):
|
||||
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
|
||||
expected = {
|
||||
"revision": "model-rev",
|
||||
"tokenizer_revision": "tokenizer-rev",
|
||||
"code_revision": "code-rev",
|
||||
"seed": 123,
|
||||
"download_dir": str(tmp_path / "downloads"),
|
||||
}
|
||||
cli_overrides = {f"stage_0_{name}" if stage_scoped else name: value for name, value in expected.items()}
|
||||
legacy_stages, omni_config = _legacy_and_typed_stages(
|
||||
pipeline,
|
||||
deploy,
|
||||
model,
|
||||
cli_overrides,
|
||||
)
|
||||
|
||||
legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model)
|
||||
typed_args = build_engine_args_dict_from_omni_stage_config(
|
||||
omni_config.stage_by_id(0),
|
||||
model,
|
||||
)
|
||||
|
||||
assert {name: legacy_args[name] for name in expected} == expected
|
||||
assert {name: typed_args[name] for name in expected} == expected
|
||||
|
||||
|
||||
def test_typed_engine_args_preserve_deploy_subdirectory_precedence(tmp_path):
|
||||
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
|
||||
stage_model = tmp_path / "stage-model"
|
||||
(stage_model / "override-model").mkdir()
|
||||
(stage_model / "override-tokenizer").mkdir()
|
||||
deploy.stages[0].engine_extras.update(
|
||||
{
|
||||
"model_subdir": "override-model",
|
||||
"tokenizer_subdir": "override-tokenizer",
|
||||
}
|
||||
)
|
||||
legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
|
||||
|
||||
typed_stage = omni_config.stage_by_id(0)
|
||||
legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model)
|
||||
typed_args = build_engine_args_dict_from_omni_stage_config(typed_stage, model)
|
||||
|
||||
assert typed_stage.model_config.model_subdir == "override-model"
|
||||
assert typed_stage.model_config.tokenizer_subdir == "override-tokenizer"
|
||||
assert typed_args["model"] == legacy_args["model"] == str(stage_model / "override-model")
|
||||
assert typed_args["tokenizer"] == legacy_args["tokenizer"] == str(stage_model / "override-tokenizer")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("cli_overrides", "expected"),
|
||||
[
|
||||
({}, {"need_recv_cache": True}),
|
||||
(
|
||||
{"stage_0_omni_kv_config": {"need_send_cache": False}},
|
||||
{"need_send_cache": False},
|
||||
),
|
||||
],
|
||||
ids=["topology-over-deploy", "cli-over-topology"],
|
||||
)
|
||||
def test_typed_engine_args_preserve_legacy_omni_kv_precedence(tmp_path, cli_overrides, expected):
|
||||
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
|
||||
thinker = replace(
|
||||
pipeline.stages[0],
|
||||
omni_kv_config={"need_recv_cache": True},
|
||||
)
|
||||
pipeline = replace(pipeline, stages=(thinker, *pipeline.stages[1:]))
|
||||
legacy_stages, omni_config = _legacy_and_typed_stages(
|
||||
pipeline,
|
||||
deploy,
|
||||
model,
|
||||
cli_overrides,
|
||||
)
|
||||
|
||||
legacy_args = build_legacy_engine_args_dict(legacy_stages[0], model)
|
||||
typed_stage = omni_config.stage_by_id(0)
|
||||
typed_args = build_engine_args_dict_from_omni_stage_config(typed_stage, model)
|
||||
|
||||
assert typed_stage.connector_config.omni_kv_config == expected
|
||||
assert typed_args["omni_kv_config"] == legacy_args["omni_kv_config"] == expected
|
||||
|
||||
|
||||
def test_build_engine_args_dict_preserves_legacy_api(tmp_path):
|
||||
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
|
||||
legacy_stages, _ = _legacy_and_typed_stages(pipeline, deploy, model)
|
||||
connector_spec = {"name": "SharedMemoryConnector", "extra": {}}
|
||||
|
||||
compatibility_args = build_engine_args_dict(
|
||||
copy.deepcopy(legacy_stages[0]),
|
||||
model,
|
||||
stage_connector_spec=connector_spec,
|
||||
cli_tokenizer="/external/tokenizer",
|
||||
)
|
||||
explicitly_legacy_args = build_legacy_engine_args_dict(
|
||||
copy.deepcopy(legacy_stages[0]),
|
||||
model,
|
||||
stage_connector_spec=connector_spec,
|
||||
cli_tokenizer="/external/tokenizer",
|
||||
)
|
||||
|
||||
assert compatibility_args == explicitly_legacy_args
|
||||
|
||||
|
||||
def test_legacy_engine_args_drop_none_tp_without_mutating_stage_config():
|
||||
stage_config = types.SimpleNamespace(
|
||||
stage_id=0,
|
||||
stage_type="llm",
|
||||
engine_args={"tensor_parallel_size": None},
|
||||
default_sampling_params={},
|
||||
)
|
||||
|
||||
engine_args = build_legacy_engine_args_dict(
|
||||
stage_config,
|
||||
model="test-model",
|
||||
)
|
||||
|
||||
assert "tensor_parallel_size" not in engine_args
|
||||
assert stage_config.engine_args == {"tensor_parallel_size": None}
|
||||
|
||||
|
||||
def test_typed_engine_args_own_rocm_attention_default(monkeypatch, tmp_path):
|
||||
from vllm_omni import platforms
|
||||
|
||||
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
|
||||
legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
|
||||
monkeypatch.setattr(platforms.current_omni_platform, "is_rocm", lambda: True)
|
||||
aiter_module = types.ModuleType("vllm._aiter_ops")
|
||||
setattr(
|
||||
aiter_module,
|
||||
"rocm_aiter_ops",
|
||||
types.SimpleNamespace(is_enabled=lambda: False),
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "vllm._aiter_ops", aiter_module)
|
||||
|
||||
legacy_args = build_legacy_engine_args_dict(legacy_stages[1], model)
|
||||
typed_args = build_engine_args_dict_from_omni_stage_config(
|
||||
omni_config.stage_by_id(1),
|
||||
model,
|
||||
)
|
||||
|
||||
assert "attention_backend" not in legacy_args
|
||||
assert typed_args["attention_backend"] == "TRITON_ATTN"
|
||||
|
||||
|
||||
def test_typed_ming_image_engine_args_defer_diffusion_batch_default():
|
||||
pipeline = resolve_pipeline_config("ming_flash_omni_image")
|
||||
assert pipeline is not None
|
||||
deploy = load_deploy_config(_DEPLOY_DIR / "ming_flash_omni_image.yaml")
|
||||
legacy_stages, omni_config = _legacy_and_typed_stages(
|
||||
pipeline,
|
||||
deploy,
|
||||
model="/tmp",
|
||||
)
|
||||
|
||||
legacy_args = build_legacy_engine_args_dict(legacy_stages[1], model="/tmp")
|
||||
typed_args = build_engine_args_dict_from_omni_stage_config(
|
||||
omni_config.stage_by_id(1),
|
||||
model="/tmp",
|
||||
)
|
||||
typed_backend_args = {name: value for name, value in typed_args.items() if name in _DIFFUSION_BACKEND_FIELDS}
|
||||
|
||||
assert "max_num_seqs" not in legacy_args
|
||||
assert "max_num_seqs" not in typed_args
|
||||
assert OmniDiffusionConfig(**typed_backend_args).max_num_seqs == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
|
||||
def test_typed_engine_args_match_current_registry_backend_semantics(model_type):
|
||||
pipeline = resolve_pipeline_config(model_type)
|
||||
if pipeline is None:
|
||||
pytest.skip(f"Pipeline {model_type!r} requires an HF config to resolve")
|
||||
|
||||
deploy = (
|
||||
load_deploy_config(_DEPLOY_DIR / pipeline.default_deploy_config_name)
|
||||
if pipeline.default_deploy_config_name is not None
|
||||
else DeployConfig()
|
||||
)
|
||||
legacy_stages, omni_config = _legacy_and_typed_stages(
|
||||
pipeline,
|
||||
deploy,
|
||||
model="/tmp",
|
||||
)
|
||||
|
||||
for legacy_stage in legacy_stages:
|
||||
stage_id = legacy_stage.stage_id
|
||||
legacy_args = build_legacy_engine_args_dict(legacy_stage, model="/tmp")
|
||||
typed_args = build_engine_args_dict_from_omni_stage_config(
|
||||
omni_config.stage_by_id(stage_id),
|
||||
model="/tmp",
|
||||
)
|
||||
if legacy_stage.stage_type == StageType.DIFFUSION:
|
||||
backend_fields = _DIFFUSION_BACKEND_FIELDS
|
||||
backend_config_cls = OmniDiffusionConfig
|
||||
else:
|
||||
backend_fields = _LLM_BACKEND_FIELDS
|
||||
backend_config_cls = OmniEngineArgs
|
||||
|
||||
missing_fields = (legacy_args.keys() & backend_fields) - typed_args.keys()
|
||||
assert not missing_fields, f"{model_type} stage {stage_id} lost backend fields: {sorted(missing_fields)}"
|
||||
|
||||
legacy_effective_args = _effective_backend_values(backend_config_cls, legacy_args)
|
||||
typed_effective_args = _effective_backend_values(backend_config_cls, typed_args)
|
||||
assert typed_effective_args == legacy_effective_args, (
|
||||
f"{model_type} stage {stage_id} changed effective backend arguments"
|
||||
)
|
||||
Reference in new issue
Block a user