diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 1deb7946c..66113a5f9 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -94,12 +94,14 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne | `max_num_batched_tokens` | int \| null | optional | `null` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. | | `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. | | `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. | -| `devices` | str \| null | optional | `null` | Device list assigned to this stage. | +| `devices` | str \| null | optional | `null` | Device list assigned to this stage. The number of device ids must equal this stage's local world size (`tensor_parallel_size` × local data-parallel size × `pipeline_parallel_size`, or `num_replicas` × that product for a replica pool); a mismatch fails early — see the note below. | | `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_`; values are names registered under top-level `connectors:`. | | `input_connectors` | dict \| null | optional | `null` | Keyed by `from_stage_`; values are names registered under top-level `connectors:`. | | `default_sampling_params` | dict \| null | optional | `null` | Baseline sampling params. Deep-merged with pipeline `sampling_constraints` (pipeline wins). | | `engine_extras` | dict | optional | `{}` | Catch-all for engine fields not listed above; deep-merged across overlays and forwarded to the stage engine. | +**Note:** a stage's `devices` count must equal its local world size (`tensor_parallel_size` × `data_parallel_size_local` × `pipeline_parallel_size`, falling back to global `data_parallel_size` when the local size is unset), or `num_replicas` × that product for a replica pool. A mismatch fails early and names the offending stage. A top-level `--tensor-parallel-size` is broadcast to every stage, so it can make a single-GPU stage violate this contract ([issue #5003](gh-issue:5003)); fix that case with `--stage-overrides` (set `tensor_parallel_size` and `devices` together per stage) or set TP only on the multi-GPU stage. + ### Connector schema Each entry under top-level `connectors:` follows this shape: diff --git a/tests/engine/test_stage_device_layout.py b/tests/engine/test_stage_device_layout.py new file mode 100644 index 000000000..0473dc341 --- /dev/null +++ b/tests/engine/test_stage_device_layout.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Regression tests for issue #5003: a per-stage world size that its assigned +``devices`` cannot satisfy must fail early in ``build_vllm_config`` with a clear +message, rather than surfacing as an opaque worker-side ``local rank ... out of +bounds`` assertion. + +Root cause: a top-level ``--tensor-parallel-size`` is broadcast to every stage, +but each stage's ``devices`` is not adjusted, so a stage can end up with e.g. +tensor_parallel_size=4 while still holding a single-GPU deploy default. Without +``--strategy-config`` the strategy-path device check never runs. +""" + +import json +import re +import types +from unittest import mock + +import pytest + +from vllm_omni.engine import stage_init_utils +from vllm_omni.engine.stage_init_utils import ( + _check_stage_device_layout, + build_vllm_config, + compute_replica_layout, + get_stage_devices_per_replica, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _stage(stage_id, devices, num_replicas=1, engine_args=None): + return types.SimpleNamespace( + stage_id=stage_id, + stage_type="llm", + engine_args=engine_args or {}, + runtime=types.SimpleNamespace(devices=devices, num_replicas=num_replicas), + ) + + +def test_tp_broadcast_without_devices_fails_early(): + """stage0 gets tensor_parallel_size=4 (broadcast) but only 1 device -> clear error.""" + stage = _stage(0, devices="0") + engine_args = { + "tensor_parallel_size": 4, + "data_parallel_size": 1, + "pipeline_parallel_size": 1, + } + with pytest.raises(ValueError) as excinfo: + _check_stage_device_layout(stage, engine_args) + msg = str(excinfo.value) + # Message names the stage, the mismatch, and the actionable workaround. + assert "Stage 0" in msg + assert "1 device" in msg and "4" in msg + assert "--stage-overrides" in msg + + +def test_workaround_example_sets_tp_on_every_stage(): + """The JSON example in the error message must not reproduce the bug: every + stage entry it suggests has to set ``tensor_parallel_size`` (a single-GPU + stage that only overrides ``devices`` would still inherit the broadcast + tp and crash again — the original #5003 failure mode).""" + stage = _stage(0, devices="0") + with pytest.raises(ValueError) as excinfo: + _check_stage_device_layout(stage, {"tensor_parallel_size": 4}) + msg = str(excinfo.value) + + match = re.search(r"\{.*\}", msg, re.DOTALL) + assert match, f"no JSON example found in error message: {msg}" + example = json.loads(match.group(0)) + + assert example, "example must contain at least one stage" + for stage_id, overrides in example.items(): + assert "tensor_parallel_size" in overrides, ( + f"stage {stage_id} in the suggested workaround omits " + f"tensor_parallel_size, which reproduces #5003: {overrides}" + ) + + +def test_consistent_tp_and_devices_pass(): + """tensor_parallel_size=4 with 4 assigned devices is valid.""" + stage = _stage(0, devices="0,1,2,3") + _check_stage_device_layout( + stage, + {"tensor_parallel_size": 4, "data_parallel_size": 1, "pipeline_parallel_size": 1}, + ) + + +def test_single_gpu_stage_passes(): + """A TP=1 stage on a single device (talker/code2wav default) is valid.""" + stage = _stage(1, devices="1") + _check_stage_device_layout(stage, {"tensor_parallel_size": 1}) + + +def test_missing_devices_is_skipped(): + """No explicit devices -> vLLM assigns them; nothing to validate here.""" + stage = _stage(0, devices=None) + _check_stage_device_layout(stage, {"tensor_parallel_size": 4}) + + +def test_replica_pool_layout_passes(): + """Pool mode: num_replicas=2 x tp=2 => 4 devices is a valid pool shape.""" + stage = _stage(0, devices="0,1,2,3", num_replicas=2) + _check_stage_device_layout(stage, {"tensor_parallel_size": 2}) + + +def test_multinode_dp_uses_local_width_for_device_validation(): + """One local DP engine needs one device even when global DP spans four nodes.""" + stage = _stage(0, devices="0") + _check_stage_device_layout( + stage, + { + "tensor_parallel_size": 1, + "data_parallel_size": 4, + "data_parallel_size_local": 1, + "pipeline_parallel_size": 1, + }, + ) + + +def test_zero_local_dp_skips_local_device_validation(): + """A head process with no local engines does not own a local device layout.""" + stage = _stage(0, devices="0") + _check_stage_device_layout( + stage, + { + "tensor_parallel_size": 4, + "data_parallel_size": 4, + "data_parallel_size_local": 0, + "pipeline_parallel_size": 1, + }, + ) + + +@pytest.mark.parametrize( + ("devices", "expected"), + [ + ("0,1", ["0,1", "2,3"]), + ("0,1,2,3", ["0,1", "2,3"]), + ], +) +def test_replica_guard_and_splitter_share_local_world_size(devices, expected): + """PP and local DP contribute to the same per-replica size in both paths.""" + engine_args = { + "tensor_parallel_size": 1, + "data_parallel_size": 4, + "data_parallel_size_local": 1, + "pipeline_parallel_size": 2, + } + stage = _stage(0, devices=devices, num_replicas=2, engine_args=engine_args) + + _check_stage_device_layout(stage, engine_args) + assert get_stage_devices_per_replica(stage) == 2 + _, replica_devices_map = compute_replica_layout([stage]) + assert replica_devices_map == {0: expected} + + +def test_non_tp_layout_error_uses_generic_guidance(): + """A PP mismatch must not be reported as a top-level TP broadcast.""" + stage = _stage(0, devices="0") + with pytest.raises(ValueError) as excinfo: + _check_stage_device_layout( + stage, + { + "tensor_parallel_size": 1, + "data_parallel_size": 1, + "pipeline_parallel_size": 2, + }, + ) + + msg = str(excinfo.value) + assert "local world size" in msg + assert "top-level --tensor-parallel-size" not in msg + + +def test_build_vllm_config_fails_before_engine_config_on_mismatch(): + """Pin the production wiring: ``build_vllm_config`` must run the device-layout + guard *before* ``create_engine_config``. If the guard call is dropped or moved + after engine-config creation, #5003 regresses while the unit tests above stay + green — so assert here that a mismatched stage raises and that neither + ``create_engine_config`` nor ``Executor.get_class`` is reached.""" + stage = _stage(0, devices="0") + with ( + mock.patch.object(stage_init_utils.OmniEngineArgs, "create_engine_config") as create_engine_config, + mock.patch.object(stage_init_utils.Executor, "get_class") as get_class, + ): + with pytest.raises(ValueError) as excinfo: + build_vllm_config( + stage, + model="dummy-model", + engine_args_dict={ + "tensor_parallel_size": 4, + "data_parallel_size": 1, + "pipeline_parallel_size": 1, + }, + ) + msg = str(excinfo.value) + assert "Stage 0" in msg + assert "--stage-overrides" in msg + # Guard fired early: worker/config construction was never reached. + create_engine_config.assert_not_called() + get_class.assert_not_called() + + +def test_build_vllm_config_proceeds_on_consistent_layout(): + """The guard must not false-positive: a consistent single-GPU stage flows past + it and reaches ``create_engine_config`` / ``Executor.get_class`` as usual.""" + stage = _stage(1, devices="1") + fake_config = types.SimpleNamespace( + quant_config=None, + model_config=types.SimpleNamespace(hf_config=types.SimpleNamespace()), + ) + sentinel_executor = object() + with ( + mock.patch.object( + stage_init_utils.OmniEngineArgs, "create_engine_config", return_value=fake_config + ) as create_engine_config, + mock.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor), + mock.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant), + ): + vllm_config, executor_class = build_vllm_config( + stage, + model="dummy-model", + engine_args_dict={"tensor_parallel_size": 1}, + ) + create_engine_config.assert_called_once() + assert vllm_config is fake_config + assert executor_class is sentinel_executor diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 263804fd6..f24eda0a1 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Stage initialization helpers for vLLM-Omni multi-stage runtime. @@ -565,7 +568,7 @@ def _maybe_set_qwen3_omni_moe_env(engine_args_dict: dict[str, Any]) -> None: def split_devices_for_replicas( devices_str: str | None, num_replicas: int, - tp_size: int, + devices_per_replica: int, stage_id: int, ) -> list[str | None]: """Split a devices string into per-replica subsets. @@ -578,19 +581,19 @@ def split_devices_for_replicas( When ``num_replicas`` is 1, returns ``[devices_str]`` unchanged. Otherwise, two YAML shapes are accepted: - 1. **Legacy / pool mode** — ``len(devices) == num_replicas * tp_size``: + 1. **Legacy / pool mode** — ``len(devices) == num_replicas * devices_per_replica``: the string enumerates the full per-stage pool. Each replica gets - ``tp_size`` consecutive entries. The values are logical indices + ``devices_per_replica`` consecutive entries. The values are logical indices into the launcher's ``CUDA_VISIBLE_DEVICES``. ``split_devices_for_replicas("1,2,3,4", 2, 2, 1) → ["1,2", "3,4"]`` - 2. **Template mode** — ``len(devices) == tp_size``: the YAML declares + 2. **Template mode** — ``len(devices) == devices_per_replica``: the YAML declares a single per-replica template (the same shape one replica would - use), and is **dp-independent**. Each replica r gets the offsets - ``[r*tp_size + a for a in template]`` of the launcher's + use), and is **replica-count-independent**. Each replica r gets the offsets + ``[r*devices_per_replica + a for a in template]`` of the launcher's ``CUDA_VISIBLE_DEVICES``. The template's entries must lie in - ``[0, tp_size)``. + ``[0, devices_per_replica)``. ``split_devices_for_replicas("0,1", 2, 2, 1) → ["0,1", "2,3"]`` ``split_devices_for_replicas("0,1", 4, 2, 1) → ["0,1", "2,3", "4,5", "6,7"]`` @@ -613,27 +616,29 @@ def split_devices_for_replicas( device_list = [d.strip() for d in devices_str.split(",") if d.strip()] - if len(device_list) == num_replicas * tp_size: - return [",".join(device_list[r * tp_size : (r + 1) * tp_size]) for r in range(num_replicas)] + if len(device_list) == num_replicas * devices_per_replica: + return [ + ",".join(device_list[r * devices_per_replica : (r + 1) * devices_per_replica]) for r in range(num_replicas) + ] - if len(device_list) == tp_size: + if len(device_list) == devices_per_replica: try: offsets = [int(a) for a in device_list] except ValueError as e: raise ValueError(f"Stage {stage_id}: template-mode devices must be ints, got {devices_str!r}") from e - bad = [a for a in offsets if not (0 <= a < tp_size)] + bad = [a for a in offsets if not (0 <= a < devices_per_replica)] if bad: raise ValueError( f"Stage {stage_id}: template-mode device offset(s) {bad} " - f"out of range [0, {tp_size}); devices={devices_str!r}" + f"out of range [0, {devices_per_replica}); devices={devices_str!r}" ) - return [",".join(str(r * tp_size + a) for a in offsets) for r in range(num_replicas)] + return [",".join(str(r * devices_per_replica + a) for a in offsets) for r in range(num_replicas)] raise ValueError( f"Stage {stage_id}: devices={devices_str!r} has {len(device_list)} id(s); " - f"need either {tp_size} (template, dp-independent) or " - f"{num_replicas * tp_size} (pool / legacy). " - f"num_replicas={num_replicas}, tensor_parallel_size={tp_size}." + f"need either {devices_per_replica} (per-replica template) or " + f"{num_replicas * devices_per_replica} (pool / legacy). " + f"num_replicas={num_replicas}, devices_per_replica={devices_per_replica}." ) @@ -645,9 +650,32 @@ def get_stage_tp_size(stage_cfg: Any) -> int: return int(getattr(engine_args, "tensor_parallel_size", 1) or 1) -def get_stage_devices_per_replica(stage_cfg: Any) -> int: +def _get_local_llm_parallel_sizes( + stage_cfg: Any, + engine_args: Any | None = None, +) -> tuple[int, int, int]: + """Return ``(tp, local_dp, pp)`` for one local LLM replica. + + ``data_parallel_size`` is cluster-wide, whereas ``runtime.devices`` is + local to this process. Prefer an explicitly resolved + ``data_parallel_size_local`` (including zero for a head process that owns + no local engines), and only fall back to the global DP width when it is + unset. + """ + if engine_args is None: + engine_args = getattr(stage_cfg, "engine_args", {}) + tp_size = int(_get_attr_or_item(engine_args, "tensor_parallel_size", 1) or 1) + pp_size = int(_get_attr_or_item(engine_args, "pipeline_parallel_size", 1) or 1) + local_dp_size = _get_attr_or_item(engine_args, "data_parallel_size_local", None) + if local_dp_size is None: + local_dp_size = _get_attr_or_item(engine_args, "data_parallel_size", 1) + return tp_size, int(local_dp_size if local_dp_size is not None else 1), pp_size + + +def get_stage_devices_per_replica(stage_cfg: Any, engine_args: Any | None = None) -> int: """Return the number of devices consumed by one replica of *stage_cfg*.""" - engine_args = getattr(stage_cfg, "engine_args", {}) + if engine_args is None: + engine_args = getattr(stage_cfg, "engine_args", {}) if getattr(stage_cfg, "stage_type", "llm") == "diffusion": parallel_config = _get_attr_or_item(engine_args, "parallel_config") if parallel_config is None: @@ -664,14 +692,15 @@ def get_stage_devices_per_replica(stage_cfg: Any) -> int: except Exception: return 1 - return get_stage_tp_size(stage_cfg) + tp_size, local_dp_size, pp_size = _get_local_llm_parallel_sizes(stage_cfg, engine_args) + return tp_size * max(1, local_dp_size) * pp_size def compute_replica_layout( stage_configs: Sequence[Any], *, allow_zero: bool = False, -) -> tuple[list[int], dict[int, list[str]]]: +) -> tuple[list[int], dict[int, list[str | None]]]: """Compute per-stage replica counts and device assignments. Args: @@ -820,7 +849,7 @@ def _project_upstream_config_fields( field_map: Mapping[str, str], ) -> dict[str, Any]: """Project every explicit upstream input, including newly added fields.""" - explicit_fields = getattr(config, "_omni_explicit_fields", frozenset()) + explicit_fields: frozenset[str] = getattr(config, "_omni_explicit_fields", frozenset()) unprojected_fields = explicit_fields - frozenset(field_map) if unprojected_fields: names = ", ".join(sorted(unprojected_fields)) @@ -838,9 +867,9 @@ def _project_omni_stage_engine_args( """Read backend inputs from one structured stage config.""" engine_args: dict[str, Any] = {} is_diffusion = isinstance(stage_config, VllmOmniDiffusionStageConfig) - if is_diffusion: - engine_args.update(_project_omni_config_fields(stage_config.diffusion_config)) + diffusion_stage = cast(VllmOmniDiffusionStageConfig, stage_config) + engine_args.update(_project_omni_config_fields(diffusion_stage.diffusion_config)) for config, excluded_fields in ( ( @@ -1112,6 +1141,74 @@ def build_engine_args_dict_from_omni_stage_config( ) +def _count_stage_devices(devices: Any) -> int | None: + if devices is None: + return None + if isinstance(devices, (list, tuple)): + return len(devices) + values = [device for device in str(devices).split(",") if device.strip()] + return len(values) or None + + +def _check_stage_device_layout(stage_config: Any, engine_args_dict: dict[str, Any]) -> None: + """Fail early when a stage's world size cannot fit its assigned ``devices``. + + Re-runs :func:`check_device_layout` (normally only reached on the + ``--strategy-config`` path) against the fully resolved per-stage layout, so + an inconsistent ``tensor_parallel_size`` vs ``devices`` (issue #5003) is + reported here with a clear message instead of surfacing later as an opaque + worker-side ``local rank ... out of bounds`` assertion. + """ + from vllm_omni.config.composable_parallel import StrategyApplyError, check_device_layout + + runtime = getattr(stage_config, "runtime", None) + devices = _get_attr_or_item(runtime, "devices", None) if runtime is not None else None + if devices is None: + # No explicit placement -> vLLM assigns devices itself; nothing to check. + return + + num_replicas = _get_attr_or_item(runtime, "num_replicas", 1) if runtime is not None else 1 + stage_id = getattr(stage_config, "stage_id", "?") + tp_size, local_dp_size, pp_size = _get_local_llm_parallel_sizes(stage_config, engine_args_dict) + if local_dp_size == 0: + # This process hosts no local DP engines, so its local device list does + # not describe the cluster-wide DP layout and must not be validated. + return + + try: + check_device_layout( + devices, + tensor_parallel_size=tp_size, + data_parallel_size=local_dp_size, + pipeline_parallel_size=pp_size, + num_replicas=int(num_replicas or 1), + role=f"stage-{stage_id}", + ) + except StrategyApplyError as e: + message = ( + f"Stage {stage_id}: device layout is inconsistent — {e} " + "Set devices and the per-stage TP, local DP, PP, and replica counts " + "so the declared device count matches the local world size." + ) + device_count = _count_stage_devices(devices) + world_without_tp = local_dp_size * pp_size + valid_without_tp = {world_without_tp, int(num_replicas or 1) * world_without_tp} + if tp_size > 1 and device_count in valid_without_tp: + message += ( + " This layout is consistent with issue #5003: a top-level " + "--tensor-parallel-size is applied to every stage, but each stage's " + "`devices` is not adjusted automatically. Pass --stage-overrides " + "to set tensor_parallel_size and devices together on every stage, " + "so single-GPU stages get tensor_parallel_size=1, e.g. " + '\'{"0": {"tensor_parallel_size": 4, "devices": "0,1,2,3"}, ' + '"1": {"tensor_parallel_size": 1, "devices": "0"}, ' + '"2": {"tensor_parallel_size": 1, "devices": "1"}}\'. ' + "Or omit the top-level --tensor-parallel-size and set it only in " + "stage-0's override." + ) + raise ValueError(message) from e + + def build_vllm_config( stage_config: Any, model: str, @@ -1147,6 +1244,16 @@ def build_vllm_config( omni_engine_args = OmniEngineArgs(**filtered_engine_args_dict) + # Guard against a per-stage world size that its assigned ``devices`` cannot + # satisfy (issue #5003). A top-level ``--tensor-parallel-size`` is broadcast + # to every stage, but ``devices`` is not, so a stage can end up with e.g. + # tensor_parallel_size=4 while still holding a single-GPU deploy default. + # Without --strategy-config the strategy-path device check never runs, so + # the mismatch used to surface only as an opaque worker-side assertion + # ("DP adjusted local rank N is out of bounds for M devices."). Re-run the + # same check here, before workers spawn, to fail early with a clear message. + _check_stage_device_layout(stage_config, filtered_engine_args_dict) + # Multi-stage pipelines (qwen3_tts code2wav, etc.) set max_model_len # larger than HF max_position_embeddings by design. vLLM's validator # rejects that without the env flag. @@ -1473,6 +1580,7 @@ def build_diffusion_config( num_devices_per_stage = od_config.parallel_config.world_size device_control_env = current_omni_platform.device_control_env_var visible_devices_str = os.environ.get(device_control_env) if device_control_env else None + physical_devices: list[str | int] if visible_devices_str: physical_devices = [device.strip() for device in visible_devices_str.split(",") if device.strip()] else: