From ff6f0da271bd1c3658c3e46e4e2cdc4db17219a5 Mon Sep 17 00:00:00 2001 From: Kushanam <42385577+kushanam@users.noreply.github.com> Date: Fri, 10 Jul 2026 22:29:24 -0500 Subject: [PATCH] [Core] Composable parallel: Phase 1 strategy overlay (TP/DP/stage-replica) (RFC #4084) (#4281) Signed-off-by: Kushan A. Signed-off-by: kushanam Signed-off-by: Kushan A Signed-off-by: kahmadian Co-authored-by: kushanam Co-authored-by: kahmadian --- docs/configuration/composable_parallel.md | 239 +++++++++++ .../offline_inference/qwen2_5_omni/README.md | 21 + .../offline_inference/qwen2_5_omni/end2end.py | 33 ++ .../qwen2_5_omni/strategy_stage_replica.yaml | 22 + .../qwen2_5_omni/strategy_tp2.yaml | 19 + .../config/composable_parallel/test_apply.py | 154 +++++++ .../test_factory_integration.py | 141 +++++++ .../composable_parallel/test_spec_types.py | 42 ++ .../test_strategy_loader.py | 97 +++++ .../composable_parallel/test_translator.py | 195 +++++++++ .../test_async_omni_engine_stage_init.py | 6 +- tests/engine/test_strategy_lb_policy.py | 63 +++ .../test_async_omni_diffusion_config.py | 2 +- .../test_resolve_dreamzero_config.py | 2 +- tests/entrypoints/test_serve.py | 109 ++++- tests/entrypoints/test_utils.py | 7 +- tests/test_config_factory.py | 5 +- tests/utils/test_tracking_parser.py | 10 +- .../config/composable_parallel/__init__.py | 104 +++++ .../config/composable_parallel/aggregation.py | 126 ++++++ vllm_omni/config/composable_parallel/apply.py | 264 ++++++++++++ .../config/composable_parallel/routing.py | 138 +++++++ vllm_omni/config/composable_parallel/spec.py | 136 +++++++ .../composable_parallel/strategy_loader.py | 149 +++++++ .../config/composable_parallel/translator.py | 378 ++++++++++++++++++ vllm_omni/config/config_factory.py | 153 ++++++- vllm_omni/engine/arg_utils.py | 3 + vllm_omni/engine/async_omni_engine.py | 55 ++- vllm_omni/entrypoints/cli/serve.py | 23 +- vllm_omni/entrypoints/utils.py | 77 +++- 30 files changed, 2722 insertions(+), 51 deletions(-) create mode 100644 docs/configuration/composable_parallel.md create mode 100644 examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml create mode 100644 examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml create mode 100644 tests/config/composable_parallel/test_apply.py create mode 100644 tests/config/composable_parallel/test_factory_integration.py create mode 100644 tests/config/composable_parallel/test_spec_types.py create mode 100644 tests/config/composable_parallel/test_strategy_loader.py create mode 100644 tests/config/composable_parallel/test_translator.py create mode 100644 tests/engine/test_strategy_lb_policy.py create mode 100644 vllm_omni/config/composable_parallel/__init__.py create mode 100644 vllm_omni/config/composable_parallel/aggregation.py create mode 100644 vllm_omni/config/composable_parallel/apply.py create mode 100644 vllm_omni/config/composable_parallel/routing.py create mode 100644 vllm_omni/config/composable_parallel/spec.py create mode 100644 vllm_omni/config/composable_parallel/strategy_loader.py create mode 100644 vllm_omni/config/composable_parallel/translator.py diff --git a/docs/configuration/composable_parallel.md b/docs/configuration/composable_parallel.md new file mode 100644 index 000000000..400fbab38 --- /dev/null +++ b/docs/configuration/composable_parallel.md @@ -0,0 +1,239 @@ +# Composable parallel strategies + +In vLLM-Omni, the `composable_parallel` system layers a declarative *parallel strategy* on top of the existing stage configs. A strategy YAML maps each pipeline stage (by its `model_stage` name) to a stack of axis declarations — one per parallelism dimension — that the engine translates into per-stage parallel sizing (`tensor_parallel_size`, `data_parallel_size`, `pipeline_parallel_size`, expert-parallel toggle, replica count) and, when relevant, an orchestrator-wide load-balancer policy. + +Strategies are intentionally narrow: each one only writes the axes it declares, leaving every other engine arg untouched. This makes a strategy file a small, reusable overlay you can swap independently of the deploy YAML (see `configuration/stage_configs.md`). + +!!! note + Composable parallel is **opt-in**. Passing `--strategy-config ` overlays the derived sizing onto the registry-merged stages *before* any CLI overrides are applied. When `--strategy-config` is omitted, stages remain exactly as their deploy YAML produces them. + +## Top-level schema reference + +A strategy file has a single top-level mapping, `strategies`, keyed by the `model_stage` name of each stage you want to overlay. Each value is a list of axis declarations (see [StrategySpec fields](#strategyspec-fields) below). + +| Field | Type | Required | Default | Description | +|-------|------|----------|---------|-------------| +| `strategies` | `dict[str, list]` | required | — | Maps `model_stage` name (e.g. `thinker`, `talker`, `code2wav`, `dit`) to a list of axis declarations. The key MUST be the stage name, not the integer `stage_id`. See `vllm_omni/config/composable_parallel/strategy_loader.py`. | + +### StrategySpec fields + +Each entry in a stage's list becomes one `StrategySpec` (see `vllm_omni/config/composable_parallel/spec.py`). The loader fills in `routing` and `aggregation` from per-kind defaults, so the file stays compact. + +| Field | Type | Required | Default | Description | +|-------|------|----------|---------|-------------| +| `axis` | str | required | — | Mesh-axis kind. One of `tp`, `dp`, `pp`, `ep`, `stage_replica` (wired today) or `sp_ulysses`, `sp_ring`, `cfg`, `vae_pp`, `hsdp`, `stage_pp`, `cp` (reserved; declaring one raises `AxisTranslationError` at translate time). See `MeshAxisKind` in `vllm_omni/config/composable_parallel/spec.py`. | +| `size` | int | required | — | Axis degree. Must be `> 0`. | +| `routing` | str | optional | per-kind default | Stateless routing policy. Accepted only for `dp` / `stage_replica`; other kinds reject it. Values: `random`, `round_robin`, `least_queue`. (`hash` is recognised but currently raises `NotImplementedError`.) See `RouteByStage` in `vllm_omni/config/composable_parallel/routing.py`. | +| `l1_owner` | str | optional | per-kind default | Overrides which layer owns request routing for this axis. `engine` for `tp`/`dp`/`pp`/`ep`, `delegated` for `stage_replica`. Recorded into `shard_extension`. | +| `name` | str | optional | axis kind | Human-readable label for this spec, used in error messages. | + +Two additional `StrategySpec` slots — `layer_hook_specs` and `kernel_specs` — let advanced authors bind layer-walker hook IDs (`LayerHookSpec`) or kernel IDs (`KernelSpec`) onto an axis. These are not exposed by the YAML loader today and are intended for programmatic construction. See `vllm_omni/config/composable_parallel/spec.py` for the dataclasses. + +### CLI flags + +| Flag | Description | +|------|-------------| +| `--strategy-config PATH` | Path to a `strategy.yaml`. Only takes effect on registry-based deploy paths (`--deploy-config` or the bundled default deploy YAML). Its derived sizing is overlaid onto the merged stages *before* CLI overrides. | +| `--stage-configs-path PATH` | **Legacy.** Loads the old `stage_args` YAML schema for un-migrated models. Passing `--strategy-config` together with this path is silently inapplicable — vLLM-Omni emits a warning and the strategy is ignored (see `vllm_omni/entrypoints/utils.py`, lines 372–380). | +| `--omni-lb-policy POLICY` | Orchestrator-wide load-balancer policy. Strategy-derived from any `stage_replica` axis; an explicit CLI value wins and must match the derived one or `AsyncOmniEngine` raises `Conflicting load-balancer policy` at construction. | + +!!! warning + `--strategy-config` cannot be used with `--stage-configs-path` (the legacy deploy path). If your model has not been migrated to the registry-based pipeline yet, the strategy is silently dropped with a warning. Migrate to a registry-based model or use `--deploy-config` to apply a strategy. + +## Sub-schemas + +### `axis` + +The set of recognised axis kinds is the `MeshAxisKind` enum in `vllm_omni/config/composable_parallel/spec.py`: + +```yaml +# Translatable today (see `_SUPPORTED_KINDS` in translator.py). +axis: tp # tensor parallel — sets tensor_parallel_size +axis: dp # engine data parallel — sets data_parallel_size (vLLM's intra-engine DP LB) +axis: pp # pipeline parallel — sets pipeline_parallel_size +axis: ep # dense expert parallel — sets enable_expert_parallel=True; size must equal tp*dp +axis: stage_replica # omni-coordinator-level replica fan-out — sets num_replicas + omni_lb_policy + +# Reserved (raise AxisTranslationError today; declared in the type system for forward compat). +axis: sp_ulysses +axis: sp_ring +axis: cfg +axis: vae_pp +axis: hsdp +axis: stage_pp +axis: cp +``` + +### `routing` + +The full closed hierarchy lives in `vllm_omni/config/composable_parallel/routing.py` (`Broadcast`, `PartitionByHash`, `ShardSequence`, `ShardSpatial`, `RouteByExpert`, `PipelineMicrobatch`, `DuplicateWithCondTag`, `RouteByStage`). For YAML authors only `RouteByStage` is exposed as a `routing:` string, and only for `dp` / `stage_replica` axes (other axes have fixed routing — `tp`/`ep` broadcast, `pp` microbatch). Recognised policies: + +```yaml +routing: random # uniform-random across replicas/DP ranks +routing: round_robin # cycle through replicas/DP ranks (the default for dp/stage_replica) +routing: least_queue # send to the replica with the smallest backlog +# routing: hash # recognised but raises NotImplementedError; no key-stable balancer yet +``` + +### `aggregation` + +Aggregation is filled in automatically by the loader (`_default_aggregation` in `vllm_omni/config/composable_parallel/strategy_loader.py`) from the axis kind, so users do not author it: + +| Axis kind | Aggregation pattern | +|-----------|---------------------| +| `tp`, `ep` | `TakeRank` (all ranks agree; pick one) | +| `dp`, ep (in some flows) | `Union` (disjoint replica results) | +| `pp` | `StitchPipeline` | +| `stage_replica` | `FanInByStage` | + +Full pattern hierarchy: `TakeRank`, `Union`, `GatherDim`, `AllGather`, `StitchSpatial`, `StitchPipeline`, `Combine`, `FanInByStage` (see `vllm_omni/config/composable_parallel/aggregation.py`). + +## Precedence + +From highest to lowest, when the same axis is touched by multiple layers: + +1. **CLI overrides** (`--stage_N_tensor_parallel_size`, `--stage-overrides`, `--omni-lb-policy`, etc.). A CLI value always wins; a warning is emitted whenever it overrides a strategy-declared axis (see `_reconcile_strategy_with_cli` in `vllm_omni/config/config_factory.py`). +2. **Strategy YAML** (`--strategy-config`). Derived sizing is written onto the merged stages with *conflict-on-explicit* semantics — if the deploy YAML already set a value to something different, application raises `StrategyApplyError`. +3. **Deploy YAML** (`--deploy-config` or the bundled `vllm_omni/deploy/.yaml`). +4. **Parser defaults.** + +After CLI overrides land, the device-layout guard re-runs (`check_device_layout`, called from `_reconcile_strategy_with_cli`) against the *effective* `tp * dp * pp * num_replicas` and the resolved `devices` string, so a `--devices` or `--tensor-parallel-size` that conflicts with the strategy's world size cannot slip past the pre-spawn check. + +## Worked example + +### Example 1 — `strategy_tp2.yaml` + +Shard the `thinker` stage with TP=2 on top of the bundled Qwen2.5-Omni deploy YAML. Routing is fixed (broadcast) for TP, so no `routing:` field is needed. + +```yaml +# examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml +strategies: + thinker: + - axis: tp + size: 2 +``` + +Launch: + +```bash +python3 end2end.py --query-type text --num-prompts 6 \ + --strategy-config strategy_tp2.yaml \ + --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}' +``` + +Effective `OmniParallelConfig` per stage after the overlay (see `OmniParallelConfig` in `vllm_omni/config/composable_parallel/translator.py`): + +| Stage (`model_stage`) | `tensor_parallel_size` | `data_parallel_size` | `pipeline_parallel_size` | `stage_replica_size` | `omni_lb_policy` | +|-----------------------|------------------------|----------------------|--------------------------|----------------------|------------------| +| `thinker` | **2** (from strategy) | 1 | 1 | 1 | — | +| `talker` | 1 | 1 | 1 | 1 | — | +| `code2wav` | 1 | 1 | 1 | 1 | — | + +### Example 2 — `strategy_stage_replica.yaml` + +Replicate the `thinker` stage across two independent engines and let omni's StagePool round-robin requests across them. + +```yaml +# examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml +strategies: + thinker: + - axis: stage_replica + size: 2 + routing: round_robin +``` + +Launch: + +```bash +python3 end2end.py --query-type text --num-prompts 6 \ + --strategy-config strategy_stage_replica.yaml \ + --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}' +``` + +Effective per-stage config and the derived orchestrator knob: + +| Stage (`model_stage`) | `tensor_parallel_size` | `stage_replica_size` | `omni_lb_policy` (derived) | +|-----------------------|------------------------|----------------------|----------------------------| +| `thinker` | 1 | **2** | **`round-robin`** | +| `talker` | 1 | 1 | — | +| `code2wav` | 1 | 1 | — | + +The derived `omni_lb_policy = round-robin` is **not** written into any stage's engine args. Instead it is surfaced on `StrategyApplyResult.omni_lb_policy` and applied once at orchestrator construction by `AsyncOmniEngine._apply_strategy_lb_policy` (see `vllm_omni/engine/async_omni_engine.py`, lines 1107–1134). If you also pass `--omni-lb-policy round-robin` the values must match; passing a different value raises `Conflicting load-balancer policy`. + +## Per-field deep-dive + +### `name` + +Default: the axis kind (e.g. `"tp"`). + +Human-readable label that appears in `AxisTranslationError` and `StrategyApplyError` messages. Set it to disambiguate two stages declaring the same axis kind when scanning logs. + +### `mesh_axis` (`axis` + `size`) + +Default: no default — both fields are required. + +`axis` chooses the parallelism dimension; `size` is its degree. Five kinds are wired end-to-end today (`tp`, `dp`, `pp`, `ep`, `stage_replica`); the remaining kinds in `MeshAxisKind` are reserved and raise `AxisTranslationError` at translate time so a typo cannot silently become a no-op (see `_SUPPORTED_KINDS` in `vllm_omni/config/composable_parallel/translator.py`). For `ep`, the size MUST equal `tensor_parallel_size * data_parallel_size`. + +### `routing` + +Default: per-kind — `RouteByStage(round_robin)` for `dp` and `stage_replica`; `PipelineMicrobatch()` for `pp`; `Broadcast()` for `tp`/`ep` and any other kind. + +User-facing string only on `dp` and `stage_replica` axes. The four recognised values are `random`, `round_robin`, `least_queue`, `hash`; `hash` is rejected today with `NotImplementedError`. See `RouteByStage` in `vllm_omni/config/composable_parallel/routing.py`. + +### `l1_owner` + +Default: per-kind — `engine` for `tp`/`dp`/`pp`/`ep`, `delegated` for `stage_replica`. + +Overrides which layer owns request routing along this axis. `engine` means vLLM realizes routing intra-engine; `delegated` means the omni coordinator owns routing across independent replicas. Set this only when overriding the per-kind default in `_DEFAULT_L1_OWNER` (`vllm_omni/config/composable_parallel/translator.py`); invalid values fail fast in `_resolve_l1_owner`. + +### `stage_replica` and `omni_lb_policy` + +`stage_replica` is *not* a vLLM world dimension; it spins up `size` independent engine replicas of one stage, coordinated by the omni `StagePool` load balancer. The axis writes both: + +- The stage's `num_replicas` (a per-stage deploy knob). +- The pipeline-wide `omni_lb_policy` string (an orchestrator-level knob). + +Because `omni_lb_policy` is read once at `AsyncOmniEngine` construction — not from per-stage engine args — only one value is allowed across the whole pipeline. If two stages declare conflicting policies, application raises (see `apply_strategy_specs` in `vllm_omni/config/composable_parallel/apply.py`, lines 253–260). To set this knob manually instead, use `--omni-lb-policy` on the CLI; see [the model stage docs](stage_configs.md) for how stage replicas plug into the orchestrator at runtime. + +### `model_stage` keys + +Default: no default — required for every entry under `strategies:`. + +Strategy keys MUST be `model_stage` *names* (strings), the same labels stage configs use for `engine_args.model_stage` (see [`model_stage` in `stage_configs.md`](stage_configs.md#engine_argsmodel_stage)). Integer `stage_id` values are not accepted; the resolver matches on `model_stage` exactly (`_resolve_stage` in `vllm_omni/config/composable_parallel/apply.py`). + +## Common pitfalls + +!!! warning "Using `stage_id` integers as keys" + Strategy keys are `model_stage` names (str), not integer stage IDs. A file that uses an integer index will raise at apply time: + + ```text + StrategyApplyError: strategy role model_stage=0 did not match any stage; + available stages: 'thinker'(id=0), 'talker'(id=1), 'code2wav'(id=2) + ``` + + Use the `model_stage` name (e.g. `thinker`) instead. + +!!! warning "Mixing CLI overrides with a strategy on the same axis" + A strategy is meant to be the single writer for the axes it declares. If a CLI override (`--stage_0_tensor_parallel_size`, `--stage-overrides '{"0": {...}}'`) sets the same axis, the **CLI value wins** and a warning is emitted from `_reconcile_strategy_with_cli`. The device-layout guard then re-runs against the effective layout, so a CLI override that breaks `tp * dp * pp * num_replicas` will still fail fast at config time. + +!!! warning "Passing `--strategy-config` against the legacy `--stage-configs-path`" + Composable parallel only applies on the registry-based deploy path. If a model still resolves through the legacy `stage_args` YAML (loaded via `--stage-configs-path`), `--strategy-config` is silently inapplicable and emits: + + ```text + --strategy-config () was provided but model '' resolves via the + legacy stage_configs YAML path, which does not support composable-parallel + strategies; the strategy is ignored. Use a registry-based model to apply it. + ``` + + Migrate the model to the new deploy schema (see `configuration/stage_configs.md`) before applying a strategy. + +!!! warning "Conflicting `omni_lb_policy` across stages" + `omni_lb_policy` is a single pipeline-wide knob. If two `stage_replica` axes in different stages derive different policies, `apply_strategy_specs` raises: + + ```text + StrategyApplyError: role 'talker' derives omni_lb_policy='least-queue-length' + but role 'thinker' already derived 'round-robin'. The load-balancer policy is + pipeline-wide; only one value is allowed. + ``` + + Make the policies match, drop one, or supply the value explicitly via `--omni-lb-policy`. diff --git a/examples/offline_inference/qwen2_5_omni/README.md b/examples/offline_inference/qwen2_5_omni/README.md index e2eae8a96..292045033 100644 --- a/examples/offline_inference/qwen2_5_omni/README.md +++ b/examples/offline_inference/qwen2_5_omni/README.md @@ -60,3 +60,24 @@ If media file paths are not provided, the script will use default assets. Suppor - `mixed_modalities`: Audio + image + video - `use_audio_in_video`: Extract audio from video - `text`: Text-only query + +### Composable parallelism (strategy configs) + +You can shard or replicate a stage with a small composable-parallel +`strategy.yaml` overlaid onto the **bundled default deploy config** -- no bespoke +deploy YAML required. Supply the strategy with `--strategy-config` and any matching +device layout via `--stage-overrides` (explicit kwargs forwarded to `Omni`). + +Both examples below require >= 3 GPUs (the thinker uses 2, talker + code2wav share a third): + +```bash +# Tensor-parallel the thinker (TP=2) +python end2end.py --query-type text --num-prompts 6 \ + --strategy-config strategy_tp2.yaml \ + --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}' + +# Replicate the thinker across 2 engines (round-robin load balancing) +python end2end.py --query-type text --num-prompts 6 \ + --strategy-config strategy_stage_replica.yaml \ + --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}' +``` diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index 8290d27b5..d8769b5f7 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -553,6 +553,39 @@ def parse_args(): default=False, help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.", ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Optional explicit deploy YAML (otherwise the bundled default is used).", + ) + parser.add_argument( + "--strategy-config", + type=str, + default=None, + help=( + "Optional composable-parallel strategy.yaml mapping role -> parallel axes " + "(e.g. tp / stage_replica). Overlaid onto the merged stage configs." + ), + ) + parser.add_argument( + "--stage-overrides", + type=str, + default=None, + help=( + "Optional JSON of per-stage overrides applied on top of the default " + 'deploy config, e.g. \'{"0": {"devices": "0,1"}}\' to give the ' + "thinker a 2-GPU pool. Lets you run a strategy on the bundled default " + "deploy config without writing a bespoke deploy YAML." + ), + ) + parser.add_argument( + "--omni-lb-policy", + type=str, + default=None, + choices=["random", "round-robin", "least-queue-length"], + help="StagePool load-balancer policy for replicated stages (orchestrator-level knob).", + ) return parser.parse_args() diff --git a/examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml b/examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml new file mode 100644 index 000000000..f0e2d9bf5 --- /dev/null +++ b/examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml @@ -0,0 +1,22 @@ +# Composable-parallel strategy for Qwen2.5-Omni. +# +# Replicate the thinker stage across 2 independent engines and let the omni +# StagePool balance requests round-robin across them. The thinker's `devices` +# entry acts as a per-replica template, so omni expands it to one GPU per replica. +# The `routing: round_robin` below auto-derives the StagePool load-balancer policy, +# so --omni-lb-policy is not needed. +# +# Rather than hand-write a bespoke deploy YAML, run on the bundled default deploy +# config and pass the matching 2-GPU device template via --stage-overrides (explicit +# kwargs): the thinker gets a two-GPU pool while talker + code2wav share a third +# card. Requires >= 3 GPUs. +# +# Usage: +# python3 end2end.py --query-type text --num-prompts 6 \ +# --strategy-config strategy_stage_replica.yaml \ +# --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}' +strategies: + thinker: + - axis: stage_replica + size: 2 + routing: round_robin diff --git a/examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml b/examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml new file mode 100644 index 000000000..02d773952 --- /dev/null +++ b/examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml @@ -0,0 +1,19 @@ +# Composable-parallel strategy for Qwen2.5-Omni: tensor-parallel the thinker. +# +# Shard the thinker stage across 2 GPUs with tensor parallelism (TP=2). TP routing +# is fixed (broadcast), so no `routing` / load-balancer policy is involved -- this +# is the simplest axis and a good starting point for the mental model. +# +# The strategy sets tensor_parallel_size=2 on the thinker. Rather than hand-write a +# bespoke deploy YAML, run on the bundled default deploy config and pass a matching +# 2-GPU device layout via --stage-overrides (explicit kwargs), so the thinker gets a +# two-GPU pool while talker + code2wav share a third card. Requires >= 3 GPUs. +# +# Usage: +# python3 end2end.py --query-type text --num-prompts 6 \ +# --strategy-config strategy_tp2.yaml \ +# --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}' +strategies: + thinker: + - axis: tp + size: 2 diff --git a/tests/config/composable_parallel/test_apply.py b/tests/config/composable_parallel/test_apply.py new file mode 100644 index 000000000..044b2cb02 --- /dev/null +++ b/tests/config/composable_parallel/test_apply.py @@ -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")], + }, + ) diff --git a/tests/config/composable_parallel/test_factory_integration.py b/tests/config/composable_parallel/test_factory_integration.py new file mode 100644 index 000000000..ea685541c --- /dev/null +++ b/tests/config/composable_parallel/test_factory_integration.py @@ -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) diff --git a/tests/config/composable_parallel/test_spec_types.py b/tests/config/composable_parallel/test_spec_types.py new file mode 100644 index 000000000..90b00cf95 --- /dev/null +++ b/tests/config/composable_parallel/test_spec_types.py @@ -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 diff --git a/tests/config/composable_parallel/test_strategy_loader.py b/tests/config/composable_parallel/test_strategy_loader.py new file mode 100644 index 000000000..01e202d05 --- /dev/null +++ b/tests/config/composable_parallel/test_strategy_loader.py @@ -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" diff --git a/tests/config/composable_parallel/test_translator.py b/tests/config/composable_parallel/test_translator.py new file mode 100644 index 000000000..b55c10e9e --- /dev/null +++ b/tests/config/composable_parallel/test_translator.py @@ -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}" diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 31850401b..ff467dc65 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -739,7 +739,7 @@ def test_resolve_stage_configs_injects_global_diffusion_attention_when_missing(m monkeypatch.setattr( engine_mod, "load_and_resolve_stage_configs", - lambda *args, **kwargs: ("dummy-config", [stage_cfg]), + lambda *args, **kwargs: ("dummy-config", [stage_cfg], None), ) _config_path, stage_configs = engine._resolve_stage_configs( @@ -772,7 +772,7 @@ def test_resolve_stage_configs_preserves_stage_diffusion_attention(monkeypatch): monkeypatch.setattr( engine_mod, "load_and_resolve_stage_configs", - lambda *args, **kwargs: ("dummy-config", [stage_cfg]), + lambda *args, **kwargs: ("dummy-config", [stage_cfg], None), ) _config_path, stage_configs = engine._resolve_stage_configs( @@ -798,7 +798,7 @@ def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stag monkeypatch.setattr( engine_mod, "load_and_resolve_stage_configs", - lambda *args, **kwargs: ("dummy-config", [stage_cfg]), + lambda *args, **kwargs: ("dummy-config", [stage_cfg], None), ) _config_path, stage_configs = engine._resolve_stage_configs( diff --git a/tests/engine/test_strategy_lb_policy.py b/tests/engine/test_strategy_lb_policy.py new file mode 100644 index 000000000..8531c2a99 --- /dev/null +++ b/tests/engine/test_strategy_lb_policy.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for strategy-derived omni_lb_policy threading into AsyncOmniEngine. + +``AsyncOmniEngine._apply_strategy_lb_policy`` only touches ``self._omni_lb_policy`` +and module-level logging, so we exercise it with a lightweight stub ``self`` to +avoid standing up a full engine (which needs GPUs/model weights). +""" + +from __future__ import annotations + +import pytest + +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Stub: + def __init__(self, current: str = "random") -> None: + self._omni_lb_policy = current + + +_apply = AsyncOmniEngine._apply_strategy_lb_policy + + +def test_no_derived_policy_is_noop(): + stub = _Stub("random") + _apply(stub, None, {}) + assert stub._omni_lb_policy == "random" + + +def test_derived_overrides_default_random(): + # "random" is the engine default == "unset", so the strategy value wins. + stub = _Stub("random") + _apply(stub, "round-robin", {}) + assert stub._omni_lb_policy == "round-robin" + + +def test_derived_overrides_when_flag_absent(): + stub = _Stub("random") + _apply(stub, "least-queue-length", {"omni_lb_policy": None}) + assert stub._omni_lb_policy == "least-queue-length" + + +def test_explicit_matching_flag_is_fine(): + stub = _Stub("round-robin") + _apply(stub, "round-robin", {"omni_lb_policy": "round-robin"}) + assert stub._omni_lb_policy == "round-robin" + + +def test_explicit_conflicting_flag_raises(): + stub = _Stub("round-robin") + with pytest.raises(ValueError, match="Conflicting load-balancer policy"): + _apply(stub, "least-queue-length", {"omni_lb_policy": "round-robin"}) + + +def test_explicit_random_is_treated_as_unset(): + # An explicit --omni-lb-policy=random is indistinguishable from the default, + # so a strategy value still overrides it (documented Phase-1 semantics). + stub = _Stub("random") + _apply(stub, "round-robin", {"omni_lb_policy": "random"}) + assert stub._omni_lb_policy == "round-robin" diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 61a420a4c..1f18d975b 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -343,7 +343,7 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo ) mocker.patch( "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", - return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage]), + return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage], None), ) engine = AsyncOmniEngine.__new__(AsyncOmniEngine) diff --git a/tests/entrypoints/test_resolve_dreamzero_config.py b/tests/entrypoints/test_resolve_dreamzero_config.py index e699edde4..4dcaddd63 100644 --- a/tests/entrypoints/test_resolve_dreamzero_config.py +++ b/tests/entrypoints/test_resolve_dreamzero_config.py @@ -32,7 +32,7 @@ def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch): lambda _model: True, ) - stage_configs = load_stage_configs_from_model("GEAR-Dreams/DreamZero-DROID") + stage_configs, _ = load_stage_configs_from_model("GEAR-Dreams/DreamZero-DROID") engine_args = stage_configs[0].engine_args assert engine_args.model_class_name == "DreamZeroPipeline" diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 6c4445306..61ffed3d5 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -9,6 +9,7 @@ import pytest from pytest_mock import MockerFixture from vllm_omni.entrypoints.cli.serve import OmniServeCommand, run_headless +from vllm_omni.entrypoints.utils import parse_stage_overrides from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -31,6 +32,20 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None: assert not explicit["async_chunk"] +def test_serve_parser_accepts_strategy_config() -> None: + """``--strategy-config`` must parse onto the ``strategy_config`` dest and be + forwarded as an explicit kwarg so the engine can overlay the strategy.""" + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + cmd = OmniServeCommand() + cmd.subparser_init(subparsers) + + argv = ["serve", "fake-model", "--omni", "--strategy-config", "/tmp/strategy.yaml"] + args = parser.parse_args(argv) + assert args.strategy_config == "/tmp/strategy.yaml" + assert args.get_explicit_kwargs_dict()["strategy_config"] == "/tmp/strategy.yaml" + + def _make_headless_args(**kwargs) -> TrackingNamespace: defaults = { "model": "fake-model", @@ -80,13 +95,97 @@ def test_run_headless_rejects_non_multiprocess_worker_backend() -> None: run_headless(args) +# --------------------------------------------------------------------------- +# --stage-overrides parsing parity (headless vs standard path) +# --------------------------------------------------------------------------- + + +def test_parse_stage_overrides_valid_json() -> None: + """A valid JSON string is parsed into the nested per-stage dict.""" + parsed = parse_stage_overrides('{"0": {"devices": "0,1"}, "1": {"devices": "2"}}') + assert parsed == {"0": {"devices": "0,1"}, "1": {"devices": "2"}} + + +def test_parse_stage_overrides_none_and_empty_return_none() -> None: + """No overrides (None / empty string) resolve to ``None``.""" + assert parse_stage_overrides(None) is None + assert parse_stage_overrides("") is None + + +def test_parse_stage_overrides_empty_dict_returns_none() -> None: + """An empty dict is falsy and must resolve to ``None``, locking in parity + with the original standard-path ``if stage_overrides_json:`` falsy check.""" + assert parse_stage_overrides({}) is None + + +def test_parse_stage_overrides_passes_through_non_str() -> None: + """An already-parsed mapping is returned unchanged (identity).""" + overrides = {"0": {"devices": "0"}} + assert parse_stage_overrides(overrides) is overrides + + +def test_parse_stage_overrides_invalid_json_raises() -> None: + """Invalid JSON raises ValueError whose message matches the standard path + verbatim: the ``--stage-overrides is not valid JSON:`` prefix AND the + ``Got: `` suffix echoing the raw input.""" + bad = "{not valid json}" + with pytest.raises(ValueError) as excinfo: + parse_stage_overrides(bad) + message = str(excinfo.value) + assert message.startswith("--stage-overrides is not valid JSON:") + assert f"Got: {bad!r}" in message + + +def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture) -> None: + """Regression: the headless path must parse ``--stage-overrides`` (a JSON + string) and forward the parsed dict to ``load_and_resolve_stage_configs``, + mirroring the standard engine path. Previously it was dropped entirely, + silently producing a different per-stage device layout.""" + captured: dict = {} + + def _fake_resolve(*args, **kwargs): + captured.update(kwargs) + # Return a stage that does NOT match stage_id=0 so run_headless stops + # right after the resolver call (we only care about how it was called). + return ("/fake/stages.yaml", [SimpleNamespace(stage_id=99)], None) + + mocker.patch( + "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + side_effect=_fake_resolve, + ) + + args = _make_headless_args( + stage_id=0, + strategy_config="/tmp/strategy.yaml", + stage_overrides='{"0": {"devices": "0,1"}, "1": {"devices": "2"}}', + ) + with pytest.raises(ValueError, match="No stage config found for stage_id=0"): + run_headless(args) + + assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}} + assert captured["strategy_config_path"] == "/tmp/strategy.yaml" + + +def test_run_headless_invalid_stage_overrides_raises(mocker: MockerFixture) -> None: + """Invalid ``--stage-overrides`` JSON in headless mode fails fast with the + shared ValueError instead of being silently ignored.""" + mocker.patch( + "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + return_value=("/fake/stages.yaml", [SimpleNamespace(stage_id=0)], None), + ) + + args = _make_headless_args(stage_id=0, stage_overrides="{not valid json}") + with pytest.raises(ValueError, match="--stage-overrides is not valid JSON"): + run_headless(args) + + def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture) -> None: """Headless looks up its assigned stage_id in the loaded deploy YAML and fails fast when the launcher's --stage-id doesn't match any entry.""" other_stage = SimpleNamespace(stage_id=99) mocker.patch( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [other_stage]), + return_value=("/fake/stages.yaml", [other_stage], None), ) args = _make_headless_args(stage_id=0) @@ -134,7 +233,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker mocker.patch( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [stage_cfg]), + return_value=("/fake/stages.yaml", [stage_cfg], None), ) mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None) @@ -210,7 +309,7 @@ def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: Mo mocker.patch( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [stage_cfg]), + return_value=("/fake/stages.yaml", [stage_cfg], None), ) mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None) @@ -269,7 +368,7 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) mocker.patch( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [stage_cfg]), + return_value=("/fake/stages.yaml", [stage_cfg], None), ) mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None) @@ -349,7 +448,7 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur mocker.patch( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [stage_cfg]), + return_value=("/fake/stages.yaml", [stage_cfg], None), ) mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None) diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 479ada963..b48ed2d91 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -318,7 +318,7 @@ class TestLoadAndResolveStageConfigs: def test_load_and_resolve_with_kwargs(self): """Ensure that dtype survives default stage creation.""" kwargs = {"dtype": torch.float32} - config_path, stage_configs = load_and_resolve_stage_configs( + config_path, stage_configs, _ = load_and_resolve_stage_configs( model="black-forest-labs/FLUX.2-klein-4B", stage_configs_path=None, kwargs=kwargs, @@ -349,10 +349,10 @@ class TestLoadAndResolveStageConfigs: ] load_stage_configs = mocker.patch( "vllm_omni.entrypoints.utils.load_stage_configs_from_model", - return_value=returned_stage_configs, + return_value=(returned_stage_configs, None), ) - config_path, stage_configs = load_and_resolve_stage_configs( + config_path, stage_configs, _ = load_and_resolve_stage_configs( model="dummy-model", stage_configs_path=str(deploy_path), kwargs={}, @@ -363,6 +363,7 @@ class TestLoadAndResolveStageConfigs: base_engine_args={}, deploy_config_path=str(deploy_path), stage_overrides=None, + strategy_config_path=None, ) assert config_path == str(deploy_path) assert len(stage_configs) == 2 diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index c1c76180a..54683e67c 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -1721,7 +1721,7 @@ class TestAuraOmniDeploy: def test_aura_omni_deploy_resolves_four_native_stages(self): pipeline_cfg = StageConfigFactory.resolve_pipeline_config("aura_omni") - stages = StageConfigFactory._create_from_registry( + stages, _ = StageConfigFactory._create_from_registry( "qwen3_tts", pipeline_cfg, cli_overrides={}, @@ -1943,11 +1943,12 @@ class TestSentinelDefaultPrecedence: model_type, Q3_OMNI_ALL_STAGES_HF_CONFIG, ) - return StageConfigFactory._create_from_registry( + stages, _ = StageConfigFactory._create_from_registry( "qwen3_omni_moe", pipeline_cfg, cli_overrides=cli_overrides, ) + return stages def test_typed_kwarg_overrides_yaml(self): stages = self._stages({"max_num_seqs": 999}) diff --git a/tests/utils/test_tracking_parser.py b/tests/utils/test_tracking_parser.py index d2e7e4034..5bee3e3b0 100644 --- a/tests/utils/test_tracking_parser.py +++ b/tests/utils/test_tracking_parser.py @@ -591,7 +591,7 @@ def test_explicit_cli_arg_reaches_runtime_overrides(mock_stages): ns = p.parse_args(["--max-num-seqs", "999"]) explicit_kwargs = ns.get_explicit_kwargs_dict() - stages = StageConfigFactory._create_from_registry( + stages, _ = StageConfigFactory._create_from_registry( _TEST_MODEL, _TEST_PIPELINE, explicit_kwargs, @@ -608,7 +608,7 @@ def test_omitted_default_not_in_runtime_overrides(mock_stages): ns = p.parse_args([]) explicit_kwargs = ns.get_explicit_kwargs_dict() - stages = StageConfigFactory._create_from_registry( + stages, _ = StageConfigFactory._create_from_registry( _TEST_MODEL, _TEST_PIPELINE, explicit_kwargs, @@ -630,7 +630,7 @@ def test_config_file_args_reach_runtime_overrides(mock_stages): ns = p.parse_args(["--config", "fake.yaml"]) explicit_kwargs = ns.get_explicit_kwargs_dict() - stages = StageConfigFactory._create_from_registry( + stages, _ = StageConfigFactory._create_from_registry( _TEST_MODEL, _TEST_PIPELINE, explicit_kwargs, @@ -648,7 +648,7 @@ def test_per_stage_override_routes_correctly(mock_stages): ns = p.parse_args(["--stage-0-gpu-memory-utilization", "0.42"]) explicit_kwargs = ns.get_explicit_kwargs_dict() - stages = StageConfigFactory._create_from_registry( + stages, _ = StageConfigFactory._create_from_registry( _TEST_MODEL, _TEST_PIPELINE, explicit_kwargs, @@ -670,7 +670,7 @@ def test_explicit_args_omitted_from_yaml(mock_stages): ns = p.parse_args(["--enforce-eager"]) explicit_kwargs = ns.get_explicit_kwargs_dict() - stages = StageConfigFactory._create_from_registry( + stages, _ = StageConfigFactory._create_from_registry( _TEST_MODEL, _TEST_PIPELINE, explicit_kwargs, diff --git a/vllm_omni/config/composable_parallel/__init__.py b/vllm_omni/config/composable_parallel/__init__.py new file mode 100644 index 000000000..499c14548 --- /dev/null +++ b/vllm_omni/config/composable_parallel/__init__.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Composable parallel strategies for vLLM-Omni. + +A small, runtime-agnostic vocabulary for declaring a stage's parallel layout +(:class:`StrategySpec` + mesh/routing/aggregation patterns), plus the logic to +translate a per-role stack of specs into concrete engine sizing and overlay it +onto merged stage configs. + +The spec/pattern *types* are deliberately pure data with no engine imports; the +translation and apply logic lives alongside them and consumes vLLM-Omni's own +config objects. +""" + +from vllm_omni.config.composable_parallel.aggregation import ( + AGGREGATION_PATTERN_VARIANTS, + AggregationConflictError, + AggregationPattern, + AllGather, + Combine, + FanInByStage, + GatherDim, + StitchPipeline, + StitchSpatial, + TakeRank, + Union, +) +from vllm_omni.config.composable_parallel.apply import ( + StrategyApplyError, + StrategyApplyResult, + apply_strategy_specs, + check_device_layout, +) +from vllm_omni.config.composable_parallel.routing import ( + ROUTING_PATTERN_VARIANTS, + Broadcast, + DuplicateWithCondTag, + PartitionByHash, + PipelineMicrobatch, + RouteByExpert, + RouteByStage, + RoutingKeyError, + RoutingPattern, + ShardSequence, + ShardSpatial, +) +from vllm_omni.config.composable_parallel.spec import ( + MESH_AXIS_KINDS, + KernelSpec, + LayerHookSpec, + MeshAxisKind, + MeshAxisSpec, + SpecMergeConflictError, + StrategySpec, +) +from vllm_omni.config.composable_parallel.translator import ( + AxisTranslationError, + OmniParallelConfig, + translate_strategy_stack, +) + +__all__ = [ + # spec types + "MeshAxisKind", + "MESH_AXIS_KINDS", + "MeshAxisSpec", + "LayerHookSpec", + "KernelSpec", + "StrategySpec", + "SpecMergeConflictError", + # routing patterns + "RoutingPattern", + "RoutingKeyError", + "Broadcast", + "PartitionByHash", + "ShardSequence", + "ShardSpatial", + "RouteByExpert", + "PipelineMicrobatch", + "DuplicateWithCondTag", + "RouteByStage", + "ROUTING_PATTERN_VARIANTS", + # aggregation patterns + "AggregationPattern", + "AggregationConflictError", + "TakeRank", + "Union", + "GatherDim", + "AllGather", + "StitchSpatial", + "StitchPipeline", + "Combine", + "FanInByStage", + "AGGREGATION_PATTERN_VARIANTS", + # translation + "OmniParallelConfig", + "translate_strategy_stack", + "AxisTranslationError", + # apply + "apply_strategy_specs", + "StrategyApplyResult", + "check_device_layout", + "StrategyApplyError", +] diff --git a/vllm_omni/config/composable_parallel/aggregation.py b/vllm_omni/config/composable_parallel/aggregation.py new file mode 100644 index 000000000..75ed8fadd --- /dev/null +++ b/vllm_omni/config/composable_parallel/aggregation.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Closed aggregation-pattern hierarchy for declarative parallel strategies. + +An aggregation pattern describes how the per-worker results of one mesh axis +are recombined into a single logical result (e.g. take one tensor-parallel +rank's output, union disjoint data-parallel shards). The hierarchy is *closed*: +subclasses may only be defined in this module. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + + +@dataclass(frozen=True) +class AggregationPattern: + """Base class for closed aggregation patterns. Subclass only within this module.""" + + def __init_subclass__(cls, **kwargs: object) -> None: + super().__init_subclass__(**kwargs) + if cls.__module__ != __name__: + raise TypeError("AggregationPattern is closed to external subclassing") + + def pattern_kind(self) -> str: + raise NotImplementedError + + +class AggregationConflictError(ValueError): + """Raised when two aggregation inputs conflict for the same request id.""" + + +@dataclass(frozen=True) +class TakeRank(AggregationPattern): + """Select one rank's result (TP — all ranks agree).""" + + rank: int = 0 + + def __post_init__(self) -> None: + if not isinstance(self.rank, int) or isinstance(self.rank, bool): + raise TypeError(f"TakeRank.rank must be int, got {type(self.rank).__name__}") + + def pattern_kind(self) -> str: + return "take_rank" + + +@dataclass(frozen=True) +class Union(AggregationPattern): + """Union of disjoint replica results (DP, HSDP).""" + + def pattern_kind(self) -> str: + return "union" + + +@dataclass(frozen=True) +class GatherDim(AggregationPattern): + """Gather along a dimension and materialize on one designated rank. + + Default designated rank is axis-rank 0 unless adapter overrides. + """ + + dim: int = 1 + + def pattern_kind(self) -> str: + return "gather_dim" + + +@dataclass(frozen=True) +class AllGather(AggregationPattern): + """All ranks materialize the same aggregate result across the axis.""" + + dim: int | None = None + + def pattern_kind(self) -> str: + return "all_gather" + + +@dataclass(frozen=True) +class StitchSpatial(AggregationPattern): + """Stitch spatial patches (VAE-PP).""" + + grid: tuple[int, int] = (1, 1) + + def pattern_kind(self) -> str: + return "stitch_spatial" + + +@dataclass(frozen=True) +class StitchPipeline(AggregationPattern): + """Stitch pipeline stage outputs (PP).""" + + def pattern_kind(self) -> str: + return "stitch_pipeline" + + +@dataclass(frozen=True) +class Combine(AggregationPattern): + """Custom reducer combination (CFG). Adapter resolves reducer_id.""" + + reducer_id: str = "default" + + def pattern_kind(self) -> str: + return "combine" + + +@dataclass(frozen=True) +class FanInByStage(AggregationPattern): + """Fan-in results routed by stage (stage replicas).""" + + def pattern_kind(self) -> str: + return "fan_in_by_stage" + + +AGGREGATION_PATTERN_VARIANTS: tuple[type[AggregationPattern], ...] = ( + TakeRank, + Union, + GatherDim, + AllGather, + StitchSpatial, + StitchPipeline, + Combine, + FanInByStage, +) + +ReducerFn = Callable[..., object] diff --git a/vllm_omni/config/composable_parallel/apply.py b/vllm_omni/config/composable_parallel/apply.py new file mode 100644 index 000000000..879c7aec3 --- /dev/null +++ b/vllm_omni/config/composable_parallel/apply.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Apply declarative parallel strategies onto merged stage configs. + +This is the *override-after-merge* seam. ``merge_pipeline_deploy`` first fuses +the pipeline topology + deploy YAML into a list of ``StageConfig`` objects; this +module then overlays the engine sizing derived from a per-role +``StrategySpec`` stack onto those stages. + +Design rules (one writer per axis): + +* A strategy only writes the axes it actually declares. If a stack has no ``tp`` + axis, this module never touches ``tensor_parallel_size`` — it does not force a + default of 1. +* **Conflict-on-explicit.** If a knob was already set explicitly by the deploy + YAML and the strategy derives a different value, we raise rather than silently + override. Equal values are a no-op; unset values are filled. +* **Pre-spawn device check.** When a stage declares ``devices``, the device + count must be consistent with the engine world size (``tp * dp * pp``) and the + replica count, so misconfigurations fail at config time instead of at spawn. + +``omni_lb_policy`` is *not* a per-stage config knob: omni reads it once at +orchestrator construction (``AsyncOmniEngine``), not from stage configs. So a +stage_replica axis's derived policy is surfaced on :class:`StrategyApplyResult` +for the caller to pass to the orchestrator; it is never silently written into a +stage where it would be ignored. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, NoReturn + +from vllm.logger import init_logger + +from vllm_omni.config.composable_parallel.spec import StrategySpec +from vllm_omni.config.composable_parallel.translator import ( + OmniParallelConfig, + translate_strategy_stack, +) + +logger = init_logger(__name__) + +# Role key -> the StageConfig fields each declared axis writes. +_ENGINE_FIELD_BY_KIND: dict[str, str] = { + "tp": "tensor_parallel_size", + "dp": "data_parallel_size", + "pp": "pipeline_parallel_size", +} + +RoleKey = str + + +class StrategyApplyError(ValueError): + """Error for applying a strategy onto stage configs. + + A single error type (rather than per-cause subclasses) keeps the public + surface small and consistent with the rest of the codebase; the specific + cause — an unmatched role, a conflict with an explicit deploy value, or a + device-count mismatch — is in the message and is logged before the raise so + it stays visible even when the type is unavailable to a caller (e.g. across a + server boundary). + """ + + +def _fail(msg: str) -> NoReturn: + """Log and raise :class:`StrategyApplyError`.""" + logger.error("[composable_parallel] %s", msg) + raise StrategyApplyError(msg) + + +@dataclass +class StrategyApplyResult: + """Outcome of applying strategies to a stage list. + + ``stages`` is the same list passed in (mutated in place and returned for + convenience). ``omni_lb_policy`` is the pipeline-wide load-balancer policy + derived from any stage_replica axes (``None`` if none declared one); the + caller wires it into the orchestrator, since omni does not read it from + stage configs. + + ``per_role_config`` is keyed by the role key the caller supplied — a + ``model_stage`` name (str). ``per_stage_config`` is the same information + keyed by the resolved integer ``stage_id`` — handy for re-validating a + stage after later layers (e.g. CLI overrides) have run. + """ + + stages: list[Any] + omni_lb_policy: str | None = None + per_role_config: dict[RoleKey, OmniParallelConfig] = field(default_factory=dict) + per_stage_config: dict[int, OmniParallelConfig] = field(default_factory=dict) + + +def _resolve_stage(stages: Sequence[Any], key: RoleKey) -> Any: + """Find the single stage matching ``key`` by ``model_stage`` name. + + Role keys are ``model_stage`` names (str) — the same human-readable names + that strategy.yaml authors use to refer to a stage (e.g. ``"thinker"``). + """ + matches = [s for s in stages if getattr(s, "model_stage", None) == key] + descriptor = f"model_stage={key!r}" + + if not matches: + available = ", ".join(f"{getattr(s, 'model_stage', '?')!r}(id={getattr(s, 'stage_id', '?')})" for s in stages) + _fail(f"strategy role {descriptor} did not match any stage; available stages: {available}") + if len(matches) > 1: + _fail(f"strategy role {descriptor} is ambiguous; it matched {len(matches)} stages") + return matches[0] + + +def _set_explicit(engine_args: dict[str, Any], field_name: str, value: Any, *, role: RoleKey) -> None: + """Write ``field_name`` with conflict-on-explicit semantics. + + ``engine_args`` is the raw mapping loaded straight from the deploy YAML + (before the engine-args dataclass is constructed), so a key is *present* only + when the YAML set it explicitly. We therefore key off membership rather than + ``.get``: an explicit ``None`` in the YAML is a real, present value and must + not be conflated with a key miss (which ``.get`` would report as ``None``). + """ + if field_name in engine_args and engine_args[field_name] != value: + existing = engine_args[field_name] + _fail( + f"role {role!r}: strategy derives {field_name}={value} but the deploy config already " + f"set {field_name}={existing}. Remove one of them (the strategy is the single writer " + "for the axes it declares)." + ) + engine_args[field_name] = value + + +def _set_num_replicas(runtime: dict[str, Any], value: int, *, role: RoleKey) -> None: + """Write ``num_replicas`` with conflict-on-explicit (treating 1 as unset).""" + existing = runtime.get("num_replicas", 1) + if existing not in (1, value): + _fail( + f"role {role!r}: strategy derives num_replicas={value} but the deploy config already " + f"set num_replicas={existing}. Remove one of them." + ) + runtime["num_replicas"] = value + + +def _parse_device_count(devices: Any) -> int | None: + """Return the number of device ids in a ``devices`` value, or None.""" + if devices is None: + return None + if isinstance(devices, (list, tuple)): + return len([d for d in devices]) + text = str(devices).strip() + if not text: + return None + return len([d for d in text.split(",") if d.strip()]) + + +def check_device_layout( + devices: Any, + *, + tensor_parallel_size: int, + data_parallel_size: int, + pipeline_parallel_size: int, + num_replicas: int, + role: RoleKey, +) -> None: + """Validate a stage's device count against its world size (× replicas in pool mode). + + Raises :class:`StrategyApplyError` when the declared ``devices`` count is + neither the per-replica world size nor the full ``num_replicas`` pool. + Exposed publicly so the resolved (post-CLI) layout can be re-checked after + later override layers, not just the strategy-derived snapshot. + """ + count = _parse_device_count(devices) + if count is None: + return + world = int(tensor_parallel_size) * int(data_parallel_size) * int(pipeline_parallel_size) + replicas = int(num_replicas or 1) + # Two valid shapes: a single per-replica template (== world) or the full + # pool across replicas (== replicas * world). This mirrors omni's own + # per-replica device splitter. + valid = {world, replicas * world} + if count not in valid: + _fail( + f"role {role!r}: declared {count} device id(s) but the strategy world size is {world} " + f"(tp={tensor_parallel_size} * dp={data_parallel_size} * pp={pipeline_parallel_size}). " + f"Provide either {world} (per-replica template) or {replicas * world} " + f"(num_replicas={replicas} pool)." + ) + + +def _check_devices(runtime: dict[str, Any], cfg: OmniParallelConfig, *, role: RoleKey) -> None: + """Pre-spawn check on the strategy-derived snapshot (see :func:`check_device_layout`).""" + check_device_layout( + runtime.get("devices"), + tensor_parallel_size=cfg.tensor_parallel_size, + data_parallel_size=cfg.data_parallel_size, + pipeline_parallel_size=cfg.pipeline_parallel_size, + num_replicas=int(runtime.get("num_replicas", 1) or 1), + role=role, + ) + + +def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> None: + engine_args = stage.yaml_engine_args + runtime = stage.yaml_runtime + + declared = set(cfg.l1_owners.keys()) + for kind, field_name in _ENGINE_FIELD_BY_KIND.items(): + if kind in declared: + _set_explicit(engine_args, field_name, getattr(cfg, field_name), role=role) + if "ep" in declared and cfg.enable_expert_parallel: + _set_explicit(engine_args, "enable_expert_parallel", True, role=role) + if "stage_replica" in declared: + _set_num_replicas(runtime, cfg.stage_replica_size, role=role) + + _check_devices(runtime, cfg, role=role) + + +def apply_strategy_specs( + stages: list[Any], + strategy_specs: Mapping[RoleKey, Sequence[StrategySpec]], +) -> StrategyApplyResult: + """Overlay per-role strategy specs onto a merged stage list. + + Args: + stages: the ``list[StageConfig]`` returned by ``merge_pipeline_deploy``. + strategy_specs: maps a role (a ``model_stage`` name, e.g. ``"thinker"``) + to that stage's stack of ``StrategySpec`` (one per declared mesh + axis). Role keys are ``model_stage`` names (str). + + Returns: + A :class:`StrategyApplyResult` with the mutated stages and the derived + pipeline-wide ``omni_lb_policy`` (if any stage_replica axis set one). + + Raises: + StrategyApplyError: a role matched zero or multiple stages, a derived + value conflicts with an explicit deploy value (including two roles + deriving different ``omni_lb_policy`` values), or a stage's device + count is inconsistent with its world size. + AxisTranslationError: the spec stack is invalid/unsupported. + NotImplementedError: the spec stack requests routing not built yet. + """ + result = StrategyApplyResult(stages=stages) + lb_policy: str | None = None + lb_owner: RoleKey | None = None + + for key, specs in strategy_specs.items(): + stage = _resolve_stage(stages, key) + cfg = translate_strategy_stack(specs) + _apply_to_stage(stage, cfg, role=key) + result.per_role_config[key] = cfg + stage_id = getattr(stage, "stage_id", None) + if stage_id is not None: + result.per_stage_config[int(stage_id)] = cfg + + if cfg.stage_replica_size > 1 and cfg.omni_lb_policy is not None: + if lb_policy is not None and lb_policy != cfg.omni_lb_policy: + _fail( + f"role {key!r} derives omni_lb_policy={cfg.omni_lb_policy!r} but role " + f"{lb_owner!r} already derived {lb_policy!r}. The load-balancer policy is " + "pipeline-wide; only one value is allowed." + ) + lb_policy = cfg.omni_lb_policy + lb_owner = key + + result.omni_lb_policy = lb_policy + return result diff --git a/vllm_omni/config/composable_parallel/routing.py b/vllm_omni/config/composable_parallel/routing.py new file mode 100644 index 000000000..38e7d4a49 --- /dev/null +++ b/vllm_omni/config/composable_parallel/routing.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Closed routing-pattern hierarchy for declarative parallel strategies. + +A routing pattern describes how a request batch is distributed across the +workers of one mesh axis (e.g. broadcast to every tensor-parallel rank, hash +partition across data-parallel ranks). The hierarchy is *closed*: subclasses +may only be defined in this module, so the set of patterns is a fixed, +exhaustively matchable contract. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Literal + + +@dataclass(frozen=True) +class RoutingPattern: + """Base class for closed routing patterns. Subclass only within this module.""" + + def __init_subclass__(cls, **kwargs: object) -> None: + super().__init_subclass__(**kwargs) + if cls.__module__ != __name__: + raise TypeError("RoutingPattern is closed to external subclassing") + + def pattern_kind(self) -> str: + raise NotImplementedError + + +class RoutingKeyError(KeyError): + """Raised when a routing key is missing from a request.""" + + def __init__(self, request_id: object, missing_key: str) -> None: + self.request_id = request_id + self.missing_key = missing_key + super().__init__(f"Missing routing key {missing_key!r} for request_id={request_id!r}") + + +@dataclass(frozen=True) +class Broadcast(RoutingPattern): + """Every worker in the group sees the whole batch (TP, dense EP).""" + + def pattern_kind(self) -> str: + return "broadcast" + + +@dataclass(frozen=True) +class PartitionByHash(RoutingPattern): + """Stable hash partition by request key (DP, HSDP).""" + + key: str = "request_id" + + def pattern_kind(self) -> str: + return "partition_by_hash" + + +@dataclass(frozen=True) +class ShardSequence(RoutingPattern): + """Shard along a sequence dimension (SP-Ulysses, SP-Ring, CP).""" + + dim: int = 1 + + def pattern_kind(self) -> str: + return "shard_sequence" + + +@dataclass(frozen=True) +class ShardSpatial(RoutingPattern): + """Shard along a spatial grid (VAE-PP).""" + + grid: tuple[int, int] = (1, 1) + + def pattern_kind(self) -> str: + return "shard_spatial" + + +@dataclass(frozen=True) +class RouteByExpert(RoutingPattern): + """Sparse MoE expert routing.""" + + router_id: str = "default" + + def pattern_kind(self) -> str: + return "route_by_expert" + + +@dataclass(frozen=True) +class PipelineMicrobatch(RoutingPattern): + """Pipeline parallelism microbatch routing (PP).""" + + microbatch_size: int = 1 + + def pattern_kind(self) -> str: + return "pipeline_microbatch" + + +@dataclass(frozen=True) +class DuplicateWithCondTag(RoutingPattern): + """Duplicate batch with a conditional tag (CFG).""" + + flag_name: str = "guidance_branch" + + def pattern_kind(self) -> str: + return "duplicate_with_cond_tag" + + +@dataclass(frozen=True) +class RouteByStage(RoutingPattern): + """Route requests by pipeline stage policy (stage replicas). + + The ``random``/``round_robin``/``least_queue`` policies describe stateless + load balancing that a runtime (e.g. an engine's built-in DP load balancer, + or the omni StagePool balancer) can realize. ``hash`` describes + deterministic, key-stable routing that a load balancer cannot honor on its + own (see ``PartitionByHash``). + """ + + routing_policy: Literal["random", "round_robin", "least_queue", "hash"] = "round_robin" + + def pattern_kind(self) -> str: + return "route_by_stage" + + +ROUTING_PATTERN_VARIANTS: tuple[type[RoutingPattern], ...] = ( + Broadcast, + PartitionByHash, + ShardSequence, + ShardSpatial, + RouteByExpert, + PipelineMicrobatch, + DuplicateWithCondTag, + RouteByStage, +) + +# Optional callable slot for Combine-style custom routing reducers in adapters. +RouterFn = Callable[..., object] diff --git a/vllm_omni/config/composable_parallel/spec.py b/vllm_omni/config/composable_parallel/spec.py new file mode 100644 index 000000000..299fc3ad3 --- /dev/null +++ b/vllm_omni/config/composable_parallel/spec.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Core declarative types for parallel-strategy specification. + +A :class:`StrategySpec` declares one parallelism scheme: the mesh axis it sizes +(kind + degree), how a batch is routed across that axis, and how per-worker +results are aggregated back. A *stack* of specs (one per axis) fully describes a +stage's parallel layout in a runtime-agnostic, data-only form that can be +translated into concrete engine sizing. + +The mesh-axis kinds below enumerate every parallelism dimension the contract can +describe. Only a subset is wired end-to-end today (see ``translator.py``): + +* **Wired** (translatable now): ``tp``, ``dp``, ``pp``, ``ep``, + ``stage_replica``. +* **Reserved** (declared in the type system but not yet translatable — the + translator raises ``AxisTranslationError`` for them): ``sp_ulysses``, + ``sp_ring``, ``cfg``, ``vae_pp``, ``hsdp``, ``stage_pp``, ``cp``. + +Reserved kinds fail fast at translation time rather than silently doing nothing, +so declaring one is an explicit "not yet" rather than a no-op. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Literal + +from vllm_omni.config.composable_parallel.aggregation import AggregationPattern +from vllm_omni.config.composable_parallel.routing import RoutingPattern + + +class MeshAxisKind(str, Enum): + """Enumerates every parallelism dimension the strategy contract can describe. + + Subclassing ``str`` keeps existing string comparisons and dict keys working + (``MeshAxisKind.TP == "tp"`` and both hash the same), while making the type + refactor-safe: a typo'd kind is an ``AttributeError`` on the enum rather than + a silently-wrong bare string. + """ + + TP = "tp" + DP = "dp" + PP = "pp" + EP = "ep" + SP_ULYSSES = "sp_ulysses" + SP_RING = "sp_ring" + CFG = "cfg" + VAE_PP = "vae_pp" + HSDP = "hsdp" + STAGE_PP = "stage_pp" + STAGE_REPLICA = "stage_replica" + CP = "cp" + + +# Bare string values, kept as a plain tuple so membership checks and error +# messages read as the kind names (``"tp"``) rather than enum reprs. ``str``-enum +# equality means both ``"tp"`` and ``MeshAxisKind.TP`` satisfy ``in`` checks. +MESH_AXIS_KINDS: tuple[str, ...] = tuple(k.value for k in MeshAxisKind) + +HookCategory = Literal[ + "linear", + "attention_pre", + "attention_post", + "ffn_pre", + "ffn_post", + "moe_dispatch", + "other", +] + +HOOK_CATEGORY_ORDER: dict[HookCategory, int] = { + "linear": 0, + "attention_pre": 1, + "attention_post": 2, + "ffn_pre": 3, + "ffn_post": 4, + "moe_dispatch": 5, + "other": 6, +} + + +class SpecMergeConflictError(ValueError): + """Raised when merged hook/kernel specs have conflicting declarations.""" + + +@dataclass(frozen=True) +class MeshAxisSpec: + """Declares one axis of a process mesh (kind + size).""" + + kind: MeshAxisKind + size: int + + def __post_init__(self) -> None: + if self.kind not in MESH_AXIS_KINDS: + raise ValueError(f"MeshAxisSpec.kind must be one of {MESH_AXIS_KINDS}, got {self.kind!r}") + if self.size <= 0: + raise ValueError(f"MeshAxisSpec.size must be > 0, got {self.size}") + + +@dataclass(frozen=True) +class LayerHookSpec: + """L2 hook slot referenced by StrategySpec; resolved by the model walker.""" + + hook_id: str + target: str | None = None + category: HookCategory = "other" + priority: int = 0 + axis_index: int = 0 + + +@dataclass(frozen=True) +class KernelSpec: + """L3 kernel slot referenced by StrategySpec; resolved by the kernel registry.""" + + kernel_id: str + target: str | None = None + category: HookCategory = "other" + priority: int = 0 + axis_index: int = 0 + group_axis_kind: MeshAxisKind | None = None + requires_collective: bool = False + + +@dataclass(frozen=True) +class StrategySpec: + """Declarative contract for one parallelism scheme (data + hook/kernel slots).""" + + name: str + mesh_axis: MeshAxisSpec + routing: RoutingPattern + aggregation: AggregationPattern + layer_hook_specs: tuple[LayerHookSpec, ...] = () + kernel_specs: tuple[KernelSpec, ...] = () + shard_extension: Mapping[str, Any] = field(default_factory=dict) diff --git a/vllm_omni/config/composable_parallel/strategy_loader.py b/vllm_omni/config/composable_parallel/strategy_loader.py new file mode 100644 index 000000000..665fbe9c2 --- /dev/null +++ b/vllm_omni/config/composable_parallel/strategy_loader.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Load per-role parallel strategies from a strategy file. + +The strategy file is a small, optional companion to the deploy config. It maps a +logical role (``model_stage``, e.g. ``thinker``) to a list of axis declarations, +each of which becomes one :class:`StrategySpec`. YAML is the only format wired up +today, but the module/symbol names stay format-agnostic so a different source can +be added without renaming the public surface:: + + strategies: + thinker: + - axis: tp + size: 2 + talker: + - axis: stage_replica + size: 2 + routing: round_robin # random | round_robin | least_queue + code2wav: + - axis: stage_replica + size: 2 + routing: round_robin + +Per-axis fields: + +* ``axis`` (required): mesh-axis kind (``tp``, ``dp``, ``pp``, ``ep``, + ``stage_replica`` are translatable today). +* ``size`` (required): axis degree (> 0). +* ``routing`` (optional): stateless policy for ``dp`` / ``stage_replica`` axes + (``random`` / ``round_robin`` / ``least_queue``). Ignored for axes whose + routing is fixed (``tp`` broadcast, ``pp`` microbatch). +* ``l1_owner`` (optional): ``engine`` or ``delegated``; overrides the default + owner for the axis (recorded in ``shard_extension``). +* ``name`` (optional): a label for the spec; defaults to the axis kind. + +Routing/aggregation patterns are filled in from sensible per-kind defaults so +the file stays compact; the translator validates the combination. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from vllm_omni.config.composable_parallel.aggregation import ( + AggregationPattern, + FanInByStage, + StitchPipeline, + TakeRank, + Union, +) +from vllm_omni.config.composable_parallel.routing import ( + Broadcast, + PipelineMicrobatch, + RouteByStage, + RoutingPattern, +) +from vllm_omni.config.composable_parallel.spec import MeshAxisSpec, StrategySpec +from vllm_omni.config.yaml_util import load_yaml_config, to_dict + +# Per-kind defaults for routing + aggregation, so the file only needs axis+size. +_ROUTING_POLICY_KINDS = frozenset({"dp", "stage_replica"}) + + +class StrategyLoadError(ValueError): + """Raised when a strategy file is malformed.""" + + +def _default_routing(kind: str, routing_policy: str | None) -> RoutingPattern: + if kind in _ROUTING_POLICY_KINDS: + return RouteByStage(routing_policy=routing_policy or "round_robin") # type: ignore[arg-type] + if kind == "pp": + return PipelineMicrobatch() + # tp, ep, and any other kind default to broadcast (the translator rejects + # kinds it cannot handle before it ever inspects routing). + return Broadcast() + + +def _default_aggregation(kind: str) -> AggregationPattern: + if kind == "stage_replica": + return FanInByStage() + if kind == "pp": + return StitchPipeline() + if kind in ("dp", "ep"): + return Union() + return TakeRank() + + +def _build_spec(role: str, entry: Mapping[str, Any]) -> StrategySpec: + if "axis" not in entry: + raise StrategyLoadError(f"role {role!r}: each strategy entry needs an 'axis' field") + if "size" not in entry: + raise StrategyLoadError(f"role {role!r}: strategy entry for axis {entry['axis']!r} needs a 'size'") + + kind = str(entry["axis"]) + try: + size = int(entry["size"]) + except (TypeError, ValueError) as exc: + raise StrategyLoadError(f"role {role!r}: axis {kind!r} size must be an integer, got {entry['size']!r}") from exc + + routing_policy = entry.get("routing") + if routing_policy is not None and kind not in _ROUTING_POLICY_KINDS: + raise StrategyLoadError( + f"role {role!r}: axis {kind!r} does not accept a 'routing' policy (only {sorted(_ROUTING_POLICY_KINDS)} do)" + ) + + shard_extension: dict[str, Any] = {} + l1_owner = entry.get("l1_owner") + if l1_owner is not None: + shard_extension["l1_owner"] = str(l1_owner) + + return StrategySpec( + name=str(entry.get("name", kind)), + mesh_axis=MeshAxisSpec(kind=kind, size=size), # type: ignore[arg-type] + routing=_default_routing(kind, routing_policy), + aggregation=_default_aggregation(kind), + shard_extension=shard_extension, + ) + + +def parse_strategy_specs(data: Mapping[str, Any]) -> dict[str, list[StrategySpec]]: + """Parse an already-loaded strategy mapping into per-role spec stacks.""" + strategies = data.get("strategies", data) + if not isinstance(strategies, Mapping): + raise StrategyLoadError("strategy file must contain a 'strategies' mapping of role -> list of axes") + + result: dict[str, list[StrategySpec]] = {} + for role, entries in strategies.items(): + if not isinstance(entries, (list, tuple)): + raise StrategyLoadError( + f"role {role!r}: expected a list of axis declarations, got {type(entries).__name__}" + ) + specs: list[StrategySpec] = [] + for entry in entries: + if not isinstance(entry, Mapping): + raise StrategyLoadError( + f"role {role!r}: each axis declaration must be a mapping, got {type(entry).__name__}" + ) + specs.append(_build_spec(str(role), dict(entry))) + result[str(role)] = specs + return result + + +def load_strategy_specs(path: str) -> dict[str, list[StrategySpec]]: + """Load and parse a strategy file into per-role spec stacks.""" + data = to_dict(load_yaml_config(path)) + if not isinstance(data, Mapping): + raise StrategyLoadError(f"strategy file {path!r} did not parse to a mapping") + return parse_strategy_specs(data) diff --git a/vllm_omni/config/composable_parallel/translator.py b/vllm_omni/config/composable_parallel/translator.py new file mode 100644 index 000000000..5c394dc9b --- /dev/null +++ b/vllm_omni/config/composable_parallel/translator.py @@ -0,0 +1,378 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Translate a ``StrategySpec`` stack into vLLM-Omni parallel sizing. + +Read a stack of strategy specs (one per mesh axis) and work out the +tensor/data/pipeline parallel sizes, the dense-EP flag, and the number of stage +replicas, plus who owns each axis's request routing. The result is a plain, +CPU-computable sizing struct (:class:`OmniParallelConfig`) keyed by the real +``OmniEngineArgs`` / ``EngineArgs`` field names, so the deploy layer can splat it +onto a stage's engine args. + +The distinction this module enforces — engine data parallelism (a true vLLM +intra-engine world dimension) vs. omni stage replicas (independent engines fanned +out by omni's coordinator) — is documented on the axis validators that depend on +it (see :func:`_validate_dp` and :func:`_stage_replica_lb_policy`). Routing we do +not support yet (key-stable / affinity routing) raises ``NotImplementedError``; +any other invalid spec raises :class:`AxisTranslationError`. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Literal, NoReturn, cast, get_args + +from vllm.logger import init_logger + +from vllm_omni.config.composable_parallel.routing import ( + Broadcast, + PartitionByHash, + PipelineMicrobatch, + RouteByStage, + RoutingPattern, +) +from vllm_omni.config.composable_parallel.spec import MeshAxisKind, StrategySpec + +logger = init_logger(__name__) + +# Who owns an axis's request routing. A closed string type so an unexpected raw +# value is caught statically and at runtime (``_VALID_L1_OWNERS`` is derived from +# it, keeping one source of truth) rather than silently flowing through. +L1Owner = Literal["delegated", "engine"] + +# DP, TP, PP, (dense) EP and stage_replica translate today. DP/TP/PP are true +# world dimensions vLLM realizes intra-engine; dense EP is a flag over the +# existing TP*DP ranks; stage_replica is an omni-coordinator-level fan-out of +# independent engine replicas (NOT a vLLM world dimension). The remaining kinds +# (sequence/context parallel, CFG, VAE pipelines, ...) need per-layer hooks or +# custom collectives and arrive in later stages. +_SUPPORTED_KINDS: tuple[MeshAxisKind, ...] = ("dp", "tp", "pp", "ep", "stage_replica") + +# Which EngineArgs field each axis kind *sizes*. EP and stage_replica are +# deliberately absent: EP is not an independent world dimension but an +# ``enable_expert_parallel`` flag whose degree equals tp*dp; stage_replica is a +# per-stage deploy ``num_replicas`` count, not an intra-engine world dimension. +_AXIS_TO_ENGINE_FIELD: dict[MeshAxisKind, str] = { + "tp": "tensor_parallel_size", + "dp": "data_parallel_size", + "pp": "pipeline_parallel_size", +} + +# Default L1 owner per axis kind. ``stage_replica`` is the only delegated axis: +# omni's coordinator/StagePool owns routing across replicas. ``dp`` is engine +# data parallelism, realized by vLLM's own intra-engine DP load balancer. +_DEFAULT_L1_OWNER: dict[MeshAxisKind, L1Owner] = { + "dp": "engine", + "tp": "engine", + "pp": "engine", + "ep": "engine", + "stage_replica": "delegated", +} + +# How a RouteByStage policy maps onto omni's load balancer policy string. These +# are the stateless options omni can actually do; note there's no "hash" here, +# because omni has no key-stable balancer. +_STAGE_POLICY_TO_OMNI_LB: dict[str, str] = { + "random": "random", + "round_robin": "round-robin", + "least_queue": "least-queue-length", +} + +_VALID_L1_OWNERS: frozenset[str] = frozenset(get_args(L1Owner)) + + +class AxisTranslationError(ValueError): + """Error for an invalid or unsupported strategy spec. + + A single error type (rather than a tree of subclasses) keeps the public + surface small and consistent with the rest of the codebase; the specific + cause is in the message and is logged before the raise so it is visible even + when the type is unavailable to a caller (e.g. across a server boundary). + + Strategies that are *valid but not built yet* — key-stable / affinity routing + today — raise ``NotImplementedError`` instead, to distinguish "we haven't + implemented this" from "your config is wrong". + """ + + +def _fail(msg: str) -> NoReturn: + """Log and raise :class:`AxisTranslationError` for an invalid/unsupported spec.""" + logger.error("[composable_parallel] %s", msg) + raise AxisTranslationError(msg) + + +def _not_implemented(msg: str) -> NoReturn: + """Log and raise ``NotImplementedError`` for a valid-but-unbuilt strategy.""" + logger.error("[composable_parallel] %s", msg) + raise NotImplementedError(msg) + + +@dataclass(frozen=True) +class OmniParallelConfig: + """Result of translating a spec stack into omni parallel sizing.""" + + tensor_parallel_size: int = 1 + data_parallel_size: int = 1 + pipeline_parallel_size: int = 1 + # Dense expert parallelism: a flag over the TP*DP ranks, not a world axis. + enable_expert_parallel: bool = False + # Number of independent omni stage replicas (the ``stage_replica`` axis). + # This is NOT a vLLM world dimension; it is the per-stage deploy + # ``num_replicas`` count. 1 means a single (un-replicated) engine. + stage_replica_size: int = 1 + # The omni StagePool LB policy string for the (delegated) stage_replica + # axis, if any. Only meaningful when ``stage_replica_size > 1``. + omni_lb_policy: str | None = None + # axis kind -> resolved L1 owner ("delegated" | "engine"). + l1_owners: Mapping[MeshAxisKind, L1Owner] = field(default_factory=dict) + + @property + def world_size(self) -> int: + # EP and stage_replica are intentionally excluded. EP reuses the TP*DP + # ranks rather than adding a dimension; stage_replica spins up separate + # engines (each its own world), not extra ranks in this engine's group. + # This product is what vLLM calls ``world_size_across_dp`` (tp * pp * dp). + return self.tensor_parallel_size * self.data_parallel_size * self.pipeline_parallel_size + + @property + def delegated_axes(self) -> tuple[MeshAxisKind, ...]: + return tuple(kind for kind, owner in self.l1_owners.items() if owner == "delegated") + + def as_engine_kwargs(self) -> dict[str, object]: + """Return per-stage kwargs keyed by real OmniEngineArgs/EngineArgs field names. + + Two derived values are intentionally *not* emitted here because they are + not per-stage engine args: + + * ``stage_replica_size`` — a per-stage deploy ``num_replicas`` knob + (``StageDeployConfig``); the deploy layer consumes it separately. + * ``omni_lb_policy`` — a pipeline-wide load-balancer policy the engine + reads once at construction (see ``StrategyApplyResult.omni_lb_policy``); + it is applied at the orchestrator level, not folded into per-stage args. + """ + kwargs: dict[str, object] = { + "tensor_parallel_size": self.tensor_parallel_size, + "data_parallel_size": self.data_parallel_size, + "pipeline_parallel_size": self.pipeline_parallel_size, + } + if self.enable_expert_parallel: + kwargs["enable_expert_parallel"] = True + return kwargs + + +def _is_affinity_dp_routing(routing: RoutingPattern) -> bool: + """True when DP routing demands key-stable (hash) placement.""" + if isinstance(routing, PartitionByHash): + return True + if isinstance(routing, RouteByStage) and routing.routing_policy == "hash": + return True + return False + + +def _resolve_l1_owner(spec: StrategySpec) -> L1Owner: + kind = spec.mesh_axis.kind + declared_owner = spec.shard_extension.get("l1_owner") + owner = declared_owner + if owner is None: + owner = _DEFAULT_L1_OWNER.get(kind, "engine") + if kind not in _DEFAULT_L1_OWNER: + logger.debug( + "[composable_parallel] axis %r has unknown kind %r with no default l1_owner; falling back to %r", + spec.name, + kind, + owner, + ) + else: + logger.debug( + "[composable_parallel] axis %r (kind %r) declared no l1_owner; using default %r", + spec.name, + kind, + owner, + ) + if owner not in _VALID_L1_OWNERS: + # Distinguish a bad value supplied via shard_extension from a bad default + # indexed out of _DEFAULT_L1_OWNER, so the source of the invalid owner is + # debuggable rather than swallowed. + source = "shard_extension" if declared_owner is not None else "_DEFAULT_L1_OWNER" + logger.debug( + "[composable_parallel] axis %r (kind %r) resolved invalid l1_owner %r from %s; valid owners are %s", + spec.name, + kind, + owner, + source, + sorted(_VALID_L1_OWNERS), + ) + _fail(f"axis {kind!r} has invalid l1_owner {owner!r}; expected one of {sorted(_VALID_L1_OWNERS)}") + return cast(L1Owner, owner) + + +def _validate_dp(spec: StrategySpec, owner: L1Owner) -> None: + """Validate an engine data-parallel axis. + + DP is realized intra-engine by vLLM's own DP load balancer, so it is + engine-owned and emits no ``omni_lb_policy`` (that string configures omni's + *replica* balancer, a different layer). Key-stable (hash) request affinity is + not something vLLM's DP LB guarantees, so it is rejected rather than silently + dropped. + """ + if owner != "engine": + _fail( + f"dp axis {spec.name!r} is engine data parallelism realized intra-engine by " + f"vLLM's DP load balancer; l1_owner must be 'engine', got {owner!r}. For " + "omni-coordinator-level request fan-out across independent replicas, use a " + "'stage_replica' axis." + ) + if _is_affinity_dp_routing(spec.routing): + _not_implemented( + f"dp axis {spec.name!r} requests key-stable (hash) routing, which vLLM's " + "intra-engine DP load balancer does not guarantee — not supported yet. Use " + "RouteByStage(random|round_robin|least_queue) for stateless DP balancing." + ) + if not isinstance(spec.routing, RouteByStage): + _fail( + f"dp axis {spec.name!r} expects RouteByStage(random|round_robin|least_queue) routing, " + f"got {type(spec.routing).__name__}" + ) + if spec.routing.routing_policy not in _STAGE_POLICY_TO_OMNI_LB: + # Recognized-but-unimplemented routing (key-stable/hash) is already + # handled above via _is_affinity_dp_routing -> _not_implemented. Anything + # left here is an unknown/invalid policy value (e.g. a typo from YAML), + # i.e. invalid input -> AxisTranslationError, not NotImplementedError. + _fail( + f"dp axis {spec.name!r} has invalid routing_policy " + f"{spec.routing.routing_policy!r}; expected one of " + f"{sorted(_STAGE_POLICY_TO_OMNI_LB)}." + ) + + +def _stage_replica_lb_policy(spec: StrategySpec, owner: L1Owner) -> str: + """Return the omni StagePool LB policy for a delegated stage_replica axis. + + A ``stage_replica`` axis is *not* a vLLM world dimension: it stands up N + independent engine replicas of one pipeline stage, coordinated by the omni + coordinator and balanced over a StagePool with stateless policies only + (random / round-robin / least-queue-length). It maps to the per-stage + ``num_replicas`` count plus the pipeline-level ``omni_lb_policy`` string, so + its ``l1_owner`` must be ``"delegated"`` (omni owns the routing). Key-stable + (hash) routing is rejected because omni has no key-stable balancer yet. + """ + if owner != "delegated": + _fail( + f"stage_replica axis {spec.name!r} must be 'delegated' to omni's StagePool load " + f"balancer; got owner {owner!r}. Replica routing is owned by omni's coordinator." + ) + + routing = spec.routing + if _is_affinity_dp_routing(routing): + _not_implemented( + f"stage_replica axis {spec.name!r} requests key-stable (hash) routing, which needs a " + "dedicated load balancer — not implemented yet. Use " + "RouteByStage(random|round_robin|least_queue) to delegate to omni's load balancer." + ) + if not isinstance(routing, RouteByStage): + _fail( + f"stage_replica axis {spec.name!r} expects RouteByStage(random|round_robin|least_queue) " + f"routing, got {type(routing).__name__}" + ) + policy = _STAGE_POLICY_TO_OMNI_LB.get(routing.routing_policy) + if policy is None: + _not_implemented( + f"stage_replica axis {spec.name!r} routing_policy {routing.routing_policy!r} has no omni LB policy" + ) + return policy + + +def _validate_tp(spec: StrategySpec, owner: L1Owner) -> None: + if not isinstance(spec.routing, Broadcast): + _fail(f"tp axis {spec.name!r} expects Broadcast routing, got {type(spec.routing).__name__}") + if owner != "engine": + _fail(f"tp axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}") + + +def _validate_pp(spec: StrategySpec, owner: L1Owner) -> None: + if not isinstance(spec.routing, PipelineMicrobatch): + _fail(f"pp axis {spec.name!r} expects PipelineMicrobatch routing, got {type(spec.routing).__name__}") + if owner != "engine": + _fail(f"pp axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}") + + +def _validate_ep(spec: StrategySpec, owner: L1Owner) -> None: + # Dense EP: every rank still sees the whole batch, experts are sharded + # across ranks inside the engine (sparse MoE all-to-all is a later stage). + if not isinstance(spec.routing, Broadcast): + _fail( + f"ep axis {spec.name!r} expects Broadcast routing (dense expert parallel), " + f"got {type(spec.routing).__name__}" + ) + if owner != "engine": + _fail(f"ep axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}") + + +def translate_strategy_stack(specs: Sequence[StrategySpec]) -> OmniParallelConfig: + """Translate a spec stack into an ``OmniParallelConfig``. + + Supported kinds: dp (engine data parallel), tp, pp, (dense) ep, and + stage_replica (omni replicas). Raises ``NotImplementedError`` for deferred + (affinity / key-stable) routing, and :class:`AxisTranslationError` for any + other invalid spec (a kind not yet translatable, a repeated kind, an owner + incompatible with the axis kind, or unsupported routing). The EP degree must + equal tensor_parallel_size * data_parallel_size. + """ + sizes: dict[str, int] = {"tensor_parallel_size": 1, "data_parallel_size": 1, "pipeline_parallel_size": 1} + owners: dict[MeshAxisKind, L1Owner] = {} + omni_lb_policy: str | None = None + stage_replica_size = 1 + enable_expert_parallel = False + ep_size: int | None = None + + for spec in specs: + kind = spec.mesh_axis.kind + if kind not in _SUPPORTED_KINDS: + _fail( + f"axis kind {kind!r} is not translatable yet (supported: {list(_SUPPORTED_KINDS)}); " + "it is designed-for and lands in a later stage" + ) + if kind in owners: + _fail(f"axis kind {kind!r} appears more than once in the spec stack") + + owner = _resolve_l1_owner(spec) + if kind == "dp": + _validate_dp(spec, owner) + elif kind == "tp": + _validate_tp(spec, owner) + elif kind == "pp": + _validate_pp(spec, owner) + elif kind == "ep": + _validate_ep(spec, owner) + enable_expert_parallel = True + ep_size = spec.mesh_axis.size + elif kind == "stage_replica": + omni_lb_policy = _stage_replica_lb_policy(spec, owner) + stage_replica_size = spec.mesh_axis.size + + if kind in _AXIS_TO_ENGINE_FIELD: + sizes[_AXIS_TO_ENGINE_FIELD[kind]] = spec.mesh_axis.size + owners[kind] = owner + + if enable_expert_parallel: + # Dense EP shards experts across exactly the TP*DP ranks, so the declared + # EP degree must match that product (it is not its own world dimension, + # and PP is excluded). An EP axis of size 1 is a degenerate no-op that + # still sets the flag; it is only valid when TP*DP == 1. + ep_ranks = sizes["tensor_parallel_size"] * sizes["data_parallel_size"] + if ep_size != ep_ranks: + _fail( + f"ep axis size {ep_size} must equal tensor_parallel_size*data_parallel_size " + f"(={ep_ranks}); dense expert parallelism shards experts across exactly those ranks" + ) + + return OmniParallelConfig( + tensor_parallel_size=sizes["tensor_parallel_size"], + data_parallel_size=sizes["data_parallel_size"], + pipeline_parallel_size=sizes["pipeline_parallel_size"], + enable_expert_parallel=enable_expert_parallel, + stage_replica_size=stage_replica_size, + omni_lb_policy=omni_lb_policy, + l1_owners=owners, + ) diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py index 4ab526a8b..1b38a3d34 100644 --- a/vllm_omni/config/config_factory.py +++ b/vllm_omni/config/config_factory.py @@ -6,6 +6,7 @@ from __future__ import annotations import dataclasses import functools +from collections.abc import Mapping from dataclasses import asdict from pathlib import Path from typing import Any @@ -33,6 +34,16 @@ from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero logger = init_logger(__name__) +# Default degree for any parallel axis / replica count that isn't set anywhere +# (CLI, deploy YAML, or pipeline default): a single, un-parallelized rank. +# TODO(composable_parallel): this "1" default and the per-axis device-layout +# fallbacks are currently re-derived in the merge, apply, and reconcile layers. +# Centralize them in one schema so the pre-spawn device guard and the engine-args +# defaults can't drift apart. This is the light slice; the full device-layout +# centralization is tracked as a follow-up. +_DEFAULT_PARALLEL_DEGREE = 1 + + class StageConfigFactory: """Factory that loads pipeline YAML and merges CLI overrides. @@ -267,14 +278,27 @@ class StageConfigFactory: trust_remote_code: bool = False, cli_overrides: dict[str, Any] | None = None, deploy_config_path: str | None = None, + strategy_specs: Mapping[Any, Any] | None = None, **deprecated_kwargs: Any, - ) -> list[StageConfig] | None: + ) -> tuple[list[StageConfig] | None, str | None]: """Load pipeline + deploy config, merge with CLI overrides. Checks OMNI_PIPELINES first, since supported models should be explicitly registered. If a model is not registered in OMNI_PIPELINES, tries to fall back to using the Transformers config & finding pipelines that have overlapping supported architectures. + + When ``strategy_specs`` is provided (a mapping of role -> list of + ``StrategySpec``), the derived parallel sizing is overlaid onto the + merged stages (see ``vllm_omni.config.composable_parallel``). This is + opt-in: omitting it leaves the existing merge path untouched. + + Returns ``(stages, omni_lb_policy)``: the merged stages (``None`` when the + model is not in the pipeline registry and the caller should fall back to + the legacy YAML path) and the strategy-derived, pipeline-wide + ``omni_lb_policy`` (``None`` when no stage_replica axis set one). The + policy is returned rather than threaded through a mutable out-param so the + engine can apply it without every intermediate call carrying the dict. """ if cli_overrides is None: cli_overrides = {} @@ -290,8 +314,9 @@ class StageConfigFactory: pipeline_cfg, cli_overrides, deploy_config_path, + strategy_specs=strategy_specs, ) - return None + return None, None @classmethod def _create_from_registry( @@ -300,12 +325,17 @@ class StageConfigFactory: pipeline_cfg: PipelineConfig, cli_overrides: dict[str, Any], deploy_config_path: str | None = None, + strategy_specs: Mapping[Any, Any] | None = None, **deprecated_kwargs: Any, - ) -> list[StageConfig]: + ) -> tuple[list[StageConfig], str | None]: """Create StageConfigs from pipeline registry + deploy YAML. Precedence: caller-typed (non-None) value > deploy YAML > StageDeployConfig dataclass default. + + Returns ``(stages, omni_lb_policy)`` — the strategy-derived pipeline-wide + load-balancer policy (``None`` when no strategy set one) travels with the + stages instead of through a mutable out-param. """ # Resolve deploy config path if deploy_config_path is None: @@ -338,12 +368,127 @@ class StageConfigFactory: stages = merge_pipeline_deploy(pipeline_cfg, deploy_cfg, cli_overrides) + # Overlay declarative parallel strategies (opt-in) before CLI overrides. + applied = cls._apply_strategy_specs(stages, strategy_specs) + explicit_overrides = {k: v for k, v in cli_overrides.items() if v is not None} for stage in stages: stage.runtime_overrides = cls._merge_cli_overrides(stage, explicit_overrides) - return stages + # Re-validate the resolved layout now that CLI overrides are on top. + cls._reconcile_strategy_with_cli(stages, applied) + + omni_lb_policy = applied.omni_lb_policy if applied is not None else None + return stages, omni_lb_policy + + @staticmethod + def _apply_strategy_specs( + stages: list[StageConfig], + strategy_specs: Mapping[Any, Any] | None, + ) -> Any: + """Overlay derived parallel sizing onto merged stages (opt-in). + + ``omni_lb_policy`` cannot be set from stage configs (omni reads it once + at orchestrator construction), so a derived non-default policy is logged + here and carried on the returned ``StrategyApplyResult`` for the caller + to hand to the orchestrator. + + Returns the ``StrategyApplyResult`` (or ``None`` when no strategy was + supplied) so the caller can re-validate the resolved layout once CLI + overrides have been merged on top, and read its ``omni_lb_policy``. + """ + if not strategy_specs: + return None + from vllm_omni.config.composable_parallel import apply_strategy_specs + + applied = apply_strategy_specs(stages, strategy_specs) + if applied.omni_lb_policy is not None: + logger.info( + "[composable_parallel] strategy derived omni_lb_policy=%r; it will be applied " + "to the orchestrator unless an explicit --omni-lb-policy was given.", + applied.omni_lb_policy, + ) + return applied + + @staticmethod + def _reconcile_strategy_with_cli( + stages: list[StageConfig], + applied: Any, + ) -> None: + """Reconcile CLI overrides applied *after* a strategy overlay. + + CLI overrides are applied last (in ``to_omegaconf``), so for any stage a + strategy touched we must (a) warn loudly when a CLI arg overrides a + strategy-declared axis — the strategy is meant to be the single writer + for the axes it declares, so a silent CLI win is surprising — and + (b) re-run the device-count check against the *effective* (post-CLI) + ``tp``/``dp``/``pp``/``num_replicas``/``devices`` so a CLI + ``--devices``/``--tensor-parallel-size``/``--num-replicas`` cannot slip + past the pre-spawn guard that exists to prevent silent OOMs. + """ + if applied is None: + return + from vllm_omni.config.composable_parallel import check_device_layout + + # Axis kind -> (engine-arg field, strategy-derived attribute on the cfg). + axis_fields = { + "tp": ("tensor_parallel_size", "tensor_parallel_size"), + "dp": ("data_parallel_size", "data_parallel_size"), + "pp": ("pipeline_parallel_size", "pipeline_parallel_size"), + } + + for stage in stages: + cfg = applied.per_stage_config.get(stage.stage_id) + if cfg is None: + continue + overrides = stage.runtime_overrides or {} + declared = set(cfg.l1_owners.keys()) + + for kind, (field_name, attr) in axis_fields.items(): + if kind in declared and overrides.get(field_name) is not None: + cli_val = overrides[field_name] + derived = getattr(cfg, attr) + if cli_val != derived: + logger.warning( + "[composable_parallel] stage %s: CLI %s=%s overrides the " + "strategy-derived %s=%s. The CLI value wins; remove one to avoid ambiguity.", + stage.stage_id, + field_name, + cli_val, + field_name, + derived, + ) + if "stage_replica" in declared and overrides.get("num_replicas") is not None: + cli_val = overrides["num_replicas"] + if cli_val != cfg.stage_replica_size: + logger.warning( + "[composable_parallel] stage %s: CLI num_replicas=%s overrides the " + "strategy-derived num_replicas=%s. The CLI value wins; remove one to avoid ambiguity.", + stage.stage_id, + cli_val, + cfg.stage_replica_size, + ) + + def _eff(field_name: str, fallback: Any) -> Any: + val = overrides.get(field_name) + return val if val is not None else fallback + + def _eff_degree(field_name: str, source: dict[str, Any]) -> int: + # Single place the per-axis "default to 1" fallback is applied, + # for both the override and the YAML-default sides (see the + # _DEFAULT_PARALLEL_DEGREE TODO). + value = _eff(field_name, source.get(field_name, _DEFAULT_PARALLEL_DEGREE)) + return int(value or _DEFAULT_PARALLEL_DEGREE) + + check_device_layout( + _eff("devices", stage.yaml_runtime.get("devices")), + tensor_parallel_size=_eff_degree("tensor_parallel_size", stage.yaml_engine_args), + data_parallel_size=_eff_degree("data_parallel_size", stage.yaml_engine_args), + pipeline_parallel_size=_eff_degree("pipeline_parallel_size", stage.yaml_engine_args), + num_replicas=_eff_degree("num_replicas", stage.yaml_runtime), + role=stage.model_stage, + ) @classmethod def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]: diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 9097237fd..2fcfd43f7 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -434,6 +434,9 @@ class OrchestratorArgs: stage_configs_path: str | None = None deploy_config: str | None = None stage_overrides: str | None = None # raw JSON string; parsed downstream + # Optional composable-parallel strategy.yaml; orchestrator reads it, overlays + # derived sizing onto merged stages, then drops it before per-stage engine args. + strategy_config: str | None = None # === Mode Switches (orchestrator reads, DeployConfig redistributes) === async_chunk: bool | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 3c898ee75..e71af65d3 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -58,7 +58,10 @@ from vllm_omni.engine.stage_runtime import ( create_stage_runtime, ) from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin -from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs +from vllm_omni.entrypoints.utils import ( + load_and_resolve_stage_configs, + parse_stage_overrides, +) from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.metrics.prometheus import OmniRequestCounter @@ -1123,11 +1126,41 @@ class AsyncOmniEngine: return result + def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None: + """Apply a strategy-derived ``omni_lb_policy`` to the engine. + + Precedence: an explicit ``--omni-lb-policy`` always wins. ``"random"`` is + the engine default and is treated as "unset" (indistinguishable from no + flag), so a strategy value overrides it. If the user explicitly passed a + non-default policy that conflicts with the strategy-derived one, raise so + the mismatch is not silently ignored. + """ + if not derived: + return + explicit = kwargs.get("omni_lb_policy") + user_set = explicit is not None and str(explicit) != "random" + if user_set: + if str(explicit) != str(derived): + raise ValueError( + f"Conflicting load-balancer policy: --omni-lb-policy={explicit!r} was given " + f"but the composable-parallel strategy derived omni_lb_policy={derived!r}. " + "Drop --omni-lb-policy to use the strategy value, or make them match." + ) + return + if self._omni_lb_policy != str(derived): + logger.info( + "[composable_parallel] applying strategy-derived omni_lb_policy=%r (was %r).", + derived, + self._omni_lb_policy, + ) + self._omni_lb_policy = str(derived) + def _resolve_stage_configs(self, model: str, kwargs: dict[str, Any]) -> tuple[str, list[Any]]: """Resolve stage configs and inject defaults shared by orchestrator/headless.""" stage_configs_path = kwargs.get("stage_configs_path", None) deploy_config_path = kwargs.pop("deploy_config", None) + strategy_config_path = kwargs.pop("strategy_config", None) stage_overrides_json = kwargs.pop("stage_overrides", None) explicit_stage_configs = kwargs.pop("stage_configs", None) if explicit_stage_configs is not None: @@ -1142,27 +1175,23 @@ class AsyncOmniEngine: base_kwargs = kwargs # Parse --stage-overrides JSON string if provided - stage_overrides = None - if stage_overrides_json: - if isinstance(stage_overrides_json, str): - try: - stage_overrides = json.loads(stage_overrides_json) - except json.JSONDecodeError as exc: - raise ValueError( - f"--stage-overrides is not valid JSON: {exc}. Got: {stage_overrides_json!r}" - ) from exc - else: - stage_overrides = stage_overrides_json + stage_overrides = parse_stage_overrides(stage_overrides_json) - config_path, stage_configs = load_and_resolve_stage_configs( + config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs( model, stage_configs_path, base_kwargs, default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs), deploy_config_path=deploy_config_path, stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, ) + # A strategy.yaml may derive a pipeline-wide load-balancer policy. It is + # an orchestrator-level knob (read once at construction), so apply it here + # rather than as a per-stage config field. + self._apply_strategy_lb_policy(strategy_lb_policy, kwargs) + # Inject diffusion LoRA-related knobs from kwargs if not present in the stage config. for cfg in stage_configs: try: diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index b1ab30dee..86576ea55 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -254,6 +254,15 @@ class OmniServeCommand(CLISubcommand): help="Path to a deploy config YAML (new format with stages/engine_args). " "Mutually exclusive with --stage-configs-path.", ) + omni_config_group.add_argument( + "--strategy-config", + type=str, + default=None, + help="Path to a composable-parallel strategy.yaml. Only applies to " + "registry-based models (e.g. qwen2_5_omni): its derived parallel sizing " + "is overlaid onto the registry-merged stages before per-stage engine " + "args are built.", + ) omni_config_group.add_argument( "--stage-overrides", type=str, @@ -766,7 +775,10 @@ def run_headless(args: TrackingNamespace) -> None: load_omni_transfer_config_for_model, prepare_engine_environment, ) - from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs + from vllm_omni.entrypoints.utils import ( + load_and_resolve_stage_configs, + parse_stage_overrides, + ) model = args.model stage_id: int | None = args.stage_id @@ -802,11 +814,18 @@ def run_headless(args: TrackingNamespace) -> None: args.replica_id, ) - config_path, stage_configs = load_and_resolve_stage_configs( + # Parse --stage-overrides (raw JSON string) exactly like the standard + # engine path (AsyncOmniEngine._resolve_stage_configs) so headless and + # standard launches resolve to the same per-stage device layout. + stage_overrides = parse_stage_overrides(args_dict.get("stage_overrides")) + + config_path, stage_configs, _ = load_and_resolve_stage_configs( model, stage_configs_path, args_dict, deploy_config_path=args_dict.get("deploy_config"), + stage_overrides=stage_overrides, + strategy_config_path=args_dict.get("strategy_config"), ) # Locate the stage config that matches stage_id. diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 5ecbaa9c9..a2c598b31 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -1,3 +1,4 @@ +import json import os import types from collections import Counter @@ -316,7 +317,8 @@ def load_stage_configs_from_model( base_engine_args: dict | None = None, deploy_config_path: str | None = None, stage_overrides: dict[str, dict[str, Any]] | None = None, -) -> list: + strategy_config_path: str | None = None, +) -> tuple[list, str | None]: """Load stage configurations from model's default config file. For models registered in the pipeline registry (new path), uses @@ -330,9 +332,15 @@ def load_stage_configs_from_model( base_engine_args: Base engine args to merge as CLI overrides. deploy_config_path: Optional explicit deploy config path. stage_overrides: Per-stage overrides from --stage-overrides. + strategy_config_path: Optional path to a composable-parallel + ``strategy.yaml`` whose derived sizing is overlaid onto the + registry-merged stages (opt-in; ignored on the legacy YAML path). Returns: - List of stage configuration dictionaries + ``(stage_configs, omni_lb_policy)``: the list of stage configuration + dictionaries plus the strategy-derived pipeline-wide ``omni_lb_policy`` + (``None`` when no strategy set one). The policy is returned rather than + written into a caller-provided mutable dict. """ if base_engine_args is None: base_engine_args = {} @@ -343,26 +351,44 @@ def load_stage_configs_from_model( for key, val in overrides.items(): cli_overrides[f"stage_{stage_id_str}_{key}"] = val - stages = StageConfigFactory.create_from_model( + strategy_specs = None + if strategy_config_path is not None: + from vllm_omni.config.composable_parallel.strategy_loader import load_strategy_specs + + strategy_specs = load_strategy_specs(strategy_config_path) + + stages, omni_lb_policy = StageConfigFactory.create_from_model( model, trust_remote_code=cli_overrides.get("trust_remote_code", False), cli_overrides=cli_overrides, deploy_config_path=deploy_config_path, + strategy_specs=strategy_specs, ) if stages is not None: # Convert StageConfig objects to OmegaConf for backward compat - return [stage.to_omegaconf() for stage in stages] + return [stage.to_omegaconf() for stage in stages], omni_lb_policy - # Legacy fallback: load from YAML + # Legacy fallback: load from YAML. A composable-parallel strategy cannot be + # applied here (it overlays onto registry-merged stages), so warn rather than + # silently dropping the operator's --strategy-config. + if strategy_config_path is not None: + logger.warning( + "--strategy-config (%s) was provided but model %r resolves via the " + "legacy stage_configs YAML path, which does not support " + "composable-parallel strategies; the strategy is ignored. Use a " + "registry-based model to apply it.", + strategy_config_path, + model, + ) stage_config_path = resolve_model_config_path(model) if stage_config_path is None: - return [] + return [], None stage_configs = load_stage_configs_from_yaml( config_path=stage_config_path, base_engine_args=base_engine_args, prefer_stage_engine_args=True, ) - return stage_configs + return stage_configs, None def load_stage_configs_from_yaml( @@ -501,6 +527,25 @@ def filter_stages( return stage_configs +def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None: + """Parse the ``--stage-overrides`` value into a per-stage override dict. + + ``value`` may be a raw JSON string (as supplied on the CLI) or an + already-parsed mapping. Returns ``None`` when no overrides are given. + + Raises: + ValueError: when ``value`` is a string that is not valid JSON. + """ + if not value: + return None + if isinstance(value, str): + try: + return json.loads(value) + except json.JSONDecodeError as exc: + raise ValueError(f"--stage-overrides is not valid JSON: {exc}. Got: {value!r}") from exc + return value + + def load_and_resolve_stage_configs( model: str, stage_configs_path: str | None, @@ -508,7 +553,8 @@ def load_and_resolve_stage_configs( default_stage_cfg_factory: Any = None, deploy_config_path: str | None = None, stage_overrides: dict[str, dict[str, Any]] | None = None, -) -> tuple[str, list]: + strategy_config_path: str | None = None, +) -> tuple[str, list, str | None]: """Load stage configurations from model or YAML file with fallback to defaults. Args: @@ -521,9 +567,13 @@ def load_and_resolve_stage_configs( Mutually exclusive with ``stage_configs_path``. stage_overrides: Per-stage overrides from ``--stage-overrides`` JSON. Keys are stage_id strings, values are dicts of overrides. + strategy_config_path: Optional path to a composable-parallel + ``strategy.yaml`` overlaid onto the registry-merged stages. Returns: - Tuple of (config_path, stage_configs) + Tuple of ``(config_path, stage_configs, omni_lb_policy)`` — the last is + the strategy-derived pipeline-wide load-balancer policy (``None`` when no + strategy set one), returned for the engine to apply. """ if stage_configs_path is not None and deploy_config_path is not None: raise ValueError( @@ -551,13 +601,15 @@ def load_and_resolve_stage_configs( stage_configs_path, ) + omni_lb_policy: str | None = None if deploy_config_path is not None: config_path = deploy_config_path - stage_configs = load_stage_configs_from_model( + stage_configs, omni_lb_policy = load_stage_configs_from_model( model, base_engine_args=kwargs, deploy_config_path=deploy_config_path, stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, ) if not stage_configs: if default_stage_cfg_factory is not None: @@ -567,10 +619,11 @@ def load_and_resolve_stage_configs( stage_configs = [] elif stage_configs_path is None: config_path = resolve_model_config_path(model) - stage_configs = load_stage_configs_from_model( + stage_configs, omni_lb_policy = load_stage_configs_from_model( model, base_engine_args=kwargs, stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, ) if not stage_configs: if default_stage_cfg_factory is not None: @@ -585,7 +638,7 @@ def load_and_resolve_stage_configs( stage_configs = filter_stages(config_path, stage_configs, kwargs) logger.debug(f"stage_configs: {stage_configs}") - return config_path, stage_configs + return config_path, stage_configs, omni_lb_policy def get_final_stage_id_for_e2e(