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