diff --git a/tests/e2e/offline_inference/test_bagel_lora.py b/tests/e2e/offline_inference/test_bagel_lora.py new file mode 100644 index 000000000..593a64047 --- /dev/null +++ b/tests/e2e/offline_inference/test_bagel_lora.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +End-to-end test for BAGEL LoRA support (Stage 1 / DiT). + +Validates that LoRA adapters are correctly loaded, applied with controllable +scale, and cleanly deactivated. Uses a synthetic rank-1 adapter targeting the +first decoder layer's QKV projection. + +Assertions: + (a) LoRA at scale=1.0 visibly changes the output (diff > 0.5) + (b) scale=2.0 produces a larger delta than scale=1.0 (linearity) + (c) The delta is bounded (diff < 80, not corrupted) + (d) Deactivating LoRA exactly restores the baseline (diff == 0) +""" + +import json +import os + +from vllm_omni.inputs.data import OmniSamplingParams +from vllm_omni.outputs import OmniRequestOutput + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + +from pathlib import Path + +import numpy as np +import pytest +import torch +from PIL import Image +from safetensors.torch import save_file + +from tests.conftest import modify_stage_config +from tests.utils import hardware_test +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.lora.request import LoRARequest +from vllm_omni.lora.utils import stable_lora_int_id + +MODEL = "ByteDance-Seed/BAGEL-7B-MoT" +BAGEL_STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml") +DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>" + + +# --------------------------------------------------------------------------- +# Helpers (reused from test_bagel_text2img.py patterns) +# --------------------------------------------------------------------------- + + +def _resolve_stage_config(config_path: str, run_level: str) -> str: + if run_level == "advanced_model": + return modify_stage_config( + config_path, + deletes={ + "stage_args": { + 0: ["engine_args.load_format"], + 1: ["engine_args.load_format"], + } + }, + ) + return config_path + + +def _configure_sampling_params(omni: Omni, num_inference_steps: int = 10) -> list[OmniSamplingParams]: + params_list = omni.default_sampling_params_list + if len(params_list) > 1: + params_list[1].num_inference_steps = num_inference_steps + params_list[1].extra_args = { + "cfg_text_scale": 4.0, + "cfg_img_scale": 1.5, + } + return params_list + + +def _extract_generated_image(omni_outputs: list[OmniRequestOutput]) -> Image.Image | None: + for req_output in omni_outputs: + if req_output.images: + return req_output.images[0] + return None + + +def _generate_bagel_image(omni: Omni) -> Image.Image: + params_list = _configure_sampling_params(omni) + params_list[1].lora_request = None + outputs = list( + omni.generate( + prompts=[{"prompt": DEFAULT_PROMPT, "modalities": ["image"]}], + sampling_params_list=params_list, + ) + ) + img = _extract_generated_image(outputs) + assert img is not None, "No image generated" + return img + + +def _generate_bagel_image_with_lora( + omni: Omni, + lora_request: LoRARequest, + lora_scale: float = 1.0, +) -> Image.Image: + params_list = _configure_sampling_params(omni) + params_list[1].lora_request = lora_request + params_list[1].lora_scale = lora_scale + outputs = list( + omni.generate( + prompts=[{"prompt": DEFAULT_PROMPT, "modalities": ["image"]}], + sampling_params_list=params_list, + ) + ) + img = _extract_generated_image(outputs) + assert img is not None, "No image generated with LoRA" + return img + + +# BAGEL uses GQA: hidden_size=3584, 28 Q heads, 4 KV heads, head_dim=128 +# QKV packed dim = 28*128 + 4*128 + 4*128 = 3584 + 512 + 512 = 4608 +_LORA_DIM = 3584 +_LORA_QKV_DIM = 4608 +_LORA_MODULE = "bagel.language_model.model.layers.0.self_attn.qkv_proj" +_LORA_RANK = 4 + + +def _make_file_lora_request(adapter_dir: Path) -> LoRARequest: + """Write synthetic adapter to disk and return a file-backed LoRARequest.""" + adapter_dir.mkdir(parents=True, exist_ok=True) + gen = torch.Generator().manual_seed(42) + lora_a = torch.randn((_LORA_RANK, _LORA_DIM), dtype=torch.float32, generator=gen) * 0.1 + lora_b = torch.randn((_LORA_QKV_DIM, _LORA_RANK), dtype=torch.float32, generator=gen) * 0.5 + save_file( + { + f"base_model.model.{_LORA_MODULE}.lora_A.weight": lora_a, + f"base_model.model.{_LORA_MODULE}.lora_B.weight": lora_b, + }, + str(adapter_dir / "adapter_model.safetensors"), + ) + (adapter_dir / "adapter_config.json").write_text( + json.dumps({"r": _LORA_RANK, "lora_alpha": _LORA_RANK, "target_modules": [_LORA_MODULE]}), + encoding="utf-8", + ) + lora_dir = str(adapter_dir) + return LoRARequest(lora_name="test_file", lora_int_id=stable_lora_int_id(lora_dir), lora_path=lora_dir) + + +# --------------------------------------------------------------------------- +# Test +# --------------------------------------------------------------------------- + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) +def test_bagel_lora_scale_and_deactivation(run_level, tmp_path): + """Validate LoRA effect, bounded perturbation, and clean deactivation.""" + config_path = _resolve_stage_config(BAGEL_STAGE_CONFIG, run_level) + omni = Omni(model=MODEL, stage_configs_path=config_path, stage_init_timeout=300) + try: + lora_request = _make_file_lora_request(tmp_path / "bagel_lora") + + # 1) Baseline (no LoRA) + baseline = _generate_bagel_image(omni) + + # 2) LoRA with scale=1.0 + img_1x = _generate_bagel_image_with_lora(omni, lora_request, lora_scale=1.0) + + # 3) LoRA with scale=2.0 + img_2x = _generate_bagel_image_with_lora(omni, lora_request, lora_scale=2.0) + + # 4) No LoRA again (deactivation) + restored = _generate_bagel_image(omni) + + baseline_arr = np.array(baseline, dtype=np.int16) + img_1x_arr = np.array(img_1x, dtype=np.int16) + img_2x_arr = np.array(img_2x, dtype=np.int16) + restored_arr = np.array(restored, dtype=np.int16) + + diff_1x = np.abs(baseline_arr - img_1x_arr).mean() + diff_2x = np.abs(baseline_arr - img_2x_arr).mean() + diff_restored = np.abs(baseline_arr - restored_arr).mean() + + # (a) Adapter has visible effect at both scales + assert diff_1x > 0.5, f"LoRA scale=1.0 had no visible effect: diff={diff_1x}" + assert diff_2x > 0.5, f"LoRA scale=2.0 had no visible effect: diff={diff_2x}" + + # (b) Different scales produce different outputs + assert not np.isclose(diff_1x, diff_2x, atol=1.0), ( + f"LoRA scale has no effect: diff_1x={diff_1x:.2f}, diff_2x={diff_2x:.2f}" + ) + + # (c) Output is not corrupted + assert diff_1x < 80, f"LoRA output looks corrupted: diff_1x={diff_1x}" + assert diff_2x < 80, f"LoRA output looks corrupted: diff_2x={diff_2x}" + + # (d) Deactivation fully restores base model + assert diff_restored == 0.0, f"Base model not restored after LoRA deactivation: diff={diff_restored}" + finally: + omni.close() diff --git a/tests/engine/test_cross_stage_lora.py b/tests/engine/test_cross_stage_lora.py new file mode 100644 index 000000000..1eccc5526 --- /dev/null +++ b/tests/engine/test_cross_stage_lora.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for cross-stage LoRA routing in the orchestrator.""" + +from __future__ import annotations + +import pytest +from vllm.lora.request import LoRARequest +from vllm.sampling_params import SamplingParams + +from vllm_omni.engine.orchestrator import build_engine_core_request_from_tokens +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestBuildEngineCoreRequestLoRA: + """Verify build_engine_core_request_from_tokens passes LoRA from params.""" + + def test_lora_extracted_from_diffusion_params(self): + lr = LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/fake") + params = OmniDiffusionSamplingParams(lora_request=lr) + + # OmniDiffusionSamplingParams is not a SamplingParams, so + # build_engine_core_request_from_tokens takes the pooling path. + # We only care that lora_request is extracted via getattr. + request = build_engine_core_request_from_tokens( + request_id="req-1", + prompt={"prompt_token_ids": [1, 2, 3]}, + params=params, + model_config=None, + ) + assert request.lora_request is lr + + def test_no_lora_on_sampling_params(self): + params = SamplingParams(max_tokens=10) + + request = build_engine_core_request_from_tokens( + request_id="req-2", + prompt={"prompt_token_ids": [1, 2, 3]}, + params=params, + model_config=None, + ) + assert request.lora_request is None diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 20dce1f0f..386b545eb 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -79,7 +79,7 @@ def build_engine_core_request_from_tokens( sampling_params=sampling_params, pooling_params=pooling_params, arrival_time=arrival_time, - lora_request=None, + lora_request=getattr(params, "lora_request", None), cache_salt=None, data_parallel_rank=None, prompt_embeds=prompt_embeds,