Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
 tzhouamandClaude Fable 5
|
4d41e84c67 | ||
|
Isotr0py
|
b95905deb1 | ||
|
|
ff6f0da271 | ||
|
NumberWan
|
fe478a95ab | ||
|
amy-why-3459
|
cbd3939c3d | ||
|
amy-why-3459
|
9a44e7e019 | ||
|
LHXuuu
|
5c3923509d | ||
|
Zheng Wengang
|
9fe26cfacb | ||
|
林鑫andyurain
|
020792e6be | ||
|
leo
|
b3a4819694 | ||
|
amy-why-3459
|
5dc787d5e8 |
No files matched your search
@@ -108,7 +108,10 @@ steps:
|
||||
depends_on: upload-merge-pipeline
|
||||
steps:
|
||||
- label: "Entrypoint Test with H100"
|
||||
timeout_in_minutes: 60
|
||||
# The suite relaunches an OmniServer per test class (~30 model loads,
|
||||
# incl. ten ~27.5 GiB Bagel loads); 60 min passed only ~1 in 3 runs
|
||||
# even on main, and vLLM 0.25 weight loading needs ~10 min more.
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
@@ -146,7 +149,9 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Entrypoint Test with L4"
|
||||
timeout_in_minutes: 30
|
||||
# 30 min left no headroom for a cold image pull (~11 min for new base
|
||||
# layers) plus per-class engine relaunches; main passed at 26/30.
|
||||
timeout_in_minutes: 45
|
||||
commands:
|
||||
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model"
|
||||
agents:
|
||||
|
||||
@@ -684,3 +684,24 @@ steps:
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · GGUF Plugin Test"
|
||||
soft_fail: true
|
||||
source_file_dependencies:
|
||||
- tests/plugin/test_gguf_diffusion.py
|
||||
- vllm_omni/diffusion/model_loader
|
||||
commands:
|
||||
- pip install "vllm-gguf-plugin>=0.0.3"
|
||||
- timeout 30m pytest -sv tests/plugin/test_gguf_diffusion.py
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
@@ -1,217 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import statistics
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import imageio.v3 as iio
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="LingBot dense parity and timing harness.")
|
||||
parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b")
|
||||
parser.add_argument("--official-repo", default="/tmp/lingbot-video")
|
||||
parser.add_argument("--output-dir", default="/tmp/lingbot_dense_parity")
|
||||
parser.add_argument("--prompt", default="a robotic arm picks up a red block")
|
||||
parser.add_argument("--height", type=int, default=192)
|
||||
parser.add_argument("--width", type=int, default=320)
|
||||
parser.add_argument("--num-frames", type=int, default=9)
|
||||
parser.add_argument("--steps", type=int, default=2)
|
||||
parser.add_argument("--guidance-scale", type=float, default=3.0)
|
||||
parser.add_argument("--shift", type=float, default=3.0)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--fps", type=int, default=24)
|
||||
parser.add_argument("--runs", type=int, default=0, help="Optional in-process steady-state native runs.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _run(cmd: list[str], cwd: Path | None = None) -> tuple[float, str]:
|
||||
start = time.perf_counter()
|
||||
proc = subprocess.run(cmd, cwd=cwd, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
|
||||
elapsed = time.perf_counter() - start
|
||||
if proc.returncode != 0:
|
||||
print(proc.stdout)
|
||||
raise subprocess.CalledProcessError(proc.returncode, proc.args, output=proc.stdout)
|
||||
return elapsed, proc.stdout
|
||||
|
||||
|
||||
def _read_video(path: Path) -> np.ndarray:
|
||||
frames = iio.imread(path, plugin="pyav")
|
||||
if frames.dtype != np.float32:
|
||||
frames = frames.astype(np.float32) / 255.0
|
||||
return frames
|
||||
|
||||
|
||||
def _compare(a: np.ndarray, b: np.ndarray) -> dict[str, Any]:
|
||||
if a.shape != b.shape:
|
||||
return {"shape_a": list(a.shape), "shape_b": list(b.shape), "shape_match": False}
|
||||
diff = a - b
|
||||
mse = float(np.mean(diff * diff))
|
||||
mae = float(np.mean(np.abs(diff)))
|
||||
psnr = float("inf") if mse == 0 else float(20.0 * np.log10(1.0 / np.sqrt(mse)))
|
||||
return {
|
||||
"shape": list(a.shape),
|
||||
"shape_match": True,
|
||||
"mae": mae,
|
||||
"mse": mse,
|
||||
"psnr": psnr,
|
||||
}
|
||||
|
||||
|
||||
def _extract_total_generation_seconds(log: str) -> float | None:
|
||||
match = re.search(r"Total generation time:\s*([0-9.]+)\s*seconds", log)
|
||||
return float(match.group(1)) if match else None
|
||||
|
||||
|
||||
def _peak_memory_mb(output: Any) -> float:
|
||||
result = output[0] if isinstance(output, list) and output else output
|
||||
return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
|
||||
|
||||
|
||||
def _benchmark_native_requests(args: argparse.Namespace) -> dict[str, Any]:
|
||||
load_start = time.perf_counter()
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
model_class_name="LingBotVideoPipeline",
|
||||
flow_shift=args.shift,
|
||||
parallel_config=DiffusionParallelConfig(),
|
||||
)
|
||||
load_seconds = time.perf_counter() - load_start
|
||||
prompt = {"prompt": args.prompt}
|
||||
latencies: list[float] = []
|
||||
peak_mb = 0.0
|
||||
|
||||
for _ in range(args.runs):
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
num_inference_steps=args.steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
generator=generator,
|
||||
output_type="pt",
|
||||
extra_args={"shift": args.shift},
|
||||
)
|
||||
start = time.perf_counter()
|
||||
output = omni.generate(prompt, sampling_params)
|
||||
latencies.append(time.perf_counter() - start)
|
||||
peak_mb = max(peak_mb, _peak_memory_mb(output))
|
||||
del output
|
||||
omni.close()
|
||||
|
||||
return {
|
||||
"load_seconds": load_seconds,
|
||||
"run_seconds": latencies,
|
||||
"mean_seconds": statistics.mean(latencies) if latencies else None,
|
||||
"median_seconds": statistics.median(latencies) if latencies else None,
|
||||
"min_seconds": min(latencies) if latencies else None,
|
||||
"max_seconds": max(latencies) if latencies else None,
|
||||
"peak_mb": peak_mb,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
official_out = output_dir / "official_diffusers.mp4"
|
||||
native_out = output_dir / "vllm_omni_native.mp4"
|
||||
|
||||
official_cmd = [
|
||||
sys.executable,
|
||||
str(Path(args.official_repo) / "scripts" / "inference.py"),
|
||||
"--backend",
|
||||
"diffusers",
|
||||
"--model_dir",
|
||||
args.model,
|
||||
"--mode",
|
||||
"t2v",
|
||||
"--prompt",
|
||||
args.prompt,
|
||||
"--output",
|
||||
str(official_out),
|
||||
"--height",
|
||||
str(args.height),
|
||||
"--width",
|
||||
str(args.width),
|
||||
"--num_frames",
|
||||
str(args.num_frames),
|
||||
"--steps",
|
||||
str(args.steps),
|
||||
"--guidance_scale",
|
||||
str(args.guidance_scale),
|
||||
"--shift",
|
||||
str(args.shift),
|
||||
"--seed",
|
||||
str(args.seed),
|
||||
"--fps",
|
||||
str(args.fps),
|
||||
"--quiet_progress",
|
||||
]
|
||||
native_cmd = [
|
||||
sys.executable,
|
||||
"examples/offline_inference/text_to_video/text_to_video_lingbot.py",
|
||||
"--model",
|
||||
args.model,
|
||||
"--prompt",
|
||||
args.prompt,
|
||||
"--output",
|
||||
str(native_out),
|
||||
"--height",
|
||||
str(args.height),
|
||||
"--width",
|
||||
str(args.width),
|
||||
"--num-frames",
|
||||
str(args.num_frames),
|
||||
"--num-inference-steps",
|
||||
str(args.steps),
|
||||
"--guidance-scale",
|
||||
str(args.guidance_scale),
|
||||
"--flow-shift",
|
||||
str(args.shift),
|
||||
"--seed",
|
||||
str(args.seed),
|
||||
"--fps",
|
||||
str(args.fps),
|
||||
]
|
||||
|
||||
official_time, official_log = _run(official_cmd, cwd=Path(args.official_repo))
|
||||
native_time, native_log = _run(native_cmd, cwd=Path.cwd())
|
||||
|
||||
result = {
|
||||
"official_seconds": official_time,
|
||||
"native_process_seconds": native_time,
|
||||
"native_request_seconds": _extract_total_generation_seconds(native_log),
|
||||
"official_output": str(official_out),
|
||||
"native_output": str(native_out),
|
||||
"accuracy": _compare(_read_video(official_out), _read_video(native_out)),
|
||||
"official_log_tail": official_log.splitlines()[-20:],
|
||||
"native_log_tail": native_log.splitlines()[-20:],
|
||||
}
|
||||
if args.runs > 0:
|
||||
result["native_steady_state"] = _benchmark_native_requests(args)
|
||||
|
||||
result_path = output_dir / "comparison.json"
|
||||
result_path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
|
||||
print(json.dumps(result, indent=2, sort_keys=True))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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`.
|
||||
@@ -123,7 +123,7 @@ After inserting the extra thinker stage, renumber the remaining stages:
|
||||
runtime:
|
||||
devices: "2"
|
||||
engine_input_source: [1]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
|
||||
sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only
|
||||
|
||||
- stage_id: 3
|
||||
runtime:
|
||||
@@ -131,7 +131,7 @@ After inserting the extra thinker stage, renumber the remaining stages:
|
||||
engine_args:
|
||||
max_num_seqs: 1
|
||||
engine_input_source: [2]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
|
||||
sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only
|
||||
```
|
||||
|
||||
Compared with the default Qwen3-Omni config:
|
||||
|
||||
@@ -475,92 +475,64 @@ thinker_hidden_states = output.multimodal_output["24"]
|
||||
|
||||
### Implementation Example
|
||||
|
||||
Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`:
|
||||
Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`. Each inter-stage edge should provide a **coherent processor set** rather than a single monolithic function:
|
||||
|
||||
| Suffix | Role | Runs when |
|
||||
|--------|------|-----------|
|
||||
| `*_full_payload` | Worker-side payload producer | `async_chunk=false`; accumulates tensors and ships via connector |
|
||||
| `*_async_chunk` | Scheduler-side streaming producer | `async_chunk=true`; emits per-chunk payloads |
|
||||
| `*_token_only` | Orchestrator placeholder builder | `async_chunk=false`; allocates downstream prompt slots only |
|
||||
|
||||
```python
|
||||
# qwen3_omni.py
|
||||
# qwen3_omni.py (Thinker → Talker, non-async path)
|
||||
|
||||
def thinker2talker(
|
||||
stage_list: list[Any],
|
||||
engine_input_source: list[int],
|
||||
def thinker2talker_token_only(
|
||||
source_outputs: list[Any],
|
||||
prompt: OmniTokensPrompt | TextPrompt | None = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
streaming_context: Any | None = None,
|
||||
) -> list[OmniTokensPrompt]:
|
||||
"""
|
||||
Process thinker outputs to create talker inputs.
|
||||
|
||||
Args:
|
||||
stage_list: List of stage objects
|
||||
engine_input_source: Source stage IDs (typically [0] for thinker)
|
||||
prompt: Original prompt data
|
||||
|
||||
Returns:
|
||||
List of OmniTokensPrompt for talker stage
|
||||
"""
|
||||
source_stage_id = engine_input_source[0]
|
||||
thinker_outputs = stage_list[source_stage_id].engine_outputs
|
||||
talker_inputs = []
|
||||
|
||||
for thinker_output in thinker_outputs:
|
||||
output = thinker_output.outputs[0]
|
||||
# Extract thinker embeddings and hidden states
|
||||
thinker_prefill_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
|
||||
thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()
|
||||
|
||||
info = {
|
||||
"thinker_prefill_embeddings": thinker_prefill_embeddings,
|
||||
"thinker_hidden_states": thinker_hidden_states,
|
||||
"thinker_sequences": thinker_output.prompt_token_ids + output.token_ids,
|
||||
"thinker_input_ids": thinker_output.prompt_token_ids,
|
||||
}
|
||||
|
||||
talker_inputs.append(
|
||||
OmniTokensPrompt(
|
||||
prompt_token_ids=[0] * computed_length,
|
||||
additional_information=info,
|
||||
multi_modal_data=None,
|
||||
)
|
||||
)
|
||||
|
||||
return talker_inputs
|
||||
"""Allocate talker prefill slots; bulk tensors arrive via the connector."""
|
||||
...
|
||||
|
||||
|
||||
def talker2code2wav(
|
||||
stage_list: list[Any],
|
||||
engine_input_source: list[int],
|
||||
prompt: OmniTokensPrompt | TextPrompt | None = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
) -> list[OmniTokensPrompt]:
|
||||
"""
|
||||
Process talker outputs to create code2wav inputs.
|
||||
"""
|
||||
source_stage_id = engine_input_source[0]
|
||||
talker_outputs = stage_list[source_stage_id].engine_outputs
|
||||
code2wav_inputs = []
|
||||
def thinker2talker_full_payload(
|
||||
transfer_manager: Any,
|
||||
pooling_output: dict[str, Any],
|
||||
request: OmniEngineCoreRequest,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Pack accumulated thinker hidden states into OmniPayload for stage-1."""
|
||||
...
|
||||
|
||||
for talker_output in talker_outputs:
|
||||
output = talker_output.outputs[0]
|
||||
# Extract codec codes
|
||||
codec_codes = (
|
||||
output.multimodal_output["code_predictor_codes"]
|
||||
.to(torch.long)
|
||||
.transpose(0, 1)
|
||||
.cpu()
|
||||
.to(torch.long)
|
||||
.reshape(-1)
|
||||
.tolist()
|
||||
)
|
||||
|
||||
code2wav_inputs.append(
|
||||
OmniTokensPrompt(
|
||||
prompt_token_ids=codec_codes,
|
||||
multi_modal_data=None,
|
||||
)
|
||||
)
|
||||
|
||||
return code2wav_inputs
|
||||
def thinker2talker_async_chunk(
|
||||
transfer_manager: Any,
|
||||
multimodal_output: OmniPayload | dict[str, Any],
|
||||
request: OmniEngineCoreRequest,
|
||||
is_finished: bool = False,
|
||||
) -> OmniPayloadStruct | None:
|
||||
"""Stream thinker rows to talker while async_chunk is enabled."""
|
||||
...
|
||||
```
|
||||
|
||||
Wire these in `pipeline.py`:
|
||||
|
||||
```python
|
||||
StagePipelineConfig(
|
||||
stage_id=0,
|
||||
custom_process_next_stage_input_func=f"{_PROC}.thinker2talker_full_payload",
|
||||
async_chunk_process_next_stage_input_func=f"{_PROC}.thinker2talker_async_chunk",
|
||||
...
|
||||
),
|
||||
StagePipelineConfig(
|
||||
stage_id=1,
|
||||
sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
|
||||
...
|
||||
),
|
||||
```
|
||||
|
||||
Do **not** add a no-suffix `thinker2talker` when a `sync_process_input_func` is already declared — `_select_processor_funcs()` always prefers the `*_token_only` hook in non-async mode, so the bare function would never run. See `docs/design/rfc_stage_input_processors_refactor.md` for the full contract.
|
||||
|
||||
## Testing
|
||||
|
||||
For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
|
||||
|
||||
@@ -46,7 +46,6 @@ th {
|
||||
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `LingBotVideoPipeline` | LingBot-Video dense T2V | `robbyant/lingbot-video-dense-1.3b` | ✅︎ | | | |
|
||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
|
||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
|
||||
|
||||
@@ -9,6 +9,15 @@ for the transformer.
|
||||
|
||||
GGUF is static quantization: the quantized weights are produced before serving.
|
||||
|
||||
!!! note
|
||||
GGUF support has migrated to OOT [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin). Make sure you have GGUF plugin installed before serving a GGUF model.
|
||||
|
||||
Before serving a GGUF model, make sure to install the [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin):
|
||||
|
||||
```bash
|
||||
uv pip install vllm-gguf-plugin>=0.0.3
|
||||
```
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | Support |
|
||||
|
||||
@@ -220,7 +220,7 @@ def parse_args() -> argparse.Namespace:
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
|
||||
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"],
|
||||
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -202,17 +202,11 @@ def parse_args() -> argparse.Namespace:
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["fp8", "int8", "gguf"],
|
||||
choices=["fp8", "int8"],
|
||||
help="Quantization method for the transformer. "
|
||||
"Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). "
|
||||
"Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8). "
|
||||
"Default: None (no quantization, uses BF16).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gguf-model",
|
||||
type=str,
|
||||
default=None,
|
||||
help=("GGUF file path or HF reference for transformer weights. Required when --quantization gguf is set."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ignored-layers",
|
||||
type=str,
|
||||
@@ -402,14 +396,7 @@ def main():
|
||||
# ignored_layers is specified so the list flows through OmniDiffusionConfig
|
||||
quant_kwargs: dict[str, Any] = {}
|
||||
ignored_layers = [s.strip() for s in args.ignored_layers.split(",") if s.strip()] if args.ignored_layers else None
|
||||
if args.quantization == "gguf":
|
||||
if not args.gguf_model:
|
||||
raise ValueError("--gguf-model is required when --quantization gguf is set.")
|
||||
quant_kwargs["quantization_config"] = {
|
||||
"method": "gguf",
|
||||
"gguf_model": args.gguf_model,
|
||||
}
|
||||
elif args.quantization and ignored_layers:
|
||||
if args.quantization and ignored_layers:
|
||||
quant_kwargs["quantization_config"] = {
|
||||
"method": args.quantization,
|
||||
"ignored_layers": ignored_layers,
|
||||
|
||||
@@ -231,7 +231,7 @@ python text_to_video.py \
|
||||
- `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
|
||||
- `--audio-sample-rate`: audio sample rate for embedded audio (when the
|
||||
pipeline returns audio; LTX-2.3 outputs 48kHz audio).
|
||||
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
|
||||
- `--quantization`: quantization method (such as `fp8` for FP8).
|
||||
- `--flow-shift`: scheduler flow_shift parameter.
|
||||
- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
|
||||
|
||||
|
||||
@@ -227,7 +227,7 @@ def parse_args() -> argparse.Namespace:
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
|
||||
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"],
|
||||
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
|
||||
)
|
||||
|
||||
|
||||
@@ -1,124 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate text-to-video with the LingBot dense preset.")
|
||||
parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b")
|
||||
parser.add_argument("--prompt", default="a robotic arm picks up a red block")
|
||||
parser.add_argument("--negative-prompt", default="")
|
||||
parser.add_argument("--output", default="lingbot_video_output.mp4")
|
||||
parser.add_argument("--height", type=int, default=192)
|
||||
parser.add_argument("--width", type=int, default=320)
|
||||
parser.add_argument("--num-frames", type=int, default=9)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=2)
|
||||
parser.add_argument("--guidance-scale", type=float, default=3.0)
|
||||
parser.add_argument("--flow-shift", type=float, default=3.0)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--fps", type=int, default=24)
|
||||
parser.add_argument("--output-type", choices=["pt", "np"], default="pt")
|
||||
parser.add_argument("--batch-cfg", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _peak_memory_mb(output: Any) -> float:
|
||||
result = output[0] if isinstance(output, list) and output else output
|
||||
return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
|
||||
|
||||
|
||||
def _extract_video(output: Any) -> Any:
|
||||
if isinstance(output, list):
|
||||
output = output[0] if output else None
|
||||
if isinstance(output, OmniRequestOutput):
|
||||
if output.final_output_type != "image":
|
||||
raise ValueError(f"Unexpected output type {output.final_output_type!r}.")
|
||||
if not output.images:
|
||||
raise ValueError("No video frames found in OmniRequestOutput.")
|
||||
return output.images[0] if len(output.images) == 1 else output.images
|
||||
return output
|
||||
|
||||
|
||||
def _to_frame_array(video: Any) -> np.ndarray:
|
||||
if isinstance(video, torch.Tensor):
|
||||
tensor = video.detach().cpu()
|
||||
if tensor.dim() == 5:
|
||||
tensor = tensor[0] if tensor.shape[1] not in (3, 4) else tensor[0].permute(1, 2, 3, 0)
|
||||
elif tensor.dim() == 4 and tensor.shape[0] in (3, 4):
|
||||
tensor = tensor.permute(1, 2, 3, 0)
|
||||
if tensor.is_floating_point():
|
||||
if float(tensor.min()) < 0.0:
|
||||
tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
|
||||
else:
|
||||
tensor = tensor.clamp(0, 1)
|
||||
return tensor.float().numpy()
|
||||
if isinstance(video, np.ndarray):
|
||||
array = video[0] if video.ndim == 5 else video
|
||||
if np.issubdtype(array.dtype, np.integer):
|
||||
return array.astype(np.float32) / 255.0
|
||||
return array.astype(np.float32, copy=False)
|
||||
if isinstance(video, list):
|
||||
return np.stack([_to_frame_array(frame) for frame in video], axis=0)
|
||||
return np.asarray(video, dtype=np.float32)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
model_class_name="LingBotVideoPipeline",
|
||||
flow_shift=args.flow_shift,
|
||||
parallel_config=DiffusionParallelConfig(),
|
||||
)
|
||||
|
||||
prompt: dict[str, str] = {"prompt": args.prompt}
|
||||
if args.negative_prompt:
|
||||
prompt["negative_prompt"] = args.negative_prompt
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
generator=generator,
|
||||
output_type=args.output_type,
|
||||
extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift},
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
output = omni.generate(prompt, sampling_params)
|
||||
elapsed = time.perf_counter() - start
|
||||
print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)")
|
||||
peak_mb = _peak_memory_mb(output)
|
||||
if peak_mb:
|
||||
print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
|
||||
|
||||
video = _to_frame_array(_extract_video(output))
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
export_to_video(list(video), str(output_path), fps=args.fps)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
omni.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -48,7 +48,6 @@ recipes/
|
||||
| [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
|
||||
| [`MammothModa2/MammothModa2-Preview.md`](./MammothModa2/MammothModa2-Preview.md) | Text-to-image with the shared offline image example (AR → DiT) | 1x L40S 48GB / 1x ≥40GB GPU |
|
||||
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
|
||||
| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense text-to-video serving | 1x L20X / A100 / H100 |
|
||||
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
|
||||
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
|
||||
| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
|
||||
|
||||
@@ -1,143 +0,0 @@
|
||||
# LingBot-Video
|
||||
|
||||
> Native dense text-to-video serving for `robbyant/lingbot-video-dense-1.3b`
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Robbyant
|
||||
- Model: `robbyant/lingbot-video-dense-1.3b`
|
||||
- Task: Text-to-video generation
|
||||
- Mode: Offline generation and online serving with the OpenAI-compatible `/v1/videos` API
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want to run the dense LingBot-Video checkpoint with
|
||||
vLLM-Omni's native pipeline. The runtime path does not import the upstream
|
||||
`lingbot_video` Python package; it loads the checkpoint components directly
|
||||
with the in-tree `LingBotVideoPipeline`, dense DiT transformer, shared FlowUniPC
|
||||
scheduler, Qwen3-VL text encoder, and Wan VAE.
|
||||
|
||||
This recipe is intentionally text-to-video only. Text-to-image, image-conditioned
|
||||
video, TI2V, MoE checkpoints, and fused-expert kernels are not covered here.
|
||||
|
||||
## References
|
||||
|
||||
- Dense checkpoint: <https://huggingface.co/robbyant/lingbot-video-dense-1.3b>
|
||||
- Upstream project: <https://github.com/Robbyant/lingbot-video>
|
||||
- Related offline example: [`examples/offline_inference/text_to_video/text_to_video_lingbot.py`](../../examples/offline_inference/text_to_video/text_to_video_lingbot.py)
|
||||
- Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
This recipe documents the CUDA single-GPU dense checkpoint path. Multi-GPU
|
||||
parallelism, Cache-DiT, CPU offload, and the LingBot MoE path are not validated
|
||||
for this model in this PR.
|
||||
|
||||
## GPU
|
||||
|
||||
### 1 x NVIDIA L20X / A100 / H100
|
||||
|
||||
The dense 1.3B checkpoint has been smoke-tested on a single NVIDIA L20X at a
|
||||
small validation shape (`192x320`, 9 frames, 2 steps). Use a larger GPU or lower
|
||||
resolution/frame count if your target workload runs out of memory.
|
||||
|
||||
#### Offline T2V
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 \
|
||||
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
|
||||
--model robbyant/lingbot-video-dense-1.3b \
|
||||
--prompt "a robotic arm picks up a red block" \
|
||||
--output lingbot_t2v.mp4 \
|
||||
--height 192 \
|
||||
--width 320 \
|
||||
--num-frames 9 \
|
||||
--num-inference-steps 2 \
|
||||
--guidance-scale 3.0 \
|
||||
--flow-shift 3.0 \
|
||||
--seed 42 \
|
||||
--fps 24
|
||||
```
|
||||
|
||||
#### Online Serving
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 \
|
||||
vllm serve robbyant/lingbot-video-dense-1.3b \
|
||||
--omni \
|
||||
--model-class-name LingBotVideoPipeline \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
After the server is ready, submit a text-to-video job:
|
||||
|
||||
```bash
|
||||
create_response=$(curl -s http://localhost:8091/v1/videos \
|
||||
-F "model=robbyant/lingbot-video-dense-1.3b" \
|
||||
-F "prompt=a robotic arm picks up a red block" \
|
||||
-F "width=320" \
|
||||
-F "height=192" \
|
||||
-F "num_frames=9" \
|
||||
-F "fps=24" \
|
||||
-F "num_inference_steps=2" \
|
||||
-F "guidance_scale=3.0" \
|
||||
-F "flow_shift=3.0" \
|
||||
-F "seed=42")
|
||||
|
||||
video_id=$(echo "${create_response}" | jq -r '.id')
|
||||
while true; do
|
||||
status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
|
||||
if [ "${status}" = "completed" ]; then
|
||||
break
|
||||
fi
|
||||
if [ "${status}" = "failed" ]; then
|
||||
curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
|
||||
exit 1
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
|
||||
curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o lingbot_t2v.mp4
|
||||
```
|
||||
|
||||
## Key Parameters
|
||||
|
||||
| Parameter | Suggested smoke value | Notes |
|
||||
|-----------|-----------------------|-------|
|
||||
| `height` | `192` | Must be a multiple of 16 |
|
||||
| `width` | `320` | Must be a multiple of 16 |
|
||||
| `num_frames` | `9` | Must be `1` or `4n + 1`; this PR validates T2V with video outputs |
|
||||
| `num_inference_steps` | `2` | Use more steps for quality sweeps |
|
||||
| `guidance_scale` | `3.0` | CFG is active when this is greater than `1.0` |
|
||||
| `flow_shift` | `3.0` | Scheduler flow-shift; aliases the pipeline's internal `shift` |
|
||||
| `negative_prompt` | model default | Optional text describing artifacts to avoid |
|
||||
| `fps` | `24` | Output MP4 frame rate |
|
||||
|
||||
## Validation
|
||||
|
||||
Local dense smoke run:
|
||||
|
||||
- Shape: 9 frames at `192x320`
|
||||
- Steps: 2
|
||||
- Request generation time: `0.2923s`
|
||||
- Peak reserved GPU memory: `14548 MiB`
|
||||
|
||||
Local parity harness against the upstream repository:
|
||||
|
||||
- Shape: `[9, 192, 320, 3]`
|
||||
- MAE: `0.0065238`
|
||||
- MSE: `0.00006650`
|
||||
- PSNR: `41.77 dB`
|
||||
- Native request time: `0.2875s`
|
||||
|
||||
Do not treat these as production benchmarks; they are draft-PR validation
|
||||
numbers for the small dense smoke configuration.
|
||||
|
||||
## Known Limitations
|
||||
|
||||
- T2V only. T2I, I2V, and TI2V are not claimed by this PR.
|
||||
- Dense checkpoint only. MoE/fused-expert support belongs in a separate PR.
|
||||
- No Cache-DiT, TeaCache, tensor parallelism, sequence parallelism, CFG
|
||||
parallelism, HSDP, CPU offload, or VAE patch parallelism validation yet.
|
||||
- Only one request per LingBot pipeline batch is currently supported.
|
||||
@@ -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}"
|
||||
@@ -751,3 +751,14 @@ def test_diffusion_config_field_classification_covers_current_fields():
|
||||
"distributed_executor_backend",
|
||||
} <= omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS
|
||||
assert "prompt_file_path" in omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS
|
||||
|
||||
|
||||
def test_diffusion_config_projection_keeps_mapping_quantization_config_serializable():
|
||||
quantization_config = {
|
||||
"method": "example_quant",
|
||||
"weights": "weights.bin",
|
||||
}
|
||||
|
||||
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config)
|
||||
|
||||
assert cfg.quantization_config == quantization_config
|
||||
@@ -1,34 +0,0 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_lingbot_video_single_device",
|
||||
"description": "Single-device LingBot dense T2V baseline at 320x192, 9 frames, 2 steps.",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "robbyant/lingbot-video-dense-1.3b",
|
||||
"serve_args": {
|
||||
"model-class-name": "LingBotVideoPipeline",
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "320x192_frames9_steps2",
|
||||
"dataset": "random",
|
||||
"task": "t2v",
|
||||
"backend": "v1/videos",
|
||||
"width": 320,
|
||||
"height": 192,
|
||||
"num-frames": 9,
|
||||
"fps": 24,
|
||||
"num-inference-steps": 2,
|
||||
"num-prompts": 3,
|
||||
"max-concurrency": 1,
|
||||
"warmup-requests": 1,
|
||||
"warmup-concurrency": 1,
|
||||
"warmup-num-inference-steps": 2,
|
||||
"enable-negative-prompt": true,
|
||||
"skip-performance-assertion": true
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -3,7 +3,10 @@
|
||||
"test_name": "test_voxcpm2",
|
||||
"server_params": {
|
||||
"model": "openbmb/VoxCPM2",
|
||||
"trust_remote_code": true
|
||||
"trust_remote_code": true,
|
||||
"extra_cli_args": [
|
||||
"--trust-remote-code"
|
||||
]
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
|
||||
@@ -1,112 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for the Flux2-Klein GGUF adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
|
||||
from vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein import (
|
||||
Flux2KleinGGUFAdapter,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_od_config(
|
||||
model_class_name: str = "Flux2Pipeline",
|
||||
model_type: str = "flux",
|
||||
):
|
||||
return SimpleNamespace(
|
||||
model_class_name=model_class_name,
|
||||
tf_model_config={"model_type": model_type},
|
||||
)
|
||||
|
||||
|
||||
def _make_source(prefix: str = "", subfolder: str = "transformer"):
|
||||
return SimpleNamespace(prefix=prefix, subfolder=subfolder)
|
||||
|
||||
|
||||
def test_flux2_klein_adapter_selected_for_flux_family():
|
||||
adapter = get_gguf_adapter(
|
||||
"dummy.gguf",
|
||||
object(),
|
||||
_make_source(),
|
||||
_make_od_config(),
|
||||
)
|
||||
|
||||
assert isinstance(adapter, Flux2KleinGGUFAdapter)
|
||||
|
||||
|
||||
def test_flux2_klein_adapter_matches_flux_model_type():
|
||||
assert Flux2KleinGGUFAdapter.is_compatible(
|
||||
_make_od_config(model_class_name="OtherPipeline", model_type="flux-dev"),
|
||||
object(),
|
||||
_make_source(),
|
||||
)
|
||||
|
||||
|
||||
def test_flux2_klein_adapter_renames_core_projection_paths(monkeypatch: pytest.MonkeyPatch):
|
||||
import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
flux_module,
|
||||
"gguf_quant_weights_iterator",
|
||||
lambda _path: iter(
|
||||
[
|
||||
("img_in.qweight", torch.ones((1, 1))),
|
||||
("time_in.in_layer.qweight_type", torch.tensor(1)),
|
||||
("double_blocks.0.img_attn.norm.query_norm.weight", torch.full((1, 1), 2.0)),
|
||||
("final_layer.linear.qweight", torch.full((1, 1), 3.0)),
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
adapter = Flux2KleinGGUFAdapter(
|
||||
"dummy.gguf",
|
||||
object(),
|
||||
_make_source(),
|
||||
_make_od_config(),
|
||||
)
|
||||
|
||||
weights = list(adapter.weights_iterator())
|
||||
names = [name for name, _ in weights]
|
||||
|
||||
assert "x_embedder.qweight" in names
|
||||
assert "time_guidance_embed.timestep_embedder.linear_1.qweight_type" in names
|
||||
assert "transformer_blocks.0.attn.norm_q.weight" in names
|
||||
assert "proj_out.qweight" in names
|
||||
|
||||
|
||||
def test_flux2_klein_adapter_swaps_final_adaln_shift_and_scale(monkeypatch: pytest.MonkeyPatch):
|
||||
import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
flux_module,
|
||||
"gguf_quant_weights_iterator",
|
||||
lambda _path: iter(
|
||||
[
|
||||
(
|
||||
"final_layer.adaLN_modulation.1.scale",
|
||||
torch.tensor([1.0, 2.0, 3.0, 4.0]),
|
||||
)
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
adapter = Flux2KleinGGUFAdapter(
|
||||
"dummy.gguf",
|
||||
object(),
|
||||
_make_source(),
|
||||
_make_od_config(),
|
||||
)
|
||||
|
||||
weights = list(adapter.weights_iterator())
|
||||
|
||||
assert len(weights) == 1
|
||||
assert weights[0][0] == "norm_out.linear.weight"
|
||||
assert torch.equal(weights[0][1], torch.tensor([3.0, 4.0, 1.0, 2.0]))
|
||||
@@ -1,184 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
|
||||
from vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image import (
|
||||
QwenImageGGUFAdapter,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_od_config(
|
||||
model_class_name: str = "QwenImagePipeline",
|
||||
model_type: str = "qwen_image",
|
||||
):
|
||||
return SimpleNamespace(
|
||||
model_class_name=model_class_name,
|
||||
tf_model_config={"model_type": model_type},
|
||||
)
|
||||
|
||||
|
||||
def _make_source(prefix: str = "", subfolder: str = "transformer"):
|
||||
return SimpleNamespace(prefix=prefix, subfolder=subfolder)
|
||||
|
||||
|
||||
class _FakeTransformer:
|
||||
def __init__(self) -> None:
|
||||
self._params = {
|
||||
"img_in.qweight": 1,
|
||||
"img_in.qweight_type": 1,
|
||||
"txt_in.qweight": 1,
|
||||
"txt_in.qweight_type": 1,
|
||||
"proj_out.qweight": 1,
|
||||
"proj_out.qweight_type": 1,
|
||||
"time_text_embed.timestep_embedder.linear_1.qweight": 1,
|
||||
"time_text_embed.timestep_embedder.linear_1.qweight_type": 1,
|
||||
"time_text_embed.timestep_embedder.linear_2.qweight": 1,
|
||||
"time_text_embed.timestep_embedder.linear_2.qweight_type": 1,
|
||||
"transformer_blocks.0.img_mod.1.qweight": 1,
|
||||
"transformer_blocks.0.img_mod.1.qweight_type": 1,
|
||||
"transformer_blocks.0.txt_mod.1.qweight": 1,
|
||||
"transformer_blocks.0.txt_mod.1.qweight_type": 1,
|
||||
"transformer_blocks.0.attn.to_qkv.qweight": 1,
|
||||
"transformer_blocks.0.attn.to_qkv.qweight_type": 1,
|
||||
"transformer_blocks.0.attn.add_kv_proj.qweight": 1,
|
||||
"transformer_blocks.0.attn.add_kv_proj.qweight_type": 1,
|
||||
"transformer_blocks.0.attn.to_out.weight": 1,
|
||||
"transformer_blocks.0.img_mlp.net.0.proj.qweight": 1,
|
||||
"transformer_blocks.0.img_mlp.net.0.proj.qweight_type": 1,
|
||||
}
|
||||
|
||||
def named_parameters(self):
|
||||
return list(self._params.items())
|
||||
|
||||
def named_buffers(self):
|
||||
return []
|
||||
|
||||
def get_submodule(self, _prefix: str):
|
||||
return self
|
||||
|
||||
|
||||
def test_qwen_adapter_selected_for_qwen_image_family():
|
||||
adapter = get_gguf_adapter(
|
||||
"dummy.gguf",
|
||||
_FakeTransformer(),
|
||||
_make_source(),
|
||||
_make_od_config(),
|
||||
)
|
||||
assert isinstance(adapter, QwenImageGGUFAdapter)
|
||||
|
||||
|
||||
def test_qwen_adapter_matches_multiple_pipeline_variants():
|
||||
for model_class_name in (
|
||||
"QwenImagePipeline",
|
||||
"QwenImageEditPipeline",
|
||||
"QwenImageEditPlusPipeline",
|
||||
"QwenImageLayeredPipeline",
|
||||
):
|
||||
assert QwenImageGGUFAdapter.is_compatible(
|
||||
_make_od_config(model_class_name=model_class_name),
|
||||
_FakeTransformer(),
|
||||
_make_source(),
|
||||
)
|
||||
|
||||
|
||||
def test_qwen_adapter_preserves_split_projection_names(monkeypatch: pytest.MonkeyPatch):
|
||||
import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
qwen_image_module,
|
||||
"gguf_quant_weights_iterator",
|
||||
lambda _path: iter(
|
||||
[
|
||||
("transformer_blocks.0.attn.to_q.qweight_type", 1),
|
||||
("transformer_blocks.0.attn.to_q.qweight", 2),
|
||||
("transformer_blocks.0.attn.to_k.qweight_type", 3),
|
||||
("transformer_blocks.0.attn.to_k.qweight", 4),
|
||||
("transformer_blocks.0.attn.to_out.0.qweight_type", 5),
|
||||
("transformer_blocks.0.attn.to_out.0.qweight", 6),
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
adapter = QwenImageGGUFAdapter(
|
||||
"dummy.gguf",
|
||||
_FakeTransformer(),
|
||||
_make_source(),
|
||||
_make_od_config(),
|
||||
)
|
||||
|
||||
weights = list(adapter.weights_iterator())
|
||||
|
||||
assert ("transformer_blocks.0.attn.to_q.qweight_type", 1) in weights
|
||||
assert ("transformer_blocks.0.attn.to_q.qweight", 2) in weights
|
||||
assert ("transformer_blocks.0.attn.to_k.qweight_type", 3) in weights
|
||||
assert ("transformer_blocks.0.attn.to_k.qweight", 4) in weights
|
||||
assert ("transformer_blocks.0.attn.to_out.0.qweight_type", 5) in weights
|
||||
assert ("transformer_blocks.0.attn.to_out.0.qweight", 6) in weights
|
||||
|
||||
|
||||
def test_qwen_adapter_keeps_top_level_quantized_weights(monkeypatch: pytest.MonkeyPatch):
|
||||
import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
qwen_image_module,
|
||||
"gguf_quant_weights_iterator",
|
||||
lambda _path: iter(
|
||||
[
|
||||
("img_in.qweight_type", 1),
|
||||
("img_in.qweight", 2),
|
||||
("time_text_embed.timestep_embedder.linear_1.qweight_type", 3),
|
||||
("time_text_embed.timestep_embedder.linear_1.qweight", 4),
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
adapter = QwenImageGGUFAdapter(
|
||||
"dummy.gguf",
|
||||
_FakeTransformer(),
|
||||
_make_source(),
|
||||
_make_od_config(),
|
||||
)
|
||||
|
||||
weights = list(adapter.weights_iterator())
|
||||
|
||||
assert ("img_in.qweight_type", 1) in weights
|
||||
assert ("img_in.qweight", 2) in weights
|
||||
assert ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3) in weights
|
||||
assert ("time_text_embed.timestep_embedder.linear_1.qweight", 4) in weights
|
||||
|
||||
|
||||
def test_qwen_adapter_keeps_modulation_quantized_weights(monkeypatch: pytest.MonkeyPatch):
|
||||
import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
qwen_image_module,
|
||||
"gguf_quant_weights_iterator",
|
||||
lambda _path: iter(
|
||||
[
|
||||
("transformer_blocks.0.img_mod.1.qweight_type", 1),
|
||||
("transformer_blocks.0.img_mod.1.qweight", 2),
|
||||
("transformer_blocks.0.txt_mod.1.qweight_type", 3),
|
||||
("transformer_blocks.0.txt_mod.1.qweight", 4),
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
adapter = QwenImageGGUFAdapter(
|
||||
"dummy.gguf",
|
||||
_FakeTransformer(),
|
||||
_make_source(),
|
||||
_make_od_config(),
|
||||
)
|
||||
|
||||
weights = list(adapter.weights_iterator())
|
||||
|
||||
assert ("transformer_blocks.0.img_mod.1.qweight_type", 1) in weights
|
||||
assert ("transformer_blocks.0.img_mod.1.qweight", 2) in weights
|
||||
assert ("transformer_blocks.0.txt_mod.1.qweight_type", 3) in weights
|
||||
assert ("transformer_blocks.0.txt_mod.1.qweight", 4) in weights
|
||||
@@ -1,83 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for the Z-Image GGUF adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
|
||||
from vllm_omni.diffusion.model_loader.gguf_adapters.z_image import (
|
||||
ZImageGGUFAdapter,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_od_config(
|
||||
model_class_name: str = "ZImagePipeline",
|
||||
model_type: str = "z_image",
|
||||
):
|
||||
return SimpleNamespace(
|
||||
model_class_name=model_class_name,
|
||||
tf_model_config={"model_type": model_type},
|
||||
)
|
||||
|
||||
|
||||
def _make_source(prefix: str = "", subfolder: str = "transformer"):
|
||||
return SimpleNamespace(prefix=prefix, subfolder=subfolder)
|
||||
|
||||
|
||||
def test_z_image_adapter_selected_for_z_image_family():
|
||||
adapter = get_gguf_adapter(
|
||||
"dummy.gguf",
|
||||
object(),
|
||||
_make_source(),
|
||||
_make_od_config(),
|
||||
)
|
||||
|
||||
assert isinstance(adapter, ZImageGGUFAdapter)
|
||||
|
||||
|
||||
def test_z_image_adapter_matches_model_type_variants():
|
||||
for model_type in ("z_image", "zimage", "z-image"):
|
||||
assert ZImageGGUFAdapter.is_compatible(
|
||||
_make_od_config(model_class_name="OtherPipeline", model_type=model_type),
|
||||
object(),
|
||||
_make_source(),
|
||||
)
|
||||
|
||||
|
||||
def test_z_image_adapter_renames_known_gguf_tensor_paths(monkeypatch: pytest.MonkeyPatch):
|
||||
import vllm_omni.diffusion.model_loader.gguf_adapters.z_image as z_image_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
z_image_module,
|
||||
"gguf_quant_weights_iterator",
|
||||
lambda _path: iter(
|
||||
[
|
||||
("model.diffusion_model.final_layer.qweight", torch.ones((1, 1))),
|
||||
("model.diffusion_model.x_embedder.qweight_type", torch.tensor(1)),
|
||||
("transformer_blocks.0.attention.out.weight", torch.full((1, 1), 2.0)),
|
||||
("transformer_blocks.0.attention.qkv.qweight", torch.full((1, 1), 3.0)),
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
adapter = ZImageGGUFAdapter(
|
||||
"dummy.gguf",
|
||||
object(),
|
||||
_make_source(),
|
||||
_make_od_config(),
|
||||
)
|
||||
|
||||
weights = list(adapter.weights_iterator())
|
||||
names = [name for name, _ in weights]
|
||||
|
||||
assert "all_final_layer.2-1.qweight" in names
|
||||
assert "all_x_embedder.2-1.qweight_type" in names
|
||||
assert "transformer_blocks.0.attention.to_out.0.weight" in names
|
||||
assert "transformer_blocks.0.attention.to_qkv.qweight" in names
|
||||
@@ -16,7 +16,6 @@ from vllm.config.load import LoadConfig
|
||||
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
|
||||
from vllm_omni.diffusion.models.helios import HeliosPipeline
|
||||
from vllm_omni.diffusion.registry import initialize_model
|
||||
|
||||
@@ -111,27 +110,6 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
|
||||
loader.load_weights(model)
|
||||
|
||||
|
||||
def test_qwen_model_class_selects_qwen_gguf_adapter():
|
||||
od_config = type(
|
||||
"Config",
|
||||
(),
|
||||
{
|
||||
"model_class_name": "QwenImagePipeline",
|
||||
"tf_model_config": {"model_type": "qwen_image"},
|
||||
},
|
||||
)()
|
||||
source = DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path="dummy",
|
||||
subfolder="transformer",
|
||||
revision=None,
|
||||
prefix="transformer.",
|
||||
)
|
||||
|
||||
adapter = get_gguf_adapter("dummy.gguf", object(), source, od_config)
|
||||
|
||||
assert adapter.__class__.__name__ == "QwenImageGGUFAdapter"
|
||||
|
||||
|
||||
class _ConfigAwareModel(nn.Module):
|
||||
def __init__(self, *, od_config):
|
||||
super().__init__()
|
||||
@@ -186,7 +164,6 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
loader.load_weights = lambda model: None # type: ignore[assignment]
|
||||
loader._process_weights_after_loading = lambda model, target_device: None # type: ignore[assignment]
|
||||
loader._is_gguf_quantization = lambda: False # type: ignore[assignment]
|
||||
|
||||
monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
|
||||
monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))
|
||||
|
||||
@@ -1,259 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for GGUF-specific DiffusersPipelineLoader behavior."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_module
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _DummyModel(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.transformer = nn.Linear(2, 2, bias=True)
|
||||
self.vae = nn.Linear(2, 2, bias=False)
|
||||
self.register_buffer("transformer_buffer", torch.ones(1))
|
||||
self.calls: list[list[str]] = []
|
||||
self.weights_sources = [
|
||||
DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path="dummy",
|
||||
subfolder="transformer",
|
||||
revision=None,
|
||||
prefix="transformer.",
|
||||
fall_back_to_pt=True,
|
||||
),
|
||||
DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path="dummy",
|
||||
subfolder="vae",
|
||||
revision=None,
|
||||
prefix="vae.",
|
||||
fall_back_to_pt=True,
|
||||
),
|
||||
]
|
||||
|
||||
def load_weights(self, weights):
|
||||
loadable = dict(self.named_parameters())
|
||||
loadable.update(dict(self.named_buffers()))
|
||||
seen: list[str] = []
|
||||
loaded: set[str] = set()
|
||||
for name, tensor in weights:
|
||||
seen.append(name)
|
||||
if name in loadable:
|
||||
target = loadable[name]
|
||||
target.data.copy_(tensor.to(dtype=target.dtype))
|
||||
loaded.add(name)
|
||||
self.calls.append(seen)
|
||||
return loaded
|
||||
|
||||
|
||||
def _make_loader() -> DiffusersPipelineLoader:
|
||||
loader = object.__new__(DiffusersPipelineLoader)
|
||||
loader.load_config = SimpleNamespace(
|
||||
download_dir="cache-dir",
|
||||
ignore_patterns=["*.tmp"],
|
||||
)
|
||||
loader.od_config = None
|
||||
loader.quant_config = None
|
||||
loader.counter_before_loading_weights = 0.0
|
||||
loader.counter_after_loading_weights = 0.0
|
||||
return loader
|
||||
|
||||
|
||||
def test_is_gguf_quantization_accepts_dict_config():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = {
|
||||
"method": "gguf",
|
||||
"gguf_model": "weights.gguf",
|
||||
}
|
||||
|
||||
assert loader._is_gguf_quantization() is True
|
||||
|
||||
|
||||
def test_is_gguf_quantization_rejects_non_gguf_dict_config():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = {
|
||||
"method": "fp8",
|
||||
"gguf_model": "weights.gguf",
|
||||
}
|
||||
|
||||
assert loader._is_gguf_quantization() is False
|
||||
|
||||
|
||||
def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = {"method": "gguf"}
|
||||
|
||||
with pytest.raises(ValueError, match="gguf_model"):
|
||||
loader._is_gguf_quantization()
|
||||
|
||||
|
||||
def test_is_gguf_quantization_accepts_object_config():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = SimpleNamespace(
|
||||
get_name=lambda: "gguf",
|
||||
gguf_model="weights.gguf",
|
||||
)
|
||||
|
||||
assert loader._is_gguf_quantization() is True
|
||||
|
||||
|
||||
def test_is_gguf_quantization_uses_fallback_object_without_get_name():
|
||||
loader = _make_loader()
|
||||
loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
|
||||
|
||||
assert loader._is_gguf_quantization() is True
|
||||
|
||||
|
||||
def test_get_model_loadable_names_collects_parameters_and_buffers():
|
||||
loader = _make_loader()
|
||||
model = _DummyModel()
|
||||
|
||||
names = loader._get_model_loadable_names(model)
|
||||
|
||||
assert "transformer.weight" in names
|
||||
assert "transformer.bias" in names
|
||||
assert "transformer_buffer" in names
|
||||
|
||||
|
||||
def test_resolve_gguf_model_path_returns_local_file(tmp_path):
|
||||
loader = _make_loader()
|
||||
gguf_file = tmp_path / "model.gguf"
|
||||
gguf_file.write_bytes(b"gguf")
|
||||
|
||||
resolved = loader._resolve_gguf_model_path(str(gguf_file), revision=None)
|
||||
|
||||
assert resolved == str(gguf_file)
|
||||
|
||||
|
||||
def test_resolve_gguf_model_path_downloads_explicit_gguf_filename(monkeypatch: pytest.MonkeyPatch):
|
||||
loader = _make_loader()
|
||||
|
||||
monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False)
|
||||
monkeypatch.setattr(
|
||||
loader_module,
|
||||
"hf_hub_download",
|
||||
lambda repo_id, filename, revision, cache_dir: f"{cache_dir}/{repo_id}/{filename}",
|
||||
)
|
||||
|
||||
resolved = loader._resolve_gguf_model_path("owner/repo/model-Q4_0.gguf", revision="main")
|
||||
|
||||
assert resolved == "cache-dir/owner/repo/model-Q4_0.gguf"
|
||||
|
||||
|
||||
def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.MonkeyPatch):
|
||||
loader = _make_loader()
|
||||
|
||||
monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False)
|
||||
monkeypatch.setattr(
|
||||
loader_module,
|
||||
"download_gguf",
|
||||
lambda repo_id, quant_type, cache_dir, revision, ignore_patterns: f"{cache_dir}/{repo_id}/{quant_type}.gguf",
|
||||
)
|
||||
|
||||
resolved = loader._resolve_gguf_model_path("owner/repo:Q4_0", revision="main")
|
||||
|
||||
assert resolved == "cache-dir/owner/repo/Q4_0.gguf"
|
||||
|
||||
|
||||
def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch):
|
||||
loader = _make_loader()
|
||||
loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
|
||||
loader.od_config = SimpleNamespace(revision="main")
|
||||
source = DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path="dummy",
|
||||
subfolder="transformer",
|
||||
revision=None,
|
||||
prefix="transformer.",
|
||||
)
|
||||
|
||||
class _Adapter:
|
||||
def weights_iterator(self):
|
||||
yield "weight", torch.ones((2, 2))
|
||||
yield "bias", torch.zeros(2)
|
||||
|
||||
monkeypatch.setattr(loader, "_resolve_gguf_model_path", lambda gguf_model, revision: "resolved.gguf")
|
||||
monkeypatch.setattr(loader_module, "get_gguf_adapter", lambda gguf_file, model, source, od_config: _Adapter())
|
||||
|
||||
weights = list(loader._get_gguf_weights_iterator(source, object()))
|
||||
|
||||
assert weights[0][0] == "transformer.weight"
|
||||
assert torch.equal(weights[0][1], torch.ones((2, 2)))
|
||||
assert weights[1][0] == "transformer.bias"
|
||||
assert torch.equal(weights[1][1], torch.zeros(2))
|
||||
|
||||
|
||||
def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch):
|
||||
loader = _make_loader()
|
||||
model = _DummyModel()
|
||||
|
||||
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
|
||||
monkeypatch.setattr(
|
||||
loader,
|
||||
"_get_gguf_weights_iterator",
|
||||
lambda source, model: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
|
||||
)
|
||||
|
||||
def _hf_iter(source):
|
||||
if source.subfolder == "transformer":
|
||||
return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))])
|
||||
return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))])
|
||||
|
||||
monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter)
|
||||
monkeypatch.setattr(
|
||||
loader,
|
||||
"_get_expected_parameter_names",
|
||||
lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
|
||||
)
|
||||
|
||||
loaded = loader._load_weights_with_gguf(model)
|
||||
|
||||
assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
|
||||
assert model.calls[0] == ["transformer.weight"]
|
||||
assert model.calls[1] == ["transformer.bias"]
|
||||
assert model.calls[2] == ["vae.weight"]
|
||||
|
||||
|
||||
def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch):
|
||||
loader = _make_loader()
|
||||
model = _DummyModel()
|
||||
|
||||
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
|
||||
monkeypatch.setattr(
|
||||
loader,
|
||||
"_get_gguf_weights_iterator",
|
||||
lambda source, model: iter(
|
||||
[
|
||||
("transformer.weight", torch.ones((2, 2), dtype=torch.float32)),
|
||||
("transformer.bias", torch.zeros(2, dtype=torch.float32)),
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
hf_calls: list[str] = []
|
||||
|
||||
def _hf_iter(source):
|
||||
hf_calls.append(source.subfolder or "")
|
||||
if source.subfolder == "transformer":
|
||||
return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))])
|
||||
return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))])
|
||||
|
||||
monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter)
|
||||
monkeypatch.setattr(
|
||||
loader,
|
||||
"_get_expected_parameter_names",
|
||||
lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
|
||||
)
|
||||
|
||||
loaded = loader._load_weights_with_gguf(model)
|
||||
|
||||
assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
|
||||
assert hf_calls == ["vae"]
|
||||
@@ -1,2 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -1,117 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _tiny_transformer(**overrides):
|
||||
from vllm_omni.diffusion.models.lingbot_video import LingBotVideoTransformer3DModel
|
||||
|
||||
config = {
|
||||
"patch_size": (1, 1, 1),
|
||||
"in_channels": 2,
|
||||
"out_channels": 2,
|
||||
"hidden_size": 16,
|
||||
"num_attention_heads": 1,
|
||||
"depth": 0,
|
||||
"intermediate_size": 32,
|
||||
"text_dim": 8,
|
||||
"freq_dim": 8,
|
||||
"axes_dims": (4, 4, 8),
|
||||
"axes_lens": (32, 32, 32),
|
||||
}
|
||||
config.update(overrides)
|
||||
return LingBotVideoTransformer3DModel(**config)
|
||||
|
||||
|
||||
def test_joint_position_ids_video_then_text_order():
|
||||
from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import make_joint_position_ids
|
||||
|
||||
positions = make_joint_position_ids(text_len=3, grid_t=1, grid_h=2, grid_w=2, device=torch.device("cpu"))
|
||||
|
||||
assert positions.shape == (7, 3)
|
||||
assert positions[:4, 0].tolist() == [4, 4, 4, 4]
|
||||
assert positions[:4, 1:].tolist() == [[0, 0], [0, 1], [1, 0], [1, 1]]
|
||||
assert positions[4:].tolist() == [[1, 0, 0], [2, 0, 0], [3, 0, 0]]
|
||||
|
||||
|
||||
def test_tiny_transformer_depth_zero_forward_shape():
|
||||
model = _tiny_transformer()
|
||||
hidden_states = torch.randn(1, 2, 1, 2, 2)
|
||||
timestep = torch.tensor([300.0])
|
||||
encoder_hidden_states = torch.randn(1, 3, 8)
|
||||
encoder_attention_mask = torch.ones(1, 3, dtype=torch.long)
|
||||
|
||||
with torch.no_grad():
|
||||
out = model(
|
||||
hidden_states,
|
||||
timestep,
|
||||
encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
assert out.shape == hidden_states.shape
|
||||
assert torch.isfinite(out).all()
|
||||
|
||||
|
||||
def test_packed_attention_uses_sdpa_fallback_without_flash_varlen(monkeypatch):
|
||||
from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module
|
||||
|
||||
monkeypatch.setattr(module, "flash_attn_varlen_func_v3", None)
|
||||
attn = module.LingBotVideoAttention(
|
||||
hidden_size=8,
|
||||
num_heads=2,
|
||||
norm_eps=1e-6,
|
||||
qkv_bias=False,
|
||||
out_bias=False,
|
||||
)
|
||||
captured = {}
|
||||
|
||||
def fake_sdpa_forward(query, key, value, attn_metadata):
|
||||
captured["mask"] = attn_metadata.attn_mask
|
||||
return torch.zeros_like(query)
|
||||
|
||||
monkeypatch.setattr(attn.attn.sdpa_fallback, "forward", fake_sdpa_forward)
|
||||
x = torch.randn(1, 5, 8)
|
||||
rotary = torch.ones(1, 5, 2, dtype=torch.complex64)
|
||||
packed_indices = {
|
||||
"cu_seqlens_kv": torch.tensor([0, 2, 5], dtype=torch.int32),
|
||||
"max_seqlen_in_batch_kv": 3,
|
||||
"attention_mask": module._packed_block_attention_mask([2, 3], x.device),
|
||||
}
|
||||
|
||||
out = attn(x, rotary, packed_indices=packed_indices)
|
||||
|
||||
assert out.shape == x.shape
|
||||
mask = captured["mask"]
|
||||
assert mask.shape == (1, 1, 5, 5)
|
||||
assert mask[0, 0, :2, :2].all()
|
||||
assert mask[0, 0, 2:, 2:].all()
|
||||
assert not mask[0, 0, :2, 2:].any()
|
||||
assert not mask[0, 0, 2:, :2].any()
|
||||
|
||||
|
||||
def test_tiny_transformer_rejects_invalid_rope_dims():
|
||||
from vllm_omni.diffusion.models.lingbot_video import LingBotVideoTransformer3DModel
|
||||
|
||||
with pytest.raises(AssertionError, match="head_dim"):
|
||||
LingBotVideoTransformer3DModel(
|
||||
hidden_size=16,
|
||||
num_attention_heads=1,
|
||||
axes_dims=(4, 4, 4),
|
||||
depth=0,
|
||||
)
|
||||
|
||||
|
||||
def test_transformer_to_keeps_sensitive_modules_in_fp32():
|
||||
model = _tiny_transformer()
|
||||
|
||||
model.to(dtype=torch.bfloat16)
|
||||
|
||||
assert model.patch_embedder.weight.dtype == torch.bfloat16
|
||||
assert model.time_embedder.linear_1.weight.dtype == torch.float32
|
||||
assert model.norm_out_modulation[1].weight.dtype == torch.float32
|
||||
@@ -1,195 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _make_pipeline():
|
||||
from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import LingBotVideoPipeline
|
||||
|
||||
pipeline = object.__new__(LingBotVideoPipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.default_negative_prompt = "default negative"
|
||||
pipeline.od_config = SimpleNamespace(flow_shift=None)
|
||||
return pipeline
|
||||
|
||||
|
||||
def _make_request_batch(prompt, **sampling_overrides):
|
||||
sampling = OmniDiffusionSamplingParams(**sampling_overrides)
|
||||
request = OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id="lingbot-test")
|
||||
return DiffusionRequestBatch([request])
|
||||
|
||||
|
||||
def test_lingbot_video_pipeline_import_and_registry():
|
||||
from vllm_omni.diffusion.models.lingbot_video import (
|
||||
LingBotVideoPipeline,
|
||||
LingBotVideoTransformer3DModel,
|
||||
get_lingbot_video_post_process_func,
|
||||
)
|
||||
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS
|
||||
|
||||
assert LingBotVideoPipeline is not None
|
||||
assert LingBotVideoTransformer3DModel is not None
|
||||
assert get_lingbot_video_post_process_func is not None
|
||||
assert _DIFFUSION_MODELS["LingBotVideoPipeline"] == (
|
||||
"lingbot_video",
|
||||
"pipeline_lingbot_video",
|
||||
"LingBotVideoPipeline",
|
||||
)
|
||||
assert _DIFFUSION_POST_PROCESS_FUNCS["LingBotVideoPipeline"] == "get_lingbot_video_post_process_func"
|
||||
|
||||
|
||||
def test_component_discovery_declarations():
|
||||
from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline
|
||||
|
||||
assert LingBotVideoPipeline._dit_modules == ["transformer"]
|
||||
assert LingBotVideoPipeline._encoder_modules == ["text_encoder"]
|
||||
assert LingBotVideoPipeline._vae_modules == ["vae"]
|
||||
assert LingBotVideoPipeline.supports_step_execution is False
|
||||
|
||||
|
||||
def test_extra_body_params_include_video_flow_shift_alias():
|
||||
from vllm_omni.model_extras import get_extra_body_params
|
||||
|
||||
params = get_extra_body_params("LingBotVideoPipeline")
|
||||
assert "shift" in params
|
||||
assert "flow_shift" in params
|
||||
assert "batch_cfg" in params
|
||||
|
||||
|
||||
def test_extract_prompt_string_and_mapping():
|
||||
from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import _extract_prompt
|
||||
|
||||
string_req = SimpleNamespace(prompt="a robot waves")
|
||||
assert _extract_prompt(string_req) == ("a robot waves", None)
|
||||
|
||||
mapping_req = SimpleNamespace(prompt={"prompt": "a robot walks", "negative_prompt": "low quality"})
|
||||
assert _extract_prompt(mapping_req) == ("a robot walks", "low quality")
|
||||
|
||||
with pytest.raises(TypeError, match="Unsupported LingBot prompt type"):
|
||||
_extract_prompt(SimpleNamespace(prompt=["not", "supported"]))
|
||||
|
||||
|
||||
def test_check_inputs_accepts_t2v_shapes_and_rejects_invalid_values():
|
||||
from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline
|
||||
|
||||
LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=1)
|
||||
LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=9)
|
||||
|
||||
with pytest.raises(ValueError, match="num_frames"):
|
||||
LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=8)
|
||||
with pytest.raises(ValueError, match="height"):
|
||||
LingBotVideoPipeline.check_inputs(height=190, width=320, num_frames=9)
|
||||
|
||||
|
||||
def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
|
||||
from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
|
||||
|
||||
frames = torch.ones(2, 4, 4, 3)
|
||||
post = get_lingbot_video_post_process_func(SimpleNamespace())
|
||||
|
||||
assert post(frames, SimpleNamespace(output_type="pt")) is frames
|
||||
array = post(frames, SimpleNamespace(output_type="np"))
|
||||
assert isinstance(array, np.ndarray)
|
||||
assert array.shape == (2, 4, 4, 3)
|
||||
|
||||
|
||||
def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
|
||||
pipeline = _make_pipeline()
|
||||
calls = []
|
||||
|
||||
def fake_generate(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return torch.zeros(9, 192, 320, 3)
|
||||
|
||||
pipeline._generate = fake_generate
|
||||
req = _make_request_batch(
|
||||
{"prompt": "a robot arm", "negative_prompt": "bad hands"},
|
||||
height=192,
|
||||
width=320,
|
||||
num_frames=9,
|
||||
num_inference_steps=2,
|
||||
guidance_scale=3.0,
|
||||
seed=42,
|
||||
output_type="pt",
|
||||
extra_args={"flow_shift": 4.0, "batch_cfg": True},
|
||||
)
|
||||
|
||||
out = pipeline.forward(req)
|
||||
|
||||
assert torch.equal(out.output, torch.zeros(9, 192, 320, 3))
|
||||
assert len(calls) == 1
|
||||
call = calls[0]
|
||||
assert call["prompt"] == "a robot arm"
|
||||
assert call["negative_prompt"] == "bad hands"
|
||||
assert call["height"] == 192
|
||||
assert call["width"] == 320
|
||||
assert call["num_frames"] == 9
|
||||
assert call["num_inference_steps"] == 2
|
||||
assert call["guidance_scale"] == 3.0
|
||||
assert call["shift"] == 4.0
|
||||
assert call["batch_cfg"] is True
|
||||
assert call["output_type"] == "pt"
|
||||
|
||||
|
||||
def test_forward_prefers_shift_over_flow_shift_and_defaults_negative_prompt():
|
||||
pipeline = _make_pipeline()
|
||||
calls = []
|
||||
|
||||
def fake_generate(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return torch.zeros(5, 192, 192, 3)
|
||||
|
||||
pipeline._generate = fake_generate
|
||||
req = _make_request_batch(
|
||||
"a robot arm",
|
||||
height=192,
|
||||
width=192,
|
||||
num_frames=5,
|
||||
num_inference_steps=2,
|
||||
guidance_scale=1.0,
|
||||
seed=42,
|
||||
extra_args={"shift": 5.0, "flow_shift": 4.0},
|
||||
)
|
||||
|
||||
pipeline.forward(req)
|
||||
|
||||
assert calls[0]["shift"] == 5.0
|
||||
assert calls[0]["negative_prompt"] == "default negative"
|
||||
|
||||
|
||||
def test_forward_rejects_multi_request_batches():
|
||||
pipeline = _make_pipeline()
|
||||
first = OmniDiffusionRequest(
|
||||
prompt="first",
|
||||
sampling_params=OmniDiffusionSamplingParams(seed=1),
|
||||
request_id="lingbot-first",
|
||||
)
|
||||
second = OmniDiffusionRequest(
|
||||
prompt="second",
|
||||
sampling_params=OmniDiffusionSamplingParams(seed=2),
|
||||
request_id="lingbot-second",
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="only supports one request"):
|
||||
pipeline.forward(DiffusionRequestBatch([first, second]))
|
||||
|
||||
|
||||
def test_load_weights_rejects_external_weight_stream():
|
||||
pipeline = _make_pipeline()
|
||||
|
||||
assert pipeline.load_weights([]) == set()
|
||||
with pytest.raises(RuntimeError, match="components are loaded directly"):
|
||||
pipeline.load_weights([("transformer.weight", torch.zeros(1))])
|
||||
@@ -154,21 +154,6 @@ def test_component_config_with_default():
|
||||
assert config.resolve("vae.encoder") is None
|
||||
|
||||
|
||||
def test_gguf_config():
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
|
||||
|
||||
config = build_quant_config(
|
||||
{
|
||||
"method": "gguf",
|
||||
"gguf_model": "path/to/model.gguf",
|
||||
}
|
||||
)
|
||||
assert isinstance(config, DiffusionGGUFConfig)
|
||||
assert config.gguf_model == "path/to/model.gguf"
|
||||
assert config.get_name() == "gguf"
|
||||
|
||||
|
||||
def test_integration_string():
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
@@ -236,7 +221,7 @@ def test_transformer_config_auto_detects_modelopt_fp8():
|
||||
def test_supported_methods_includes_vllm():
|
||||
from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
|
||||
|
||||
for method in ["fp8", "gguf", "awq", "gptq", "bitsandbytes", "modelopt"]:
|
||||
for method in ["fp8", "awq", "gptq", "bitsandbytes", "modelopt"]:
|
||||
assert method in SUPPORTED_QUANTIZATION_METHODS, f"{method} missing"
|
||||
|
||||
|
||||
|
||||
@@ -1,128 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for diffusion GGUF quantization config and linear method."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
|
||||
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.gguf_config import (
|
||||
UNQUANTIZED_TYPES,
|
||||
DiffusionGGUFConfig,
|
||||
DiffusionGGUFLinearMethod,
|
||||
dequant_gemm_gguf,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_gguf_config_creation_and_delegation():
|
||||
config = DiffusionGGUFConfig(
|
||||
gguf_model="weights.gguf",
|
||||
unquantized_modules=["proj_out"],
|
||||
)
|
||||
|
||||
assert config.gguf_model == "weights.gguf"
|
||||
assert config.unquantized_modules == ["proj_out"]
|
||||
assert config.get_name() == "gguf"
|
||||
|
||||
|
||||
def test_build_quant_config_builds_gguf_config():
|
||||
config = build_quant_config(
|
||||
"gguf",
|
||||
gguf_model="repo/model-Q4_0.gguf",
|
||||
unquantized_modules=["norm_out"],
|
||||
)
|
||||
|
||||
assert isinstance(config, DiffusionGGUFConfig)
|
||||
assert config.gguf_model == "repo/model-Q4_0.gguf"
|
||||
assert config.unquantized_modules == ["norm_out"]
|
||||
|
||||
|
||||
def test_gguf_vllm_config_returns_diffusion_linear_method_for_linear_layers():
|
||||
linear = object.__new__(LinearBase)
|
||||
method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(linear, "transformer.img_in")
|
||||
|
||||
assert isinstance(method, DiffusionGGUFLinearMethod)
|
||||
|
||||
|
||||
def test_gguf_vllm_config_respects_unquantized_modules():
|
||||
linear = object.__new__(LinearBase)
|
||||
method = DiffusionGGUFConfig(unquantized_modules=["proj_out"]).get_quant_method(linear, "transformer.proj_out")
|
||||
|
||||
assert isinstance(method, UnquantizedLinearMethod)
|
||||
|
||||
|
||||
def test_gguf_vllm_config_returns_none_for_non_linear_layers():
|
||||
method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(torch.nn.LayerNorm(4), "norm")
|
||||
|
||||
assert method is None
|
||||
|
||||
|
||||
def test_dequant_gemm_gguf_uses_plain_matmul_for_unquantized_types():
|
||||
qweight_type = next(iter(UNQUANTIZED_TYPES))
|
||||
x = torch.tensor([[1.0, 2.0]], dtype=torch.float32)
|
||||
qweight = torch.tensor([[3.0, 4.0], [5.0, 6.0]], dtype=torch.float32)
|
||||
|
||||
out = dequant_gemm_gguf(x, qweight, qweight_type)
|
||||
|
||||
assert torch.allclose(out, x @ qweight.T)
|
||||
|
||||
|
||||
def test_diffusion_gguf_linear_method_applies_bias_on_unquantized_weight():
|
||||
qweight_type = next(iter(UNQUANTIZED_TYPES))
|
||||
method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
|
||||
layer = SimpleNamespace(
|
||||
qweight=torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32),
|
||||
qweight_type=SimpleNamespace(weight_type=qweight_type),
|
||||
)
|
||||
x = torch.tensor([[[1.0, 0.5]]], dtype=torch.float32)
|
||||
bias = torch.tensor([0.25, -0.5], dtype=torch.float32)
|
||||
|
||||
out = method.apply(layer, x, bias)
|
||||
|
||||
expected = x @ layer.qweight.T + bias
|
||||
assert torch.allclose(out, expected)
|
||||
|
||||
|
||||
def test_diffusion_gguf_linear_method_concatenates_sharded_outputs():
|
||||
qweight_type = next(iter(UNQUANTIZED_TYPES))
|
||||
method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
|
||||
qweight = torch.nn.Parameter(
|
||||
torch.tensor(
|
||||
[
|
||||
[1.0, 0.0],
|
||||
[0.0, 1.0],
|
||||
[2.0, 0.0],
|
||||
[0.0, 2.0],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
qweight.shard_id = ["left", "right"]
|
||||
qweight.shard_offset_map = {
|
||||
"left": (0, 2, 2),
|
||||
"right": (2, 4, 2),
|
||||
}
|
||||
layer = SimpleNamespace(
|
||||
qweight=qweight,
|
||||
qweight_type=SimpleNamespace(
|
||||
shard_weight_type={
|
||||
"left": qweight_type,
|
||||
"right": qweight_type,
|
||||
}
|
||||
),
|
||||
)
|
||||
x = torch.tensor([[1.5, 2.0]], dtype=torch.float32)
|
||||
|
||||
out = method.apply(layer, x)
|
||||
|
||||
expected_q = x @ qweight[:2].T
|
||||
expected_k = x @ qweight[2:4].T
|
||||
assert torch.allclose(out, torch.cat([expected_q, expected_k], dim=-1))
|
||||
@@ -99,7 +99,6 @@ def test_candidate_model_can_point_to_offline_checkpoint_without_online_quantiza
|
||||
candidate_quantization_config_json=None,
|
||||
candidate_ignored_layers=None,
|
||||
ignored_layers=None,
|
||||
candidate_gguf_model=None,
|
||||
candidate_load_format="default",
|
||||
)
|
||||
|
||||
|
||||
@@ -34,12 +34,10 @@ def get_cuda_graph_config():
|
||||
"stages": {
|
||||
0: {
|
||||
"max_num_seqs": 1,
|
||||
"gpu_memory_utilization": 0.2,
|
||||
"enforce_eager": True,
|
||||
"async_scheduling": False,
|
||||
},
|
||||
1: {
|
||||
"gpu_memory_utilization": 0.2,
|
||||
"enforce_eager": True,
|
||||
"async_scheduling": False,
|
||||
},
|
||||
|
||||
@@ -1,61 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Online serving smoke for ``robbyant/lingbot-video-dense-1.3b``.
|
||||
|
||||
The LingBot dense PR supports text-to-video only, so this module exercises the
|
||||
native ``/v1/videos`` path with a small 9-frame request.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
MODEL = "robbyant/lingbot-video-dense-1.3b"
|
||||
PROMPT = "a robotic arm picks up a red block"
|
||||
NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
|
||||
|
||||
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
|
||||
|
||||
def _get_diffusion_feature_cases(model: str):
|
||||
return [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=["--model-class-name", "LingBotVideoPipeline"],
|
||||
),
|
||||
id="default",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
|
||||
def test_text_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"form_data": {
|
||||
"model": omni_server.model,
|
||||
"prompt": PROMPT,
|
||||
"negative_prompt": NEGATIVE_PROMPT,
|
||||
"height": 192,
|
||||
"width": 320,
|
||||
"num_frames": 9,
|
||||
"fps": 24,
|
||||
"num_inference_steps": 2,
|
||||
"guidance_scale": 3.0,
|
||||
"flow_shift": 3.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
openai_client.send_video_diffusion_request(request_config)
|
||||
@@ -1,82 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
L4 expansion coverage for ``robbyant/lingbot-video-dense-1.3b``.
|
||||
|
||||
This is intentionally T2V-only. T2I, I2V, TI2V, MoE, and multi-GPU feature
|
||||
rows should be added in follow-up PRs when those paths are implemented.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL = "robbyant/lingbot-video-dense-1.3b"
|
||||
PROMPT = "a robotic arm picks up a red block"
|
||||
NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
|
||||
|
||||
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
|
||||
|
||||
|
||||
def _get_diffusion_feature_cases(model: str):
|
||||
return [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=["--model-class-name", "LingBotVideoPipeline"],
|
||||
),
|
||||
id="default",
|
||||
marks=SINGLE_CARD_FEATURE_MARKS,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
|
||||
def test_cfg_off(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"form_data": {
|
||||
"model": omni_server.model,
|
||||
"prompt": PROMPT,
|
||||
"height": 192,
|
||||
"width": 192,
|
||||
"num_frames": 9,
|
||||
"fps": 24,
|
||||
"num_inference_steps": 2,
|
||||
"guidance_scale": 1.0,
|
||||
"flow_shift": 3.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
openai_client.send_video_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
|
||||
def test_batch_cfg_extra_params(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"form_data": {
|
||||
"model": omni_server.model,
|
||||
"prompt": PROMPT,
|
||||
"negative_prompt": NEGATIVE_PROMPT,
|
||||
"height": 192,
|
||||
"width": 320,
|
||||
"num_frames": 9,
|
||||
"fps": 24,
|
||||
"num_inference_steps": 2,
|
||||
"guidance_scale": 3.0,
|
||||
"flow_shift": 3.0,
|
||||
"seed": 42,
|
||||
"extra_params": json.dumps({"batch_cfg": True}, separators=(",", ":")),
|
||||
},
|
||||
}
|
||||
openai_client.send_video_diffusion_request(request_config)
|
||||
@@ -410,6 +410,44 @@ def test_audio_in_video_002(omni_server, openai_client) -> None:
|
||||
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Regression: ``use_audio_in_video`` with default video loader sampling (no
|
||||
``media_io_kwargs.video.fps``) after a prior request that pinned fps.
|
||||
|
||||
Before the fix the second request could partial-hit mm-cache (audio cached,
|
||||
video re-processed without audio) and fail with ``StopIteration`` in
|
||||
``replace_multimodal_special_tokens``.
|
||||
"""
|
||||
video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 128, embed_audio=True)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(
|
||||
system_prompt=get_system_prompt(),
|
||||
video_data_url=video_data_url,
|
||||
content_text="Reply with one word: OK.",
|
||||
)
|
||||
|
||||
base_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"stream": False,
|
||||
"use_audio_in_video": True,
|
||||
"modalities": ["text"],
|
||||
}
|
||||
|
||||
# Prime mm-cache with explicit loader fps (historical failure mode).
|
||||
openai_client.send_omni_request(
|
||||
{
|
||||
**base_config,
|
||||
"extra_body": {"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}}},
|
||||
}
|
||||
)
|
||||
|
||||
# Default loader sampling: must succeed without explicit media_io fps.
|
||||
openai_client.send_omni_request(dict(base_config))
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_one_word_prompt_001(omni_server, openai_client) -> None:
|
||||
|
||||
@@ -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"
|
||||
@@ -39,7 +39,8 @@ from vllm_omni.entrypoints.openai.serving_speech import (
|
||||
OmniOpenAIServingSpeech,
|
||||
_create_wav_header,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest, SpeechServingContext
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter
|
||||
from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
|
||||
FISH_TEXT_ONLY_SYSTEM_PROMPT,
|
||||
build_fish_voice_clone_prompt_ids,
|
||||
@@ -973,6 +974,29 @@ class TestTTSMethods:
|
||||
)
|
||||
)
|
||||
|
||||
def test_upload_ming_voice_embedding_wrong_dims_rejected_without_replacing_existing(self, speech_server):
|
||||
"""Ming embedding uploads must be 192-dim before replacing stored voices."""
|
||||
speech_server._tts_model_type = "ming_tts"
|
||||
existing = {
|
||||
"name": "bad_emb_voice",
|
||||
"file_path": "/tmp/voice_samples/bad_emb_voice.safetensors",
|
||||
"mime_type": "application/x-safetensors",
|
||||
"embedding_source": "direct",
|
||||
"embedding_dim": 192,
|
||||
}
|
||||
speech_server.uploaded_speakers = {"bad_emb_voice": existing.copy()}
|
||||
|
||||
with pytest.raises(ValueError, match="Ming speaker embedding must have 192 dims, got 191"):
|
||||
asyncio.run(
|
||||
speech_server.upload_voice_embedding(
|
||||
embedding_json=json.dumps([0.0] * 191),
|
||||
consent="consent",
|
||||
name="bad_emb_voice",
|
||||
)
|
||||
)
|
||||
|
||||
assert speech_server.uploaded_speakers["bad_emb_voice"] == existing
|
||||
|
||||
def test_base_task_requires_ref_audio_or_speaker_embedding(self, speech_server):
|
||||
"""Base task without ref_audio or speaker_embedding is rejected."""
|
||||
req = OpenAICreateSpeechRequest(input="Hello", task_type="Base")
|
||||
@@ -1557,6 +1581,86 @@ class TestTTSMethods:
|
||||
assert params["x_vector_only_mode"] == [True]
|
||||
assert "ref_audio" not in params
|
||||
|
||||
def test_ming_adapter_uploaded_direct_embedding_uses_embedding_not_audio(
|
||||
self, speech_server, mocker: MockerFixture
|
||||
):
|
||||
"""Ming uploaded direct embeddings are loaded into prompt construction."""
|
||||
speech_server.uploaded_speakers = {
|
||||
"emb_voice": {
|
||||
"name": "emb_voice",
|
||||
"file_path": "/tmp/voice_samples/emb_voice.safetensors",
|
||||
"mime_type": "application/x-safetensors",
|
||||
"embedding_source": "direct",
|
||||
"embedding_dim": 192,
|
||||
}
|
||||
}
|
||||
fake_embedding = [0.1] * 192
|
||||
mock_get_emb = mocker.patch.object(
|
||||
speech_server, "_get_uploaded_speaker_embedding", return_value=fake_embedding
|
||||
)
|
||||
mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data")
|
||||
mock_prompt = mocker.patch.object(
|
||||
speech_server,
|
||||
"_build_ming_dense_prompt",
|
||||
return_value={"additional_information": {"speaker_count": 1}},
|
||||
)
|
||||
|
||||
adapter = MingTTSAdapter(SpeechServingContext(server=speech_server))
|
||||
req = OpenAICreateSpeechRequest(input="Hello", voice="emb_voice")
|
||||
prepared = asyncio.run(adapter.build(req, [], False))
|
||||
|
||||
mock_get_emb.assert_called_once_with("emb_voice")
|
||||
mock_get_audio.assert_not_called()
|
||||
mock_prompt.assert_called_once_with(req, ref_audio_data=None)
|
||||
prompt_request = mock_prompt.call_args.args[0]
|
||||
assert prompt_request.speaker_embedding == fake_embedding
|
||||
assert len(prompt_request.speaker_embedding) == 192
|
||||
assert prepared.prompt["additional_information"]["speaker_count"] == 1
|
||||
assert prepared.model_type == "ming_tts"
|
||||
|
||||
def test_ming_adapter_uploaded_audio_path_preserves_ref_text(self, speech_server, mocker: MockerFixture):
|
||||
"""Ming uploaded audio voices still resolve audio and stored ref_text."""
|
||||
speech_server.uploaded_speakers = {
|
||||
"audio_voice": {
|
||||
"name": "audio_voice",
|
||||
"file_path": "/tmp/voice_samples/audio_voice.wav",
|
||||
"mime_type": "audio/wav",
|
||||
"ref_text": "Reference transcript.",
|
||||
}
|
||||
}
|
||||
ref_audio_source = "data:audio/wav;base64,ZmFrZWF1ZGlv"
|
||||
ref_audio_data = ([0.0, 0.1], 16000)
|
||||
fake_embedding = [0.2] * 192
|
||||
mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data", return_value=ref_audio_source)
|
||||
mock_get_emb = mocker.patch.object(speech_server, "_get_uploaded_speaker_embedding")
|
||||
mocker.patch.object(
|
||||
speech_server,
|
||||
"_resolve_ref_audio",
|
||||
new=mocker.AsyncMock(return_value=ref_audio_data),
|
||||
)
|
||||
mock_extract = mocker.patch.object(
|
||||
speech_server,
|
||||
"_extract_ming_speaker_embeddings_from_ref_audio",
|
||||
return_value=[fake_embedding],
|
||||
)
|
||||
mock_prompt = mocker.patch.object(
|
||||
speech_server,
|
||||
"_build_ming_dense_prompt",
|
||||
return_value={"additional_information": {"speaker_count": 1}},
|
||||
)
|
||||
|
||||
adapter = MingTTSAdapter(SpeechServingContext(server=speech_server))
|
||||
req = OpenAICreateSpeechRequest(input="Hello", voice="audio_voice")
|
||||
prepared = asyncio.run(adapter.build(req, [], False))
|
||||
|
||||
mock_get_audio.assert_called_once_with("audio_voice")
|
||||
mock_get_emb.assert_not_called()
|
||||
mock_extract.assert_called_once_with([ref_audio_data])
|
||||
mock_prompt.assert_called_once_with(req, ref_audio_data=ref_audio_data)
|
||||
assert req.ref_text == "Reference transcript."
|
||||
assert req.speaker_embedding == fake_embedding
|
||||
assert prepared.model_type == "ming_tts"
|
||||
|
||||
# ── regression: full flow from issue #1603 ──
|
||||
|
||||
def test_regression_1603_speaker_key_with_uploaded_audio_voice(self, speech_server, mocker: MockerFixture):
|
||||
|
||||
@@ -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)
|
||||
@@ -359,3 +359,15 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
|
||||
|
||||
assert not hasattr(stage_configs[0].engine_args, "additional_config")
|
||||
assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}}
|
||||
|
||||
|
||||
def test_default_stage_config_includes_quantization_config():
|
||||
"""Ensure structured quantization_config survives default diffusion-stage creation."""
|
||||
quantization_config = {
|
||||
"method": "example_quant",
|
||||
"weights": "weights.bin",
|
||||
}
|
||||
|
||||
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"quantization_config": quantization_config})[0]
|
||||
|
||||
assert stage_cfg["engine_args"]["quantization_config"] == quantization_config
|
||||
@@ -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
|
||||
|
||||
@@ -18,8 +18,11 @@ import types
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
# Direct file import to avoid vllm_omni.__init__ patch dependencies.
|
||||
_MODELS = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
@@ -542,3 +545,92 @@ class TestCodePredictorWrapperConfig:
|
||||
wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=True),
|
||||
)
|
||||
assert graph_wrapper._prefix_graphs_enabled is True
|
||||
|
||||
|
||||
class TestGumbelMaxSampling:
|
||||
"""Sanity tests for the shared Gumbel-max sampling helper."""
|
||||
|
||||
@staticmethod
|
||||
def _sample(logits: torch.Tensor, generator=None) -> torch.Tensor:
|
||||
common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
|
||||
return common_mod.CodePredictorWrapper._sample_codes_gumbel(logits, generator=generator)
|
||||
|
||||
def test_gumbel_max_picks_unique_high_logit(self, loaded_target_classes) -> None:
|
||||
"""If one logit is dominant, Gumbel-max must always pick it."""
|
||||
torch.manual_seed(0)
|
||||
bsz, vocab = 4, 64
|
||||
logits = torch.full((bsz, vocab), -1e4)
|
||||
logits[:, 7] = 0.0 # only finite entry
|
||||
out = self._sample(logits)
|
||||
assert out.shape == (bsz, 1)
|
||||
assert torch.equal(out.squeeze(-1), torch.full((bsz,), 7, dtype=torch.long))
|
||||
|
||||
def test_gumbel_max_handles_single_finite_entry_rows(self, loaded_target_classes) -> None:
|
||||
"""Degenerate rows with one surviving finite entry should still return
|
||||
a valid index instead of tripping multinomial-style asserts."""
|
||||
bsz, vocab = 2, 8
|
||||
logits = torch.full((bsz, vocab), float("-inf"))
|
||||
logits[0, 3] = 0.0
|
||||
logits[1, 5] = 0.0
|
||||
out = self._sample(logits)
|
||||
assert out[0].item() == 3
|
||||
assert out[1].item() == 5
|
||||
|
||||
def test_gumbel_max_generator_makes_results_reproducible(self, loaded_target_classes) -> None:
|
||||
"""Same ``Generator.manual_seed`` should produce identical samples."""
|
||||
torch.manual_seed(0)
|
||||
logits = torch.randn(8, 32)
|
||||
g1 = torch.Generator()
|
||||
g1.manual_seed(123)
|
||||
g2 = torch.Generator()
|
||||
g2.manual_seed(123)
|
||||
g3 = torch.Generator()
|
||||
g3.manual_seed(456)
|
||||
s1 = self._sample(logits, generator=g1)
|
||||
s2 = self._sample(logits, generator=g2)
|
||||
s3 = self._sample(logits, generator=g3)
|
||||
assert torch.equal(s1, s2)
|
||||
# Different seeds must (almost surely) yield a different sequence.
|
||||
assert not torch.equal(s1, s3)
|
||||
|
||||
def test_gumbel_max_per_row_generators_match_individual_rows(self, loaded_target_classes) -> None:
|
||||
"""A per-row generator list should preserve row independence."""
|
||||
torch.manual_seed(11)
|
||||
logits = torch.randn(3, 16)
|
||||
seeds = [101, 202, 303]
|
||||
batched_generators = []
|
||||
expected_rows = []
|
||||
for row, seed in enumerate(seeds):
|
||||
row_generator = torch.Generator()
|
||||
row_generator.manual_seed(seed)
|
||||
batched_generators.append(row_generator)
|
||||
|
||||
solo_generator = torch.Generator()
|
||||
solo_generator.manual_seed(seed)
|
||||
expected_rows.append(self._sample(logits[row : row + 1], generator=solo_generator))
|
||||
|
||||
batched = self._sample(logits, generator=batched_generators)
|
||||
expected = torch.cat(expected_rows, dim=0)
|
||||
assert torch.equal(batched, expected)
|
||||
|
||||
def test_gumbel_max_rejects_generator_length_mismatch(self, loaded_target_classes) -> None:
|
||||
logits = torch.randn(2, 8)
|
||||
generators = [torch.Generator()]
|
||||
with pytest.raises(ValueError, match="Expected 2 per-row generators"):
|
||||
self._sample(logits, generator=generators)
|
||||
|
||||
def test_gumbel_max_distribution_matches_softmax(self, loaded_target_classes) -> None:
|
||||
"""Empirical frequency under Gumbel-max should be close to the
|
||||
softmax probability mass. Loose tolerance because we draw a
|
||||
finite number of samples."""
|
||||
torch.manual_seed(7)
|
||||
vocab = 16
|
||||
logits = torch.randn(vocab) * 2.0
|
||||
target = F.softmax(logits, dim=-1, dtype=torch.float32)
|
||||
|
||||
n_draws = 20000
|
||||
batched_logits = logits.unsqueeze(0).expand(n_draws, vocab).contiguous()
|
||||
samples = self._sample(batched_logits).squeeze(-1)
|
||||
empirical = torch.bincount(samples, minlength=vocab).float() / n_draws
|
||||
# L1 distance under 0.03 with 20k draws is comfortable for vocab=16.
|
||||
assert (empirical - target).abs().sum().item() < 0.05
|
||||
+2
-13
@@ -8,7 +8,6 @@ import torch
|
||||
|
||||
from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
|
||||
talker2code2wav_async_chunk,
|
||||
text2flow,
|
||||
text2flow_full_payload,
|
||||
text2flow_token_only,
|
||||
)
|
||||
@@ -49,13 +48,13 @@ def _transfer_manager(
|
||||
)
|
||||
|
||||
|
||||
def test_text2flow_supports_batched_source_outputs():
|
||||
def test_text2flow_token_only_supports_batched_source_outputs():
|
||||
source_outputs = [
|
||||
_source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[8, 9]])}),
|
||||
_source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[6, 7]])}),
|
||||
]
|
||||
|
||||
outputs = text2flow(source_outputs=source_outputs, prompt=None)
|
||||
outputs = text2flow_token_only(source_outputs=source_outputs, prompt=None)
|
||||
|
||||
assert len(outputs) == 2
|
||||
assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
|
||||
@@ -64,16 +63,6 @@ def test_text2flow_supports_batched_source_outputs():
|
||||
assert outputs[1]["additional_information"]["ids"]["prompt"] == [20, 21]
|
||||
|
||||
|
||||
def test_text2flow_strips_reference_speech_prefix_from_cumulative_ids():
|
||||
source_outputs = [
|
||||
_source_output("req-0", [10, 11], [8, 9, 1, 2, 3], {"speech_token": torch.tensor([[8, 9]])}),
|
||||
]
|
||||
|
||||
outputs = text2flow(source_outputs=source_outputs, prompt=None)
|
||||
|
||||
assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
|
||||
|
||||
|
||||
def test_text2flow_token_only_strips_reference_speech_prefix_from_cumulative_ids():
|
||||
source_outputs = [
|
||||
_source_output(
|
||||
|
||||
@@ -7,7 +7,6 @@ import torch
|
||||
|
||||
from vllm_omni.model_executor.stage_input_processors.indextts2 import (
|
||||
_strip_stop_token,
|
||||
talker2s2mel,
|
||||
talker2s2mel_full_payload,
|
||||
talker2s2mel_token_only,
|
||||
)
|
||||
@@ -108,118 +107,6 @@ def test_strip_stop_token_all_stop():
|
||||
assert lat.shape[1] == 0
|
||||
|
||||
|
||||
# ── talker2s2mel ──
|
||||
|
||||
|
||||
def test_talker2s2mel_basic():
|
||||
mel_codes = torch.tensor([10, 20, 30, STOP_MEL_TOKEN])
|
||||
latent = torch.randn(4, LATENT_DIM)
|
||||
s_ref = torch.randn(1, 128)
|
||||
ref_mel = torch.randn(1, 80, 200)
|
||||
style = torch.randn(1, 256)
|
||||
|
||||
outputs = talker2s2mel(
|
||||
source_outputs=[
|
||||
_make_talker_output(
|
||||
mel_codes=mel_codes,
|
||||
latent=latent,
|
||||
s_ref=s_ref,
|
||||
ref_mel=ref_mel,
|
||||
style=style,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert len(outputs) == 1
|
||||
info = outputs[0]["additional_information"]
|
||||
assert info["mel_codes"].shape == (1, 3)
|
||||
assert info["latent"].shape == (1, 3, LATENT_DIM)
|
||||
assert info["code_lens"].tolist() == [3]
|
||||
assert "S_ref" in info
|
||||
assert "ref_mel" in info
|
||||
assert "style" in info
|
||||
assert outputs[0]["prompt_token_ids"] == [0]
|
||||
|
||||
|
||||
def test_talker2s2mel_skips_unfinished():
|
||||
mel_codes = torch.tensor([10, 20])
|
||||
latent = torch.randn(2, LATENT_DIM)
|
||||
|
||||
outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, finished=False)])
|
||||
|
||||
assert len(outputs) == 0
|
||||
|
||||
|
||||
def test_talker2s2mel_skips_empty_mel():
|
||||
outputs = talker2s2mel(
|
||||
source_outputs=[
|
||||
_make_talker_output(
|
||||
mel_codes=torch.tensor([]),
|
||||
latent=torch.randn(0, LATENT_DIM),
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert len(outputs) == 0
|
||||
|
||||
|
||||
def test_talker2s2mel_missing_optional_meta():
|
||||
mel_codes = torch.tensor([10, 20])
|
||||
latent = torch.randn(2, LATENT_DIM)
|
||||
|
||||
outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent)])
|
||||
|
||||
assert len(outputs) == 1
|
||||
info = outputs[0]["additional_information"]
|
||||
assert "S_ref" not in info
|
||||
assert "ref_mel" not in info
|
||||
assert "style" not in info
|
||||
|
||||
|
||||
def test_talker2s2mel_batch():
|
||||
out0 = _make_talker_output(
|
||||
mel_codes=torch.tensor([1, 2, STOP_MEL_TOKEN]),
|
||||
latent=torch.randn(3, LATENT_DIM),
|
||||
s_ref=torch.randn(1, 64),
|
||||
)
|
||||
out1 = _make_talker_output(
|
||||
mel_codes=torch.tensor([5, 6, 7, 8]),
|
||||
latent=torch.randn(4, LATENT_DIM),
|
||||
ref_mel=torch.randn(1, 80, 100),
|
||||
)
|
||||
|
||||
outputs = talker2s2mel(source_outputs=[out0, out1])
|
||||
|
||||
assert len(outputs) == 2
|
||||
assert outputs[0]["additional_information"]["code_lens"].tolist() == [2]
|
||||
assert outputs[1]["additional_information"]["code_lens"].tolist() == [4]
|
||||
|
||||
|
||||
def test_talker2s2mel_tensors_on_cpu():
|
||||
device = torch.device("cpu")
|
||||
mel_codes = torch.tensor([10, 20], device=device)
|
||||
latent = torch.randn(2, LATENT_DIM, device=device)
|
||||
s_ref = torch.randn(1, 64, device=device)
|
||||
|
||||
outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, s_ref=s_ref)])
|
||||
|
||||
info = outputs[0]["additional_information"]
|
||||
assert info["latent"].device.type == "cpu"
|
||||
assert info["mel_codes"].device.type == "cpu"
|
||||
assert info["S_ref"].device.type == "cpu"
|
||||
|
||||
|
||||
def test_talker2s2mel_no_multimodal_dict():
|
||||
output = SimpleNamespace(
|
||||
finished=True,
|
||||
outputs=[SimpleNamespace(multimodal_output="not_a_dict")],
|
||||
)
|
||||
|
||||
results = talker2s2mel(source_outputs=[output])
|
||||
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
# ── full-payload connector path ──
|
||||
|
||||
|
||||
|
||||
@@ -11,59 +11,40 @@ import torch
|
||||
from vllm_omni.model_executor.stage_input_processors import mimo_audio as sip
|
||||
from vllm_omni.model_executor.stage_input_processors.mimo_audio import (
|
||||
MAX_CODE2WAV_TOKENS,
|
||||
llm2code2wav,
|
||||
llm2code2wav_full_payload,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_source_outputs(codec_codes: torch.Tensor, request_id: str = "req-0"):
|
||||
output = SimpleNamespace(multimodal_output={"codes": {"audio": codec_codes}})
|
||||
"""Build minimal source_outputs with one talker output carrying codec codes."""
|
||||
talker_output = SimpleNamespace(outputs=[output], request_id=request_id)
|
||||
return [talker_output]
|
||||
|
||||
|
||||
def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog):
|
||||
"""Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated, not passed through."""
|
||||
# prepend_and_flatten_colmajor produces 36 ids per (8, 4) codec frame:
|
||||
# pad adds one row -> (9, 4) per frame, permuted and flattened.
|
||||
# Pick enough frames to comfortably exceed the cap.
|
||||
def test_llm2code2wav_full_payload_truncates_when_flat_exceeds_max(caplog):
|
||||
"""Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated."""
|
||||
frames = (MAX_CODE2WAV_TOKENS // 36) + 100
|
||||
codec_codes = torch.ones(frames, 1, 8, 4, dtype=torch.long)
|
||||
request = SimpleNamespace(request_id="req-long")
|
||||
|
||||
source_outputs = _make_source_outputs(codec_codes, request_id="req-long")
|
||||
|
||||
# Attach caplog's handler directly to the module logger so the warning is
|
||||
# captured regardless of propagation (vllm's logger configuration can
|
||||
# interact badly with caplog.at_level's default root-handler path).
|
||||
target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.mimo_audio")
|
||||
target_logger.addHandler(caplog.handler)
|
||||
prev_level = target_logger.level
|
||||
target_logger.setLevel(logging.WARNING)
|
||||
try:
|
||||
prompts = llm2code2wav(source_outputs)
|
||||
payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request)
|
||||
finally:
|
||||
target_logger.removeHandler(caplog.handler)
|
||||
target_logger.setLevel(prev_level)
|
||||
|
||||
assert len(prompts) == 1
|
||||
assert len(prompts[0]["prompt_token_ids"]) == MAX_CODE2WAV_TOKENS
|
||||
assert any("truncating" in rec.getMessage() for rec in caplog.records), (
|
||||
f"Expected a 'truncating' warning; captured records: {[r.getMessage() for r in caplog.records]}"
|
||||
)
|
||||
assert payload is not None
|
||||
assert len(payload["codes"]["audio"]) == MAX_CODE2WAV_TOKENS
|
||||
|
||||
|
||||
def test_llm2code2wav_short_sequence_unchanged():
|
||||
"""Short codec sequences are returned without truncation."""
|
||||
def test_llm2code2wav_full_payload_short_sequence_unchanged():
|
||||
codec_codes = torch.ones(4, 1, 8, 4, dtype=torch.long)
|
||||
source_outputs = _make_source_outputs(codec_codes, request_id="req-short")
|
||||
request = SimpleNamespace(request_id="req-short")
|
||||
|
||||
prompts = llm2code2wav(source_outputs)
|
||||
payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request)
|
||||
|
||||
assert len(prompts) == 1
|
||||
# 4 frames + 1 pad row, flattened col-major → well below the cap
|
||||
assert 0 < len(prompts[0]["prompt_token_ids"]) <= MAX_CODE2WAV_TOKENS
|
||||
assert payload is not None
|
||||
assert 0 < len(payload["codes"]["audio"]) <= MAX_CODE2WAV_TOKENS
|
||||
|
||||
|
||||
def test_llm2code2wav_truncation_boundary_constant_matches_yaml():
|
||||
|
||||
@@ -163,7 +163,7 @@ def test_talker2code2wav_full_payload_filters_by_output_token_ids() -> None:
|
||||
|
||||
assert payload is not None
|
||||
assert payload["codes"]["audio"] == [10, 20, 11, 21, 12, 22]
|
||||
assert payload["code_predictor_codes"] == payload["codes"]["audio"]
|
||||
assert "code_predictor_codes" not in payload
|
||||
|
||||
|
||||
def test_talker2code2wav_full_payload_drops_count_matched_terminal_row() -> None:
|
||||
@@ -204,7 +204,7 @@ def test_talker2code2wav_full_payload_drops_rows_aligned_to_non_codec_ids() -> N
|
||||
|
||||
assert payload is not None
|
||||
assert payload["codes"]["audio"] == [0, 0, 0]
|
||||
assert payload["code_predictor_codes"] == payload["codes"]["audio"]
|
||||
assert "code_predictor_codes" not in payload
|
||||
|
||||
|
||||
def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
|
||||
@@ -224,7 +224,7 @@ def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
|
||||
|
||||
assert payload is not None
|
||||
assert payload["codes"]["audio"] == [0, 7, 0, 8, 0, 9]
|
||||
assert payload["code_predictor_codes"] == payload["codes"]["audio"]
|
||||
assert "code_predictor_codes" not in payload
|
||||
|
||||
|
||||
def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
|
||||
@@ -247,7 +247,6 @@ def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
|
||||
assert payload["ids"]["all"] == [151644, 872, 3]
|
||||
assert payload["embed"]["prefill"].device.type == "cpu"
|
||||
assert payload["hidden_states"]["output"].device.type == "cpu"
|
||||
assert payload["next_stage_prompt_len"] > 0
|
||||
assert payload["embed"]["prefill"].shape[0] == 2
|
||||
assert payload["hidden_states"]["output"].shape[0] == 2
|
||||
|
||||
|
||||
@@ -12,9 +12,9 @@ from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
|
||||
max_ic_for_chunk_size,
|
||||
)
|
||||
from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
|
||||
talker2code2wav,
|
||||
talker2code2wav_async_chunk,
|
||||
talker2code2wav_full_payload,
|
||||
talker2code2wav_token_only,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -368,7 +368,15 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
|
||||
assert rid in tm.request_payload
|
||||
|
||||
|
||||
def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
|
||||
def test_non_async_token_only_sizes_placeholder_for_ref_and_audio_frames():
|
||||
"""``talker2code2wav_token_only`` only allocates placeholder prompt slots.
|
||||
|
||||
After the connector refactor, actual codec flattening (ref prepend +
|
||||
codebook-major layout) is performed by ``talker2code2wav_full_payload`` on
|
||||
the worker data plane. The orchestrator hook still derives
|
||||
``left_context_size`` from stage-0 multimodal_output so Code2Wav can trim
|
||||
reference frames once the connector payload arrives.
|
||||
"""
|
||||
ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
|
||||
audio_codes = torch.tensor(
|
||||
[
|
||||
@@ -387,12 +395,38 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
|
||||
engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
|
||||
)
|
||||
|
||||
prompts = talker2code2wav(stage.engine_outputs)
|
||||
prompts = talker2code2wav_token_only(stage.engine_outputs)
|
||||
|
||||
assert len(prompts) == 1
|
||||
prompt = prompts[0]
|
||||
assert prompt["additional_information"] == {"meta": {"left_context_size": 2}}
|
||||
assert prompt["prompt_token_ids"] == [
|
||||
# 2 ref frames + 2 valid audio frames (zero row filtered), 4 quantizers.
|
||||
assert prompt["prompt_token_ids"] == [0] * (_Q * (2 + 2))
|
||||
|
||||
|
||||
def test_full_payload_prepends_ref_code_and_flattens_codebook_major():
|
||||
"""Worker producer is authoritative for ref prepend + codec flatten."""
|
||||
ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
|
||||
audio_codes = torch.tensor(
|
||||
[
|
||||
[0, 0, 0, 0],
|
||||
[1, 2, 3, 4],
|
||||
[5, 6, 7, 8],
|
||||
],
|
||||
dtype=torch.long,
|
||||
)
|
||||
pooling_output = {
|
||||
"codes.audio": audio_codes,
|
||||
"codes.ref": ref_code,
|
||||
"meta.ref_code_len": 2,
|
||||
}
|
||||
request = SimpleNamespace(request_id="r", output_token_ids=list(range(3)))
|
||||
|
||||
payload = talker2code2wav_full_payload(transfer_manager=None, pooling_output=pooling_output, request=request)
|
||||
|
||||
assert payload is not None
|
||||
assert payload["meta"]["left_context_size"] == 2
|
||||
assert payload["codes"]["audio"].tolist() == [
|
||||
9,
|
||||
8,
|
||||
1,
|
||||
@@ -433,7 +467,7 @@ def test_non_async_processor_filters_out_of_range_codec_values():
|
||||
engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
|
||||
)
|
||||
|
||||
prompts = talker2code2wav(stage.engine_outputs)
|
||||
prompts = talker2code2wav_token_only(stage.engine_outputs)
|
||||
|
||||
assert len(prompts) == 1
|
||||
prompt = prompts[0]
|
||||
|
||||
@@ -8,61 +8,12 @@ import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.stage_input_processors.voxtral_tts import (
|
||||
generator2tokenizer,
|
||||
generator2tokenizer_async_chunk,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
# ---- Helpers for generator2tokenizer (non-streaming) ----
|
||||
|
||||
|
||||
def _make_stage(engine_outputs):
|
||||
return SimpleNamespace(engine_outputs=engine_outputs)
|
||||
|
||||
|
||||
def _make_output(audio_tensors):
|
||||
"""Build a single generator output with multimodal_output codes.audio."""
|
||||
return SimpleNamespace(
|
||||
outputs=[SimpleNamespace(multimodal_output={"codes": {"audio": audio_tensors}})],
|
||||
)
|
||||
|
||||
|
||||
# ---- Tests for generator2tokenizer ----
|
||||
|
||||
|
||||
class TestGenerator2Tokenizer:
|
||||
def test_single_output_single_chunk(self):
|
||||
"""Single output with one audio tensor produces correct prompt_token_ids."""
|
||||
audio = torch.tensor([10, 20, 30, 40])
|
||||
stage = _make_stage([_make_output([audio])])
|
||||
|
||||
result = generator2tokenizer(stage.engine_outputs)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["prompt_token_ids"] == [10, 20, 30, 40]
|
||||
assert result[0]["multi_modal_data"] is None
|
||||
|
||||
def test_2d_audio_tensors_flattened(self):
|
||||
"""2D audio tensors (e.g., multi-codebook frames) are flattened."""
|
||||
audio = torch.tensor([[1, 2, 3], [4, 5, 6]]) # (2, 3)
|
||||
stage = _make_stage([_make_output([audio])])
|
||||
|
||||
result = generator2tokenizer(stage.engine_outputs)
|
||||
|
||||
assert result[0]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6]
|
||||
|
||||
def test_no_outputs_yet_raises(self):
|
||||
"""Stage with engine_outputs=None raises RuntimeError."""
|
||||
stage = _make_stage(engine_outputs=None)
|
||||
with pytest.raises(TypeError):
|
||||
generator2tokenizer(stage.engine_outputs)
|
||||
|
||||
|
||||
# ---- Helpers for generator2tokenizer_async_chunk (streaming) ----
|
||||
|
||||
|
||||
def _req(external_req_id: str, *, finished: bool):
|
||||
return SimpleNamespace(
|
||||
external_req_id=external_req_id,
|
||||
|
||||
@@ -66,7 +66,6 @@ EXCLUDED_MODELS = [
|
||||
"AudioXPipeline",
|
||||
"HunyuanVideo15Pipeline",
|
||||
"HunyuanVideo15ImageToVideoPipeline",
|
||||
"LingBotVideoPipeline",
|
||||
"MagiHumanPipeline",
|
||||
"OmniVoicePipeline",
|
||||
"OmniVoice",
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
from vllm.config import CompilationConfig, CompilationMode, DeviceConfig, VllmConfig
|
||||
|
||||
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _vllm_config(*, backend: str = "inductor", mode: CompilationMode) -> VllmConfig:
|
||||
return VllmConfig(
|
||||
device_config=DeviceConfig(device="cpu"),
|
||||
compilation_config=CompilationConfig(backend=backend, mode=mode),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mode",
|
||||
[CompilationMode.NONE, CompilationMode.VLLM_COMPILE, CompilationMode.STOCK_TORCH_COMPILE],
|
||||
)
|
||||
def test_cuda_default_ir_op_priority_prefers_vllm_c_when_inductor_backend(mode: CompilationMode) -> None:
|
||||
"""Regression for #4964: inductor-active configs must not switch to native-only."""
|
||||
priority = CudaOmniPlatform.get_default_ir_op_priority(_vllm_config(mode=mode))
|
||||
assert priority.rms_norm == ["vllm_c", "native"]
|
||||
assert priority.fused_add_rms_norm == ["vllm_c", "native"]
|
||||
|
||||
|
||||
def test_cuda_default_ir_op_priority_with_oink(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import vllm.envs as envs
|
||||
|
||||
monkeypatch.setattr(envs, "VLLM_USE_OINK_OPS", True)
|
||||
priority = CudaOmniPlatform.get_default_ir_op_priority(
|
||||
_vllm_config(mode=CompilationMode.VLLM_COMPILE),
|
||||
)
|
||||
assert priority.rms_norm == ["oink", "vllm_c", "native"]
|
||||
assert priority.fused_add_rms_norm == ["oink", "vllm_c", "native"]
|
||||
@@ -0,0 +1 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
@@ -0,0 +1,154 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""End-to-end tests for GGUF quantization on diffusion models.
|
||||
|
||||
Validates that GGUF-quantized diffusion models generate valid images and
|
||||
use less peak GPU memory than BF16 baseline.
|
||||
|
||||
Requires vllm-omni to be installed alongside the plugin.
|
||||
|
||||
Usage:
|
||||
pytest tests/diffusion/test_gguf_memory.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import NamedTuple
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.helpers.env import DeviceMemoryMonitor
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
class DiffusionGGUFTestConfig(NamedTuple):
|
||||
artifact_prefix: str
|
||||
gguf_model: str
|
||||
hf_model: str
|
||||
min_cosine_similarity: float = 0.95
|
||||
|
||||
|
||||
Z_IMAGE_CONFIG = DiffusionGGUFTestConfig(
|
||||
artifact_prefix="zimage",
|
||||
hf_model="Tongyi-MAI/Z-Image-Turbo",
|
||||
gguf_model="unsloth/Z-Image-Turbo-GGUF:Q4_0",
|
||||
)
|
||||
|
||||
|
||||
FLUX_CONFIG = DiffusionGGUFTestConfig(
|
||||
artifact_prefix="flux2_klein",
|
||||
hf_model="black-forest-labs/FLUX.2-klein-4B",
|
||||
gguf_model="unsloth/FLUX.2-klein-4B-GGUF:Q8_0",
|
||||
)
|
||||
|
||||
|
||||
def _image_cosine_similarity(hf_image, gguf_image) -> float:
|
||||
hf_tensor = torch.as_tensor(np.array(hf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten()
|
||||
gguf_tensor = torch.as_tensor(np.array(gguf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten()
|
||||
return torch.nn.functional.cosine_similarity(hf_tensor, gguf_tensor, dim=0).item()
|
||||
|
||||
|
||||
def _generate_single_stage_image(
|
||||
model: str,
|
||||
height: int = 256,
|
||||
width: int = 256,
|
||||
num_inference_steps: int = 20,
|
||||
seed: int = 42,
|
||||
**extra_kwargs,
|
||||
) -> tuple[list, float]:
|
||||
"""Generate an image with a single-stage diffusion model.
|
||||
|
||||
Returns (images, peak_memory_gib).
|
||||
"""
|
||||
omni_kwargs = dict(extra_kwargs)
|
||||
|
||||
memory_monitor = DeviceMemoryMonitor(device_index=0)
|
||||
memory_monitor.start()
|
||||
try:
|
||||
with OmniRunner(model, **omni_kwargs) as runner:
|
||||
generator = torch.Generator(
|
||||
device=current_omni_platform.device_type,
|
||||
).manual_seed(seed)
|
||||
outputs = runner.omni.generate(
|
||||
"a photo of a cat sitting on a laptop keyboard",
|
||||
OmniDiffusionSamplingParams(
|
||||
height=height,
|
||||
width=width,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=0.0,
|
||||
generator=generator,
|
||||
),
|
||||
)
|
||||
finally:
|
||||
memory_monitor.stop()
|
||||
peak_mem = memory_monitor.peak_used_mb / 1024
|
||||
|
||||
first_output = outputs[0]
|
||||
assert first_output.final_output_type == "image"
|
||||
if hasattr(first_output, "images") and first_output.images:
|
||||
images = first_output.images
|
||||
else:
|
||||
assert hasattr(first_output, "request_output") and first_output.request_output
|
||||
request_output = first_output.request_output
|
||||
if isinstance(request_output, list):
|
||||
req_out = request_output[0]
|
||||
else:
|
||||
req_out = request_output
|
||||
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
|
||||
images = req_out.images
|
||||
assert len(images) >= 1
|
||||
assert images[0].width == width
|
||||
assert images[0].height == height
|
||||
|
||||
return images, peak_mem
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"})
|
||||
@pytest.mark.full_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.parametrize("model", [Z_IMAGE_CONFIG, FLUX_CONFIG], ids=["Z-Image-Turbo", "FLUX.2-klein"])
|
||||
def test_single_stage_diffusion_gguf(model: DiffusionGGUFTestConfig, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Z-Image-Turbo GGUF generates valid images and uses less memory than BF16."""
|
||||
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
|
||||
# BF16 baseline
|
||||
hf_images, mem_bf16 = _generate_single_stage_image(
|
||||
model=model.hf_model,
|
||||
)
|
||||
|
||||
# GGUF
|
||||
images, mem_gguf = _generate_single_stage_image(
|
||||
model=model.hf_model,
|
||||
diffusion_quantization_config={
|
||||
"method": "gguf",
|
||||
"gguf_model": model.gguf_model,
|
||||
},
|
||||
)
|
||||
|
||||
hf_image_path = f"test_{model.artifact_prefix}_hf.png"
|
||||
gguf_image_path = f"test_{model.artifact_prefix}_gguf.png"
|
||||
assert len(hf_images) >= 1
|
||||
hf_images[0].save(hf_image_path)
|
||||
assert len(images) >= 1
|
||||
images[0].save(gguf_image_path)
|
||||
print(f"Saved HF image: {hf_image_path}")
|
||||
print(f"Saved GGUF image: {gguf_image_path}")
|
||||
|
||||
image_cosine_similarity = _image_cosine_similarity(hf_images[0], images[0])
|
||||
print(f"{model.artifact_prefix} image cosine similarity: {image_cosine_similarity:.4f}")
|
||||
assert image_cosine_similarity >= model.min_cosine_similarity, (
|
||||
f"GGUF image cosine similarity ({image_cosine_similarity:.4f}) should be >= {model.min_cosine_similarity:.4f}"
|
||||
)
|
||||
|
||||
print(f"{model.artifact_prefix} BF16 peak VRAM delta: {mem_bf16:.2f} GiB")
|
||||
print(f"{model.artifact_prefix} GGUF peak VRAM delta: {mem_gguf:.2f} GiB")
|
||||
reduction = (mem_bf16 - mem_gguf) / mem_bf16 * 100
|
||||
print(f"VRAM reduction: {reduction:.1f}%")
|
||||
assert mem_gguf < mem_bf16, f"GGUF ({mem_gguf:.2f} GiB) should use less VRAM than BF16 ({mem_bf16:.2f} GiB)"
|
||||
@@ -1166,7 +1166,10 @@ class TestQwen3OmniPipeline:
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.input_sources == (0,)
|
||||
assert s.sampling_constraints["stop_token_ids"] == [2150]
|
||||
assert s.custom_process_input_func is not None
|
||||
# thinker2talker was removed: sync_process_input_func always wins.
|
||||
assert s.custom_process_input_func is None
|
||||
assert s.sync_process_input_func is not None
|
||||
assert s.sync_process_input_func.endswith("thinker2talker_token_only")
|
||||
assert s.custom_process_next_stage_input_func is not None
|
||||
|
||||
def test_code2wav(self):
|
||||
@@ -1180,7 +1183,7 @@ class TestQwen3OmniPipeline:
|
||||
assert isinstance(s, StagePipelineConfig)
|
||||
assert s.execution_type == StageExecutionType.LLM_GENERATION
|
||||
assert s.final_output_type == "audio"
|
||||
assert s.custom_process_input_func is not None
|
||||
assert s.engine_output_type == "audio"
|
||||
|
||||
|
||||
class TestQwen2_5OmniPipeline:
|
||||
@@ -1718,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={},
|
||||
@@ -1940,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})
|
||||
|
||||
@@ -56,3 +56,26 @@ def test_warning_contains_version_strings():
|
||||
msg = str(record[0].message)
|
||||
assert "0.19.0rc2.dev21" in msg
|
||||
assert "0.18.0" in msg
|
||||
|
||||
|
||||
def test_omni_plugin_loader_invokes_vllm_general_plugins_first(monkeypatch):
|
||||
import vllm.plugins as vllm_plugins
|
||||
|
||||
import vllm_omni.plugins as omni_plugins
|
||||
|
||||
calls = []
|
||||
monkeypatch.setattr(omni_plugins, "omni_plugins_loaded", False)
|
||||
monkeypatch.setattr(vllm_plugins, "load_general_plugins", lambda: calls.append("vllm"))
|
||||
monkeypatch.setattr(
|
||||
omni_plugins,
|
||||
"load_omni_plugins_by_group",
|
||||
lambda group: {"omni": lambda: calls.append(group)},
|
||||
)
|
||||
|
||||
omni_plugins.load_omni_general_plugins()
|
||||
|
||||
assert calls == ["vllm", omni_plugins.OMNI_DEFAULT_PLUGINS_GROUP]
|
||||
|
||||
omni_plugins.load_omni_general_plugins()
|
||||
|
||||
assert calls == ["vllm", omni_plugins.OMNI_DEFAULT_PLUGINS_GROUP]
|
||||
@@ -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]]:
|
||||
|
||||
@@ -551,7 +551,7 @@ class _DiffusionConfigProjection:
|
||||
elif isinstance(self.quantization_config, str):
|
||||
self.quantization_config = build_quant_config(self.quantization_config)
|
||||
elif isinstance(self.quantization_config, Mapping):
|
||||
self.quantization_config = build_quant_config(dict(self.quantization_config))
|
||||
self.quantization_config = dict(self.quantization_config)
|
||||
else:
|
||||
raise TypeError(
|
||||
"quantization_config must be str, dict, QuantizationConfig, or None, "
|
||||
|
||||
@@ -47,7 +47,6 @@ stages:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
repetition_penalty: 1.05
|
||||
|
||||
- stage_id: 2
|
||||
@@ -67,7 +66,6 @@ stages:
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
repetition_penalty: 1.1
|
||||
|
||||
platforms:
|
||||
|
||||
@@ -35,7 +35,7 @@ connectors:
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_seqs: 10
|
||||
max_num_seqs: 64
|
||||
gpu_memory_utilization: 0.3
|
||||
trust_remote_code: true
|
||||
# Prefix caching on by default. The HS opt-out (talker only consumes the
|
||||
@@ -43,9 +43,9 @@ stages:
|
||||
# PR make per-step bookkeeping negligible, and any text repetition across
|
||||
# requests (same voice or shared ref_audio) is a net throughput win.
|
||||
# Disable explicitly if memory headroom is tight on a single-GPU deploy.
|
||||
enable_prefix_caching: true
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: true
|
||||
max_num_batched_tokens: 512
|
||||
max_num_batched_tokens: 32768
|
||||
max_model_len: 4096
|
||||
devices: "0"
|
||||
output_connectors:
|
||||
@@ -67,7 +67,7 @@ stages:
|
||||
|
||||
- stage_id: 1
|
||||
# Tuned for Base voice clone; CustomVoice / VoiceDesign are TTFA-optimal at 1.
|
||||
max_num_seqs: 10
|
||||
max_num_seqs: 64
|
||||
gpu_memory_utilization: 0.3
|
||||
enforce_eager: false
|
||||
trust_remote_code: true
|
||||
@@ -95,3 +95,8 @@ platforms:
|
||||
stages:
|
||||
- stage_id: 0
|
||||
enforce_eager: false
|
||||
compilation_config:
|
||||
cudagraph_mode: PIECEWISE
|
||||
- stage_id: 1
|
||||
compilation_config:
|
||||
cudagraph_mode: PIECEWISE
|
||||
@@ -11,7 +11,6 @@ from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
import torch
|
||||
from huggingface_hub import hf_hub_download
|
||||
from torch import nn
|
||||
from vllm.config import ModelConfig
|
||||
from vllm.config.load import LoadConfig
|
||||
@@ -36,7 +35,6 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp
|
||||
from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
|
||||
get_checkpoint_adapter,
|
||||
)
|
||||
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
|
||||
from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
|
||||
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
|
||||
from vllm_omni.diffusion.registry import initialize_model
|
||||
@@ -381,10 +379,7 @@ class DiffusersPipelineLoader:
|
||||
# DiffusionPipeline.from_pretrained() internally; it does
|
||||
# not use our native customized pipeline classes.
|
||||
cast(DiffusersAdapterPipeline, model).load_weights()
|
||||
elif self._is_gguf_quantization():
|
||||
self._load_weights_with_gguf(model)
|
||||
else:
|
||||
# Quantization does not happen in `load_weights` but after it
|
||||
self.load_weights(model)
|
||||
|
||||
# Process weights after loading for quantization (e.g., FP8 online quantization)
|
||||
@@ -454,7 +449,7 @@ class DiffusersPipelineLoader:
|
||||
def _is_expected_quantized_weight(name: str) -> bool:
|
||||
"""Return True if *name* is a quantization-specific parameter.
|
||||
|
||||
Quantization methods (GPTQ, AWQ, FP8, GGUF, Autoround, etc.) create extra
|
||||
Quantization methods (GPTQ, AWQ, FP8, Autoround, etc.) create extra
|
||||
parameters that have no counterpart in an unquantized checkpoint.
|
||||
These are expected to be absent and should not trigger a load error.
|
||||
"""
|
||||
@@ -467,8 +462,6 @@ class DiffusersPipelineLoader:
|
||||
".weight_scale",
|
||||
".weight_scale_inv",
|
||||
".input_scale",
|
||||
# GGUF
|
||||
".qweight_type",
|
||||
# INT8 (weight_scale already covered above)
|
||||
)
|
||||
return name.endswith(_QUANTIZED_WEIGHT_SUFFIXES)
|
||||
@@ -500,116 +493,6 @@ class DiffusersPipelineLoader:
|
||||
if unexpected_missing:
|
||||
raise ValueError(f"Following weights were not initialized from checkpoint: {unexpected_missing}")
|
||||
|
||||
def _get_gguf_model_from_config(self) -> str | None:
|
||||
"""Extract gguf_model from quant_config (handles both dict and object)."""
|
||||
if self.quant_config is None:
|
||||
return None
|
||||
if isinstance(self.quant_config, dict):
|
||||
return self.quant_config.get("gguf_model")
|
||||
return getattr(self.quant_config, "gguf_model", None)
|
||||
|
||||
def _is_gguf_quantization(self) -> bool:
|
||||
"""Check whether or not this pipeline loader is pointing at a GGUF config."""
|
||||
from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
|
||||
|
||||
maybe_gguf_model = self._get_gguf_model_from_config()
|
||||
is_gguf = (
|
||||
isinstance(self.quant_config, DiffusionGGUFConfig)
|
||||
or (isinstance(self.quant_config, dict) and self.quant_config.get("method") == "gguf")
|
||||
or (hasattr(self.quant_config, "get_name") and self.quant_config.get_name() == "gguf")
|
||||
)
|
||||
if is_gguf:
|
||||
if not maybe_gguf_model:
|
||||
raise ValueError("GGUF quantization requires gguf_model")
|
||||
return True
|
||||
|
||||
if isinstance(self.quant_config, dict):
|
||||
return False
|
||||
|
||||
return bool(maybe_gguf_model)
|
||||
|
||||
def _is_transformer_source(self, source: "ComponentSource") -> bool:
|
||||
if source.subfolder == "transformer":
|
||||
return True
|
||||
return source.prefix.startswith("transformer.")
|
||||
|
||||
def _get_model_loadable_names(self, model: nn.Module) -> set[str]:
|
||||
# Avoid model.state_dict() here because GGUF uses UninitializedParameter
|
||||
# which raises during detach(). Collect names directly.
|
||||
return {name for name, _ in model.named_parameters()} | {name for name, _ in model.named_buffers()}
|
||||
|
||||
def _resolve_gguf_model_path(self, gguf_model: str, revision: str | None) -> str:
|
||||
if os.path.isfile(gguf_model):
|
||||
return gguf_model
|
||||
# repo_id/filename.gguf
|
||||
if "/" in gguf_model and gguf_model.endswith(".gguf"):
|
||||
repo_id, filename = gguf_model.rsplit("/", 1)
|
||||
return hf_hub_download(
|
||||
repo_id=repo_id,
|
||||
filename=filename,
|
||||
revision=revision,
|
||||
cache_dir=self.load_config.download_dir,
|
||||
)
|
||||
# repo_id:quant_type
|
||||
if "/" in gguf_model and ":" in gguf_model:
|
||||
repo_id, quant_type = gguf_model.rsplit(":", 1)
|
||||
return download_gguf(
|
||||
repo_id,
|
||||
quant_type,
|
||||
cache_dir=self.load_config.download_dir,
|
||||
revision=revision,
|
||||
ignore_patterns=self.load_config.ignore_patterns,
|
||||
)
|
||||
raise ValueError(
|
||||
f"Unrecognized GGUF reference: {gguf_model!r} (expected local file, "
|
||||
"<repo_id>/<filename>.gguf, or <repo_id>:<quant_type>)"
|
||||
)
|
||||
|
||||
def _get_gguf_weights_iterator(
|
||||
self,
|
||||
source: "ComponentSource",
|
||||
model: nn.Module,
|
||||
) -> Generator[tuple[str, torch.Tensor], None, None]:
|
||||
gguf_model = self._get_gguf_model_from_config()
|
||||
if gguf_model is None:
|
||||
raise ValueError("GGUF quantization requires quantization_config.gguf_model")
|
||||
gguf_file = self._resolve_gguf_model_path(gguf_model, self.od_config.revision)
|
||||
adapter = get_gguf_adapter(gguf_file, model, source, self.od_config)
|
||||
weights_iter = adapter.weights_iterator()
|
||||
return ((source.prefix + name, tensor) for (name, tensor) in weights_iter)
|
||||
|
||||
def _load_weights_with_gguf(self, model: nn.Module) -> set[str]:
|
||||
sources = self._get_weight_sources(model)
|
||||
loaded: set[str] = set()
|
||||
loadable_names: set[str] | None = None
|
||||
|
||||
for source in sources:
|
||||
if self._is_transformer_source(source):
|
||||
loaded |= model.load_weights(self._get_gguf_weights_iterator(source, model))
|
||||
|
||||
# GGUF checkpoints can be transformer-only or partially quantized.
|
||||
# Only fall back to HF if this source still has missing loadable weights.
|
||||
loadable_names = loadable_names or self._get_model_loadable_names(model)
|
||||
has_missing_for_source = any(
|
||||
name.startswith(source.prefix) and name not in loaded for name in loadable_names
|
||||
)
|
||||
if not has_missing_for_source:
|
||||
continue
|
||||
|
||||
hf_iter = self._get_weights_iterator(source)
|
||||
hf_iter = (
|
||||
(name, tensor) for (name, tensor) in hf_iter if name in loadable_names and name not in loaded
|
||||
)
|
||||
loaded |= model.load_weights(hf_iter)
|
||||
else:
|
||||
loaded |= model.load_weights(self._get_weights_iterator(source))
|
||||
|
||||
weights_to_load = self._get_expected_parameter_names(model)
|
||||
weights_not_loaded = weights_to_load - loaded
|
||||
if weights_not_loaded:
|
||||
raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}")
|
||||
return loaded
|
||||
|
||||
def _init_from_load_format(
|
||||
self,
|
||||
load_format: str,
|
||||
|
||||
@@ -1,47 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from .base import GGUFAdapter
|
||||
from .flux2_klein import Flux2KleinGGUFAdapter
|
||||
from .qwen_image import QwenImageGGUFAdapter
|
||||
from .z_image import ZImageGGUFAdapter
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import (
|
||||
DiffusersPipelineLoader,
|
||||
)
|
||||
|
||||
|
||||
def get_gguf_adapter(
|
||||
gguf_file: str,
|
||||
model: torch.nn.Module,
|
||||
source: DiffusersPipelineLoader.ComponentSource,
|
||||
od_config: OmniDiffusionConfig,
|
||||
) -> GGUFAdapter:
|
||||
adapter_classes = (QwenImageGGUFAdapter, ZImageGGUFAdapter, Flux2KleinGGUFAdapter)
|
||||
for adapter_cls in adapter_classes:
|
||||
if adapter_cls.is_compatible(od_config, model, source):
|
||||
return adapter_cls(gguf_file, model, source, od_config)
|
||||
model_type = None
|
||||
if od_config.tf_model_config is not None:
|
||||
model_type = od_config.tf_model_config.get("model_type")
|
||||
supported = ", ".join(cls.__name__ for cls in adapter_classes)
|
||||
raise ValueError(
|
||||
"No GGUF adapter matched diffusion model "
|
||||
f"(model_class_name={od_config.model_class_name!r}, model_type={model_type!r}). "
|
||||
f"Supported adapters: {supported}."
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"GGUFAdapter",
|
||||
"Flux2KleinGGUFAdapter",
|
||||
"QwenImageGGUFAdapter",
|
||||
"ZImageGGUFAdapter",
|
||||
"get_gguf_adapter",
|
||||
]
|
||||
@@ -1,109 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Generator
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
try:
|
||||
import gguf
|
||||
except ImportError:
|
||||
# Optional dependency: only required for GGUF-quantized diffusion models.
|
||||
# Keep this import lazy so a missing ``gguf`` does not break importing the
|
||||
# diffusion model loader (and thus inspection/loading of every other
|
||||
# diffusion model that imports through this package).
|
||||
gguf = None
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import (
|
||||
DiffusersPipelineLoader,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MappedTensor:
|
||||
name: str
|
||||
tensor: Any
|
||||
tensor_type: Any
|
||||
row_slice: slice | None = None
|
||||
swap_scale_shift: bool = False
|
||||
|
||||
|
||||
class GGUFAdapter(ABC):
|
||||
"""Base class for model-specific GGUF adapters."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
gguf_file: str,
|
||||
model: torch.nn.Module,
|
||||
source: DiffusersPipelineLoader.ComponentSource,
|
||||
od_config: OmniDiffusionConfig,
|
||||
) -> None:
|
||||
self.gguf_file = gguf_file
|
||||
self.model = model
|
||||
self.source = source
|
||||
self.od_config = od_config
|
||||
|
||||
@staticmethod
|
||||
def is_compatible(
|
||||
od_config: OmniDiffusionConfig,
|
||||
model: torch.nn.Module,
|
||||
source: DiffusersPipelineLoader.ComponentSource,
|
||||
) -> bool:
|
||||
return False
|
||||
|
||||
@abstractmethod
|
||||
def weights_iterator(self) -> Generator[tuple[str, torch.Tensor], None, None]:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
# FIXME(Isotr0py): Sync implemnentation with upstream vLLM?
|
||||
def gguf_quant_weights_iterator(gguf_file: str) -> Generator[tuple[str, torch.Tensor]]:
|
||||
"""
|
||||
Iterate over the quant weights in the model gguf files and convert
|
||||
them to torch tensors.
|
||||
Be careful of the order of yielding weight types and weights data,
|
||||
we have to yield all weight types first before yielding any weights.
|
||||
Otherwise it would cause issue when loading weights with for packed
|
||||
layer with different quant types.
|
||||
"""
|
||||
|
||||
if gguf is None:
|
||||
raise ImportError(
|
||||
"The 'gguf' package is required to load GGUF-quantized diffusion models. "
|
||||
"Install it with `pip install gguf`."
|
||||
)
|
||||
reader = gguf.GGUFReader(gguf_file)
|
||||
|
||||
for tensor in reader.tensors:
|
||||
weight_type = tensor.tensor_type
|
||||
name = tensor.name
|
||||
|
||||
if weight_type.name not in ("F32", "F16"):
|
||||
weight_type_name = name.replace("weight", "qweight_type")
|
||||
weight_type = torch.tensor(weight_type)
|
||||
yield weight_type_name, weight_type
|
||||
|
||||
for tensor in reader.tensors:
|
||||
weight = tensor.data
|
||||
weight_type = tensor.tensor_type
|
||||
name = tensor.name
|
||||
if weight_type.name not in ("F32", "F16"):
|
||||
name = name.replace("weight", "qweight")
|
||||
if weight_type.name == "BF16" and tensor.data.dtype == np.uint8:
|
||||
# BF16 is currently the only "quantization" type that isn't
|
||||
# actually quantized but is read as a raw byte tensor.
|
||||
# Reinterpret as `torch.bfloat16` tensor.
|
||||
weight = weight.view(np.uint16)
|
||||
if reader.byte_order == "S":
|
||||
# GGUF endianness != system endianness
|
||||
weight = weight.byteswap()
|
||||
param = torch.tensor(weight).view(torch.bfloat16)
|
||||
else:
|
||||
param = torch.tensor(weight)
|
||||
yield name, param
|
||||
@@ -1,98 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
|
||||
from .base import GGUFAdapter, gguf_quant_weights_iterator
|
||||
|
||||
FLUX2_TRANSFORMER_KEYS_RENAME_DICT = {
|
||||
"single_blocks.": "single_transformer_blocks.",
|
||||
# Image and text input projections
|
||||
"img_in": "x_embedder",
|
||||
"txt_in": "context_embedder",
|
||||
# Timestep and guidance embeddings
|
||||
"time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1",
|
||||
"time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2",
|
||||
"guidance_in.in_layer": "time_guidance_embed.guidance_embedder.linear_1",
|
||||
"guidance_in.out_layer": "time_guidance_embed.guidance_embedder.linear_2",
|
||||
# Modulation parameters
|
||||
"double_stream_modulation_img.lin": "double_stream_modulation_img.linear",
|
||||
"double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear",
|
||||
"single_stream_modulation.lin": "single_stream_modulation.linear",
|
||||
# Final output layer
|
||||
# "final_layer.adaLN_modulation.1": "norm_out.linear", # Handle separately since we need to swap mod params
|
||||
"final_layer.linear": "proj_out",
|
||||
}
|
||||
|
||||
FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP = {
|
||||
"final_layer.adaLN_modulation.1": "norm_out.linear",
|
||||
}
|
||||
|
||||
FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP = {
|
||||
"double_blocks.": "transformer_blocks.",
|
||||
# Handle fused QKV projections separately as we need to break into Q, K, V projections
|
||||
"img_attn.norm.query_norm": "attn.norm_q",
|
||||
"img_attn.norm.key_norm": "attn.norm_k",
|
||||
"img_attn.proj": "attn.to_out.0",
|
||||
"img_mlp.0": "ff.linear_in",
|
||||
"img_mlp.2": "ff.linear_out",
|
||||
"txt_attn.norm.query_norm": "attn.norm_added_q",
|
||||
"txt_attn.norm.key_norm": "attn.norm_added_k",
|
||||
"txt_attn.proj": "attn.to_add_out",
|
||||
"txt_mlp.0": "ff_context.linear_in",
|
||||
"txt_mlp.2": "ff_context.linear_out",
|
||||
# Additional for fuse qkv
|
||||
"img_attn.qkv": "attn.to_qkv",
|
||||
"txt_attn.qkv": "attn.add_kv_proj",
|
||||
}
|
||||
|
||||
FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP = {
|
||||
"linear1": "attn.to_qkv_mlp_proj",
|
||||
"linear2": "attn.to_out",
|
||||
"norm.query_norm": "attn.norm_q",
|
||||
"norm.key_norm": "attn.norm_k",
|
||||
}
|
||||
|
||||
|
||||
class Flux2KleinGGUFAdapter(GGUFAdapter):
|
||||
"""GGUF adapter for Flux2-Klein models with qkv splitting and adaLN swap."""
|
||||
|
||||
@staticmethod
|
||||
def is_compatible(od_config, model: torch.nn.Module, source) -> bool:
|
||||
model_class = od_config.model_class_name or ""
|
||||
if model_class.startswith("Flux2"):
|
||||
return True
|
||||
cfg = od_config.tf_model_config
|
||||
if cfg is not None:
|
||||
model_type = str(cfg.get("model_type", "")).lower()
|
||||
if model_type.startswith("flux"):
|
||||
return True
|
||||
return False
|
||||
|
||||
gguf_to_hf_mapper = WeightsMapper(
|
||||
# double_stream_modulation
|
||||
orig_to_new_prefix=FLUX2_TRANSFORMER_KEYS_RENAME_DICT | FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP,
|
||||
orig_to_new_substr=FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP | FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP,
|
||||
)
|
||||
|
||||
def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]:
|
||||
def custom_weights_adapter(
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
) -> Iterable[tuple[str, torch.Tensor]]:
|
||||
for name, weight in weights:
|
||||
# Handle the special case for adaLN modulation parameters that require swapping shift and scale
|
||||
if name.endswith(".scale"):
|
||||
name = name.replace(".scale", ".weight")
|
||||
if name == "norm_out.linear.weight":
|
||||
shift, scale = weight.chunk(2, dim=0)
|
||||
weight = torch.cat([scale, shift], dim=0)
|
||||
yield name, weight
|
||||
else:
|
||||
yield name, weight
|
||||
|
||||
weights = gguf_quant_weights_iterator(self.gguf_file)
|
||||
weights = self.gguf_to_hf_mapper.apply(weights)
|
||||
yield from custom_weights_adapter(weights)
|
||||
@@ -1,27 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
|
||||
from .base import GGUFAdapter, gguf_quant_weights_iterator
|
||||
|
||||
|
||||
class QwenImageGGUFAdapter(GGUFAdapter):
|
||||
"""GGUF adapter for the Qwen-Image transformer family."""
|
||||
|
||||
@staticmethod
|
||||
def is_compatible(od_config, model: torch.nn.Module, source) -> bool:
|
||||
model_class = od_config.model_class_name or ""
|
||||
if model_class.startswith("QwenImage"):
|
||||
return True
|
||||
cfg = od_config.tf_model_config
|
||||
if cfg is not None:
|
||||
model_type = str(cfg.get("model_type", "")).lower()
|
||||
if model_type.startswith("qwen_image"):
|
||||
return True
|
||||
return False
|
||||
|
||||
def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]:
|
||||
yield from gguf_quant_weights_iterator(self.gguf_file)
|
||||
@@ -1,43 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
|
||||
from .base import GGUFAdapter, gguf_quant_weights_iterator
|
||||
|
||||
Z_IMAGE_KEYS_RENAME_DICT = {
|
||||
"final_layer.": "all_final_layer.2-1.",
|
||||
"x_embedder.": "all_x_embedder.2-1.",
|
||||
".attention.qkv": ".attention.to_qkv",
|
||||
".attention.k_norm": ".attention.norm_k",
|
||||
".attention.q_norm": ".attention.norm_q",
|
||||
".attention.out": ".attention.to_out.0",
|
||||
"model.diffusion_model.": "",
|
||||
}
|
||||
|
||||
|
||||
class ZImageGGUFAdapter(GGUFAdapter):
|
||||
"""GGUF adapter for Z-Image models with QKV/FFN shard support."""
|
||||
|
||||
@staticmethod
|
||||
def is_compatible(od_config, model: torch.nn.Module, source) -> bool:
|
||||
model_class = od_config.model_class_name or ""
|
||||
if model_class.startswith("ZImage"):
|
||||
return True
|
||||
cfg = od_config.tf_model_config
|
||||
if cfg is not None:
|
||||
model_type = str(cfg.get("model_type", "")).lower()
|
||||
if model_type in {"z_image", "zimage", "z-image"}:
|
||||
return True
|
||||
return False
|
||||
|
||||
gguf_to_hf_mapper = WeightsMapper(
|
||||
orig_to_new_substr=Z_IMAGE_KEYS_RENAME_DICT,
|
||||
)
|
||||
|
||||
def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]:
|
||||
weights = gguf_quant_weights_iterator(self.gguf_file)
|
||||
yield from self.gguf_to_hf_mapper.apply(weights)
|
||||
@@ -1034,7 +1034,7 @@ class Flux2Transformer2DModel(nn.Module):
|
||||
name = original_name
|
||||
if name not in params_dict and ".to_out.0." in name:
|
||||
name = name.replace(".to_out.0.", ".to_out.")
|
||||
# Some GGUF checkpoints include quantized tensors for modules that
|
||||
# Some quantized checkpoints include tensors for modules that
|
||||
# are intentionally left unquantized in this model.
|
||||
param = params_dict.get(name)
|
||||
if param is None:
|
||||
|
||||
@@ -1044,7 +1044,7 @@ class Flux2Transformer2DModel(nn.Module):
|
||||
name = original_name
|
||||
if name not in params_dict and ".to_out.0." in name:
|
||||
name = name.replace(".to_out.0.", ".to_out.")
|
||||
# Some GGUF checkpoints include quantized tensors for modules that
|
||||
# Some quantized checkpoints include tensors for modules that
|
||||
# are intentionally left unquantized in this model.
|
||||
param = params_dict.get(name)
|
||||
if param is None:
|
||||
|
||||
@@ -1,16 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import (
|
||||
LingBotVideoTransformer3DModel,
|
||||
)
|
||||
from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import (
|
||||
LingBotVideoPipeline,
|
||||
get_lingbot_video_post_process_func,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"LingBotVideoPipeline",
|
||||
"LingBotVideoTransformer3DModel",
|
||||
"get_lingbot_video_post_process_func",
|
||||
]
|
||||
@@ -1,702 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# Adapted from LingBot-Video (https://github.com/Robbyant/lingbot-video).
|
||||
|
||||
import math
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.utils.fa import (
|
||||
flash_attn_varlen_func as flash_attn_varlen_func_v3,
|
||||
)
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
|
||||
LINGBOT_VIDEO_FP32_MODULES = (
|
||||
"time_embedder",
|
||||
"time_modulation",
|
||||
"scale_shift_table",
|
||||
"norm",
|
||||
"norm1",
|
||||
"norm2",
|
||||
"norm_q",
|
||||
"norm_k",
|
||||
"norm_post_attn",
|
||||
"norm_post_ffn",
|
||||
"norm_out",
|
||||
"norm_out_modulation",
|
||||
"router",
|
||||
)
|
||||
|
||||
|
||||
def should_keep_in_fp32(name: str) -> bool:
|
||||
return any(module_name in name.split(".") for module_name in LINGBOT_VIDEO_FP32_MODULES)
|
||||
|
||||
|
||||
def _all_to_all_split_cat(
|
||||
local_input: torch.Tensor,
|
||||
scatter_dim: int,
|
||||
gather_dim: int,
|
||||
group: dist.ProcessGroup,
|
||||
) -> torch.Tensor:
|
||||
world_size = dist.get_world_size(group)
|
||||
input_list = [tensor.contiguous() for tensor in torch.tensor_split(local_input, world_size, scatter_dim)]
|
||||
output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)]
|
||||
dist.all_to_all(output_list, input_list, group=group)
|
||||
return torch.cat(output_list, dim=gather_dim).contiguous()
|
||||
|
||||
|
||||
class LingBotVideoRMSNorm(nn.Module):
|
||||
"""RMSNorm with fp32 accumulation."""
|
||||
|
||||
def __init__(self, dim: int, eps: float = 1e-6):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
self.variance_epsilon = eps
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
input_dtype = hidden_states.dtype
|
||||
hidden_states = hidden_states.to(torch.float32)
|
||||
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
||||
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
||||
return (self.weight * hidden_states).to(input_dtype)
|
||||
|
||||
|
||||
def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
|
||||
"""Apply complex RoPE to `(B, S, H, D)` attention tensors."""
|
||||
with torch.amp.autocast("cuda", enabled=False):
|
||||
x_c = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
|
||||
out = torch.view_as_real(x_c * freqs_cis.unsqueeze(2)).flatten(3)
|
||||
return out.type_as(x)
|
||||
|
||||
|
||||
class LingBotVideoRotaryEmbedding(nn.Module):
|
||||
"""Complex64 RoPE table indexed by position ids."""
|
||||
|
||||
def __init__(self, axes_dims: tuple[int, ...], axes_lens: tuple[int, ...], theta: float):
|
||||
super().__init__()
|
||||
self.axes_dims = tuple(axes_dims)
|
||||
self.axes_lens = list(axes_lens)
|
||||
self.theta = theta
|
||||
self.freqs_cis = None
|
||||
|
||||
@staticmethod
|
||||
def precompute_freqs_cis(dim: tuple[int, ...], end: tuple[int, ...], theta: float):
|
||||
freqs_cis = []
|
||||
for d, e in zip(dim, end):
|
||||
freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d))
|
||||
timestep = torch.arange(e, device=freqs.device, dtype=torch.float64)
|
||||
freqs = torch.outer(timestep, freqs).float()
|
||||
freqs_cis.append(torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64))
|
||||
return freqs_cis
|
||||
|
||||
def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
|
||||
# position_ids: (S, 3) int -> (S, head_dim/2) complex64
|
||||
device = position_ids.device
|
||||
max_vals = position_ids.max(dim=0).values.tolist()
|
||||
needs_rebuild = self.freqs_cis is None or any(
|
||||
max_val >= axis_len for max_val, axis_len in zip(max_vals, self.axes_lens)
|
||||
)
|
||||
if needs_rebuild:
|
||||
for i in range(len(self.axes_lens)):
|
||||
if max_vals[i] >= self.axes_lens[i]:
|
||||
self.axes_lens[i] = int(max_vals[i] * 1.5) + 1
|
||||
self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, tuple(self.axes_lens), theta=self.theta)
|
||||
self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
|
||||
elif self.freqs_cis[0].device != device:
|
||||
self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
|
||||
|
||||
return torch.cat([self.freqs_cis[i][position_ids[:, i]] for i in range(len(self.axes_dims))], dim=-1)
|
||||
|
||||
|
||||
def make_joint_position_ids(text_len: int, grid_t: int, grid_h: int, grid_w: int, device: torch.device) -> torch.Tensor:
|
||||
"""3D positions in [video; text] order. Text t-axis is 1..text_len; video t-axis starts at text_len+1.
|
||||
|
||||
Matches patchify_and_embed: cap start (1,0,0); vision start (cap_len+1,0,0);
|
||||
freqs ordered with x first and cap second (same order as cat_interleave).
|
||||
"""
|
||||
tt = torch.arange(grid_t, device=device, dtype=torch.int32) + (text_len + 1)
|
||||
hh = torch.arange(grid_h, device=device, dtype=torch.int32)
|
||||
ww = torch.arange(grid_w, device=device, dtype=torch.int32)
|
||||
grid = torch.stack(torch.meshgrid(tt, hh, ww, indexing="ij"), dim=-1).flatten(0, 2)
|
||||
text_t = torch.arange(text_len, device=device, dtype=torch.int32) + 1
|
||||
text_pos = torch.stack([text_t, torch.zeros_like(text_t), torch.zeros_like(text_t)], dim=-1)
|
||||
return torch.cat([grid, text_pos], dim=0) # (Nx + L, 3)
|
||||
|
||||
|
||||
def _cat_interleave(
|
||||
a: torch.Tensor,
|
||||
len_a: list[int],
|
||||
b: torch.Tensor,
|
||||
len_b: list[int],
|
||||
) -> torch.Tensor:
|
||||
a_split = torch.split(a, len_a, dim=1)
|
||||
b_split = torch.split(b, len_b, dim=1)
|
||||
blocks: list[torch.Tensor] = []
|
||||
for x_part, text_part in zip(a_split, b_split):
|
||||
blocks.extend([x_part, text_part])
|
||||
return torch.cat(blocks, dim=1)
|
||||
|
||||
|
||||
def _packed_block_attention_mask(sample_seq_lens: list[int], device: torch.device) -> torch.Tensor:
|
||||
total_seq_len = sum(sample_seq_lens)
|
||||
mask = torch.zeros(
|
||||
1,
|
||||
1,
|
||||
total_seq_len,
|
||||
total_seq_len,
|
||||
dtype=torch.bool,
|
||||
device=device,
|
||||
)
|
||||
start = 0
|
||||
for seq_len in sample_seq_lens:
|
||||
end = start + seq_len
|
||||
mask[:, :, start:end, start:end] = True
|
||||
start = end
|
||||
return mask
|
||||
|
||||
|
||||
class LingBotVideoTextEmbedder(nn.Module):
|
||||
"""Matches CondProjection: RMSNorm(text_dim, eps=1e-6 fixed) -> Linear-SiLU-Linear."""
|
||||
|
||||
def __init__(self, text_dim: int, hidden_size: int):
|
||||
super().__init__()
|
||||
self.norm = LingBotVideoRMSNorm(text_dim, eps=1e-6)
|
||||
self.linear_1 = nn.Linear(text_dim, hidden_size, bias=True)
|
||||
self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.norm(x)
|
||||
return self.linear_2(F.silu(self.linear_1(x)))
|
||||
|
||||
|
||||
class LingBotVideoAttention(nn.Module):
|
||||
def __init__(self, hidden_size, num_heads, norm_eps, qkv_bias, out_bias):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = hidden_size // num_heads
|
||||
self.to_q = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
|
||||
self.to_k = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
|
||||
self.to_v = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
|
||||
self.norm_q = LingBotVideoRMSNorm(self.head_dim, norm_eps)
|
||||
self.norm_k = LingBotVideoRMSNorm(self.head_dim, norm_eps)
|
||||
self.to_out = nn.Linear(hidden_size, hidden_size, bias=out_bias)
|
||||
self.attn = Attention(
|
||||
num_heads=num_heads,
|
||||
head_size=self.head_dim,
|
||||
softmax_scale=1.0 / math.sqrt(self.head_dim),
|
||||
causal=False,
|
||||
num_kv_heads=num_heads,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
rotary_emb,
|
||||
attention_mask=None,
|
||||
packed_indices: dict[str, torch.Tensor] | None = None,
|
||||
parallel_config=None,
|
||||
):
|
||||
B, S, _ = x.shape
|
||||
if os.environ.get("LINGBOT_FUSED_QKV_LINEAR") == "1":
|
||||
weight = torch.cat(
|
||||
(self.to_q.weight, self.to_k.weight, self.to_v.weight),
|
||||
dim=0,
|
||||
)
|
||||
bias = None
|
||||
if self.to_q.bias is not None:
|
||||
bias = torch.cat(
|
||||
(self.to_q.bias, self.to_k.bias, self.to_v.bias),
|
||||
dim=0,
|
||||
)
|
||||
qkv = F.linear(x, weight, bias)
|
||||
q, k, v = qkv.view(B, S, 3, self.num_heads, self.head_dim).unbind(2)
|
||||
else:
|
||||
q = self.to_q(x).unflatten(2, (self.num_heads, self.head_dim))
|
||||
k = self.to_k(x).unflatten(2, (self.num_heads, self.head_dim))
|
||||
v = self.to_v(x).unflatten(2, (self.num_heads, self.head_dim))
|
||||
q = apply_rotary_emb(self.norm_q(q), rotary_emb)
|
||||
k = apply_rotary_emb(self.norm_k(k), rotary_emb)
|
||||
if packed_indices is None:
|
||||
out = self.attn(q, k, v, attn_metadata=AttentionMetadata(attn_mask=attention_mask))
|
||||
else:
|
||||
packed_attention_mask = packed_indices.get("attention_mask")
|
||||
if packed_attention_mask is not None and parallel_config is None:
|
||||
out = self.attn.sdpa_fallback.forward(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
AttentionMetadata(attn_mask=packed_attention_mask),
|
||||
)
|
||||
return self.to_out(out.flatten(2, 3).type_as(x))
|
||||
if flash_attn_varlen_func_v3 is None:
|
||||
raise RuntimeError(
|
||||
"A flash attention varlen function is required for packed context parallel attention."
|
||||
)
|
||||
if parallel_config is None:
|
||||
result = flash_attn_varlen_func_v3(
|
||||
q=q.reshape(-1, self.num_heads, self.head_dim),
|
||||
k=k.reshape(-1, self.num_heads, self.head_dim),
|
||||
v=v.reshape(-1, self.num_heads, self.head_dim),
|
||||
cu_seqlens_q=packed_indices["cu_seqlens_kv"],
|
||||
cu_seqlens_k=packed_indices["cu_seqlens_kv"],
|
||||
max_seqlen_q=packed_indices["max_seqlen_in_batch_kv"],
|
||||
max_seqlen_k=packed_indices["max_seqlen_in_batch_kv"],
|
||||
causal=False,
|
||||
)
|
||||
out = result[0] if isinstance(result, tuple) else result
|
||||
out = out.reshape(B, S, self.num_heads, self.head_dim)
|
||||
else:
|
||||
group = parallel_config.context_parallel_config._ulysses_mesh.get_group()
|
||||
world_size = dist.get_world_size(group)
|
||||
local_heads = self.num_heads // world_size
|
||||
q_global = _all_to_all_split_cat(
|
||||
q.reshape(B, S, self.num_heads * self.head_dim),
|
||||
scatter_dim=2,
|
||||
gather_dim=1,
|
||||
group=group,
|
||||
).view(B, S * world_size, local_heads, self.head_dim)
|
||||
k_global = _all_to_all_split_cat(
|
||||
k.reshape(B, S, self.num_heads * self.head_dim),
|
||||
scatter_dim=2,
|
||||
gather_dim=1,
|
||||
group=group,
|
||||
).view(B, S * world_size, local_heads, self.head_dim)
|
||||
v_global = _all_to_all_split_cat(
|
||||
v.reshape(B, S, self.num_heads * self.head_dim),
|
||||
scatter_dim=2,
|
||||
gather_dim=1,
|
||||
group=group,
|
||||
).view(B, S * world_size, local_heads, self.head_dim)
|
||||
q_flat = q_global.reshape(-1, local_heads, self.head_dim)
|
||||
k_flat = k_global.reshape(-1, local_heads, self.head_dim)
|
||||
v_flat = v_global.reshape(-1, local_heads, self.head_dim)
|
||||
result = flash_attn_varlen_func_v3(
|
||||
q=q_flat,
|
||||
k=k_flat,
|
||||
v=v_flat,
|
||||
cu_seqlens_q=packed_indices["cu_seqlens_kv"],
|
||||
cu_seqlens_k=packed_indices["cu_seqlens_kv"],
|
||||
max_seqlen_q=packed_indices["max_seqlen_in_batch_kv"],
|
||||
max_seqlen_k=packed_indices["max_seqlen_in_batch_kv"],
|
||||
causal=False,
|
||||
)
|
||||
out_global = result[0] if isinstance(result, tuple) else result
|
||||
out_global = out_global.reshape(B, S * world_size, local_heads * self.head_dim)
|
||||
out = _all_to_all_split_cat(
|
||||
out_global,
|
||||
scatter_dim=1,
|
||||
gather_dim=2,
|
||||
group=group,
|
||||
).view(B, S, self.num_heads, self.head_dim)
|
||||
return self.to_out(out.flatten(2, 3).type_as(x))
|
||||
|
||||
|
||||
class LingBotVideoMLP(nn.Module):
|
||||
def __init__(self, hidden_size, intermediate_size):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
|
||||
self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
|
||||
self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class LingBotVideoBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
num_attention_heads,
|
||||
intermediate_size,
|
||||
norm_eps,
|
||||
qkv_bias,
|
||||
out_bias,
|
||||
num_experts,
|
||||
num_experts_per_tok,
|
||||
moe_intermediate_size,
|
||||
decoder_sparse_step,
|
||||
mlp_only_layers,
|
||||
n_shared_experts,
|
||||
score_func,
|
||||
norm_topk_prob,
|
||||
n_group,
|
||||
topk_group,
|
||||
routed_scaling_factor,
|
||||
layer_idx: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.layer_idx = layer_idx
|
||||
h = hidden_size
|
||||
self.scale_shift_table = nn.Parameter(torch.zeros(1, 6 * h))
|
||||
self.norm1 = LingBotVideoRMSNorm(h, norm_eps)
|
||||
self.attn = LingBotVideoAttention(h, num_attention_heads, norm_eps, qkv_bias, out_bias)
|
||||
self.norm_post_attn = LingBotVideoRMSNorm(h, norm_eps)
|
||||
self.norm2 = LingBotVideoRMSNorm(h, norm_eps)
|
||||
if num_experts > 0:
|
||||
raise NotImplementedError(
|
||||
"LingBotVideoTransformer3DModel in vLLM-Omni currently supports "
|
||||
"the dense checkpoint only. MoE/fused-expert support belongs in "
|
||||
"the follow-up MoE PR."
|
||||
)
|
||||
self.ffn = LingBotVideoMLP(h, intermediate_size)
|
||||
self.norm_post_ffn = LingBotVideoRMSNorm(h, norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
temb6,
|
||||
rotary_emb,
|
||||
attention_mask=None,
|
||||
moe_padding_mask=None,
|
||||
packed_indices: dict[str, torch.Tensor] | None = None,
|
||||
parallel_config=None,
|
||||
):
|
||||
expected_tokens = x.shape[0] * x.shape[1]
|
||||
if temb6.ndim != 2 or temb6.shape[0] != expected_tokens:
|
||||
raise ValueError(
|
||||
"LingBotVideoBlock expects token-level temb6 with shape "
|
||||
f"(B*S, 6D); got {tuple(temb6.shape)} for hidden states {tuple(x.shape)}."
|
||||
)
|
||||
# AdaLN modulation and normalization stay in fp32 for the sensitive path.
|
||||
mod = temb6.view(x.shape[0], x.shape[1], -1) + self.scale_shift_table.unsqueeze(0)
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=-1)
|
||||
gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
|
||||
scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
|
||||
|
||||
# AdaLN modulation / norms run in fp32 (sensitive path); cast to the bulk
|
||||
# compute dtype only at the bf16 Linear boundary. This replaces the old
|
||||
# ambient autocast, which rounded Linear inputs to bf16 at the same point.
|
||||
bulk_dtype = self.attn.to_q.weight.dtype
|
||||
attn_in = (self.norm1(x) * scale_msa + shift_msa).to(bulk_dtype)
|
||||
attn_out = self.attn(
|
||||
attn_in,
|
||||
rotary_emb,
|
||||
attention_mask,
|
||||
packed_indices=packed_indices,
|
||||
parallel_config=parallel_config,
|
||||
)
|
||||
x = x + (gate_msa * self.norm_post_attn(attn_out)).to(x.dtype)
|
||||
|
||||
ffn_in = (self.norm2(x) * scale_mlp + shift_mlp).to(bulk_dtype)
|
||||
ffn_out = self.ffn(ffn_in)
|
||||
ffn_normed = self.norm_post_ffn(ffn_out)
|
||||
x = x + (gate_mlp * ffn_normed).to(x.dtype)
|
||||
return x
|
||||
|
||||
|
||||
class LingBotVideoTransformer3DModel(ModelMixin, ConfigMixin):
|
||||
_supports_gradient_checkpointing = False
|
||||
_repeated_blocks = ["LingBotVideoBlock"]
|
||||
_layerwise_offload_blocks_attr = "blocks"
|
||||
_no_split_modules = ["LingBotVideoBlock"]
|
||||
_keep_in_fp32_modules = list(LINGBOT_VIDEO_FP32_MODULES)
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
|
||||
if dtype is None or dtype == torch.float32:
|
||||
return super().to(*args, **kwargs)
|
||||
|
||||
dtype_is_floating = torch.is_floating_point(torch.empty((), dtype=dtype))
|
||||
if not dtype_is_floating:
|
||||
return super().to(*args, **kwargs)
|
||||
|
||||
if device is not None:
|
||||
super().to(device=device, non_blocking=non_blocking)
|
||||
|
||||
for name, param in self.named_parameters():
|
||||
if not torch.is_floating_point(param):
|
||||
continue
|
||||
target_dtype = torch.float32 if should_keep_in_fp32(name) else dtype
|
||||
param.data = param.data.to(dtype=target_dtype, non_blocking=non_blocking)
|
||||
if param.grad is not None:
|
||||
param.grad.data = param.grad.data.to(dtype=target_dtype, non_blocking=non_blocking)
|
||||
|
||||
for name, buffer in self.named_buffers():
|
||||
if not torch.is_floating_point(buffer):
|
||||
continue
|
||||
target_dtype = torch.float32 if should_keep_in_fp32(name) else dtype
|
||||
buffer.data = buffer.data.to(dtype=target_dtype, non_blocking=non_blocking)
|
||||
|
||||
return self
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
patch_size: tuple[int, int, int] = (1, 2, 2),
|
||||
in_channels: int = 16,
|
||||
out_channels: int = 16,
|
||||
hidden_size: int = 2048,
|
||||
num_attention_heads: int = 16,
|
||||
depth: int = 24,
|
||||
intermediate_size: int = 6144,
|
||||
text_dim: int = 2560,
|
||||
freq_dim: int = 256,
|
||||
norm_eps: float = 1e-6,
|
||||
rope_theta: float = 256.0,
|
||||
axes_dims: tuple[int, int, int] = (32, 48, 48),
|
||||
axes_lens: tuple[int, int, int] = (8192, 1024, 1024),
|
||||
qkv_bias: bool = False,
|
||||
out_bias: bool = True,
|
||||
patch_embed_bias: bool = True,
|
||||
timestep_mlp_bias: bool = True,
|
||||
num_experts: int = 0,
|
||||
num_experts_per_tok: int = 8,
|
||||
moe_intermediate_size: int = 512,
|
||||
decoder_sparse_step: int = 1,
|
||||
mlp_only_layers: tuple[int, ...] = (),
|
||||
n_shared_experts: int | None = None,
|
||||
score_func: str = "sigmoid",
|
||||
norm_topk_prob: bool = True,
|
||||
n_group: int | None = None,
|
||||
topk_group: int | None = None,
|
||||
routed_scaling_factor: float = 1.0,
|
||||
):
|
||||
super().__init__()
|
||||
head_dim = hidden_size // num_attention_heads
|
||||
assert head_dim == sum(axes_dims), f"head_dim {head_dim} != sum(axes_dims) {sum(axes_dims)}"
|
||||
mlp_only_layers = tuple(mlp_only_layers)
|
||||
|
||||
self.patch_embedder = nn.Linear(in_channels * math.prod(patch_size), hidden_size, bias=patch_embed_bias)
|
||||
self.time_proj = Timesteps(freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
|
||||
self.time_embedder = TimestepEmbedding(freq_dim, hidden_size, act_fn="silu", sample_proj_bias=timestep_mlp_bias)
|
||||
self.time_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size))
|
||||
self.text_embedder = LingBotVideoTextEmbedder(text_dim, hidden_size)
|
||||
self.rope = LingBotVideoRotaryEmbedding(axes_dims, axes_lens, rope_theta)
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
LingBotVideoBlock(
|
||||
hidden_size=hidden_size,
|
||||
num_attention_heads=num_attention_heads,
|
||||
intermediate_size=intermediate_size,
|
||||
norm_eps=norm_eps,
|
||||
qkv_bias=qkv_bias,
|
||||
out_bias=out_bias,
|
||||
num_experts=num_experts,
|
||||
num_experts_per_tok=num_experts_per_tok,
|
||||
moe_intermediate_size=moe_intermediate_size,
|
||||
decoder_sparse_step=decoder_sparse_step,
|
||||
mlp_only_layers=mlp_only_layers,
|
||||
n_shared_experts=n_shared_experts,
|
||||
score_func=score_func,
|
||||
norm_topk_prob=norm_topk_prob,
|
||||
n_group=n_group,
|
||||
topk_group=topk_group,
|
||||
routed_scaling_factor=routed_scaling_factor,
|
||||
layer_idx=i,
|
||||
)
|
||||
for i in range(depth)
|
||||
]
|
||||
)
|
||||
self.norm_out = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps)
|
||||
self.norm_out_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size))
|
||||
self.proj_out = nn.Linear(hidden_size, math.prod(patch_size) * out_channels)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor, # (B, C, T, H, W)
|
||||
timestep: torch.Tensor, # (B,) in [0, 1000](= sigma*1000)
|
||||
encoder_hidden_states: torch.Tensor, # (B, L, text_dim)
|
||||
encoder_attention_mask: torch.Tensor | None = None, # (B, L) 1=valid
|
||||
return_dict: bool = True,
|
||||
):
|
||||
B, C, T, H, W = hidden_states.shape
|
||||
pF, pH, pW = self.config.patch_size
|
||||
gt, gh, gw = T // pF, H // pH, W // pW
|
||||
n_video = gt * gh * gw
|
||||
L = encoder_hidden_states.shape[1]
|
||||
device = hidden_states.device
|
||||
if encoder_attention_mask is not None:
|
||||
text_lens = encoder_attention_mask.sum(dim=-1).long()
|
||||
else:
|
||||
text_lens = torch.full((B,), L, dtype=torch.long, device=device)
|
||||
text_lens_list = [int(v) for v in text_lens.detach().cpu().tolist()]
|
||||
packed_batch = B > 1
|
||||
|
||||
# patchify: token order (f h w), feature order (pf ph pw c) -- matches patchify_and_embed
|
||||
patch_tokens = hidden_states.reshape(B, C, gt, pF, gh, pH, gw, pW)
|
||||
patch_tokens = patch_tokens.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(
|
||||
B,
|
||||
n_video,
|
||||
pF * pH * pW * C,
|
||||
)
|
||||
if packed_batch:
|
||||
x = torch.cat(
|
||||
[self.patch_embedder(patch_tokens[i : i + 1]) for i in range(B)],
|
||||
dim=1,
|
||||
)
|
||||
else:
|
||||
x = self.patch_embedder(patch_tokens)
|
||||
|
||||
if packed_batch:
|
||||
text_parts = [
|
||||
self.text_embedder(encoder_hidden_states[i : i + 1, : text_lens_list[i], :]) for i in range(B)
|
||||
]
|
||||
text = torch.cat(text_parts, dim=1)
|
||||
joint = _cat_interleave(
|
||||
x,
|
||||
[n_video] * B,
|
||||
text,
|
||||
text_lens_list,
|
||||
)
|
||||
else:
|
||||
text = self.text_embedder(encoder_hidden_states)
|
||||
joint = torch.cat([x, text], dim=1) # [video; text]
|
||||
joint_seq_len = joint.shape[1]
|
||||
|
||||
# Per-sample RoPE: video t-axis start = real text length of this sample + 1
|
||||
rotary_parts = [self.rope(make_joint_position_ids(text_lens_list[i], gt, gh, gw, device)) for i in range(B)]
|
||||
if packed_batch:
|
||||
rotary = torch.cat(rotary_parts, dim=0).unsqueeze(0)
|
||||
else:
|
||||
rotary = torch.stack(rotary_parts, dim=0) # (B, S, head_dim/2) complex64
|
||||
|
||||
parallel_config = getattr(self, "_parallel_config", None)
|
||||
use_packed_attention = parallel_config is not None
|
||||
|
||||
attention_mask = None
|
||||
moe_padding_mask = None
|
||||
packed_indices = None
|
||||
has_padding = encoder_attention_mask is not None and bool((text_lens < L).any())
|
||||
if packed_batch or use_packed_attention:
|
||||
sample_seq_lens = [n_video + text_len for text_len in text_lens_list]
|
||||
cu_seqlens = torch.zeros(B + 1, device=device, dtype=torch.int32)
|
||||
cu_seqlens[1:] = torch.cumsum(
|
||||
torch.tensor(sample_seq_lens, device=device, dtype=torch.int32),
|
||||
dim=0,
|
||||
)
|
||||
packed_indices = {
|
||||
"cu_seqlens_kv": cu_seqlens,
|
||||
"max_seqlen_in_batch_kv": max(sample_seq_lens),
|
||||
}
|
||||
if packed_batch and not use_packed_attention:
|
||||
packed_indices["attention_mask"] = _packed_block_attention_mask(sample_seq_lens, device)
|
||||
has_padding = False
|
||||
if has_padding:
|
||||
key_mask = torch.cat(
|
||||
[torch.ones(B, n_video, dtype=torch.bool, device=device), encoder_attention_mask.bool()],
|
||||
dim=1,
|
||||
)
|
||||
attention_mask = key_mask[:, None, None, :] # (B,1,1,S) -> SDPA broadcast
|
||||
moe_padding_mask = key_mask.reshape(-1).float() # (B*S,)
|
||||
packed_cp = packed_indices is not None and parallel_config is not None
|
||||
padding_size = 0
|
||||
if packed_cp:
|
||||
cp_config = parallel_config.context_parallel_config
|
||||
cp_world_size = int(getattr(cp_config, "ulysses_degree", getattr(cp_config, "_world_size", 1)))
|
||||
padding_size = (cp_world_size - (joint_seq_len % cp_world_size)) % cp_world_size
|
||||
if padding_size:
|
||||
joint = torch.cat(
|
||||
[
|
||||
joint,
|
||||
torch.zeros(
|
||||
joint.shape[0],
|
||||
padding_size,
|
||||
joint.shape[2],
|
||||
device=joint.device,
|
||||
dtype=joint.dtype,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
rotary = torch.cat(
|
||||
[
|
||||
rotary,
|
||||
torch.zeros(
|
||||
rotary.shape[0],
|
||||
padding_size,
|
||||
rotary.shape[2],
|
||||
device=rotary.device,
|
||||
dtype=rotary.dtype,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
if packed_indices is None:
|
||||
raise RuntimeError("packed_indices must be initialized for packed context parallel.")
|
||||
packed_indices["cu_seqlens_kv"] = torch.cat(
|
||||
[
|
||||
packed_indices["cu_seqlens_kv"],
|
||||
packed_indices["cu_seqlens_kv"][-1:] + padding_size,
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
packed_indices["max_seqlen_in_batch_kv"] = max(
|
||||
int(packed_indices["max_seqlen_in_batch_kv"]),
|
||||
int(padding_size),
|
||||
)
|
||||
joint_seq_len = joint.shape[1]
|
||||
|
||||
timestep_for_embed = timestep.float()
|
||||
timestep_proj = self.time_proj(timestep_for_embed)
|
||||
t_emb = self.time_embedder(timestep_proj) # (B, D)
|
||||
if packed_batch:
|
||||
temb_input = torch.cat(
|
||||
[t_emb[i : i + 1].unsqueeze(1).expand(1, n_video + text_lens_list[i], -1) for i in range(B)],
|
||||
dim=1,
|
||||
)
|
||||
if padding_size:
|
||||
temb_input = torch.cat(
|
||||
[
|
||||
temb_input,
|
||||
torch.zeros(
|
||||
temb_input.shape[0],
|
||||
padding_size,
|
||||
temb_input.shape[2],
|
||||
device=temb_input.device,
|
||||
dtype=temb_input.dtype,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
temb6 = self.time_modulation(temb_input.reshape(joint_seq_len, -1))
|
||||
temb6 = temb6.reshape(1, joint_seq_len, -1)
|
||||
else:
|
||||
temb_input = t_emb.unsqueeze(1).expand(B, joint_seq_len, -1) # (B, S, D)
|
||||
temb6 = self.time_modulation(temb_input.reshape(B * joint_seq_len, -1))
|
||||
temb6 = temb6.reshape(B, joint_seq_len, -1) # (B, S, 6D)
|
||||
|
||||
temb6 = temb6.reshape(temb6.shape[0] * temb6.shape[1], -1)
|
||||
|
||||
for block in self.blocks:
|
||||
joint = block(
|
||||
joint,
|
||||
temb6,
|
||||
rotary,
|
||||
attention_mask,
|
||||
moe_padding_mask,
|
||||
packed_indices=packed_indices,
|
||||
parallel_config=parallel_config,
|
||||
)
|
||||
final_mod = self.norm_out_modulation(temb_input.reshape(joint.shape[0] * joint.shape[1], -1))
|
||||
shift, scale = final_mod.reshape(joint.shape[0], joint.shape[1], -1).chunk(2, dim=-1)
|
||||
final_hidden = self.norm_out(joint) * (1.0 + scale) + shift
|
||||
projected = self.proj_out(final_hidden.to(self.proj_out.weight.dtype))
|
||||
if packed_cp:
|
||||
if padding_size:
|
||||
projected = projected[:, :-padding_size, :]
|
||||
if packed_batch:
|
||||
split_lengths: list[int] = []
|
||||
for text_len in text_lens_list:
|
||||
split_lengths.extend([n_video, text_len])
|
||||
parts = torch.split(projected, split_lengths, dim=1)
|
||||
x = torch.cat(parts[::2], dim=1).reshape(B, n_video, -1)
|
||||
else:
|
||||
x = projected[:, :n_video]
|
||||
|
||||
# unpatchify (matches the rearrange in postprocess)
|
||||
Cout = self.config.out_channels
|
||||
x = x.reshape(B, gt, gh, gw, pF, pH, pW, Cout)
|
||||
x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(B, Cout, T, H, W)
|
||||
|
||||
if not return_dict:
|
||||
return (x,)
|
||||
return Transformer2DModelOutput(sample=x)
|
||||
@@ -1,695 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# Adapted from LingBot-Video (https://github.com/Robbyant/lingbot-video).
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Iterable, Mapping
|
||||
from contextlib import nullcontext
|
||||
from typing import Any, ClassVar
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from diffusers import AutoencoderKLWan
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from torch import nn
|
||||
from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
|
||||
from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import LingBotVideoTransformer3DModel
|
||||
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
||||
from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
|
||||
TOKEN_LENGTH = 37698
|
||||
HIDDEN_STATE_SKIP_LAYER = 0
|
||||
LOW_NOISE_TAIL_V1_DEFAULT_STEPS = 2
|
||||
|
||||
PROMPT_TEMPLATE = (
|
||||
"<|im_start|>system\nGiven a user input that may include a text prompt alone, "
|
||||
"a text prompt with an image reference, or a text prompt with a video reference "
|
||||
'or a video reference alone, generate an "Enhanced prompt" that provides detailed '
|
||||
"visual descriptions suitable for video generation. Evaluate the level of detail "
|
||||
"in the user's input: if it is simple, enrich it by adding specifics about colors, "
|
||||
"shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal "
|
||||
"progression, and spatial relationships to create vivid, concrete, and temporally "
|
||||
"coherent scenes to create vivid and concrete scenes. Please generate only the "
|
||||
"enhanced description for the prompt below and avoid including any additional "
|
||||
"commentary or evaluations:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
DEFAULT_NEGATIVE_PROMPT = (
|
||||
'{"universal_negative": {"visual_quality": ["low quality", "worst quality", "blurry", '
|
||||
'"pixelated", "jpeg artifacts", "low resolution", "unstable color", "color flicker", '
|
||||
'"underexposed", "overexposed", "invisible subject", "subject hidden in darkness"], '
|
||||
'"artistic_style": ["painting", "illustration", "drawing", "cartoon", "3d render", '
|
||||
'"cgi", "sketch", "digital art"], "composition_and_content": ["text", "watermark", '
|
||||
'"signature", "logo", "subtitles", "pillarboxed", "side bars", '
|
||||
'"portrait image in landscape frame"], "temporal_and_motion_stability": ["flickering", '
|
||||
'"jittery", "motion blur", "temporal inconsistency", "warping", "morphing", '
|
||||
'"incoherent motion", "unnatural movement", "static object with sudden jump", '
|
||||
'"frame-to-frame inconsistency"], "material_and_structure": ["plastic-like glass", '
|
||||
'"unrealistic texture", "deformed bottle", "liquid freezing improperly", '
|
||||
'"distorted reflections"]}}'
|
||||
)
|
||||
|
||||
|
||||
def _dtype_from_name(value: Any, default: torch.dtype) -> torch.dtype:
|
||||
if isinstance(value, torch.dtype):
|
||||
return value
|
||||
if value is None:
|
||||
return default
|
||||
normalized = str(value).lower()
|
||||
if normalized in {"bf16", "bfloat16", "torch.bfloat16"}:
|
||||
return torch.bfloat16
|
||||
if normalized in {"fp16", "float16", "torch.float16"}:
|
||||
return torch.float16
|
||||
if normalized in {"fp32", "float32", "torch.float32"}:
|
||||
return torch.float32
|
||||
raise ValueError(f"Unsupported LingBot dtype: {value!r}.")
|
||||
|
||||
|
||||
def _extract_prompt(req: OmniDiffusionRequest) -> tuple[str, str | None]:
|
||||
prompt_obj = req.prompt
|
||||
if isinstance(prompt_obj, str):
|
||||
return prompt_obj, None
|
||||
if isinstance(prompt_obj, Mapping):
|
||||
prompt = prompt_obj.get("prompt", "")
|
||||
negative_prompt = prompt_obj.get("negative_prompt")
|
||||
return str(prompt), None if negative_prompt is None else str(negative_prompt)
|
||||
raise TypeError(f"Unsupported LingBot prompt type: {type(prompt_obj)!r}.")
|
||||
|
||||
|
||||
def _module_dtype(module: torch.nn.Module) -> torch.dtype:
|
||||
try:
|
||||
return next(module.parameters()).dtype
|
||||
except StopIteration:
|
||||
return torch.float32
|
||||
|
||||
|
||||
def _module_device(module: torch.nn.Module) -> torch.device:
|
||||
try:
|
||||
return next(module.parameters()).device
|
||||
except StopIteration:
|
||||
return torch.device("cpu")
|
||||
|
||||
|
||||
def _transformer_timestep(timestep: torch.Tensor, transformer_dtype: torch.dtype) -> torch.Tensor:
|
||||
sigma = timestep.float() / 1000.0
|
||||
if transformer_dtype in {torch.bfloat16, torch.float16}:
|
||||
sigma = sigma.to(transformer_dtype)
|
||||
return (sigma * 1000.0).float()
|
||||
|
||||
|
||||
def _transformer_autocast(device: torch.device, transformer_dtype: torch.dtype):
|
||||
if device.type != "cuda" or transformer_dtype not in {torch.bfloat16, torch.float16}:
|
||||
return nullcontext()
|
||||
return torch.autocast(device_type="cuda", dtype=transformer_dtype)
|
||||
|
||||
|
||||
def _group_global_rank(group: Any | None, group_rank: int) -> int:
|
||||
if group is None:
|
||||
return group_rank
|
||||
get_global_rank = getattr(dist, "get_global_rank", None)
|
||||
if get_global_rank is None:
|
||||
return group_rank
|
||||
return int(get_global_rank(group, group_rank))
|
||||
|
||||
|
||||
def _validate_refiner_sigmas(sigmas: np.ndarray, t_thresh: float | None = None) -> np.ndarray:
|
||||
arr = np.asarray(list(sigmas), dtype=np.float64)
|
||||
if arr.ndim != 1 or arr.size == 0:
|
||||
raise ValueError("refiner sigma schedule must be a non-empty 1D list")
|
||||
if not np.all(np.isfinite(arr)):
|
||||
raise ValueError("refiner sigma schedule contains non-finite values")
|
||||
if np.any(arr < 0.0) or np.any(arr > 1.0):
|
||||
raise ValueError(f"refiner sigma schedule values must be in [0, 1], got {arr.tolist()}")
|
||||
if arr.size > 1 and not np.all(np.diff(arr) < 0.0):
|
||||
raise ValueError(f"refiner sigma schedule must be strictly descending, got {arr.tolist()}")
|
||||
if t_thresh is not None and abs(float(arr[0]) - float(t_thresh)) > 1e-6:
|
||||
raise ValueError(f"refiner sigma schedule must start at t_thresh={float(t_thresh)}, got {float(arr[0])}")
|
||||
return arr
|
||||
|
||||
|
||||
def _compute_refiner_sigmas(
|
||||
*,
|
||||
sigma_max: float,
|
||||
sigma_min: float,
|
||||
num_inference_steps: int,
|
||||
shift: float,
|
||||
t_thresh: float | None,
|
||||
tail_steps: int = 0,
|
||||
) -> np.ndarray | None:
|
||||
if t_thresh is None:
|
||||
return None
|
||||
t_value = float(t_thresh)
|
||||
if not (0.0 < t_value <= 1.0):
|
||||
raise ValueError(f"refiner t_thresh must lie in (0, 1], got {t_value}")
|
||||
steps = int(num_inference_steps)
|
||||
if steps < 1:
|
||||
raise ValueError(f"num_inference_steps must be >= 1, got {steps}")
|
||||
tail = int(tail_steps or 0)
|
||||
if tail < 0:
|
||||
raise ValueError(f"refiner_sigma_tail_steps must be >= 0, got {tail}")
|
||||
|
||||
base = np.linspace(float(sigma_max), float(sigma_min), steps + 1).copy()[:-1]
|
||||
shift_value = float(shift)
|
||||
shifted = shift_value * base / (1.0 + (shift_value - 1.0) * base)
|
||||
eps = 1e-6
|
||||
sigmas = shifted[shifted <= t_value + eps]
|
||||
if sigmas.size == 0 or abs(float(sigmas[0]) - t_value) > eps:
|
||||
sigmas = np.concatenate([[t_value], sigmas])
|
||||
if tail > 0:
|
||||
start = float(sigmas[-1])
|
||||
stop = min(float(sigma_min), start)
|
||||
extra = np.linspace(start, stop, tail + 2, dtype=np.float64)[1:-1]
|
||||
sigmas = np.concatenate([sigmas, extra])
|
||||
return _validate_refiner_sigmas(sigmas, t_value).astype(np.float32)
|
||||
|
||||
|
||||
def _pad_prompt_embeds(
|
||||
embeds: torch.Tensor,
|
||||
mask: torch.Tensor,
|
||||
target_length: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if embeds.shape[0] != 1:
|
||||
raise ValueError(f"batched CFG helper expects batch=1 inputs, got {embeds.shape[0]}")
|
||||
if embeds.shape[1] > target_length:
|
||||
raise ValueError(f"cannot pad length {embeds.shape[1]} down to {target_length}")
|
||||
pad_len = target_length - embeds.shape[1]
|
||||
if pad_len == 0:
|
||||
return embeds, mask
|
||||
embed_pad = torch.zeros(embeds.shape[0], pad_len, embeds.shape[2], dtype=embeds.dtype, device=embeds.device)
|
||||
mask_pad = torch.zeros(mask.shape[0], pad_len, dtype=mask.dtype, device=mask.device)
|
||||
return torch.cat([embeds, embed_pad], dim=1), torch.cat([mask, mask_pad], dim=1)
|
||||
|
||||
|
||||
def _batch_cfg_prompt_inputs(
|
||||
prompt_embeds: torch.Tensor,
|
||||
prompt_mask: torch.Tensor,
|
||||
negative_embeds: torch.Tensor,
|
||||
negative_mask: torch.Tensor,
|
||||
*,
|
||||
null_cond_clone_zero: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if null_cond_clone_zero:
|
||||
zero_negative = torch.zeros_like(prompt_embeds)
|
||||
return (
|
||||
torch.cat([prompt_embeds, zero_negative], dim=0),
|
||||
torch.cat([prompt_mask, prompt_mask.clone()], dim=0),
|
||||
)
|
||||
|
||||
target_length = max(int(prompt_embeds.shape[1]), int(negative_embeds.shape[1]))
|
||||
prompt_padded, prompt_mask_padded = _pad_prompt_embeds(prompt_embeds, prompt_mask, target_length)
|
||||
negative_padded, negative_mask_padded = _pad_prompt_embeds(negative_embeds, negative_mask, target_length)
|
||||
return (
|
||||
torch.cat([prompt_padded, negative_padded], dim=0),
|
||||
torch.cat([prompt_mask_padded, negative_mask_padded], dim=0),
|
||||
)
|
||||
|
||||
|
||||
def get_lingbot_video_post_process_func(od_config: OmniDiffusionConfig):
|
||||
del od_config
|
||||
|
||||
def post_process_func(frames: torch.Tensor, sampling_params=None):
|
||||
output_type = getattr(sampling_params, "output_type", None) or "pt"
|
||||
if output_type == "np" and isinstance(frames, torch.Tensor):
|
||||
return frames.float().cpu().numpy()
|
||||
return frames
|
||||
|
||||
return post_process_func
|
||||
|
||||
|
||||
class LingBotVideoPipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery):
|
||||
"""Native vLLM-Omni entry for the dense LingBot-Video checkpoint.
|
||||
|
||||
This PR supports the dense checkpoint only. MoE and fused-expert kernels are
|
||||
intentionally left for the follow-up MoE PR.
|
||||
"""
|
||||
|
||||
supports_step_execution: ClassVar[bool] = False
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_encoder"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae"]
|
||||
|
||||
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
del prefix
|
||||
self.od_config = od_config
|
||||
self.device = get_local_device()
|
||||
self.vae_scale_factor_temporal = 4
|
||||
self.vae_scale_factor_spatial = 8
|
||||
self.token_length = TOKEN_LENGTH
|
||||
self.hidden_state_skip_layer = HIDDEN_STATE_SKIP_LAYER
|
||||
self.prompt_template = PROMPT_TEMPLATE
|
||||
self._crop_start: int | None = None
|
||||
|
||||
model = od_config.model
|
||||
local_files_only = os.path.exists(model)
|
||||
dtype = getattr(od_config, "dtype", torch.bfloat16)
|
||||
model_config = getattr(od_config, "model_config", None) or {}
|
||||
transformer_dtype = _dtype_from_name(model_config.get("transformer_dtype"), dtype)
|
||||
text_encoder_dtype = _dtype_from_name(model_config.get("text_encoder_dtype"), dtype)
|
||||
vae_dtype = _dtype_from_name(model_config.get("vae_dtype"), torch.float32)
|
||||
|
||||
transformer_subfolder = str(model_config.get("transformer_subfolder", "transformer"))
|
||||
text_encoder_subfolder = str(model_config.get("text_encoder_subfolder", "text_encoder"))
|
||||
processor_subfolder = str(model_config.get("processor_subfolder", "processor"))
|
||||
vae_subfolder = str(model_config.get("vae_subfolder", "vae"))
|
||||
scheduler_subfolder = str(model_config.get("scheduler_subfolder", "scheduler"))
|
||||
|
||||
self.transformer = LingBotVideoTransformer3DModel.from_pretrained(
|
||||
model,
|
||||
subfolder=transformer_subfolder,
|
||||
torch_dtype=transformer_dtype,
|
||||
local_files_only=local_files_only,
|
||||
).to(self.device)
|
||||
text_encoder_kwargs: dict[str, Any] = {
|
||||
"dtype": text_encoder_dtype,
|
||||
"local_files_only": local_files_only,
|
||||
}
|
||||
self.text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
|
||||
model,
|
||||
subfolder=text_encoder_subfolder,
|
||||
**text_encoder_kwargs,
|
||||
).to(self.device)
|
||||
self.processor = Qwen3VLProcessor.from_pretrained(
|
||||
model,
|
||||
subfolder=processor_subfolder,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
self.vae = AutoencoderKLWan.from_pretrained(
|
||||
model,
|
||||
subfolder=vae_subfolder,
|
||||
torch_dtype=vae_dtype,
|
||||
local_files_only=local_files_only,
|
||||
).to(self.device)
|
||||
self.scheduler = FlowUniPCMultistepScheduler.from_pretrained(
|
||||
model,
|
||||
subfolder=scheduler_subfolder,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
self.set_progress_bar_config(disable=bool(model_config.get("quiet_progress", True)))
|
||||
self.default_negative_prompt = DEFAULT_NEGATIVE_PROMPT
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
|
||||
super().to(*args, **kwargs)
|
||||
if device is not None:
|
||||
self.device = torch.device(device)
|
||||
self.transformer.to(device=self.device, non_blocking=non_blocking)
|
||||
self.text_encoder.to(device=self.device, non_blocking=non_blocking)
|
||||
self.vae.to(device=self.device, non_blocking=non_blocking)
|
||||
if dtype is not None:
|
||||
self.transformer.to(dtype=dtype, non_blocking=non_blocking)
|
||||
return self
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
consumed = list(weights)
|
||||
if consumed:
|
||||
raise RuntimeError(
|
||||
f"{self.__class__.__name__}.load_weights received {len(consumed)} weight tensors; "
|
||||
"LingBot components are loaded directly from subfolders during __init__."
|
||||
)
|
||||
return {name for name, _ in self.named_parameters()}
|
||||
|
||||
@staticmethod
|
||||
def check_inputs(height: int, width: int, num_frames: int) -> None:
|
||||
if num_frames != 1 and (num_frames - 1) % 4 != 0:
|
||||
raise ValueError(f"`num_frames` must be 1 or 4n+1, got {num_frames}.")
|
||||
if height % 16 != 0 or width % 16 != 0:
|
||||
raise ValueError(f"`height` and `width` must be multiples of 16, got {height}x{width}.")
|
||||
|
||||
@staticmethod
|
||||
def apply_text_to_template(text: str, template: str = PROMPT_TEMPLATE) -> str:
|
||||
return template.format(text)
|
||||
|
||||
def _compute_crop_start(self) -> int:
|
||||
if self._crop_start is None:
|
||||
marker = "<|USER_INPUT_MARKER|>"
|
||||
marked = self.prompt_template.format(marker)
|
||||
marker_pos = marked.find(marker)
|
||||
if marker_pos < 0:
|
||||
self._crop_start = 0
|
||||
else:
|
||||
prefix = self.processor(
|
||||
text=marked[:marker_pos],
|
||||
images=None,
|
||||
videos=None,
|
||||
return_tensors="pt",
|
||||
)
|
||||
self._crop_start = int(prefix["input_ids"].shape[1])
|
||||
return self._crop_start
|
||||
|
||||
def _build_prompt_inputs(self, prompt: str | list[str]):
|
||||
prompts = [prompt] if isinstance(prompt, str) else list(prompt)
|
||||
texts = [self.apply_text_to_template(text, self.prompt_template) for text in prompts]
|
||||
return self.processor(
|
||||
text=texts,
|
||||
images=None,
|
||||
videos=None,
|
||||
do_resize=False,
|
||||
truncation=True,
|
||||
max_length=self.token_length,
|
||||
padding="longest",
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def encode_prompt(
|
||||
self,
|
||||
prompt: str | list[str],
|
||||
*,
|
||||
device: str | torch.device | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
device = torch.device(device) if device is not None else self.device
|
||||
inputs = self._build_prompt_inputs(prompt).to(device)
|
||||
outputs = self.text_encoder(
|
||||
**inputs,
|
||||
output_hidden_states=self.hidden_state_skip_layer is not None,
|
||||
)
|
||||
if self.hidden_state_skip_layer is not None:
|
||||
prompt_embeds = outputs.hidden_states[-(self.hidden_state_skip_layer + 1)]
|
||||
else:
|
||||
prompt_embeds = outputs.last_hidden_state
|
||||
|
||||
prompt_mask = inputs["attention_mask"]
|
||||
crop_start = self._compute_crop_start()
|
||||
if crop_start > 0:
|
||||
prompt_embeds = prompt_embeds[:, crop_start:]
|
||||
prompt_mask = prompt_mask[:, crop_start:]
|
||||
|
||||
if prompt_embeds.shape[0] == 1:
|
||||
true_len = int(prompt_mask[0].sum().item())
|
||||
prompt_embeds = prompt_embeds[:, :true_len]
|
||||
prompt_mask = prompt_mask[:, :true_len]
|
||||
return prompt_embeds, prompt_mask
|
||||
|
||||
def prepare_latents(
|
||||
self,
|
||||
num_frames: int,
|
||||
height: int,
|
||||
width: int,
|
||||
generator: torch.Generator | None,
|
||||
latents: torch.Tensor | None,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
|
||||
latent_height = height // self.vae_scale_factor_spatial
|
||||
latent_width = width // self.vae_scale_factor_spatial
|
||||
shape = (1, self.transformer.config.in_channels, latent_frames, latent_height, latent_width)
|
||||
if latents is None:
|
||||
return randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
|
||||
if tuple(latents.shape) != shape:
|
||||
raise ValueError(f"`latents` shape must be {shape}, got {tuple(latents.shape)}.")
|
||||
return latents.to(device=device, dtype=torch.float32)
|
||||
|
||||
def _dit_latent_to_vae(self, latents: torch.Tensor) -> torch.Tensor:
|
||||
mean = torch.tensor(self.vae.config.latents_mean, device=latents.device, dtype=torch.float32)
|
||||
std_inv = 1.0 / torch.tensor(self.vae.config.latents_std, device=latents.device, dtype=torch.float32)
|
||||
mean = mean.view(1, -1, 1, 1, 1)
|
||||
std_inv = std_inv.view(1, -1, 1, 1, 1)
|
||||
return latents.float() / std_inv + mean
|
||||
|
||||
@torch.no_grad()
|
||||
def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
|
||||
vae_device = _module_device(self.vae)
|
||||
vae_dtype = _module_dtype(self.vae)
|
||||
vae_latents = self._dit_latent_to_vae(latents).to(device=vae_device, dtype=torch.float32)
|
||||
if vae_latents.ndim == 5:
|
||||
vae_latents = vae_latents.contiguous(memory_format=torch.channels_last_3d)
|
||||
autocast_dtype = (
|
||||
vae_dtype if vae_device.type == "cuda" and vae_dtype in {torch.bfloat16, torch.float16} else None
|
||||
)
|
||||
with torch.autocast("cuda", dtype=autocast_dtype or torch.bfloat16, enabled=autocast_dtype is not None):
|
||||
decoded = self.vae.decode(vae_latents)
|
||||
frames = decoded[0] if isinstance(decoded, tuple) else decoded.sample
|
||||
frames = frames.float().clamp_(-1, 1)
|
||||
frames = (frames + 1.0) / 2.0
|
||||
frames = frames.permute(0, 2, 3, 4, 1).cpu()
|
||||
return frames[0]
|
||||
|
||||
@torch.no_grad()
|
||||
def _generate(
|
||||
self,
|
||||
*,
|
||||
prompt: str,
|
||||
negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
|
||||
height: int = 480,
|
||||
width: int = 480,
|
||||
num_frames: int = 81,
|
||||
num_inference_steps: int = 40,
|
||||
guidance_scale: float = 6.0,
|
||||
shift: float = 3.0,
|
||||
generator: torch.Generator | None = None,
|
||||
latents: torch.Tensor | None = None,
|
||||
prompt_embeds: torch.Tensor | None = None,
|
||||
prompt_mask: torch.Tensor | None = None,
|
||||
negative_prompt_embeds: torch.Tensor | None = None,
|
||||
negative_prompt_mask: torch.Tensor | None = None,
|
||||
output_type: str = "pt",
|
||||
cfg_parallel_group: Any | None = None,
|
||||
batch_cfg: bool = False,
|
||||
null_cond_clone_zero: bool = False,
|
||||
t_thresh: float | None = None,
|
||||
refiner_sigma_tail_steps: int = LOW_NOISE_TAIL_V1_DEFAULT_STEPS,
|
||||
offload_vae_during_denoise: bool = False,
|
||||
**extra_args,
|
||||
) -> torch.Tensor:
|
||||
del extra_args
|
||||
self.check_inputs(height, width, num_frames)
|
||||
device = self.device
|
||||
do_cfg = guidance_scale > 1.0
|
||||
effective_batch_cfg = bool(batch_cfg)
|
||||
cfg_parallel = cfg_parallel_group is not None
|
||||
cfg_parallel_rank = 0
|
||||
if cfg_parallel:
|
||||
if not dist.is_available() or not dist.is_initialized():
|
||||
raise ValueError("`cfg_parallel_group` requires an initialized process group.")
|
||||
if effective_batch_cfg:
|
||||
raise ValueError("`cfg_parallel_group` and `batch_cfg` are mutually exclusive.")
|
||||
if not do_cfg:
|
||||
raise ValueError("CFG parallel requires `guidance_scale > 1.0`.")
|
||||
cfg_parallel_rank = dist.get_rank(cfg_parallel_group)
|
||||
cfg_parallel_world_size = dist.get_world_size(cfg_parallel_group)
|
||||
if cfg_parallel_world_size != 2:
|
||||
raise ValueError(f"CFG parallel currently requires exactly 2 ranks, got {cfg_parallel_world_size}.")
|
||||
|
||||
if prompt_embeds is not None:
|
||||
if prompt_mask is None:
|
||||
raise ValueError("`prompt_mask` is required when `prompt_embeds` is provided.")
|
||||
prompt_embeds = prompt_embeds.to(device=device)
|
||||
prompt_mask = prompt_mask.to(device=device)
|
||||
if negative_prompt_embeds is not None:
|
||||
if negative_prompt_mask is None:
|
||||
raise ValueError("`negative_prompt_mask` is required when `negative_prompt_embeds` is provided.")
|
||||
negative_prompt_embeds = negative_prompt_embeds.to(device=device)
|
||||
negative_prompt_mask = negative_prompt_mask.to(device=device)
|
||||
|
||||
negative_embeds = None
|
||||
negative_mask = None
|
||||
if cfg_parallel and cfg_parallel_rank == 1:
|
||||
if negative_prompt_embeds is not None:
|
||||
negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask
|
||||
else:
|
||||
negative_embeds, negative_mask = self.encode_prompt(negative_prompt, device=device)
|
||||
prompt_embeds = prompt_mask = None
|
||||
else:
|
||||
if prompt_embeds is None:
|
||||
prompt_embeds, prompt_mask = self.encode_prompt(prompt, device=device)
|
||||
if do_cfg and not cfg_parallel:
|
||||
if null_cond_clone_zero:
|
||||
negative_embeds = torch.zeros_like(prompt_embeds)
|
||||
negative_mask = prompt_mask.clone()
|
||||
elif negative_prompt_embeds is not None:
|
||||
negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask
|
||||
else:
|
||||
negative_embeds, negative_mask = self.encode_prompt(negative_prompt, device=device)
|
||||
|
||||
latents = self.prepare_latents(num_frames, height, width, generator, latents, device)
|
||||
sigmas = _compute_refiner_sigmas(
|
||||
sigma_max=float(self.scheduler.sigma_max),
|
||||
sigma_min=float(self.scheduler.sigma_min),
|
||||
num_inference_steps=num_inference_steps,
|
||||
shift=shift,
|
||||
t_thresh=t_thresh,
|
||||
tail_steps=refiner_sigma_tail_steps,
|
||||
)
|
||||
if sigmas is None:
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
|
||||
else:
|
||||
self.scheduler.set_timesteps(int(sigmas.shape[0]), device=device, sigmas=sigmas, shift=1.0)
|
||||
|
||||
transformer_dtype = _module_dtype(self.transformer)
|
||||
vae_restore_device: torch.device | None = None
|
||||
vae_offloaded = False
|
||||
if offload_vae_during_denoise and output_type != "latent":
|
||||
vae_device = _module_device(self.vae)
|
||||
if vae_device.type == "cuda":
|
||||
self.vae.to("cpu")
|
||||
torch.accelerator.empty_cache()
|
||||
vae_restore_device = vae_device
|
||||
vae_offloaded = True
|
||||
|
||||
cfg_latent_src = _group_global_rank(cfg_parallel_group, 0)
|
||||
cfg_uncond_src = _group_global_rank(cfg_parallel_group, 1)
|
||||
for timestep in self.progress_bar(self.scheduler.timesteps):
|
||||
if cfg_parallel:
|
||||
dist.broadcast(latents, src=cfg_latent_src, group=cfg_parallel_group)
|
||||
timestep_batch = _transformer_timestep(timestep, transformer_dtype).expand(1).to(device)
|
||||
latent_model_input = latents
|
||||
if cfg_parallel:
|
||||
if cfg_parallel_rank == 0:
|
||||
branch_embeds = prompt_embeds
|
||||
branch_mask = prompt_mask
|
||||
else:
|
||||
branch_embeds = negative_embeds
|
||||
branch_mask = negative_mask
|
||||
if branch_embeds is None:
|
||||
raise RuntimeError("CFG branch embeddings were not initialized.")
|
||||
with _transformer_autocast(device, transformer_dtype):
|
||||
branch_noise_pred = self.transformer(
|
||||
latent_model_input,
|
||||
timestep_batch,
|
||||
branch_embeds.to(transformer_dtype),
|
||||
encoder_attention_mask=branch_mask,
|
||||
return_dict=False,
|
||||
)[0].float()
|
||||
if cfg_parallel_rank == 0:
|
||||
noise_pred = branch_noise_pred
|
||||
noise_pred_uncond = torch.empty_like(noise_pred)
|
||||
else:
|
||||
noise_pred_uncond = branch_noise_pred
|
||||
dist.broadcast(noise_pred_uncond, src=cfg_uncond_src, group=cfg_parallel_group)
|
||||
if cfg_parallel_rank != 0:
|
||||
continue
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
|
||||
else:
|
||||
if prompt_embeds is None:
|
||||
raise RuntimeError("Prompt embeddings were not initialized.")
|
||||
prompt_model_input = prompt_embeds.to(transformer_dtype)
|
||||
if do_cfg and effective_batch_cfg:
|
||||
if negative_embeds is None or negative_mask is None:
|
||||
raise RuntimeError("Negative embeddings were not initialized for CFG.")
|
||||
cfg_embeds, cfg_mask = _batch_cfg_prompt_inputs(
|
||||
prompt_model_input,
|
||||
prompt_mask,
|
||||
negative_embeds.to(transformer_dtype),
|
||||
negative_mask,
|
||||
null_cond_clone_zero=False,
|
||||
)
|
||||
cfg_latents = torch.cat([latent_model_input, latent_model_input], dim=0)
|
||||
cfg_timesteps = torch.cat([timestep_batch, timestep_batch], dim=0)
|
||||
with _transformer_autocast(device, transformer_dtype):
|
||||
noise_batched = self.transformer(
|
||||
cfg_latents,
|
||||
cfg_timesteps,
|
||||
cfg_embeds,
|
||||
encoder_attention_mask=cfg_mask,
|
||||
return_dict=False,
|
||||
)[0].float()
|
||||
noise_pred, noise_pred_uncond = noise_batched.chunk(2, dim=0)
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
|
||||
else:
|
||||
with _transformer_autocast(device, transformer_dtype):
|
||||
noise_pred = self.transformer(
|
||||
latent_model_input,
|
||||
timestep_batch,
|
||||
prompt_model_input,
|
||||
encoder_attention_mask=prompt_mask,
|
||||
return_dict=False,
|
||||
)[0].float()
|
||||
|
||||
if do_cfg and not effective_batch_cfg:
|
||||
if negative_embeds is None or negative_mask is None:
|
||||
raise RuntimeError("Negative embeddings were not initialized for CFG.")
|
||||
with _transformer_autocast(device, transformer_dtype):
|
||||
noise_pred_uncond = self.transformer(
|
||||
latent_model_input,
|
||||
timestep_batch,
|
||||
negative_embeds.to(transformer_dtype),
|
||||
encoder_attention_mask=negative_mask,
|
||||
return_dict=False,
|
||||
)[0].float()
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
|
||||
|
||||
latents = self.scheduler.step(noise_pred, timestep, latents, return_dict=False, generator=generator)[0]
|
||||
|
||||
if cfg_parallel:
|
||||
dist.barrier(group=cfg_parallel_group)
|
||||
if cfg_parallel_rank != 0:
|
||||
return latents if output_type == "latent" else []
|
||||
|
||||
if output_type == "latent":
|
||||
return latents
|
||||
if output_type in {"pt", "np"}:
|
||||
if vae_offloaded and vae_restore_device is not None:
|
||||
self.vae.to(device=vae_restore_device)
|
||||
torch.accelerator.empty_cache()
|
||||
return self._decode_latents(latents)
|
||||
raise ValueError(f"Unsupported output_type: {output_type}")
|
||||
|
||||
@torch.inference_mode()
|
||||
def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
|
||||
if req.num_reqs != 1:
|
||||
raise ValueError(f"LingBotVideoPipeline only supports one request per batch, got {req.num_reqs}.")
|
||||
request = req.requests[0]
|
||||
prompt, prompt_negative = _extract_prompt(request)
|
||||
sampling = request.sampling_params
|
||||
extra_args = dict(sampling.extra_args or {})
|
||||
|
||||
generator = sampling.generator
|
||||
if isinstance(generator, list):
|
||||
generator = generator[0] if generator else None
|
||||
if generator is None and sampling.seed is not None:
|
||||
generator = torch.Generator(device=self.device).manual_seed(int(sampling.seed))
|
||||
|
||||
height = sampling.height if sampling.height is not None else extra_args.pop("height", 480)
|
||||
width = sampling.width if sampling.width is not None else extra_args.pop("width", 480)
|
||||
num_frames = sampling.num_frames or extra_args.pop("num_frames", 81)
|
||||
num_inference_steps = (
|
||||
sampling.num_inference_steps
|
||||
if sampling.num_inference_steps is not None
|
||||
else extra_args.pop("num_inference_steps", 40)
|
||||
)
|
||||
guidance_scale = (
|
||||
sampling.guidance_scale
|
||||
if sampling.guidance_scale_provided or sampling.guidance_scale > 0
|
||||
else extra_args.pop("guidance_scale", 6.0)
|
||||
)
|
||||
shift = extra_args.pop(
|
||||
"shift",
|
||||
extra_args.pop("flow_shift", getattr(self.od_config, "flow_shift", None) or 3.0),
|
||||
)
|
||||
negative_prompt = extra_args.pop("negative_prompt", prompt_negative or self.default_negative_prompt)
|
||||
output_type = (
|
||||
sampling.output_type or getattr(self.od_config, "output_type", None) or extra_args.pop("output_type", "pt")
|
||||
)
|
||||
if output_type not in {"pt", "np", "latent"}:
|
||||
output_type = "pt"
|
||||
sampling.output_type = output_type
|
||||
|
||||
frames = self._generate(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=guidance_scale,
|
||||
shift=shift,
|
||||
generator=generator,
|
||||
latents=sampling.latents,
|
||||
output_type=output_type,
|
||||
batch_cfg=bool(extra_args.pop("batch_cfg", False)),
|
||||
null_cond_clone_zero=bool(extra_args.pop("null_cond_clone_zero", False)),
|
||||
offload_vae_during_denoise=bool(extra_args.pop("offload_vae_during_denoise", False)),
|
||||
t_thresh=extra_args.pop("t_thresh", None),
|
||||
refiner_sigma_tail_steps=int(extra_args.pop("refiner_sigma_tail_steps", LOW_NOISE_TAIL_V1_DEFAULT_STEPS)),
|
||||
)
|
||||
return DiffusionOutput(output=frames)
|
||||
@@ -17,7 +17,6 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
from diffusers.models.normalization import AdaLayerNormContinuous
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
@@ -555,8 +554,8 @@ class QwenImageCrossAttention(nn.Module):
|
||||
self.query_num_heads = self.to_qkv.num_heads
|
||||
self.kv_num_heads = self.to_qkv.num_kv_heads
|
||||
|
||||
self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
|
||||
self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
|
||||
self.norm_q = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
|
||||
self.norm_k = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
|
||||
|
||||
self.inner_dim = out_dim if out_dim is not None else head_dim * self.total_num_heads
|
||||
|
||||
@@ -593,8 +592,8 @@ class QwenImageCrossAttention(nn.Module):
|
||||
prefix=f"{prefix}.to_out.0",
|
||||
)
|
||||
|
||||
self.norm_added_q = RMSNorm(head_dim, eps=eps)
|
||||
self.norm_added_k = RMSNorm(head_dim, eps=eps)
|
||||
self.norm_added_q = nn.RMSNorm(head_dim, eps=eps)
|
||||
self.norm_added_k = nn.RMSNorm(head_dim, eps=eps)
|
||||
|
||||
self.attn = Attention(
|
||||
num_heads=self.query_num_heads,
|
||||
@@ -1007,7 +1006,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
|
||||
quant_config=quant_config,
|
||||
)
|
||||
|
||||
self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
|
||||
self.txt_norm = nn.RMSNorm(joint_attention_dim, eps=1e-6)
|
||||
|
||||
# Entry projections (image/text) are kept full precision —
|
||||
# small sensitive layers at the network boundary (see #2728).
|
||||
|
||||
@@ -256,11 +256,6 @@ _DIFFUSION_MODELS = {
|
||||
"pipeline_hunyuan_video_1_5_i2v",
|
||||
"HunyuanVideo15I2VPipeline",
|
||||
),
|
||||
"LingBotVideoPipeline": (
|
||||
"lingbot_video",
|
||||
"pipeline_lingbot_video",
|
||||
"LingBotVideoPipeline",
|
||||
),
|
||||
"MagiHumanPipeline": (
|
||||
"magi_human",
|
||||
"pipeline_magi_human",
|
||||
@@ -531,7 +526,6 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
|
||||
"Flux2Pipeline": "get_flux2_post_process_func",
|
||||
"HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func",
|
||||
"HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func",
|
||||
"LingBotVideoPipeline": "get_lingbot_video_post_process_func",
|
||||
"MagiHumanPipeline": "get_magi_human_post_process_func",
|
||||
"OmniVoicePipeline": "get_omnivoice_post_process_func",
|
||||
"DreamIDOmniPipeline": "get_dreamid_omni_post_process_func",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -1040,6 +1043,7 @@ class AsyncOmniEngine:
|
||||
"enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
|
||||
"num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
|
||||
"quantization": kwargs.get("quantization", None),
|
||||
"quantization_config": kwargs.get("quantization_config", None),
|
||||
"diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None),
|
||||
"diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None),
|
||||
"diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None),
|
||||
@@ -1123,11 +1127,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 +1176,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:
|
||||
@@ -1204,7 +1234,7 @@ class AsyncOmniEngine:
|
||||
kwargs.get("diffusion_attention_config"),
|
||||
attention_backend=kwargs.get("diffusion_attention_backend"),
|
||||
)
|
||||
quantization_config = kwargs.get("diffusion_quantization_config")
|
||||
quantization_config = kwargs.get("diffusion_quantization_config") or kwargs.get("quantization_config")
|
||||
if quantization_config is not None:
|
||||
if (
|
||||
not hasattr(cfg.engine_args, "quantization_config")
|
||||
|
||||
@@ -230,7 +230,7 @@ class StageEngineCoreClientBase(StageClientBase):
|
||||
|
||||
async def add_request_async(self, request: EngineCoreRequest) -> None:
|
||||
"""Add request to the stage engine core."""
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[%s] stage-%s [rep-%s] add request: %s",
|
||||
self.__class__.__name__,
|
||||
self.stage_id,
|
||||
|
||||
@@ -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,
|
||||
@@ -469,7 +478,7 @@ class OmniServeCommand(CLISubcommand):
|
||||
default=None,
|
||||
help=(
|
||||
"JSON string for diffusion quantization_config. "
|
||||
'Example: \'{"method":"gguf","gguf_model":"/path/to/model.gguf"}\'.'
|
||||
'Example: \'{"method":"fp8","activation_scheme":"dynamic"}\'.'
|
||||
),
|
||||
)
|
||||
omni_config_group.add_argument(
|
||||
@@ -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.
|
||||
|
||||
@@ -1395,7 +1395,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
emb_dim = len(embedding)
|
||||
if self._tts_model_type == "ming_tts":
|
||||
if emb_dim != 192:
|
||||
logger.warning("speaker_embedding has %d dimensions; Ming dense expects 192", emb_dim)
|
||||
raise ValueError(f"Ming speaker embedding must have 192 dims, got {emb_dim}")
|
||||
else:
|
||||
dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_dim)
|
||||
if dim_err is not None:
|
||||
|
||||
@@ -25,9 +25,18 @@ class MingTTSAdapter(ARTTSAdapter):
|
||||
ref_audio_source = request.ref_audio
|
||||
voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
|
||||
if ref_audio_source is None and voice_lower in server.uploaded_speakers:
|
||||
ref_audio_source = server._get_uploaded_audio_data(request.voice)
|
||||
if request.ref_text is None:
|
||||
request.ref_text = server.uploaded_speakers[voice_lower].get("ref_text")
|
||||
speaker_info = server.uploaded_speakers[voice_lower]
|
||||
if speaker_info.get("embedding_source") == "direct":
|
||||
if request.speaker_embedding is None:
|
||||
request.speaker_embedding = server._get_uploaded_speaker_embedding(request.voice)
|
||||
if request.speaker_embedding is None:
|
||||
raise ValueError(f"Speaker embedding for uploaded voice '{request.voice}' is missing")
|
||||
else:
|
||||
ref_audio_source = server._get_uploaded_audio_data(request.voice)
|
||||
if not ref_audio_source:
|
||||
raise ValueError(f"Audio file for uploaded voice '{request.voice}' is missing")
|
||||
if request.ref_text is None:
|
||||
request.ref_text = speaker_info.get("ref_text")
|
||||
ref_audio_data = None
|
||||
if isinstance(ref_audio_source, list):
|
||||
ref_audio_data = await server._resolve_ref_audio_many(ref_audio_source)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -28,6 +28,9 @@ from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_GeneratorLike = torch.Generator | Sequence[torch.Generator | None] | None
|
||||
_UNIFORM_EPS = 1e-20
|
||||
|
||||
|
||||
# ===================================================================
|
||||
# HF-numerics-compatible layers for code predictor
|
||||
@@ -605,6 +608,30 @@ class CodePredictorWrapper(nn.Module):
|
||||
values.add(parsed)
|
||||
return values
|
||||
|
||||
@staticmethod
|
||||
def _normalize_generators(
|
||||
generator: _GeneratorLike, batch_size: int
|
||||
) -> torch.Generator | list[torch.Generator | None] | None:
|
||||
if generator is None or isinstance(generator, torch.Generator):
|
||||
return generator
|
||||
|
||||
row_generators = list(generator)
|
||||
if len(row_generators) != batch_size:
|
||||
raise ValueError(f"Expected {batch_size} per-row generators, but got {len(row_generators)}.")
|
||||
return row_generators
|
||||
|
||||
@classmethod
|
||||
def _sample_codes_gumbel(cls, logits: torch.Tensor, generator: _GeneratorLike = None) -> torch.Tensor:
|
||||
"""Sample ``logits`` via Gumbel-max with optional per-row generators."""
|
||||
row_generators = cls._normalize_generators(generator, int(logits.shape[0]))
|
||||
u = torch.empty_like(logits, dtype=torch.float32)
|
||||
if isinstance(row_generators, list):
|
||||
for row, row_generator in enumerate(row_generators):
|
||||
u[row : row + 1].uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generator)
|
||||
else:
|
||||
u.uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generators)
|
||||
return (logits.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True)
|
||||
|
||||
def _prefix_seq_lens(self, max_seq: int) -> list[int]:
|
||||
all_seq_lens = list(range(2, max_seq))
|
||||
if not self._prefix_graph_seq_lens:
|
||||
@@ -780,6 +807,7 @@ class CodePredictorWrapper(nn.Module):
|
||||
bsz = int(layer0_code.shape[0])
|
||||
if generators is not None and len(generators) != bsz:
|
||||
raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}")
|
||||
sample_generator: _GeneratorLike = generators if generators is not None else generator
|
||||
num_groups = self._num_groups
|
||||
device = layer0_code.device
|
||||
|
||||
@@ -856,9 +884,20 @@ class CodePredictorWrapper(nn.Module):
|
||||
|
||||
logits = lm_heads[step - 1](hidden_out[:bsz, step, :])
|
||||
|
||||
# Sample next code
|
||||
# Sample next code via Gumbel-max.
|
||||
#
|
||||
# ``argmax_i(logits_i + Gumbel_i)`` with
|
||||
# ``Gumbel_i = -log(-log(u_i)), u_i ~ Uniform(0, 1)`` is
|
||||
# distributionally identical to sampling from ``softmax(logits)``.
|
||||
# In this file the motivations are practical rather than graph
|
||||
# related: it is measurably cheaper than ``softmax + multinomial``
|
||||
# on the B x 2048 shapes used here, it stays well-defined for
|
||||
# degenerate masked rows with a surviving finite entry (and is more
|
||||
# defensive than ``multinomial`` around fully-masked/NaN inputs),
|
||||
# and the helper below can honor either one batch generator or one
|
||||
# generator per seeded row.
|
||||
if stored_mode:
|
||||
# "stored" mode: top-k -> top-p -> softmax -> multinomial
|
||||
# "stored" mode: top-k -> top-p -> Gumbel-max
|
||||
if s_top_k > 0:
|
||||
topk_vals, _ = logits.topk(s_top_k, dim=-1)
|
||||
logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf"))
|
||||
@@ -869,17 +908,15 @@ class CodePredictorWrapper(nn.Module):
|
||||
remove_mask = (cumulative_probs - sorted_probs) >= s_top_p
|
||||
sorted_logits[remove_mask] = float("-inf")
|
||||
logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)
|
||||
probs = F.softmax(logits, dim=-1, dtype=torch.float32)
|
||||
code = self._multinomial(probs, generator, generators)
|
||||
code = self._sample_codes_gumbel(logits, generator=sample_generator)
|
||||
else:
|
||||
# "per_call" mode: temperature-scaled + top-k
|
||||
# "per_call" mode: temperature-scaled + top-k -> Gumbel-max
|
||||
if use_sampling:
|
||||
scaled = logits * inv_temperature
|
||||
if top_k > 0:
|
||||
topk_vals, _ = scaled.topk(top_k, dim=-1)
|
||||
scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
|
||||
probs = F.softmax(scaled, dim=-1, dtype=torch.float32)
|
||||
code = self._multinomial(probs, generator, generators)
|
||||
code = self._sample_codes_gumbel(scaled, generator=sample_generator)
|
||||
else:
|
||||
code = logits.argmax(dim=-1, keepdim=True)
|
||||
|
||||
|
||||
@@ -5,9 +5,8 @@
|
||||
Stage 0: Talker — text prompt → speech tokens (LLM autoregressive).
|
||||
Stage 1: Code2Wav — flow-matching decoder → acoustic features → waveform.
|
||||
* ``sync_process_input_func`` (``text2flow_token_only``) runs when
|
||||
``deploy.async_chunk=false``: stage 1 builds full-sequence flow input
|
||||
from the emitted speech tokens. ``text2flow`` is the
|
||||
``custom_process_input_func`` variant.
|
||||
``deploy.async_chunk=false``: stage 1 allocates placeholder slots while
|
||||
bulk tensors arrive via ``text2flow_full_payload`` on the connector.
|
||||
* ``async_chunk_process_next_stage_input_func`` runs when
|
||||
``deploy.async_chunk=true``: stage 0 streams codec chunks to stage 1
|
||||
through the shared-memory connector.
|
||||
@@ -47,7 +46,6 @@ COSYVOICE3_PIPELINE = PipelineConfig(
|
||||
final_output=True,
|
||||
final_output_type="audio",
|
||||
engine_output_type="latent",
|
||||
custom_process_input_func=f"{_PROC}.text2flow",
|
||||
sync_process_input_func=f"{_PROC}.text2flow_token_only",
|
||||
),
|
||||
),
|
||||
|
||||
@@ -56,7 +56,6 @@ MIMO_AUDIO_PIPELINE = PipelineConfig(
|
||||
final_output=True,
|
||||
final_output_type="audio",
|
||||
engine_output_type="audio",
|
||||
custom_process_input_func=f"{_PROC}.llm2code2wav",
|
||||
sync_process_input_func=f"{_PROC}.llm2code2wav_token_only",
|
||||
sampling_constraints={"detokenize": False},
|
||||
),
|
||||
|
||||
@@ -17,7 +17,6 @@ from vllm.logger import init_logger
|
||||
from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP
|
||||
from vllm.model_executor.models.qwen2_5_omni_thinker import (
|
||||
Qwen2_5OmniConditionalGenerationMixin,
|
||||
Qwen2_5OmniThinkerProcessingInfo,
|
||||
)
|
||||
from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
|
||||
from vllm.model_executor.models.vision import (
|
||||
@@ -39,6 +38,7 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
|
||||
Qwen2_5OmniThinkerDummyInputsBuilder,
|
||||
Qwen2_5OmniThinkerMultiModalProcessor,
|
||||
Qwen2_5OmniThinkerProcessingInfo,
|
||||
)
|
||||
from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
@@ -4,6 +4,7 @@ from collections.abc import Iterable, Iterator, Mapping
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import PIL.Image as PILImage
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
|
||||
@@ -26,16 +27,21 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
|
||||
from vllm.model_executor.models.qwen2_5_omni_thinker import (
|
||||
Qwen2_5OmniAudioFeatureInputs,
|
||||
Qwen2_5OmniThinkerDummyInputsBuilder,
|
||||
Qwen2_5OmniThinkerProcessingInfo,
|
||||
check_interleaved_audio_video,
|
||||
merge_interleaved_embeddings,
|
||||
)
|
||||
from vllm.model_executor.models.qwen2_5_omni_thinker import (
|
||||
Qwen2_5OmniConditionalGenerationMixin as Qwen2_5OmniConditionalGenerationMixinBase,
|
||||
)
|
||||
from vllm.model_executor.models.qwen2_5_omni_thinker import (
|
||||
Qwen2_5OmniThinkerMultiModalDataParser as _Qwen2_5OmniThinkerMultiModalDataParserBase,
|
||||
)
|
||||
from vllm.model_executor.models.qwen2_5_omni_thinker import (
|
||||
Qwen2_5OmniThinkerMultiModalProcessor as _Qwen2_5OmniThinkerMultiModalProcessorBase,
|
||||
)
|
||||
from vllm.model_executor.models.qwen2_5_omni_thinker import (
|
||||
Qwen2_5OmniThinkerProcessingInfo as _Qwen2_5OmniThinkerProcessingInfoBase,
|
||||
)
|
||||
from vllm.model_executor.models.qwen2_5_vl import (
|
||||
Qwen2_5_VisionTransformer,
|
||||
Qwen2_5_VLImageEmbeddingInputs,
|
||||
@@ -56,12 +62,13 @@ from vllm.multimodal.inputs import (
|
||||
MultiModalFeatureSpec,
|
||||
MultiModalKwargsItems,
|
||||
)
|
||||
from vllm.multimodal.parse import MultiModalDataItems
|
||||
from vllm.multimodal.parse import MultiModalDataItems, VideoProcessorItems
|
||||
from vllm.multimodal.processing.processor import (
|
||||
MultiModalPromptUpdates,
|
||||
PlaceholderFeaturesInfo,
|
||||
)
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.utils.collection_utils import is_list_of
|
||||
|
||||
from vllm_omni.quantization.component_config import (
|
||||
resolve_encoder_quant_config,
|
||||
@@ -74,11 +81,143 @@ except (ImportError, ModuleNotFoundError):
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _presampled_videos_hf_kwargs(
|
||||
mm_data: Mapping[str, object],
|
||||
mm_kwargs: Mapping[str, object],
|
||||
) -> Mapping[str, object]:
|
||||
"""Adjust HF video kwargs for videos pre-sampled by vLLM's video loader.
|
||||
|
||||
When ``video_metadata`` is present (emitted by
|
||||
``Qwen2_5OmniVideoProcessorItems``), the frames were already sampled
|
||||
according to ``media_io_kwargs``, so the HF processor must not re-sample
|
||||
them, and it needs the sampled fps (instead of its default) to compute
|
||||
``video_second_per_grid`` correctly. Otherwise the audio/video temporal
|
||||
alignment breaks under ``use_audio_in_video=True``.
|
||||
"""
|
||||
video_metadata = mm_data.get("video_metadata")
|
||||
if not video_metadata:
|
||||
return mm_kwargs
|
||||
|
||||
mm_kwargs = dict(mm_kwargs)
|
||||
videos_kwargs = dict(mm_kwargs.get("videos_kwargs") or {})
|
||||
videos_kwargs["do_sample_frames"] = False
|
||||
|
||||
def _compute_sampled_video_fps(metadata) -> float | None:
|
||||
duration = getattr(metadata, "duration", None)
|
||||
indices = getattr(metadata, "frames_indices", None)
|
||||
if not duration or not indices or duration <= 0:
|
||||
return None
|
||||
return len(indices) / float(duration)
|
||||
|
||||
# An explicit user-provided fps takes precedence.
|
||||
if "fps" not in videos_kwargs:
|
||||
fps_values = [_compute_sampled_video_fps(m) for m in video_metadata]
|
||||
# HF accepts a single fps per call and uses it for every video's
|
||||
# video_second_per_grid.
|
||||
known_fps = [fps for fps in fps_values if fps is not None]
|
||||
unique_fps = set(known_fps)
|
||||
if len(unique_fps) == 1:
|
||||
videos_kwargs["fps"] = known_fps[0]
|
||||
elif len(unique_fps) > 1:
|
||||
logger.warning(
|
||||
f"Mixed sampled FPS {sorted(unique_fps)} in one request; HF accepts a single fps, using {known_fps[0]}."
|
||||
)
|
||||
videos_kwargs["fps"] = known_fps[0]
|
||||
|
||||
mm_kwargs["videos_kwargs"] = videos_kwargs
|
||||
return mm_kwargs
|
||||
|
||||
|
||||
class Qwen2_5OmniVideoProcessorItems(VideoProcessorItems):
|
||||
"""Video items that carry the loader's ``(frames, metadata)`` tuples.
|
||||
|
||||
The tuples are kept in ``data`` so that mm hashing and cache-miss
|
||||
re-parsing retain the metadata; ``get_processor_data`` unpacks them into
|
||||
the ``videos`` + ``video_metadata`` arguments of the HF processor.
|
||||
"""
|
||||
|
||||
def get_processor_data(self) -> Mapping[str, object]:
|
||||
from transformers.video_utils import VideoMetadata
|
||||
|
||||
videos = [frames for frames, _ in self.data]
|
||||
video_metadata = [
|
||||
VideoMetadata(**{k: v for k, v in metadata.items() if k != "do_sample_frames"}) for _, metadata in self.data
|
||||
]
|
||||
return {"videos": videos, "video_metadata": video_metadata}
|
||||
|
||||
|
||||
class Qwen2_5OmniThinkerMultiModalDataParser(_Qwen2_5OmniThinkerMultiModalDataParserBase):
|
||||
def _parse_video_data(self, data):
|
||||
if data is None or isinstance(data, dict) or self.is_embeddings(data):
|
||||
return super()._parse_video_data(data)
|
||||
|
||||
# Normalize to a list of per-video items (same as the base parser).
|
||||
if (is_list_of(data, PILImage.Image) and len(data) > 0) or (
|
||||
isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4
|
||||
):
|
||||
data_items = [data]
|
||||
elif isinstance(data, (np.ndarray, torch.Tensor)):
|
||||
data_items = [elem for elem in data]
|
||||
elif isinstance(data, tuple) and len(data) == 2:
|
||||
data_items = [data]
|
||||
else:
|
||||
data_items = data
|
||||
|
||||
videos_with_metadata = [self._get_video_with_metadata(item) for item in data_items]
|
||||
|
||||
# Metadata is optional: videos fetched by vLLM's media connector
|
||||
# arrive as (frames, metadata) tuples, while plain arrays (e.g.
|
||||
# offline `multi_modal_data` or dummy profiling inputs) carry no
|
||||
# metadata and parse as before.
|
||||
if any(metadata is None for _, metadata in videos_with_metadata):
|
||||
return super()._parse_video_data(data)
|
||||
|
||||
return Qwen2_5OmniVideoProcessorItems(
|
||||
videos_with_metadata,
|
||||
metadata=[metadata for _, metadata in videos_with_metadata],
|
||||
)
|
||||
|
||||
|
||||
class Qwen2_5OmniThinkerProcessingInfo(_Qwen2_5OmniThinkerProcessingInfoBase):
|
||||
def get_data_parser(self):
|
||||
feature_extractor = self.get_feature_extractor()
|
||||
|
||||
return Qwen2_5OmniThinkerMultiModalDataParser(
|
||||
spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size,
|
||||
target_sr=feature_extractor.sampling_rate,
|
||||
target_channels=self.get_target_channels(),
|
||||
expected_hidden_size=self._get_expected_hidden_size(),
|
||||
)
|
||||
|
||||
|
||||
class Qwen2_5OmniThinkerMultiModalProcessor(
|
||||
_Qwen2_5OmniThinkerMultiModalProcessorBase,
|
||||
):
|
||||
"""Override to fix use_audio_in_video detection when mm cache returns None."""
|
||||
|
||||
def _cached_apply_hf_processor(self, inputs, timing_ctx):
|
||||
# If use_audio_in_video, process video and audio together; otherwise,
|
||||
# skip cache to avoid errors.
|
||||
# Prevents partial cache hits from causing processing failures.
|
||||
if inputs.hf_processor_mm_kwargs.get("use_audio_in_video"):
|
||||
return self._apply_hf_processor(inputs, timing_ctx)
|
||||
return super()._cached_apply_hf_processor(inputs, timing_ctx)
|
||||
|
||||
def _call_hf_processor(
|
||||
self,
|
||||
prompt: str,
|
||||
mm_data: Mapping[str, object],
|
||||
mm_kwargs: Mapping[str, object],
|
||||
tok_kwargs: Mapping[str, object],
|
||||
):
|
||||
mm_kwargs = _presampled_videos_hf_kwargs(mm_data, mm_kwargs)
|
||||
return super()._call_hf_processor(
|
||||
prompt=prompt,
|
||||
mm_data=mm_data,
|
||||
mm_kwargs=mm_kwargs,
|
||||
tok_kwargs=tok_kwargs,
|
||||
)
|
||||
|
||||
def _maybe_apply_prompt_updates(
|
||||
self,
|
||||
mm_items: MultiModalDataItems,
|
||||
|
||||
@@ -51,7 +51,6 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
|
||||
input_sources=(0,),
|
||||
hf_config_name="talker_config",
|
||||
engine_output_type="latent",
|
||||
custom_process_input_func=f"{_PROC}.thinker2talker",
|
||||
sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
|
||||
custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_full_payload"),
|
||||
async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
|
||||
@@ -69,7 +68,6 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
|
||||
final_output_type="audio",
|
||||
hf_config_name="thinker_config",
|
||||
engine_output_type="audio",
|
||||
custom_process_input_func=f"{_PROC}.talker2code2wav",
|
||||
sampling_constraints={"detokenize": True},
|
||||
),
|
||||
),
|
||||
|
||||
Loaded 100 of 121 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user