Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: kushanam <kushanam@users.noreply.github.com> Signed-off-by: Kushan A <kahmadian@nvidia.com> Signed-off-by: kahmadian <kahmadian@users.noreply.github.com> Co-authored-by: kushanam <kushanam@users.noreply.github.com> Co-authored-by: kahmadian <kahmadian@users.noreply.github.com>
This commit is contained in:
30 files changed
+2722
-51
No files matched your search
@@ -0,0 +1,154 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for applying strategy specs onto merged stage configs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.config.composable_parallel import (
|
||||
Broadcast,
|
||||
FanInByStage,
|
||||
MeshAxisSpec,
|
||||
RouteByStage,
|
||||
StrategyApplyError,
|
||||
StrategySpec,
|
||||
TakeRank,
|
||||
apply_strategy_specs,
|
||||
)
|
||||
from vllm_omni.config.stage_config import StageConfig
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _tp(size: int) -> StrategySpec:
|
||||
return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank())
|
||||
|
||||
|
||||
def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec:
|
||||
return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage())
|
||||
|
||||
|
||||
def _stage(stage_id: int, model_stage: str, engine_args=None, runtime=None) -> StageConfig:
|
||||
return StageConfig(
|
||||
stage_id=stage_id,
|
||||
model_stage=model_stage,
|
||||
yaml_engine_args=dict(engine_args or {}),
|
||||
yaml_runtime=dict(runtime or {"num_replicas": 1}),
|
||||
)
|
||||
|
||||
|
||||
def _qwen_stages() -> list[StageConfig]:
|
||||
return [
|
||||
_stage(0, "thinker"),
|
||||
_stage(1, "talker"),
|
||||
_stage(2, "code2wav"),
|
||||
]
|
||||
|
||||
|
||||
def test_apply_tp_by_role():
|
||||
stages = _qwen_stages()
|
||||
apply_strategy_specs(stages, {"thinker": [_tp(2)]})
|
||||
assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
|
||||
# untouched roles keep their config
|
||||
assert "tensor_parallel_size" not in stages[1].yaml_engine_args
|
||||
|
||||
|
||||
def test_apply_by_model_stage():
|
||||
stages = _qwen_stages()
|
||||
apply_strategy_specs(stages, {"thinker": [_tp(2)]})
|
||||
assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
|
||||
|
||||
|
||||
def test_apply_stage_replica_sets_num_replicas_and_surfaces_lb():
|
||||
stages = _qwen_stages()
|
||||
result = apply_strategy_specs(stages, {"talker": [_stage_replica(2, "round_robin")]})
|
||||
assert stages[1].yaml_runtime["num_replicas"] == 2
|
||||
assert result.omni_lb_policy == "round-robin"
|
||||
|
||||
|
||||
def test_only_declared_axes_are_written():
|
||||
stages = _qwen_stages()
|
||||
# strategy declares only stage_replica -> tp must not be forced.
|
||||
apply_strategy_specs(stages, {"talker": [_stage_replica(2)]})
|
||||
assert "tensor_parallel_size" not in stages[1].yaml_engine_args
|
||||
|
||||
|
||||
def test_conflict_on_explicit_tp():
|
||||
stages = _qwen_stages()
|
||||
stages[0].yaml_engine_args["tensor_parallel_size"] = 4
|
||||
with pytest.raises(StrategyApplyError):
|
||||
apply_strategy_specs(stages, {"thinker": [_tp(2)]})
|
||||
|
||||
|
||||
def test_equal_explicit_value_is_noop():
|
||||
stages = _qwen_stages()
|
||||
stages[0].yaml_engine_args["tensor_parallel_size"] = 2
|
||||
apply_strategy_specs(stages, {"thinker": [_tp(2)]})
|
||||
assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
|
||||
|
||||
|
||||
def test_explicit_none_conflicts_with_derived_value():
|
||||
# An explicit YAML ``null`` (``tensor_parallel_size: null``) is a *present*
|
||||
# value, not a missing key, so a strategy deriving a non-None size must raise
|
||||
# rather than silently clobber the explicit None.
|
||||
stages = _qwen_stages()
|
||||
stages[0].yaml_engine_args["tensor_parallel_size"] = None
|
||||
with pytest.raises(StrategyApplyError):
|
||||
apply_strategy_specs(stages, {"thinker": [_tp(2)]})
|
||||
|
||||
|
||||
def test_missing_key_is_filled_without_conflict():
|
||||
# A genuinely absent key (never set in the YAML) is filled by the strategy
|
||||
# and must NOT raise — the contrast case to an explicit None.
|
||||
stages = _qwen_stages()
|
||||
assert "tensor_parallel_size" not in stages[0].yaml_engine_args
|
||||
apply_strategy_specs(stages, {"thinker": [_tp(2)]})
|
||||
assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
|
||||
|
||||
|
||||
def test_num_replicas_conflict():
|
||||
stages = _qwen_stages()
|
||||
stages[1].yaml_runtime["num_replicas"] = 3
|
||||
with pytest.raises(StrategyApplyError):
|
||||
apply_strategy_specs(stages, {"talker": [_stage_replica(2)]})
|
||||
|
||||
|
||||
def test_device_count_ok_template():
|
||||
stages = _qwen_stages()
|
||||
stages[0].yaml_runtime["devices"] = "0,1"
|
||||
apply_strategy_specs(stages, {"thinker": [_tp(2)]})
|
||||
assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
|
||||
|
||||
|
||||
def test_device_count_ok_pool():
|
||||
# tp=2 -> world=2; 2 replicas -> pool of 4 device ids is valid.
|
||||
stages = _qwen_stages()
|
||||
stages[1].yaml_runtime["devices"] = "0,1,2,3"
|
||||
apply_strategy_specs(stages, {"talker": [_tp(2), _stage_replica(2)]})
|
||||
assert stages[1].yaml_runtime["num_replicas"] == 2
|
||||
|
||||
|
||||
def test_device_count_mismatch():
|
||||
stages = _qwen_stages()
|
||||
stages[0].yaml_runtime["devices"] = "0,1,2"
|
||||
with pytest.raises(StrategyApplyError):
|
||||
apply_strategy_specs(stages, {"thinker": [_tp(2)]})
|
||||
|
||||
|
||||
def test_unknown_role_raises():
|
||||
stages = _qwen_stages()
|
||||
with pytest.raises(StrategyApplyError):
|
||||
apply_strategy_specs(stages, {"nonexistent": [_tp(2)]})
|
||||
|
||||
|
||||
def test_conflicting_lb_policy_across_roles():
|
||||
stages = _qwen_stages()
|
||||
with pytest.raises(StrategyApplyError):
|
||||
apply_strategy_specs(
|
||||
stages,
|
||||
{
|
||||
"talker": [_stage_replica(2, "round_robin")],
|
||||
"code2wav": [_stage_replica(2, "least_queue")],
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,141 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Integration tests: strategy overlay on the real merge_pipeline_deploy seam."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.config.composable_parallel import (
|
||||
Broadcast,
|
||||
FanInByStage,
|
||||
MeshAxisSpec,
|
||||
RouteByStage,
|
||||
StrategyApplyError,
|
||||
StrategySpec,
|
||||
TakeRank,
|
||||
apply_strategy_specs,
|
||||
)
|
||||
from vllm_omni.config.config_factory import StageConfigFactory
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
||||
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_DEPLOY = Path(__file__).parents[3] / "vllm_omni" / "deploy" / "qwen2_5_omni.yaml"
|
||||
|
||||
|
||||
def _tp(size: int) -> StrategySpec:
|
||||
return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank())
|
||||
|
||||
|
||||
def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec:
|
||||
return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage())
|
||||
|
||||
|
||||
def _qwen_stages():
|
||||
pipeline = OMNI_PIPELINES["qwen2_5_omni"]
|
||||
deploy = load_deploy_config(_DEPLOY)
|
||||
return merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
|
||||
def _stage(stages, role):
|
||||
return next(s for s in stages if s.model_stage == role)
|
||||
|
||||
|
||||
def test_overlay_tp_on_thinker():
|
||||
stages = _qwen_stages()
|
||||
# The bundled deploy pins the thinker to one GPU; a TP=2 strategy needs a
|
||||
# matching 2-GPU layout (mirrors what a TP2 deploy would declare).
|
||||
_stage(stages, "thinker").yaml_runtime["devices"] = "0,1"
|
||||
result = apply_strategy_specs(stages, {"thinker": [_tp(2)]})
|
||||
assert _stage(result.stages, "thinker").yaml_engine_args["tensor_parallel_size"] == 2
|
||||
|
||||
|
||||
def test_device_guard_rejects_tp_on_single_gpu_deploy():
|
||||
# The bundled deploy pins the thinker to a single GPU, so the pre-spawn
|
||||
# device check must refuse a TP=2 strategy on it.
|
||||
stages = _qwen_stages()
|
||||
with pytest.raises(StrategyApplyError):
|
||||
apply_strategy_specs(stages, {"thinker": [_tp(2)]})
|
||||
|
||||
|
||||
def test_overlay_stage_replica_on_talker():
|
||||
stages = _qwen_stages()
|
||||
result = apply_strategy_specs(stages, {"talker": [_stage_replica(2, "round_robin")]})
|
||||
assert _stage(result.stages, "talker").yaml_runtime["num_replicas"] == 2
|
||||
assert result.omni_lb_policy == "round-robin"
|
||||
|
||||
|
||||
def test_overlay_mixed_roles():
|
||||
stages = _qwen_stages()
|
||||
_stage(stages, "thinker").yaml_runtime["devices"] = "0,1"
|
||||
result = apply_strategy_specs(
|
||||
stages,
|
||||
{
|
||||
"thinker": [_tp(2)],
|
||||
"talker": [_stage_replica(2, "round_robin")],
|
||||
"code2wav": [_stage_replica(2, "round_robin")],
|
||||
},
|
||||
)
|
||||
by_role = {s.model_stage: s for s in result.stages}
|
||||
assert by_role["thinker"].yaml_engine_args["tensor_parallel_size"] == 2
|
||||
assert by_role["talker"].yaml_runtime["num_replicas"] == 2
|
||||
assert by_role["code2wav"].yaml_runtime["num_replicas"] == 2
|
||||
assert result.omni_lb_policy == "round-robin"
|
||||
|
||||
|
||||
def _resolved(stages, role):
|
||||
"""Return a role's fully-resolved (post-CLI) OmegaConf stage config."""
|
||||
stage = next(s for s in stages if s.model_stage == role)
|
||||
return stage.to_omegaconf()
|
||||
|
||||
|
||||
def test_device_check_survives_cli_override():
|
||||
# Strategy replicates the talker (1-GPU template -> valid at apply time),
|
||||
# but a CLI --stage_1_devices with 3 ids must NOT slip past the device
|
||||
# guard: effective world=1, replicas=2 admits only 1 or 2 device ids.
|
||||
pipeline_cfg = OMNI_PIPELINES["qwen2_5_omni"]
|
||||
with pytest.raises(StrategyApplyError):
|
||||
StageConfigFactory._create_from_registry(
|
||||
"qwen2_5_omni",
|
||||
pipeline_cfg,
|
||||
cli_overrides={"stage_1_devices": "0,1,2"},
|
||||
strategy_specs={"talker": [_stage_replica(2, "round_robin")]},
|
||||
)
|
||||
|
||||
|
||||
def test_cli_overrides_strategy_with_warning():
|
||||
# Strategy derives num_replicas=2 for the talker; a CLI override to 3 wins
|
||||
# (it is the most explicit user action) but must be surfaced loudly rather
|
||||
# than silently. The resulting layout (1 template device) stays valid.
|
||||
#
|
||||
# vLLM's logger sets propagate=False, so attach a handler directly to it
|
||||
# rather than relying on pytest's caplog (which listens on the root logger).
|
||||
messages: list[str] = []
|
||||
|
||||
class _Capture(logging.Handler):
|
||||
def emit(self, record: logging.LogRecord) -> None:
|
||||
messages.append(record.getMessage())
|
||||
|
||||
log = logging.getLogger("vllm_omni.config.config_factory")
|
||||
handler = _Capture(level=logging.WARNING)
|
||||
log.addHandler(handler)
|
||||
try:
|
||||
pipeline_cfg = OMNI_PIPELINES["qwen2_5_omni"]
|
||||
stages, _ = StageConfigFactory._create_from_registry(
|
||||
"qwen2_5_omni",
|
||||
pipeline_cfg,
|
||||
cli_overrides={"stage_1_num_replicas": 3},
|
||||
strategy_specs={"talker": [_stage_replica(2, "round_robin")]},
|
||||
)
|
||||
finally:
|
||||
log.removeHandler(handler)
|
||||
|
||||
# CLI value wins in the resolved config.
|
||||
assert _resolved(stages, "talker").runtime.num_replicas == 3
|
||||
# ...and the override was warned about, naming the conflicting field.
|
||||
assert any("num_replicas" in m and "overrides the strategy-derived" in m for m in messages)
|
||||
@@ -0,0 +1,42 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the vendored spec/pattern data types."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.config.composable_parallel.aggregation import AggregationPattern, TakeRank
|
||||
from vllm_omni.config.composable_parallel.routing import RoutingPattern
|
||||
from vllm_omni.config.composable_parallel.spec import MeshAxisSpec
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_mesh_axis_rejects_bad_kind():
|
||||
with pytest.raises(ValueError):
|
||||
MeshAxisSpec(kind="not_a_kind", size=2) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_mesh_axis_rejects_nonpositive_size():
|
||||
with pytest.raises(ValueError):
|
||||
MeshAxisSpec(kind="tp", size=0)
|
||||
|
||||
|
||||
def test_routing_pattern_is_closed():
|
||||
with pytest.raises(TypeError):
|
||||
|
||||
class RogueRouting(RoutingPattern): # noqa: F811
|
||||
pass
|
||||
|
||||
|
||||
def test_aggregation_pattern_is_closed():
|
||||
with pytest.raises(TypeError):
|
||||
|
||||
class RogueAggregation(AggregationPattern): # noqa: F811
|
||||
pass
|
||||
|
||||
|
||||
def test_take_rank_type_guard():
|
||||
with pytest.raises(TypeError):
|
||||
TakeRank(rank=True) # bool is not a valid rank
|
||||
@@ -0,0 +1,97 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for parsing/loading strategy files."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.config.composable_parallel.routing import Broadcast, RouteByStage
|
||||
from vllm_omni.config.composable_parallel.strategy_loader import (
|
||||
StrategyLoadError,
|
||||
load_strategy_specs,
|
||||
parse_strategy_specs,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_parse_basic():
|
||||
data = {
|
||||
"strategies": {
|
||||
"thinker": [{"axis": "tp", "size": 2}],
|
||||
"talker": [{"axis": "stage_replica", "size": 2, "routing": "round_robin"}],
|
||||
}
|
||||
}
|
||||
specs = parse_strategy_specs(data)
|
||||
assert set(specs) == {"thinker", "talker"}
|
||||
|
||||
(tp_spec,) = specs["thinker"]
|
||||
assert tp_spec.mesh_axis.kind == "tp"
|
||||
assert tp_spec.mesh_axis.size == 2
|
||||
assert isinstance(tp_spec.routing, Broadcast)
|
||||
|
||||
(sr_spec,) = specs["talker"]
|
||||
assert sr_spec.mesh_axis.kind == "stage_replica"
|
||||
assert isinstance(sr_spec.routing, RouteByStage)
|
||||
assert sr_spec.routing.routing_policy == "round_robin"
|
||||
|
||||
|
||||
def test_parse_without_strategies_key():
|
||||
# A bare role mapping (no top-level "strategies") is also accepted.
|
||||
specs = parse_strategy_specs({"thinker": [{"axis": "tp", "size": 4}]})
|
||||
assert specs["thinker"][0].mesh_axis.size == 4
|
||||
|
||||
|
||||
def test_l1_owner_goes_to_shard_extension():
|
||||
specs = parse_strategy_specs({"talker": [{"axis": "stage_replica", "size": 2, "l1_owner": "delegated"}]})
|
||||
assert specs["talker"][0].shard_extension["l1_owner"] == "delegated"
|
||||
|
||||
|
||||
def test_missing_axis_raises():
|
||||
with pytest.raises(StrategyLoadError):
|
||||
parse_strategy_specs({"thinker": [{"size": 2}]})
|
||||
|
||||
|
||||
def test_missing_size_raises():
|
||||
with pytest.raises(StrategyLoadError):
|
||||
parse_strategy_specs({"thinker": [{"axis": "tp"}]})
|
||||
|
||||
|
||||
def test_routing_on_non_policy_axis_raises():
|
||||
with pytest.raises(StrategyLoadError):
|
||||
parse_strategy_specs({"thinker": [{"axis": "tp", "size": 2, "routing": "random"}]})
|
||||
|
||||
|
||||
def test_bad_size_raises():
|
||||
with pytest.raises(StrategyLoadError):
|
||||
parse_strategy_specs({"thinker": [{"axis": "tp", "size": "two"}]})
|
||||
|
||||
|
||||
def test_entries_must_be_list():
|
||||
with pytest.raises(StrategyLoadError):
|
||||
parse_strategy_specs({"thinker": {"axis": "tp", "size": 2}})
|
||||
|
||||
|
||||
def test_non_mapping_entry_raises():
|
||||
# A list whose elements are not mappings (e.g. a bare string) must raise a
|
||||
# StrategyLoadError, not an opaque TypeError from dict(entry).
|
||||
with pytest.raises(StrategyLoadError):
|
||||
parse_strategy_specs({"thinker": ["tp"]})
|
||||
|
||||
|
||||
def test_load_from_file(tmp_path):
|
||||
path = tmp_path / "strategy.yaml"
|
||||
path.write_text(
|
||||
"strategies:\n"
|
||||
" thinker:\n"
|
||||
" - axis: tp\n"
|
||||
" size: 2\n"
|
||||
" talker:\n"
|
||||
" - axis: stage_replica\n"
|
||||
" size: 2\n"
|
||||
" routing: least_queue\n"
|
||||
)
|
||||
specs = load_strategy_specs(str(path))
|
||||
assert specs["thinker"][0].mesh_axis.size == 2
|
||||
assert specs["talker"][0].routing.routing_policy == "least_queue"
|
||||
@@ -0,0 +1,195 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for StrategySpec -> OmniParallelConfig translation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.config.composable_parallel import (
|
||||
AxisTranslationError,
|
||||
Broadcast,
|
||||
FanInByStage,
|
||||
MeshAxisSpec,
|
||||
PartitionByHash,
|
||||
PipelineMicrobatch,
|
||||
RouteByStage,
|
||||
StrategySpec,
|
||||
TakeRank,
|
||||
Union,
|
||||
translate_strategy_stack,
|
||||
)
|
||||
from vllm_omni.config.composable_parallel.aggregation import StitchPipeline
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _tp(size: int) -> StrategySpec:
|
||||
return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank())
|
||||
|
||||
|
||||
def _dp(size: int, policy: str = "round_robin") -> StrategySpec:
|
||||
return StrategySpec("dp", MeshAxisSpec("dp", size), RouteByStage(policy), Union())
|
||||
|
||||
|
||||
def _pp(size: int) -> StrategySpec:
|
||||
return StrategySpec("pp", MeshAxisSpec("pp", size), PipelineMicrobatch(), StitchPipeline())
|
||||
|
||||
|
||||
def _ep(size: int) -> StrategySpec:
|
||||
return StrategySpec("ep", MeshAxisSpec("ep", size), Broadcast(), Union())
|
||||
|
||||
|
||||
def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec:
|
||||
return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage())
|
||||
|
||||
|
||||
def test_tp_only():
|
||||
cfg = translate_strategy_stack([_tp(2)])
|
||||
assert cfg.tensor_parallel_size == 2
|
||||
assert cfg.data_parallel_size == 1
|
||||
assert cfg.world_size == 2
|
||||
assert cfg.stage_replica_size == 1
|
||||
assert cfg.l1_owners == {"tp": "engine"}
|
||||
|
||||
|
||||
def test_dp_is_engine_owned():
|
||||
cfg = translate_strategy_stack([_dp(2)])
|
||||
assert cfg.data_parallel_size == 2
|
||||
assert cfg.world_size == 2
|
||||
# engine DP does not emit an omni LB policy (vLLM's internal DP LB is used).
|
||||
assert cfg.omni_lb_policy is None
|
||||
assert cfg.as_engine_kwargs().get("omni_lb_policy") is None
|
||||
assert cfg.l1_owners["dp"] == "engine"
|
||||
|
||||
|
||||
def test_tp_times_dp_world_size():
|
||||
cfg = translate_strategy_stack([_tp(2), _dp(2)])
|
||||
assert cfg.world_size == 4
|
||||
assert cfg.tensor_parallel_size == 2
|
||||
assert cfg.data_parallel_size == 2
|
||||
|
||||
|
||||
def test_stage_replica_is_delegated():
|
||||
cfg = translate_strategy_stack([_stage_replica(3, "least_queue")])
|
||||
# stage_replica is NOT a world dimension.
|
||||
assert cfg.world_size == 1
|
||||
assert cfg.stage_replica_size == 3
|
||||
assert cfg.omni_lb_policy == "least-queue-length"
|
||||
assert cfg.l1_owners["stage_replica"] == "delegated"
|
||||
# stage_replica_size is not an engine kwarg; it's a per-stage deploy knob.
|
||||
assert "stage_replica_size" not in cfg.as_engine_kwargs()
|
||||
|
||||
|
||||
def test_stage_replica_lb_policy_only_when_replicated():
|
||||
cfg = translate_strategy_stack([_stage_replica(1)])
|
||||
assert cfg.stage_replica_size == 1
|
||||
assert cfg.as_engine_kwargs().get("omni_lb_policy") is None
|
||||
|
||||
|
||||
def test_tp_with_stage_replica():
|
||||
cfg = translate_strategy_stack([_tp(2), _stage_replica(2)])
|
||||
assert cfg.world_size == 2 # tp only
|
||||
assert cfg.tensor_parallel_size == 2
|
||||
assert cfg.stage_replica_size == 2
|
||||
assert cfg.omni_lb_policy == "round-robin"
|
||||
|
||||
|
||||
def test_ep_must_match_tp_dp_product():
|
||||
cfg = translate_strategy_stack([_tp(2), _ep(2)])
|
||||
assert cfg.enable_expert_parallel is True
|
||||
assert cfg.as_engine_kwargs()["enable_expert_parallel"] is True
|
||||
|
||||
|
||||
def test_ep_mismatch_raises():
|
||||
with pytest.raises(Exception) as exc:
|
||||
translate_strategy_stack([_tp(2), _ep(4)])
|
||||
assert "must equal" in str(exc.value)
|
||||
|
||||
|
||||
def test_duplicate_kind_rejected():
|
||||
with pytest.raises(AxisTranslationError):
|
||||
translate_strategy_stack([_tp(2), _tp(2)])
|
||||
|
||||
|
||||
def test_unsupported_kind_rejected():
|
||||
spec = StrategySpec("sp", MeshAxisSpec("sp_ulysses", 2), Broadcast(), TakeRank())
|
||||
with pytest.raises(AxisTranslationError):
|
||||
translate_strategy_stack([spec])
|
||||
|
||||
|
||||
def test_dp_hash_routing_rejected():
|
||||
spec = StrategySpec("dp", MeshAxisSpec("dp", 2), PartitionByHash(), Union())
|
||||
with pytest.raises(NotImplementedError):
|
||||
translate_strategy_stack([spec])
|
||||
|
||||
|
||||
def test_dp_invalid_routing_policy_rejected():
|
||||
spec = StrategySpec("dp", MeshAxisSpec("dp", 2), RouteByStage("bogus"), Union())
|
||||
with pytest.raises(AxisTranslationError):
|
||||
translate_strategy_stack([spec])
|
||||
|
||||
|
||||
def test_lb_policy_never_in_engine_kwargs_even_when_replicated():
|
||||
# omni_lb_policy is pipeline-wide (applied at engine construction), not a
|
||||
# per-stage engine arg, so it must never leak into as_engine_kwargs().
|
||||
cfg = translate_strategy_stack([_tp(2), _stage_replica(2, "least_queue")])
|
||||
assert cfg.omni_lb_policy == "least-queue-length"
|
||||
assert "omni_lb_policy" not in cfg.as_engine_kwargs()
|
||||
|
||||
|
||||
def test_stage_replica_hash_routing_rejected():
|
||||
spec = StrategySpec("sr", MeshAxisSpec("stage_replica", 2), RouteByStage("hash"), FanInByStage())
|
||||
with pytest.raises(NotImplementedError):
|
||||
translate_strategy_stack([spec])
|
||||
|
||||
|
||||
def test_dp_wrong_owner_rejected():
|
||||
spec = StrategySpec(
|
||||
"dp",
|
||||
MeshAxisSpec("dp", 2),
|
||||
RouteByStage("round_robin"),
|
||||
Union(),
|
||||
shard_extension={"l1_owner": "delegated"},
|
||||
)
|
||||
with pytest.raises(AxisTranslationError):
|
||||
translate_strategy_stack([spec])
|
||||
|
||||
|
||||
def test_stage_replica_wrong_owner_rejected():
|
||||
spec = StrategySpec(
|
||||
"sr",
|
||||
MeshAxisSpec("stage_replica", 2),
|
||||
RouteByStage("round_robin"),
|
||||
FanInByStage(),
|
||||
shard_extension={"l1_owner": "engine"},
|
||||
)
|
||||
with pytest.raises(AxisTranslationError):
|
||||
translate_strategy_stack([spec])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"specs",
|
||||
[
|
||||
[_tp(2)],
|
||||
[_dp(2)],
|
||||
[_tp(2), _dp(2)],
|
||||
[_tp(2), _ep(2)],
|
||||
[_tp(2), _stage_replica(2)],
|
||||
[_pp(2)],
|
||||
],
|
||||
)
|
||||
def test_as_engine_kwargs_are_valid_omni_engine_args(specs):
|
||||
# Guard against desync: every key emitted by as_engine_kwargs() must be a
|
||||
# real OmniEngineArgs field, otherwise constructing the per-stage engine args
|
||||
# would explode at runtime. Keep this CPU-only by checking field names rather
|
||||
# than instantiating the engine.
|
||||
from vllm_omni.engine.arg_utils import OmniEngineArgs
|
||||
|
||||
valid_fields = {f.name for f in dataclasses.fields(OmniEngineArgs)}
|
||||
cfg = translate_strategy_stack(specs)
|
||||
kwargs = cfg.as_engine_kwargs()
|
||||
unknown = set(kwargs) - valid_fields
|
||||
assert not unknown, f"as_engine_kwargs emitted non-OmniEngineArgs fields: {unknown}"
|
||||
Reference in new issue
Block a user