Merge remote-tracking branch 'origin/main' into dev/vllm-align
This commit is contained in:
tzhouam committed
commit
7f54fd2f54
26 files changed
+757
-501
No files matched your search
+13
-6
@@ -20,12 +20,19 @@ steps:
|
||||
- label: "Simple Unit Test"
|
||||
depends_on: image-build
|
||||
commands:
|
||||
- pytest -v -s tests/entrypoints/
|
||||
- pytest -v -s tests/diffusion/cache/
|
||||
- pytest -v -s tests/diffusion/lora/
|
||||
- pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
|
||||
- pytest -v -s tests/worker/
|
||||
- pytest -v -s tests/distributed/omni_connectors/test_kv_flow.py
|
||||
- |
|
||||
pytest -v -s \
|
||||
tests/entrypoints/ \
|
||||
tests/diffusion/cache/ \
|
||||
tests/diffusion/lora/ \
|
||||
tests/model_executor/models/qwen2_5_omni/test_audio_length.py \
|
||||
tests/worker/ \
|
||||
tests/distributed/omni_connectors/test_kv_flow.py \
|
||||
--cov=vllm_omni \
|
||||
--cov-branch \
|
||||
--cov-report=term-missing \
|
||||
--cov-report=html \
|
||||
--cov-report=xml
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
steps:
|
||||
- label: ":docker: Build image"
|
||||
key: image-build
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
|
||||
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
- label: "Omni Model Test with H100"
|
||||
timeout_in_minutes: 180
|
||||
depends_on: image-build
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
@@ -7,12 +7,8 @@ WORKDIR ${APP_DIR}
|
||||
|
||||
COPY . .
|
||||
|
||||
# Remove this replace when the dispatch of requirements is ready
|
||||
RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
|
||||
&& sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
|
||||
|
||||
# Install vllm-omni with dev dependencies
|
||||
RUN pip install --no-cache-dir -e .
|
||||
RUN pip install --no-cache-dir -e . --no-build-isolation
|
||||
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
|
||||
@@ -7,12 +7,8 @@ WORKDIR ${APP_DIR}
|
||||
|
||||
COPY . .
|
||||
|
||||
# Remove this replace when the dispatch of requirements is ready
|
||||
RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
|
||||
&& sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
|
||||
|
||||
# Install vllm-omni with dev dependencies
|
||||
RUN pip install --no-cache-dir -e .
|
||||
RUN pip install --no-cache-dir -e . --no-build-isolation
|
||||
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
|
||||
@@ -15,11 +15,7 @@ RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
|
||||
|
||||
# Step 2: Copy vllm-omni code and install without uv
|
||||
COPY . ${COMMON_WORKDIR}/vllm-omni
|
||||
RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
|
||||
|
||||
# When we are installing onnxruntime-rocm, we need to uninstall the system-installed onnxruntime first.
|
||||
# These are the dependencies of Qwen3-TTS.
|
||||
RUN uv pip uninstall onnxruntime --system && uv pip install --no-cache-dir onnxruntime-rocm sox --system
|
||||
RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]" --no-build-isolation
|
||||
|
||||
RUN ln -sf /usr/bin/python3 /usr/bin/python
|
||||
|
||||
|
||||
@@ -26,6 +26,8 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
|
||||
|
||||
### Pre-built wheels
|
||||
|
||||
Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
|
||||
|
||||
=== "NVIDIA CUDA"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-wheels"
|
||||
|
||||
@@ -17,8 +17,6 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
|
||||
# --8<-- [start:pre-built-wheels]
|
||||
|
||||
#### Installation of vLLM
|
||||
Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
|
||||
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
```bash
|
||||
|
||||
@@ -9,10 +9,60 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
|
||||
|
||||
# --8<-- [start:pre-built-wheels]
|
||||
|
||||
#### Installation of vLLM
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
```bash
|
||||
uv pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
|
||||
```
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
|
||||
```bash
|
||||
# we need to add --no-build-isolation as the torch
|
||||
# is not obtained from pypi, we have to install using the
|
||||
# torch installed in our environment
|
||||
uv pip install vllm-omni
|
||||
|
||||
# Optional if want to run Qwen3 TTS
|
||||
uv pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
|
||||
uv pip install onnxruntime-rocm sox
|
||||
```
|
||||
|
||||
# --8<-- [end:pre-built-wheels]
|
||||
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.14.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
|
||||
```
|
||||
|
||||
The release 0.14.0 of vLLM requires ROCm 7.0 environment.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
VLLM_OMNI_TARGET_DEVICE=rocm uv pip install -e .
|
||||
# OR
|
||||
uv pip install -e . --no-build-isolation
|
||||
```
|
||||
|
||||
<details><summary>(Optional) Installation of vLLM from source</summary>
|
||||
If you want to check, modify or debug with source code of vLLM, install the library from source with the following instructions:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.14.0
|
||||
python3 -m pip install -r requirements/rocm.txt
|
||||
python3 setup.py develop
|
||||
```
|
||||
|
||||
# --8<-- [end:build-wheel-from-source]
|
||||
|
||||
# --8<-- [start:build-docker]
|
||||
|
||||
@@ -38,12 +38,9 @@ docker run --rm \
|
||||
cd /vllm-workspace
|
||||
git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git
|
||||
|
||||
# Remove this replace when the dispatch of requirements is ready
|
||||
RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
|
||||
&& sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
|
||||
|
||||
cd vllm-omni
|
||||
pip install -v -e .
|
||||
VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
|
||||
# OR pip install -v -e . --no-build-isolation
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
```
|
||||
|
||||
|
||||
@@ -49,7 +49,7 @@ The following table shows which models are currently supported by parallelism me
|
||||
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-SP | Tensor-Parallel |
|
||||
|-------|------------------|------------|---------|--------------------------|
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ❌ |
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ |
|
||||
|
||||
### Tensor Parallelism
|
||||
|
||||
|
||||
@@ -109,7 +109,12 @@ def parse_args() -> argparse.Namespace:
|
||||
choices=[1, 2],
|
||||
help="Number of GPUs used for classifier free guidance parallel size.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--tensor_parallel_size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -141,6 +146,7 @@ def main():
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
)
|
||||
|
||||
# Check if profiling is requested via environment variable
|
||||
@@ -173,7 +179,7 @@ def main():
|
||||
print(f" Inference steps: {args.num_inference_steps}")
|
||||
print(f" Frames: {args.num_frames}")
|
||||
print(
|
||||
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"
|
||||
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}"
|
||||
)
|
||||
print(f" Video size: {args.width}x{args.height}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
+13
-28
@@ -1,5 +1,9 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=61.0", "wheel"]
|
||||
requires = [
|
||||
"setuptools>=77.0.3,<81.0.0",
|
||||
"wheel",
|
||||
"setuptools-scm>=8.0",
|
||||
]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
@@ -8,16 +12,16 @@ version = "0.14.0"
|
||||
description = "A framework for efficient model inference with omni-modality models"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10,<3.14"
|
||||
license = {text = "Apache-2.0"}
|
||||
license = "Apache-2.0"
|
||||
authors = [
|
||||
{name = "vLLM-Omni Team"}
|
||||
]
|
||||
keywords = ["vllm", "multimodal", "diffusion", "transformer", "inference", "serving"]
|
||||
dynamic = ["dependencies"]
|
||||
classifiers = [
|
||||
"Development Status :: 3 - Alpha",
|
||||
"Intended Audience :: Developers",
|
||||
"Intended Audience :: Science/Research",
|
||||
"License :: OSI Approved :: Apache Software License",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
@@ -26,26 +30,10 @@ classifiers = [
|
||||
"Topic :: Software Development :: Libraries :: Python Modules",
|
||||
]
|
||||
|
||||
|
||||
dependencies = [
|
||||
# Core runtime dependencies (required for actual usage)
|
||||
"omegaconf>=2.3.0",
|
||||
"librosa>=0.11.0",
|
||||
"resampy>=0.4.3",
|
||||
"diffusers>=0.36.0",
|
||||
"accelerate==1.12.0",
|
||||
"gradio==5.50",
|
||||
"soundfile>=0.13.1",
|
||||
"cache-dit==1.2.0",
|
||||
"tqdm>=4.66.0",
|
||||
"torchsde>=0.2.6", # Required for Stable Audio scheduler
|
||||
"fa3-fwd==0.0.1", # flash attention 3, maintained by @ZJY0516
|
||||
"openai-whisper>=20250625",
|
||||
"imageio[ffmpeg]>=2.37.2",
|
||||
"onnxruntime>=1.19.0",
|
||||
"sox>=1.5.0",
|
||||
# "vllm==0.15.0", # TODO: fix the entrypoints overwrite problem
|
||||
]
|
||||
# Dependencies are now managed dynamically via setup.py based on detected hardware platform.
|
||||
# This allows automatic installation of the correct platform-specific dependencies (CUDA/ROCm/CPU/XPU/NPU)
|
||||
# without requiring extras like [cuda]. See requirements/ directory for platform-specific dependencies.
|
||||
# Note: vllm is intentionally excluded due to entrypoints overwrite issue.
|
||||
|
||||
[project.optional-dependencies]
|
||||
|
||||
@@ -155,11 +143,7 @@ python_classes = ["Test*"]
|
||||
python_functions = ["test_*"]
|
||||
addopts = [
|
||||
"--strict-markers",
|
||||
"--strict-config",
|
||||
"--cov=vllm_omni",
|
||||
"--cov-report=term-missing",
|
||||
"--cov-report=html",
|
||||
"--cov-report=xml",
|
||||
"--strict-config"
|
||||
]
|
||||
markers = [
|
||||
# ci/cd required
|
||||
@@ -204,4 +188,5 @@ extend-ignore-identifiers-re = [
|
||||
".*NOTHINK.*",
|
||||
".*nin.*",
|
||||
"Ono_Anna",
|
||||
".*cann.*",
|
||||
]
|
||||
@@ -0,0 +1,14 @@
|
||||
# Common dependencies for all platforms
|
||||
omegaconf>=2.3.0
|
||||
librosa>=0.11.0
|
||||
resampy>=0.4.3
|
||||
diffusers>=0.36.0
|
||||
accelerate==1.12.0
|
||||
gradio==5.50
|
||||
soundfile>=0.13.1
|
||||
cache-dit==1.2.0
|
||||
tqdm>=4.66.0
|
||||
torchsde>=0.2.6
|
||||
openai-whisper>=20250625
|
||||
imageio[ffmpeg]>=2.37.2
|
||||
sox>=1.5.0
|
||||
@@ -0,0 +1,2 @@
|
||||
-r common.txt
|
||||
onnxruntime>=1.23.2
|
||||
@@ -0,0 +1,3 @@
|
||||
-r common.txt
|
||||
onnxruntime>=1.23.2
|
||||
fa3-fwd==0.0.1
|
||||
@@ -0,0 +1,3 @@
|
||||
-r common.txt
|
||||
onnxruntime-cann>=1.23.2
|
||||
torchaudio==2.9.0
|
||||
@@ -0,0 +1,2 @@
|
||||
-r common.txt
|
||||
onnxruntime-rocm>=1.22.2
|
||||
@@ -0,0 +1,2 @@
|
||||
-r common.txt
|
||||
onnxruntime>=1.23.2
|
||||
@@ -0,0 +1,171 @@
|
||||
"""
|
||||
Setup script for vLLM-Omni with hardware-dependent installation.
|
||||
|
||||
This setup.py implements platform-aware dependency routing so users can run
|
||||
`pip install vllm-omni` and automatically receive the correct platform-specific
|
||||
dependencies (CUDA/ROCm/CPU/XPU/NPU) without requiring extras like `[cuda]`.
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from setuptools import setup
|
||||
|
||||
|
||||
def uninstall_onnxruntime() -> None:
|
||||
"""
|
||||
Uninstall onnxruntime package if it exists.
|
||||
|
||||
This is necessary for ROCm environments where onnxruntime may conflict
|
||||
with ROCm-specific dependencies.
|
||||
"""
|
||||
try:
|
||||
import pkg_resources
|
||||
|
||||
try:
|
||||
pkg_resources.get_distribution("onnxruntime")
|
||||
print("Found onnxruntime installed, uninstalling for ROCm compatibility...")
|
||||
subprocess.check_call(
|
||||
[sys.executable, "-m", "pip", "uninstall", "-y", "onnxruntime"],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
print("Successfully uninstalled onnxruntime")
|
||||
except pkg_resources.DistributionNotFound:
|
||||
print("onnxruntime not installed, skipping uninstall")
|
||||
except Exception as e:
|
||||
print(f"Warning: Failed to uninstall onnxruntime: {e}")
|
||||
|
||||
|
||||
def detect_target_device() -> str:
|
||||
"""
|
||||
Detect the target device for installation following RFC priority rules.
|
||||
|
||||
Priority order:
|
||||
1. VLLM_OMNI_TARGET_DEVICE environment variable (highest priority)
|
||||
2. Torch backend detection (cuda, rocm, npu, xpu)
|
||||
3. CPU fallback (default)
|
||||
|
||||
Returns:
|
||||
str: Device name ('cuda', 'rocm', 'npu', 'xpu', or 'cpu')
|
||||
"""
|
||||
# Priority 1: Explicit override via environment variable
|
||||
target_device = os.environ.get("VLLM_OMNI_TARGET_DEVICE")
|
||||
if target_device:
|
||||
valid_devices = ["cuda", "rocm", "npu", "xpu", "cpu"]
|
||||
if target_device.lower() in valid_devices:
|
||||
print(f"Using target device from VLLM_OMNI_TARGET_DEVICE: {target_device.lower()}")
|
||||
return target_device.lower()
|
||||
else:
|
||||
print(f"Warning: Invalid VLLM_OMNI_TARGET_DEVICE '{target_device}', falling back to auto-detection")
|
||||
|
||||
# Priority 2: Torch backend detection
|
||||
# This is a code path for when user is using
|
||||
# --no-build-isolation flag
|
||||
try:
|
||||
import torch
|
||||
|
||||
# Check for CUDA
|
||||
if torch.version.cuda is not None:
|
||||
print("Detected CUDA backend from torch")
|
||||
return "cuda"
|
||||
|
||||
# Check for ROCm (AMD)
|
||||
if torch.version.hip is not None:
|
||||
print("Detected ROCm backend from torch")
|
||||
uninstall_onnxruntime()
|
||||
return "rocm"
|
||||
|
||||
# Check for NPU (Ascend)
|
||||
if hasattr(torch, "npu"):
|
||||
try:
|
||||
if torch.npu.is_available():
|
||||
print("Detected NPU backend from torch")
|
||||
return "npu"
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Check for XPU (Intel)
|
||||
if hasattr(torch, "xpu"):
|
||||
try:
|
||||
if torch.xpu.is_available():
|
||||
print("Detected XPU backend from torch")
|
||||
return "xpu"
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
print("No GPU backend detected in torch, defaulting to CPU")
|
||||
return "cpu"
|
||||
|
||||
except ImportError:
|
||||
print("PyTorch not found, defaulting to CUDA installation")
|
||||
return "cuda"
|
||||
|
||||
|
||||
def load_requirements(file_path: Path) -> list[str]:
|
||||
"""
|
||||
Load requirements from a file, supporting -r directive for recursive loading.
|
||||
|
||||
Args:
|
||||
file_path: Path to the requirements file
|
||||
|
||||
Returns:
|
||||
List of requirement strings
|
||||
"""
|
||||
requirements = []
|
||||
|
||||
if not file_path.exists():
|
||||
print(f"Warning: Requirements file not found: {file_path}")
|
||||
return requirements
|
||||
|
||||
with open(file_path) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
|
||||
# Skip empty lines and comments
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
|
||||
# Handle -r directive for recursive loading
|
||||
if line.startswith("-r "):
|
||||
nested_file = line[3:].strip()
|
||||
nested_path = file_path.parent / nested_file
|
||||
requirements.extend(load_requirements(nested_path))
|
||||
else:
|
||||
requirements.append(line)
|
||||
|
||||
return requirements
|
||||
|
||||
|
||||
def get_install_requires() -> list[str]:
|
||||
"""
|
||||
Get the list of dependencies based on detected platform.
|
||||
|
||||
Returns:
|
||||
List of requirement strings for the detected platform
|
||||
"""
|
||||
device = detect_target_device()
|
||||
requirements_dir = Path(__file__).parent / "requirements"
|
||||
requirements_file = requirements_dir / f"{device}.txt"
|
||||
|
||||
print(f"Loading requirements from: {requirements_file}")
|
||||
requirements = load_requirements(requirements_file)
|
||||
|
||||
if not requirements:
|
||||
print(f"Warning: No requirements loaded for device '{device}'")
|
||||
else:
|
||||
print(f"Loaded {len(requirements)} requirements for {device}")
|
||||
|
||||
return requirements
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Get platform-specific dependencies
|
||||
install_requires = get_install_requires()
|
||||
|
||||
# Setup configuration
|
||||
setup(
|
||||
install_requires=install_requires,
|
||||
)
|
||||
@@ -60,6 +60,12 @@ def clean_gpu_memory_between_tests():
|
||||
_run_post_test_cleanup()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def log_test_name_before_test(request):
|
||||
print(f"--- Running test: {request.node.name}")
|
||||
yield
|
||||
|
||||
|
||||
def _run_pre_test_cleanup(enable_force=False):
|
||||
if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force:
|
||||
print("GPU cleanup disabled")
|
||||
|
||||
@@ -24,6 +24,7 @@ if str(REPO_ROOT) not in sys.path:
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
|
||||
|
||||
@@ -62,7 +63,7 @@ def test_cache_dit(model_name: str):
|
||||
width=width,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=0.0,
|
||||
generator=torch.Generator("cuda").manual_seed(42),
|
||||
generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
|
||||
num_outputs_per_prompt=1, # Single output for speed
|
||||
),
|
||||
)
|
||||
|
||||
@@ -47,7 +47,7 @@ def test_diffusion_model(model_name: str):
|
||||
width=width,
|
||||
num_inference_steps=2,
|
||||
guidance_scale=0.0,
|
||||
generator=torch.Generator("cuda").manual_seed(42),
|
||||
generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
|
||||
num_outputs_per_prompt=2,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -14,6 +14,7 @@ if str(REPO_ROOT) not in sys.path:
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
@@ -42,7 +43,7 @@ def test_video_diffusion_model(model_name: str):
|
||||
num_frames=num_frames,
|
||||
num_inference_steps=2,
|
||||
guidance_scale=1.0,
|
||||
generator=torch.Generator("cuda").manual_seed(42),
|
||||
generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
|
||||
),
|
||||
)
|
||||
first_output = outputs[0]
|
||||
|
||||
@@ -8,6 +8,7 @@ import concurrent.futures
|
||||
import os
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
@@ -21,7 +22,6 @@ from tests.conftest import (
|
||||
dummy_messages_from_mix_data,
|
||||
generate_synthetic_audio,
|
||||
generate_synthetic_image,
|
||||
modify_stage_config,
|
||||
)
|
||||
|
||||
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
|
||||
@@ -33,6 +33,29 @@ stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omn
|
||||
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
|
||||
|
||||
|
||||
_omni_server_lock = threading.Lock()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def omni_server(request):
|
||||
"""Start vLLM-Omni server as a subprocess with actual model weights.
|
||||
Uses session scope so the server starts only once for the entire test session.
|
||||
Multi-stage initialization can take 10-20+ minutes.
|
||||
"""
|
||||
with _omni_server_lock:
|
||||
model, stage_config_path = request.param
|
||||
|
||||
print(f"Starting OmniServer with model: {model}")
|
||||
|
||||
with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server:
|
||||
print("OmniServer started successfully")
|
||||
yield server
|
||||
print("OmniServer stopping...")
|
||||
|
||||
print("OmniServer stopped")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(omni_server):
|
||||
"""OpenAI client for the running vLLM-Omni server."""
|
||||
return openai.OpenAI(
|
||||
@@ -86,227 +109,180 @@ def get_deploy_config(deploy_type="TP1"):
|
||||
return result.get(deploy_type, result["TP1"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
def test_text_to_text_001(test_config: tuple[str, str]) -> None:
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_text_to_text_001(client: openai.OpenAI, omni_server) -> None:
|
||||
"""Test processing text, generating text output via OpenAI API."""
|
||||
model, stage_config_path = test_config
|
||||
with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
|
||||
messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
|
||||
messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
|
||||
|
||||
# Test single completion
|
||||
api_client = client(server)
|
||||
start_time = time.perf_counter()
|
||||
chat_completion = api_client.chat.completions.create(
|
||||
model=server.model, messages=messages, max_tokens=20, modalities=["text"]
|
||||
)
|
||||
# Verify E2E
|
||||
print(f"the request e2e is: {time.perf_counter() - start_time}")
|
||||
# TODO: Verify the E2E latency after confirmation baseline.
|
||||
# Test single completion
|
||||
start_time = time.perf_counter()
|
||||
chat_completion = client.chat.completions.create(
|
||||
model=omni_server.model, messages=messages, max_tokens=20, modalities=["text"]
|
||||
)
|
||||
# Verify E2E
|
||||
print(f"the request e2e is: {time.perf_counter() - start_time}")
|
||||
# TODO: Verify the E2E latency after confirmation baseline.
|
||||
|
||||
# Verify only output text
|
||||
assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
|
||||
# Verify only output text
|
||||
assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
|
||||
|
||||
# Verify text output success
|
||||
text_choice = chat_completion.choices[0]
|
||||
assert text_choice.message.content is not None, "No text output is generated"
|
||||
assert chat_completion.usage.completion_tokens <= 20, "The output length more than the requested max_tokens."
|
||||
assert "beijing" in text_choice.message.content.lower(), "The output do not contain keywords."
|
||||
# Verify text output success
|
||||
text_choice = chat_completion.choices[0]
|
||||
assert text_choice.message.content is not None, "No text output is generated"
|
||||
assert chat_completion.usage.completion_tokens <= 20, "The output length more than the requested max_tokens."
|
||||
assert "beijing" in text_choice.message.content.lower(), "The output do not contain keywords."
|
||||
|
||||
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
def test_audio_to_text_001(test_config: tuple[str, str]) -> None:
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_audio_to_text_001(client: openai.OpenAI, omni_server) -> None:
|
||||
"""Test processing text, generating text output via OpenAI API."""
|
||||
model, stage_config_path = test_config
|
||||
deploy_config = get_deploy_config()
|
||||
deploy_config[0]["default_sampling_params.ignore_eos"] = True
|
||||
stage_config_path = modify_stage_config(stage_config_path, deploy_config)
|
||||
with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
|
||||
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(1, 1)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url)
|
||||
# Test single completion
|
||||
api_client = client(server)
|
||||
start_time = time.perf_counter()
|
||||
chat_completion = api_client.chat.completions.create(
|
||||
model=server.model, messages=messages, max_tokens=200, modalities=["text"]
|
||||
)
|
||||
# Verify only output text
|
||||
assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
|
||||
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(1, 1)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url)
|
||||
# Test single completion
|
||||
start_time = time.perf_counter()
|
||||
chat_completion = client.chat.completions.create(model=omni_server.model, messages=messages, modalities=["text"])
|
||||
# Verify only output text
|
||||
assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
|
||||
|
||||
# Verify text output success
|
||||
text_choice = chat_completion.choices[0]
|
||||
assert text_choice.message.content is not None, "No text output is generated"
|
||||
assert chat_completion.usage.completion_tokens == 200, (
|
||||
"The output length differs from the requested max_tokens."
|
||||
)
|
||||
# Verify text output success
|
||||
text_choice = chat_completion.choices[0]
|
||||
assert text_choice.message.content is not None, "No text output is generated"
|
||||
assert "water" in text_choice.message.content.lower(), "The output do not contain keywords."
|
||||
|
||||
# Verify E2E
|
||||
print(f"the request e2e is: {time.perf_counter() - start_time}")
|
||||
# TODO: Verify the E2E latency after confirmation baseline.
|
||||
# Verify E2E
|
||||
print(f"the request e2e is: {time.perf_counter() - start_time}")
|
||||
# TODO: Verify the E2E latency after confirmation baseline.
|
||||
|
||||
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
def test_audio_to_text_audio_001(test_config: tuple[str, str]) -> None:
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_audio_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
|
||||
"""Test processing text, generating audio output via OpenAI API."""
|
||||
|
||||
model, stage_config_path = test_config
|
||||
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(1, 5)['base64']}"
|
||||
|
||||
messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url)
|
||||
num_concurrent_requests = get_max_batch_size()
|
||||
stage_config_path = modify_stage_config(
|
||||
stage_config_path,
|
||||
{
|
||||
"stage_args": {
|
||||
0: {"runtime.max_batch_size": num_concurrent_requests},
|
||||
1: {"runtime.max_batch_size": num_concurrent_requests},
|
||||
}
|
||||
},
|
||||
)
|
||||
with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
|
||||
audio_data_url = []
|
||||
for _ in range(5):
|
||||
audio_data_url.append(f"data:audio/wav;base64,{generate_synthetic_audio(1, 5)['base64']}")
|
||||
|
||||
messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url)
|
||||
|
||||
# Test single completion
|
||||
api_client = client(server)
|
||||
e2e_list = list()
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
|
||||
# Submit multiple completion requests concurrently
|
||||
futures = [
|
||||
executor.submit(api_client.chat.completions.create, model=server.model, messages=messages)
|
||||
for _ in range(num_concurrent_requests)
|
||||
]
|
||||
start_time = time.perf_counter()
|
||||
# Wait for all requests to complete and collect results
|
||||
chat_completions = list()
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
chat_completions.append(future.result())
|
||||
# Verify E2E
|
||||
current_e2e = time.perf_counter() - start_time
|
||||
print(f"the request e2e is: {current_e2e}")
|
||||
# TODO: Verify the E2E latency after confirmation baseline.
|
||||
e2e_list.append(current_e2e)
|
||||
|
||||
print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
|
||||
# Verify all completions succeeded
|
||||
assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded."
|
||||
for chat_completion in chat_completions:
|
||||
# Verify audio output success
|
||||
audio_message = chat_completion.choices[1].message
|
||||
audio_data = audio_message.audio.data
|
||||
assert audio_data is not None, "No audio output is generated"
|
||||
assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
|
||||
|
||||
# Verify text output success
|
||||
text_choice = chat_completion.choices[0]
|
||||
text_content = text_choice.message.content
|
||||
assert text_choice.message.content is not None, "No text output is generated"
|
||||
|
||||
# Verify text output same as audio output
|
||||
audio_content = convert_audio_to_text(audio_data)
|
||||
print(f"text content is: {text_content}")
|
||||
print(f"audio content is: {audio_content}")
|
||||
similarity = cosine_similarity_text(audio_content, text_content)
|
||||
print(f"similarity between audio and text is: {similarity}")
|
||||
assert similarity > 0.9, "The audio content is not same as the text"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
def test_image_to_text_001(test_config: tuple[str, str]) -> None:
|
||||
"""Test processing text, generating text output via OpenAI API."""
|
||||
model, stage_config_path = test_config
|
||||
deploy_config = get_deploy_config()
|
||||
stage_config_path = modify_stage_config(stage_config_path, deploy_config)
|
||||
|
||||
with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
|
||||
# Test single completion
|
||||
api_client = client(server)
|
||||
# Test single completion
|
||||
e2e_list = list()
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
|
||||
# Submit multiple completion requests concurrently
|
||||
futures = [
|
||||
executor.submit(client.chat.completions.create, model=omni_server.model, messages=messages)
|
||||
for _ in range(num_concurrent_requests)
|
||||
]
|
||||
start_time = time.perf_counter()
|
||||
chat_completion = api_client.chat.completions.create(
|
||||
model=server.model, messages=messages, max_tokens=100, modalities=["text"]
|
||||
)
|
||||
# Verify E2E
|
||||
print(f"the request e2e is: {time.perf_counter() - start_time}")
|
||||
# TODO: Verify the E2E latency after confirmation baseline.
|
||||
# Wait for all requests to complete and collect results
|
||||
chat_completions = list()
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
chat_completions.append(future.result())
|
||||
# Verify E2E
|
||||
current_e2e = time.perf_counter() - start_time
|
||||
print(f"the request e2e is: {current_e2e}")
|
||||
# TODO: Verify the E2E latency after confirmation baseline.
|
||||
e2e_list.append(current_e2e)
|
||||
|
||||
# Verify only output text
|
||||
assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
|
||||
print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
|
||||
# Verify all completions succeeded
|
||||
assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded."
|
||||
for chat_completion in chat_completions:
|
||||
# Verify audio output success
|
||||
audio_message = chat_completion.choices[1].message
|
||||
audio_data = audio_message.audio.data
|
||||
assert audio_data is not None, "No audio output is generated"
|
||||
assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
|
||||
|
||||
# Verify text output success
|
||||
text_choice = chat_completion.choices[0]
|
||||
text_content = text_choice.message.content
|
||||
assert text_choice.message.content is not None, "No text output is generated"
|
||||
|
||||
# Verify text output same as audio output
|
||||
audio_content = convert_audio_to_text(audio_data)
|
||||
print(f"text content is: {text_content}")
|
||||
print(f"audio content is: {audio_content}")
|
||||
similarity = cosine_similarity_text(audio_content, text_content)
|
||||
print(f"similarity between audio and text is: {similarity}")
|
||||
assert similarity > 0.9, "The audio content is not same as the text"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_image_to_text_001(client: openai.OpenAI, omni_server) -> None:
|
||||
"""Test processing text, generating text output via OpenAI API."""
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
|
||||
# Test single completion
|
||||
start_time = time.perf_counter()
|
||||
chat_completion = client.chat.completions.create(
|
||||
model=omni_server.model, messages=messages, max_tokens=100, modalities=["text"]
|
||||
)
|
||||
# Verify E2E
|
||||
print(f"the request e2e is: {time.perf_counter() - start_time}")
|
||||
# TODO: Verify the E2E latency after confirmation baseline.
|
||||
|
||||
# Verify only output text
|
||||
assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
|
||||
|
||||
# Verify text output success
|
||||
text_choice = chat_completion.choices[0]
|
||||
text_content = text_choice.message.content
|
||||
assert text_content is not None, "No text output is generated"
|
||||
assert chat_completion.usage.completion_tokens <= 100, "The output length more than the requested max_tokens."
|
||||
assert "square" in text_content.lower(), "The output do not contain keywords."
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_image_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
|
||||
"""Test processing text, generating audio output via OpenAI API."""
|
||||
|
||||
num_concurrent_requests = get_max_batch_size()
|
||||
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
|
||||
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
|
||||
|
||||
# Test single completion
|
||||
e2e_list = list()
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
|
||||
# Submit multiple completion requests concurrently
|
||||
futures = [
|
||||
executor.submit(
|
||||
client.chat.completions.create,
|
||||
model=omni_server.model,
|
||||
messages=messages,
|
||||
)
|
||||
for _ in range(num_concurrent_requests)
|
||||
]
|
||||
start_time = time.perf_counter()
|
||||
# Wait for all requests to complete and collect results
|
||||
chat_completions = list()
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
chat_completions.append(future.result())
|
||||
# Verify E2E
|
||||
current_e2e = time.perf_counter() - start_time
|
||||
print(f"the request e2e is: {current_e2e}")
|
||||
# TODO: Verify the E2E latency after confirmation baseline.
|
||||
e2e_list.append(current_e2e)
|
||||
|
||||
print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
|
||||
# Verify all completions succeeded
|
||||
assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded."
|
||||
for chat_completion in chat_completions:
|
||||
# Verify audio output success
|
||||
audio_message = chat_completion.choices[1].message
|
||||
audio_data = audio_message.audio.data
|
||||
assert audio_data is not None, "No audio output is generated"
|
||||
assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
|
||||
|
||||
# Verify text output success
|
||||
text_choice = chat_completion.choices[0]
|
||||
text_content = text_choice.message.content
|
||||
assert text_content is not None, "No text output is generated"
|
||||
assert chat_completion.usage.completion_tokens <= 100, "The output length more than the requested max_tokens."
|
||||
assert "square" in text_content.lower(), "The output do not contain keywords."
|
||||
|
||||
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
def test_image_to_text_audio_001(test_config: tuple[str, str]) -> None:
|
||||
"""Test processing text, generating audio output via OpenAI API."""
|
||||
|
||||
model, stage_config_path = test_config
|
||||
num_concurrent_requests = 5
|
||||
stage_config_path = modify_stage_config(
|
||||
stage_config_path,
|
||||
{
|
||||
"stage_args": {
|
||||
0: {"runtime.max_batch_size": num_concurrent_requests},
|
||||
1: {"runtime.max_batch_size": num_concurrent_requests},
|
||||
}
|
||||
},
|
||||
)
|
||||
with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "90"]) as server:
|
||||
image_data_url = []
|
||||
for _ in range(4):
|
||||
image_data_url.append(f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}")
|
||||
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
|
||||
|
||||
# Test single completion
|
||||
api_client = client(server)
|
||||
e2e_list = list()
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
|
||||
# Submit multiple completion requests concurrently
|
||||
futures = [
|
||||
executor.submit(
|
||||
api_client.chat.completions.create,
|
||||
model=server.model,
|
||||
messages=messages,
|
||||
)
|
||||
for _ in range(num_concurrent_requests)
|
||||
]
|
||||
start_time = time.perf_counter()
|
||||
# Wait for all requests to complete and collect results
|
||||
chat_completions = list()
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
chat_completions.append(future.result())
|
||||
# Verify E2E
|
||||
current_e2e = time.perf_counter() - start_time
|
||||
print(f"the request e2e is: {current_e2e}")
|
||||
# TODO: Verify the E2E latency after confirmation baseline.
|
||||
e2e_list.append(current_e2e)
|
||||
|
||||
print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
|
||||
# Verify all completions succeeded
|
||||
assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded."
|
||||
for chat_completion in chat_completions:
|
||||
# Verify audio output success
|
||||
audio_message = chat_completion.choices[1].message
|
||||
audio_data = audio_message.audio.data
|
||||
assert audio_data is not None, "No audio output is generated"
|
||||
assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
|
||||
|
||||
# Verify text output success
|
||||
text_choice = chat_completion.choices[0]
|
||||
text_content = text_choice.message.content
|
||||
assert text_content is not None, "No text output is generated"
|
||||
assert "square" in text_content.lower(), "The output do not contain keywords."
|
||||
|
||||
# Verify text output same as audio output
|
||||
audio_content = convert_audio_to_text(audio_data)
|
||||
print(f"text content is: {text_content}")
|
||||
print(f"audio content is: {audio_content}")
|
||||
similarity = cosine_similarity_text(audio_content, text_content)
|
||||
print(f"similarity between audio and text is: {similarity}")
|
||||
assert similarity > 0.9, "The audio content is not same as the text"
|
||||
# Verify text output same as audio output
|
||||
audio_content = convert_audio_to_text(audio_data)
|
||||
print(f"text content is: {text_content}")
|
||||
print(f"audio content is: {audio_content}")
|
||||
similarity = cosine_similarity_text(audio_content, text_content)
|
||||
print(f"similarity between audio and text is: {similarity}")
|
||||
assert similarity > 0.9, "The audio content is not same as the text"
|
||||
@@ -7,14 +7,19 @@ from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from diffusers.models.attention import FeedForward
|
||||
from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
from diffusers.models.normalization import FP32LayerNorm
|
||||
from vllm.distributed import (
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.conv import Conv3dLayer
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinear
|
||||
from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
@@ -51,6 +56,94 @@ def apply_rotary_emb_wan(
|
||||
return out.type_as(hidden_states)
|
||||
|
||||
|
||||
class DistributedRMSNorm(nn.Module):
|
||||
"""
|
||||
RMSNorm that computes global RMS across tensor parallel ranks.
|
||||
This ensures mathematically equivalent results to non-TP execution.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size: int, eps: float = 1e-6):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
|
||||
input_dtype = x.dtype
|
||||
x_float = x.float()
|
||||
local_sum_sq = (x_float**2).sum(dim=-1, keepdim=True)
|
||||
local_count = x.shape[-1]
|
||||
|
||||
if tp_size > 1:
|
||||
global_sum_sq = local_sum_sq.clone()
|
||||
tensor_model_parallel_all_reduce(global_sum_sq)
|
||||
global_count = local_count * tp_size
|
||||
else:
|
||||
global_sum_sq = local_sum_sq
|
||||
global_count = local_count
|
||||
|
||||
rms = torch.sqrt(global_sum_sq / global_count + self.eps)
|
||||
|
||||
output = (x_float / rms) * self.weight.float()
|
||||
return output.to(input_dtype)
|
||||
|
||||
|
||||
class ColumnParallelGELU(nn.Module):
|
||||
"""Column parallel linear with GELU activation."""
|
||||
|
||||
def __init__(self, dim_in: int, dim_out: int, *, approximate: str = "tanh", bias: bool = True):
|
||||
super().__init__()
|
||||
self.proj = ColumnParallelLinear(
|
||||
dim_in,
|
||||
dim_out,
|
||||
bias=bias,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
)
|
||||
self.approximate = approximate
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.proj(x)
|
||||
return F.gelu(x, approximate=self.approximate)
|
||||
|
||||
|
||||
class WanFeedForward(nn.Module):
|
||||
"""
|
||||
TP-enabled FeedForward network for WAN2.2.
|
||||
Replaces diffusers FeedForward with ColumnParallel + RowParallel.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
inner_dim: int,
|
||||
dim_out: int | None = None,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
dim_out = dim_out or dim
|
||||
|
||||
# ColumnParallel: scatter to each tp_rank
|
||||
self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias)
|
||||
# Placeholder for weight loading compatibility
|
||||
self.net_1 = nn.Identity()
|
||||
# RowParallel: gather from each tp_rank
|
||||
self.net_2 = RowParallelLinear(
|
||||
inner_dim,
|
||||
dim_out,
|
||||
bias=bias,
|
||||
input_is_parallel=True,
|
||||
return_bias=False,
|
||||
)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
hidden_states = self.net_0(hidden_states)
|
||||
hidden_states = self.net_1(hidden_states)
|
||||
hidden_states = self.net_2(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class WanRotaryPosEmbed(nn.Module):
|
||||
"""
|
||||
Rotary position embeddings for 3D video data (temporal + spatial dimensions).
|
||||
@@ -97,8 +190,8 @@ class WanRotaryPosEmbed(nn.Module):
|
||||
t = torch.arange(max_seq_len, dtype=freqs_dtype)
|
||||
freqs = torch.outer(t, freqs)
|
||||
# Repeat interleave for real representation
|
||||
freqs_cos = freqs.cos().repeat_interleave(2, dim=-1)
|
||||
freqs_sin = freqs.sin().repeat_interleave(2, dim=-1)
|
||||
freqs_cos = freqs.cos().float().repeat_interleave(2, dim=-1)
|
||||
freqs_sin = freqs.sin().float().repeat_interleave(2, dim=-1)
|
||||
return freqs_cos.float(), freqs_sin.float()
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
@@ -229,25 +322,30 @@ class WanSelfAttention(nn.Module):
|
||||
head_size=head_dim,
|
||||
total_num_heads=num_heads,
|
||||
bias=True,
|
||||
disable_tp=True,
|
||||
)
|
||||
|
||||
self.num_heads = self.to_qkv.num_heads
|
||||
self.num_kv_heads = self.to_qkv.num_kv_heads
|
||||
self.tp_inner_dim = self.num_heads * head_dim
|
||||
|
||||
# QK normalization using vLLM's RMSNorm
|
||||
self.norm_q = RMSNorm(self.inner_dim, eps=eps)
|
||||
self.norm_k = RMSNorm(self.inner_dim, eps=eps)
|
||||
self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
|
||||
self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
|
||||
|
||||
# Output projection
|
||||
self.to_out = nn.ModuleList(
|
||||
[
|
||||
ReplicatedLinear(self.inner_dim, dim, bias=True),
|
||||
nn.Dropout(dropout),
|
||||
]
|
||||
self.to_out = RowParallelLinear(
|
||||
self.inner_dim,
|
||||
dim,
|
||||
bias=True,
|
||||
input_is_parallel=True,
|
||||
return_bias=False,
|
||||
)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
# Unified attention layer
|
||||
self.attn = Attention(
|
||||
num_heads=num_heads,
|
||||
num_heads=self.num_heads,
|
||||
head_size=head_dim,
|
||||
num_kv_heads=self.num_kv_heads,
|
||||
softmax_scale=1.0 / (head_dim**0.5),
|
||||
causal=False,
|
||||
)
|
||||
@@ -259,16 +357,19 @@ class WanSelfAttention(nn.Module):
|
||||
) -> torch.Tensor:
|
||||
# Fused QKV projection
|
||||
qkv, _ = self.to_qkv(hidden_states)
|
||||
query, key, value = qkv.chunk(3, dim=-1)
|
||||
|
||||
q_size = self.num_heads * self.head_dim
|
||||
kv_size = self.num_kv_heads * self.head_dim
|
||||
query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
|
||||
|
||||
# Apply QK normalization
|
||||
query = self.norm_q(query)
|
||||
key = self.norm_k(key)
|
||||
|
||||
# Reshape for multi-head attention
|
||||
query = query.unflatten(2, (self.num_heads, -1))
|
||||
key = key.unflatten(2, (self.num_heads, -1))
|
||||
value = value.unflatten(2, (self.num_heads, -1))
|
||||
query = query.unflatten(2, (self.num_heads, self.head_dim))
|
||||
key = key.unflatten(2, (self.num_kv_heads, self.head_dim))
|
||||
value = value.unflatten(2, (self.num_kv_heads, self.head_dim))
|
||||
|
||||
# Apply rotary embeddings
|
||||
if rotary_emb is not None:
|
||||
@@ -282,8 +383,8 @@ class WanSelfAttention(nn.Module):
|
||||
hidden_states = hidden_states.type_as(query)
|
||||
|
||||
# Output projection
|
||||
hidden_states, _ = self.to_out[0](hidden_states)
|
||||
hidden_states = self.to_out[1](hidden_states)
|
||||
hidden_states = self.to_out(hidden_states)
|
||||
hidden_states = self.dropout(hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
@@ -312,39 +413,77 @@ class WanCrossAttention(nn.Module):
|
||||
self.kv_inner_dim = head_dim * num_heads # For cross-attention, K/V come from encoder
|
||||
|
||||
# Query projection
|
||||
self.to_q = ReplicatedLinear(dim, self.inner_dim, bias=True)
|
||||
self.to_q = ColumnParallelLinear(
|
||||
dim,
|
||||
self.inner_dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
)
|
||||
|
||||
# Separate K and V projections for cross-attention
|
||||
self.to_k = ReplicatedLinear(dim, self.kv_inner_dim, bias=True)
|
||||
self.to_v = ReplicatedLinear(dim, self.kv_inner_dim, bias=True)
|
||||
self.to_k = ColumnParallelLinear(
|
||||
dim,
|
||||
self.kv_inner_dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
)
|
||||
|
||||
self.to_v = ColumnParallelLinear(
|
||||
dim,
|
||||
self.kv_inner_dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
)
|
||||
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
self.num_heads = num_heads // tp_size
|
||||
self.tp_inner_dim = self.num_heads * head_dim
|
||||
|
||||
# QK normalization
|
||||
self.norm_q = RMSNorm(self.inner_dim, eps=eps)
|
||||
self.norm_k = RMSNorm(self.kv_inner_dim, eps=eps)
|
||||
self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
|
||||
self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
|
||||
|
||||
# Optional added KV projections for I2V (image embeddings)
|
||||
self.added_kv_proj_dim = added_kv_proj_dim
|
||||
if added_kv_proj_dim is not None:
|
||||
self.add_k_proj = ReplicatedLinear(added_kv_proj_dim, self.inner_dim, bias=True)
|
||||
self.add_v_proj = ReplicatedLinear(added_kv_proj_dim, self.inner_dim, bias=True)
|
||||
self.norm_added_k = RMSNorm(self.inner_dim, eps=eps)
|
||||
self.add_k_proj = ColumnParallelLinear(
|
||||
added_kv_proj_dim,
|
||||
self.inner_dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
)
|
||||
self.add_v_proj = ColumnParallelLinear(
|
||||
added_kv_proj_dim,
|
||||
self.inner_dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
)
|
||||
self.norm_added_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
|
||||
else:
|
||||
self.add_k_proj = None
|
||||
self.add_v_proj = None
|
||||
self.norm_added_k = None
|
||||
|
||||
# Output projection
|
||||
self.to_out = nn.ModuleList(
|
||||
[
|
||||
ReplicatedLinear(self.inner_dim, dim, bias=True),
|
||||
nn.Dropout(dropout),
|
||||
]
|
||||
self.to_out = RowParallelLinear(
|
||||
self.inner_dim,
|
||||
dim,
|
||||
bias=True,
|
||||
input_is_parallel=True,
|
||||
return_bias=False,
|
||||
)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
# Unified attention layer
|
||||
self.attn = Attention(
|
||||
num_heads=num_heads,
|
||||
num_heads=self.num_heads,
|
||||
head_size=head_dim,
|
||||
num_kv_heads=self.num_heads,
|
||||
softmax_scale=1.0 / (head_dim**0.5),
|
||||
causal=False,
|
||||
)
|
||||
@@ -363,28 +502,28 @@ class WanCrossAttention(nn.Module):
|
||||
encoder_hidden_states = encoder_hidden_states[:, image_context_length:]
|
||||
|
||||
# Query projection
|
||||
query, _ = self.to_q(hidden_states)
|
||||
query = self.to_q(hidden_states)
|
||||
query = self.norm_q(query)
|
||||
|
||||
# KV projection from encoder
|
||||
key, _ = self.to_k(encoder_hidden_states)
|
||||
value, _ = self.to_v(encoder_hidden_states)
|
||||
key = self.to_k(encoder_hidden_states)
|
||||
value = self.to_v(encoder_hidden_states)
|
||||
key = self.norm_k(key)
|
||||
|
||||
# Reshape for multi-head attention
|
||||
query = query.unflatten(2, (self.num_heads, -1))
|
||||
key = key.unflatten(2, (self.num_heads, -1))
|
||||
value = value.unflatten(2, (self.num_heads, -1))
|
||||
query = query.unflatten(2, (self.num_heads, self.head_dim))
|
||||
key = key.unflatten(2, (self.num_heads, self.head_dim))
|
||||
value = value.unflatten(2, (self.num_heads, self.head_dim))
|
||||
|
||||
# I2V: Additional attention with image embeddings
|
||||
hidden_states_img = None
|
||||
if encoder_hidden_states_img is not None:
|
||||
key_img, _ = self.add_k_proj(encoder_hidden_states_img)
|
||||
value_img, _ = self.add_v_proj(encoder_hidden_states_img)
|
||||
key_img = self.add_k_proj(encoder_hidden_states_img)
|
||||
value_img = self.add_v_proj(encoder_hidden_states_img)
|
||||
key_img = self.norm_added_k(key_img)
|
||||
|
||||
key_img = key_img.unflatten(2, (self.num_heads, -1))
|
||||
value_img = value_img.unflatten(2, (self.num_heads, -1))
|
||||
key_img = key_img.unflatten(2, (self.num_heads, self.head_dim))
|
||||
value_img = value_img.unflatten(2, (self.num_heads, self.head_dim))
|
||||
|
||||
hidden_states_img = self.attn(query, key_img, value_img)
|
||||
hidden_states_img = hidden_states_img.flatten(2, 3)
|
||||
@@ -400,8 +539,8 @@ class WanCrossAttention(nn.Module):
|
||||
hidden_states = hidden_states + hidden_states_img
|
||||
|
||||
# Output projection
|
||||
hidden_states, _ = self.to_out[0](hidden_states)
|
||||
hidden_states = self.to_out[1](hidden_states)
|
||||
hidden_states = self.to_out(hidden_states)
|
||||
hidden_states = self.dropout(hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
@@ -445,7 +584,7 @@ class WanTransformerBlock(nn.Module):
|
||||
self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
|
||||
|
||||
# 3. Feed-forward
|
||||
self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate")
|
||||
self.ffn = WanFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim)
|
||||
self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)
|
||||
|
||||
# Scale-shift table for modulation
|
||||
@@ -725,11 +864,13 @@ class WanTransformer3DModel(nn.Module):
|
||||
Returns:
|
||||
Set of parameter names that were successfully loaded.
|
||||
"""
|
||||
tp_rank = get_tensor_model_parallel_rank()
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
# Stacked params mapping for self-attention QKV fusion
|
||||
# Format: (param_name, shard_name, shard_id)
|
||||
# Note: Only fuse attn1 (self-attention), NOT attn2 (cross-attention)
|
||||
stacked_params_mapping = [
|
||||
# self-attention QKV fusion (attn1 only)
|
||||
# self-attention QKV fusion
|
||||
(".attn1.to_qkv", ".attn1.to_q", "q"),
|
||||
(".attn1.to_qkv", ".attn1.to_k", "k"),
|
||||
(".attn1.to_qkv", ".attn1.to_v", "v"),
|
||||
@@ -739,18 +880,54 @@ class WanTransformer3DModel(nn.Module):
|
||||
loaded_params: set[str] = set()
|
||||
|
||||
for name, loaded_weight in weights:
|
||||
original_name = name
|
||||
lookup_name = name
|
||||
|
||||
# Handle QKV fusion
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
if weight_name not in name:
|
||||
if weight_name not in original_name:
|
||||
continue
|
||||
name = name.replace(weight_name, param_name)
|
||||
param = params_dict[name]
|
||||
lookup_name = original_name.replace(weight_name, param_name)
|
||||
param = params_dict[lookup_name]
|
||||
weight_loader = param.weight_loader
|
||||
weight_loader(param, loaded_weight, shard_id)
|
||||
break
|
||||
else:
|
||||
param = params_dict[name]
|
||||
# diffusers: ffn.net.0.proj.weight -> our: ffn.net_0.proj.weight
|
||||
if ".ffn.net.0." in lookup_name:
|
||||
lookup_name = lookup_name.replace(".ffn.net.0.", ".ffn.net_0.")
|
||||
elif ".ffn.net.2." in lookup_name:
|
||||
lookup_name = lookup_name.replace(".ffn.net.2.", ".ffn.net_2.")
|
||||
|
||||
if ".to_out.0." in lookup_name:
|
||||
lookup_name = lookup_name.replace(".to_out.0.", ".to_out.")
|
||||
|
||||
if lookup_name not in params_dict:
|
||||
logger.warning(f"Skipping weight {original_name} -> {lookup_name}")
|
||||
continue
|
||||
|
||||
param = params_dict[lookup_name]
|
||||
|
||||
# Handle RMSNorm weights that need to be sharded for TP
|
||||
# These norms are applied after ColumnParallelLinear outputs,
|
||||
# so their weights must be sharded to match the sharded hidden dim
|
||||
if tp_size > 1 and any(
|
||||
norm_name in lookup_name
|
||||
for norm_name in [
|
||||
".attn1.norm_q.",
|
||||
".attn1.norm_k.",
|
||||
".attn2.norm_q.",
|
||||
".attn2.norm_k.",
|
||||
".attn2.norm_added_k.",
|
||||
]
|
||||
):
|
||||
shard_size = loaded_weight.shape[0] // tp_size
|
||||
loaded_weight = loaded_weight[tp_rank * shard_size : (tp_rank + 1) * shard_size]
|
||||
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, loaded_weight)
|
||||
loaded_params.add(name)
|
||||
|
||||
loaded_params.add(original_name)
|
||||
loaded_params.add(lookup_name)
|
||||
|
||||
return loaded_params
|
||||
@@ -3,7 +3,6 @@ from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
|
||||
Qwen3OmniMoeTalkerConfig,
|
||||
)
|
||||
@@ -11,12 +10,8 @@ from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
|
||||
Qwen3OmniMoeAudioEncoder,
|
||||
)
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.distributed import get_tensor_model_parallel_world_size
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY
|
||||
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
|
||||
from vllm.model_executor.layers.linear import ReplicatedLinear
|
||||
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||
from vllm.model_executor.models.interfaces import (
|
||||
MultiModalEmbeddings,
|
||||
SupportsPP,
|
||||
@@ -24,7 +19,6 @@ from vllm.model_executor.models.interfaces import (
|
||||
from vllm.model_executor.models.qwen2_5_omni_thinker import (
|
||||
Qwen2_5OmniThinkerDummyInputsBuilder,
|
||||
)
|
||||
from vllm.model_executor.models.qwen3_moe import Qwen3MoeMLP, Qwen3MoeSparseMoeBlock
|
||||
from vllm.model_executor.models.qwen3_omni_moe_thinker import Qwen3Omni_VisionTransformer
|
||||
from vllm.model_executor.models.utils import (
|
||||
AutoWeightsLoader,
|
||||
@@ -512,162 +506,14 @@ class Qwen3OmniMoeTalkerResizeMLP(nn.Module):
|
||||
return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))
|
||||
|
||||
|
||||
class Qwen3OmniMoeTalkerSharedExpertWrapper(nn.Module):
|
||||
"""
|
||||
Wrapper that combines shared_expert MLP with its sigmoid gate.
|
||||
|
||||
This matches the HuggingFace weight structure where:
|
||||
- mlp.shared_expert.{gate_proj, up_proj, down_proj}.weight
|
||||
- mlp.shared_expert_gate.weight (sibling, not child)
|
||||
|
||||
The wrapper applies: sigmoid(shared_expert_gate(x)) * shared_expert(x).
|
||||
|
||||
It also exposes the underlying shared_expert interface to keep
|
||||
compatibility with backends that split shared-expert computation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
shared_expert: Qwen3MoeMLP,
|
||||
shared_expert_gate: nn.Linear,
|
||||
):
|
||||
super().__init__()
|
||||
self._shared_expert = shared_expert
|
||||
self._shared_expert_gate = shared_expert_gate
|
||||
|
||||
@property
|
||||
def gate_up_proj(self):
|
||||
return self._shared_expert.gate_up_proj
|
||||
|
||||
@property
|
||||
def down_proj(self):
|
||||
return self._shared_expert.down_proj
|
||||
|
||||
@property
|
||||
def act_fn(self):
|
||||
return self._shared_expert.act_fn
|
||||
|
||||
def expert_gate(self, x: torch.Tensor):
|
||||
gate_out = self._shared_expert_gate(x)
|
||||
if isinstance(gate_out, tuple):
|
||||
return gate_out
|
||||
return gate_out, None
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
out = self._shared_expert(x)
|
||||
gate_out = self._shared_expert_gate(x)
|
||||
if isinstance(gate_out, tuple):
|
||||
gate_out = gate_out[0]
|
||||
gate_values = F.sigmoid(gate_out) # [batch, 1]
|
||||
return gate_values * out # Broadcasting: [batch, 1] * [batch, hidden]
|
||||
|
||||
|
||||
class Qwen3OmniMoeTalkerSparseMoeBlock(nn.Module):
|
||||
"""
|
||||
Sparse MoE block for Qwen3 Omni MoE Talker with shared expert support.
|
||||
|
||||
This block uses SharedFusedMoE to efficiently compute both routed experts
|
||||
and the shared expert, potentially overlapping computation with communication.
|
||||
|
||||
Weight structure matches HuggingFace:
|
||||
- mlp.gate.weight (router)
|
||||
- mlp.shared_expert.{gate_proj, up_proj, down_proj}.weight
|
||||
- mlp.shared_expert_gate.weight
|
||||
- mlp.experts.{0..n}.{gate_proj, up_proj, down_proj}.weight
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: Qwen3OmniMoeTalkerConfig,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
text_config = config.text_config
|
||||
self.tp_size = get_tensor_model_parallel_world_size()
|
||||
|
||||
if self.tp_size > text_config.num_experts:
|
||||
raise ValueError(
|
||||
f"Tensor parallel size {self.tp_size} is greater than the number of experts {text_config.num_experts}."
|
||||
)
|
||||
|
||||
# Router gate for selecting top-k experts
|
||||
self.gate = ReplicatedLinear(
|
||||
text_config.hidden_size,
|
||||
text_config.num_experts,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.gate",
|
||||
)
|
||||
|
||||
# Shared expert MLP (matches HF: mlp.shared_expert.*)
|
||||
if text_config.shared_expert_intermediate_size > 0:
|
||||
self.shared_expert = Qwen3MoeMLP(
|
||||
hidden_size=text_config.hidden_size,
|
||||
intermediate_size=text_config.shared_expert_intermediate_size,
|
||||
hidden_act=text_config.hidden_act,
|
||||
quant_config=quant_config,
|
||||
reduce_results=False, # Don't reduce, we'll handle it
|
||||
prefix=f"{prefix}.shared_expert",
|
||||
)
|
||||
# Shared expert gate (matches HF: mlp.shared_expert_gate.weight)
|
||||
# This is a sibling of shared_expert, not a child
|
||||
self.shared_expert_gate = torch.nn.Linear(text_config.hidden_size, 1, bias=False)
|
||||
# Create wrapper for SharedFusedMoE
|
||||
self._shared_expert_wrapper = Qwen3OmniMoeTalkerSharedExpertWrapper(
|
||||
self.shared_expert, self.shared_expert_gate
|
||||
)
|
||||
else:
|
||||
self.shared_expert = None
|
||||
self.shared_expert_gate = None
|
||||
self._shared_expert_wrapper = None
|
||||
|
||||
# Fused MoE with shared expert support
|
||||
self.experts = SharedFusedMoE(
|
||||
shared_experts=self._shared_expert_wrapper,
|
||||
num_experts=text_config.num_experts,
|
||||
top_k=text_config.num_experts_per_tok,
|
||||
hidden_size=text_config.hidden_size,
|
||||
intermediate_size=text_config.moe_intermediate_size,
|
||||
reduce_results=False, # We'll reduce manually after combining
|
||||
renormalize=text_config.norm_topk_prob,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.experts",
|
||||
)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
# NOTE: hidden_states can have either 1D or 2D shape.
|
||||
orig_shape = hidden_states.shape
|
||||
hidden_dim = hidden_states.shape[-1]
|
||||
hidden_states = hidden_states.view(-1, hidden_dim)
|
||||
|
||||
# Compute router logits
|
||||
router_logits, _ = self.gate(hidden_states)
|
||||
|
||||
# Forward through SharedFusedMoE
|
||||
# Returns (shared_out, fused_out) when shared_expert is present
|
||||
final_hidden_states = self.experts(hidden_states=hidden_states, router_logits=router_logits)
|
||||
|
||||
# Combine shared and routed expert outputs
|
||||
if self._shared_expert_wrapper is not None:
|
||||
# SharedFusedMoE returns tuple: (shared_out, fused_out)
|
||||
final_hidden_states = final_hidden_states[0] + final_hidden_states[1]
|
||||
|
||||
# Apply tensor parallel reduction if needed
|
||||
if self.tp_size > 1:
|
||||
final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(final_hidden_states)
|
||||
|
||||
return final_hidden_states.view(orig_shape)
|
||||
|
||||
|
||||
class Qwen3OmniMoeModel(Qwen3MoeLLMForCausalLM):
|
||||
"""
|
||||
Qwen3 Omni MoE Talker language model.
|
||||
|
||||
This model extends Qwen3MoeLLMForCausalLM with:
|
||||
- Shared expert support via SharedFusedMoE
|
||||
- Codec embedding instead of text embedding
|
||||
- No LM head (codec head is separate in the parent class)
|
||||
Extends Qwen3MoeLLMForCausalLM (which already uses SharedFusedMoE with
|
||||
shared-expert support) and replaces the text embedding / LM head with a
|
||||
codec embedding so the talker operates over audio-codec tokens instead
|
||||
of text tokens.
|
||||
"""
|
||||
|
||||
def __init__(self, vllm_config: VllmConfig, talker_config: Qwen3OmniMoeTalkerConfig, prefix: str):
|
||||
@@ -699,32 +545,6 @@ class Qwen3OmniMoeModel(Qwen3MoeLLMForCausalLM):
|
||||
talker_config.text_config.hidden_size,
|
||||
)
|
||||
|
||||
# Replace MoE blocks with shared expert versions
|
||||
self._replace_moe_blocks_with_shared_expert(prefix)
|
||||
|
||||
def _replace_moe_blocks_with_shared_expert(self, prefix: str) -> None:
|
||||
"""
|
||||
Replace Qwen3MoeSparseMoeBlock layers with Qwen3OmniMoeTalkerSparseMoeBlock
|
||||
that includes shared expert support via SharedFusedMoE.
|
||||
"""
|
||||
# Get compilation config to clean up registered layer names
|
||||
compilation_config = self.talker_vllm_config.compilation_config
|
||||
|
||||
for layer_idx, layer in enumerate(self.model.layers):
|
||||
# Check if this layer has a MoE block (has experts attribute)
|
||||
if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock):
|
||||
# Remove old layer registration from static_forward_context
|
||||
old_experts_prefix = f"{prefix}.model.layers.{layer_idx}.mlp.experts"
|
||||
if old_experts_prefix in compilation_config.static_forward_context:
|
||||
del compilation_config.static_forward_context[old_experts_prefix]
|
||||
|
||||
# Create new MoE block with shared expert support
|
||||
layer.mlp = Qwen3OmniMoeTalkerSparseMoeBlock(
|
||||
config=self.config,
|
||||
quant_config=self.talker_vllm_config.quant_config,
|
||||
prefix=f"{prefix}.model.layers.{layer_idx}.mlp",
|
||||
)
|
||||
|
||||
def embed_input_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
|
||||
Reference in new issue
Block a user