diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 402e62518..e02d49472 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -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: diff --git a/.buildkite/test-nightly.yaml b/.buildkite/test-nightly.yaml new file mode 100644 index 000000000..27fdc8378 --- /dev/null +++ b/.buildkite/test-nightly.yaml @@ -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 diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index ab58ebec0..c555372d9 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -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 diff --git a/docker/Dockerfile.npu.a3 b/docker/Dockerfile.npu.a3 index 17515fdb9..413ed88c3 100644 --- a/docker/Dockerfile.npu.a3 +++ b/docker/Dockerfile.npu.a3 @@ -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 diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index bbb75a196..d0af83423 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -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 diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md index 73c974280..141256dbc 100644 --- a/docs/getting_started/installation/gpu.md +++ b/docs/getting_started/installation/gpu.md @@ -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" diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index dfbfa1bf1..09323cd25 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -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 diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index 1a8ffb612..f3b893f8e 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -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 +``` + +
(Optional) Installation of vLLM from source +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] diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index ff71f4009..c473820fe 100644 --- a/docs/getting_started/installation/npu/npu.inc.md +++ b/docs/getting_started/installation/npu/npu.inc.md @@ -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 ``` diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index 6e2c18d64..f91a626e0 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -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 diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index e9dd2d085..903fc001f 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -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") diff --git a/pyproject.toml b/pyproject.toml index de35d45a3..706d0152e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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.*", ] diff --git a/requirements/common.txt b/requirements/common.txt new file mode 100644 index 000000000..19ebca949 --- /dev/null +++ b/requirements/common.txt @@ -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 diff --git a/requirements/cpu.txt b/requirements/cpu.txt new file mode 100644 index 000000000..c7732b1e1 --- /dev/null +++ b/requirements/cpu.txt @@ -0,0 +1,2 @@ +-r common.txt +onnxruntime>=1.23.2 diff --git a/requirements/cuda.txt b/requirements/cuda.txt new file mode 100644 index 000000000..823f47235 --- /dev/null +++ b/requirements/cuda.txt @@ -0,0 +1,3 @@ +-r common.txt +onnxruntime>=1.23.2 +fa3-fwd==0.0.1 diff --git a/requirements/npu.txt b/requirements/npu.txt new file mode 100644 index 000000000..955742757 --- /dev/null +++ b/requirements/npu.txt @@ -0,0 +1,3 @@ +-r common.txt +onnxruntime-cann>=1.23.2 +torchaudio==2.9.0 diff --git a/requirements/rocm.txt b/requirements/rocm.txt new file mode 100644 index 000000000..849610eec --- /dev/null +++ b/requirements/rocm.txt @@ -0,0 +1,2 @@ +-r common.txt +onnxruntime-rocm>=1.22.2 diff --git a/requirements/xpu.txt b/requirements/xpu.txt new file mode 100644 index 000000000..c7732b1e1 --- /dev/null +++ b/requirements/xpu.txt @@ -0,0 +1,2 @@ +-r common.txt +onnxruntime>=1.23.2 diff --git a/setup.py b/setup.py new file mode 100644 index 000000000..e33c848a3 --- /dev/null +++ b/setup.py @@ -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, + ) diff --git a/tests/conftest.py b/tests/conftest.py index b1d7e6cc8..84fbb91b7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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") diff --git a/tests/e2e/offline_inference/test_cache_dit.py b/tests/e2e/offline_inference/test_cache_dit.py index 18d3988e0..281ac48dc 100644 --- a/tests/e2e/offline_inference/test_cache_dit.py +++ b/tests/e2e/offline_inference/test_cache_dit.py @@ -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 ), ) diff --git a/tests/e2e/offline_inference/test_t2i_model.py b/tests/e2e/offline_inference/test_t2i_model.py index e7351f157..74c1bd1ce 100644 --- a/tests/e2e/offline_inference/test_t2i_model.py +++ b/tests/e2e/offline_inference/test_t2i_model.py @@ -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, ), ) diff --git a/tests/e2e/offline_inference/test_t2v_model.py b/tests/e2e/offline_inference/test_t2v_model.py index a378291ac..9565534bc 100644 --- a/tests/e2e/offline_inference/test_t2v_model.py +++ b/tests/e2e/offline_inference/test_t2v_model.py @@ -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] diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index 5db07fd89..9e84c7559 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -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" diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index ab92ad0c8..328a64024 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -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 diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 946694ed2..ac7d8120d 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -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,