698f716012
Signed-off-by: TaffyOfficial <2324465096@qq.com> Signed-off-by: zuiho <2324465096@qq.com> Signed-off-by: Gao Han <hgaoaf@connect.ust.hk> Co-authored-by: TaffyOfficial <2324465096@qq.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
52 lines
2.6 KiB
Python
52 lines
2.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Check the sleep/wake E2E topology before loading large BAGEL weights."""
|
|
|
|
import pytest
|
|
|
|
from tests.entrypoints.test_omni_sleep_mode import _sleep_deploy_config
|
|
from vllm_omni.config.omni_config import (
|
|
VllmOmniConfig,
|
|
VllmOmniDiffusionStageConfig,
|
|
extract_diffusion_stage_config_kwargs,
|
|
)
|
|
from vllm_omni.config.pipeline_registry import resolve_pipeline_config
|
|
from vllm_omni.config.stage_config import load_deploy_config
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.mark.parametrize("diffusion_only,tp_size", [(False, 1), (False, 2), (True, 2)])
|
|
def test_sleep_deploy_resolves_expected_stages(diffusion_only, tp_size):
|
|
deploy = load_deploy_config(_sleep_deploy_config(diffusion_only=diffusion_only, tp_size=tp_size))
|
|
pipeline = resolve_pipeline_config(deploy.pipeline or "bagel")
|
|
config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
|
|
|
|
assert len(config.stage_configs) == (1 if diffusion_only else 2)
|
|
for stage in config.stage_configs:
|
|
is_diffusion = diffusion_only or stage.stage_id == 1
|
|
assert isinstance(stage, VllmOmniDiffusionStageConfig) == is_diffusion
|
|
# Keep BAGEL's batch budget: a single multimodal item needs 8625 tokens.
|
|
assert stage.scheduler_config.max_num_batched_tokens == 32768
|
|
assert stage.model_config.enable_sleep_mode is True
|
|
assert stage.parallel_config.tensor_parallel_size == tp_size
|
|
assert stage.runtime_config.devices == ",".join(str(stage.stage_id * tp_size + rank) for rank in range(tp_size))
|
|
if not is_diffusion:
|
|
assert stage.cache_config.gpu_memory_utilization == 0.8
|
|
if is_diffusion:
|
|
engine_args = build_engine_args_dict_from_omni_stage_config(stage, "test-bagel")
|
|
terminal = OmniDiffusionConfig.from_kwargs(
|
|
**extract_diffusion_stage_config_kwargs(
|
|
engine_args, stage_id=stage.stage_id, include_engine_adapter_metadata=True
|
|
)
|
|
)
|
|
assert terminal.enable_sleep_mode is True
|
|
assert terminal.parallel_config.tensor_parallel_size == tp_size
|
|
|
|
if not diffusion_only:
|
|
assert config.stage_by_id(1).stage_pipeline_config.input_sources == (0,)
|
|
assert deploy.stages[1].input_connectors["from_stage_0"] == "shared_memory_connector"
|