From 486d77d7970e6deb88fab915e224c7659e379e10 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zhengyuan=20Su=20=28=E8=8B=8F=E6=94=BF=E6=B8=8A=29?= Date: Mon, 6 Apr 2026 19:30:52 +0800 Subject: [PATCH] [Feature] Enable LoRA adapter injection for BAGEL (#2490) Signed-off-by: Zhengyuan Su Signed-off-by: Claude Co-authored-by: Claude --- tests/diffusion/lora/conftest.py | 56 ++++ tests/diffusion/lora/test_lora_manager.py | 80 +++--- .../diffusion/models/bagel/test_bagel_lora.py | 248 ++++++++++++++++++ vllm_omni/diffusion/lora/manager.py | 9 +- .../models/bagel/bagel_transformer.py | 14 + .../model_executor/models/bagel/bagel.py | 16 ++ 6 files changed, 392 insertions(+), 31 deletions(-) create mode 100644 tests/diffusion/lora/conftest.py create mode 100644 tests/diffusion/models/bagel/test_bagel_lora.py diff --git a/tests/diffusion/lora/conftest.py b/tests/diffusion/lora/conftest.py new file mode 100644 index 000000000..8b9b1ef4d --- /dev/null +++ b/tests/diffusion/lora/conftest.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared test helpers for diffusion LoRA tests.""" + +from __future__ import annotations + +import torch +from vllm.model_executor.layers.linear import LinearBase + + +class FakeLinearBase(LinearBase): + """Minimal LinearBase stub for LoRA layer discovery.""" + + def __init__(self): + torch.nn.Module.__init__(self) + + +class DummyBaseLayerWithLoRA(torch.nn.Module): + """Fake LoRA wrapper that records set/reset/create calls.""" + + def __init__(self, base_layer: torch.nn.Module): + super().__init__() + self.base_layer = base_layer + + self.set_calls: list[ + tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor] + ] = [] + self.reset_calls: int = 0 + self.create_calls: int = 0 + + def set_lora(self, index: int, lora_a, lora_b): + assert index == 0 + self.set_calls.append((lora_a, lora_b)) + + def reset_lora(self, index: int): + assert index == 0 + self.reset_calls += 1 + + def create_lora_weights(self, max_loras, lora_config, model_config): + self.create_calls += 1 + + +def fake_replace_submodule( + root: torch.nn.Module, + module_name: str, + submodule: torch.nn.Module, + replace_calls: list[str] | None = None, +) -> None: + """Replace a submodule by traversing dotted paths correctly.""" + if replace_calls is not None: + replace_calls.append(module_name) + parts = module_name.split(".") + parent = root + for attr in parts[:-1]: + parent = getattr(parent, attr) + setattr(parent, parts[-1], submodule) diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 8d4a1487f..83ac7a114 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -7,8 +7,12 @@ import pytest import torch from vllm.lora.lora_weights import LoRALayerWeights from vllm.lora.utils import get_supported_lora_modules -from vllm.model_executor.layers.linear import LinearBase +from tests.diffusion.lora.conftest import ( + DummyBaseLayerWithLoRA, + FakeLinearBase, + fake_replace_submodule, +) from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager from vllm_omni.lora.request import LoRARequest @@ -33,35 +37,9 @@ class _DummyLoRALayer: self.reset_calls += 1 -class _FakeLinearBase(LinearBase): - def __init__(self): - torch.nn.Module.__init__(self) - - -class _DummyBaseLayerWithLoRA(torch.nn.Module): - def __init__(self, base_layer: torch.nn.Module): - super().__init__() - self.base_layer = base_layer - - self.set_calls: list[ - tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor] - ] = [] - self.reset_calls: int = 0 - self.create_calls: int = 0 - - def set_lora(self, index: int, lora_a, lora_b): - assert index == 0 - self.set_calls.append((lora_a, lora_b)) - - def reset_lora(self, index: int): - assert index == 0 - self.reset_calls += 1 - - def create_lora_weights(self, max_loras, lora_config, model_config): - # Needs to be callable for scale test when rank changes, but not - # actually used since we mock everything and check everything based - # on set calls. - self.create_calls += 1 +# Aliases for backward compatibility within this file +_FakeLinearBase = FakeLinearBase +_DummyBaseLayerWithLoRA = DummyBaseLayerWithLoRA class _DummyPipeline(torch.nn.Module): @@ -555,3 +533,45 @@ def test_lora_manager_max_rank_validation(monkeypatch, rank): req1 = _dummy_lora_request(1) with pytest.raises(ValueError): manager.add_adapter(req1) + + +def test_lora_manager_discovers_bagel_component(monkeypatch): + """Verify that _replace_layers_with_lora finds layers under 'bagel'.""" + import vllm_omni.diffusion.lora.manager as manager_mod + + monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA) + + def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs): + if isinstance(layer, _FakeLinearBase): + return _DummyBaseLayerWithLoRA(layer) + return layer + + replace_calls: list[str] = [] + + monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion) + monkeypatch.setattr( + manager_mod, + "replace_submodule", + lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls), + ) + + # Pipeline with a 'bagel' component (no 'transformer') + pipeline = torch.nn.Module() + pipeline.bagel = torch.nn.Module() + pipeline.bagel.language_model = torch.nn.Module() + pipeline.bagel.language_model.qkv_proj = _FakeLinearBase() + + manager = DiffusionLoRAManager( + pipeline=pipeline, + device=torch.device("cpu"), + dtype=torch.bfloat16, + max_cached_adapters=1, + ) + + peft_helper = type("_PH", (), {"r": 1})() + manager._replace_layers_with_lora(peft_helper) + + assert "language_model.qkv_proj" in replace_calls + assert "bagel.language_model.qkv_proj" in manager._lora_modules + # Verify the module was actually replaced in the tree (not just recorded) + assert isinstance(pipeline.bagel.language_model.qkv_proj, _DummyBaseLayerWithLoRA) diff --git a/tests/diffusion/models/bagel/test_bagel_lora.py b/tests/diffusion/models/bagel/test_bagel_lora.py new file mode 100644 index 000000000..8cb3446ed --- /dev/null +++ b/tests/diffusion/models/bagel/test_bagel_lora.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for BAGEL LoRA support across Stage 0 (Thinker) and Stage 1 (DiT).""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file + +from tests.diffusion.lora.conftest import ( + DummyBaseLayerWithLoRA, + FakeLinearBase, + fake_replace_submodule, +) +from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager +from vllm_omni.lora.request import LoRARequest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_FakeLinearBase = FakeLinearBase + + +# --------------------------------------------------------------------------- +# Stage 0 (Thinker / AR) -- packed_modules_mapping on the AR model class +# --------------------------------------------------------------------------- + + +class TestStage0ThinkerLoRA: + """Validate that OmniBagelForConditionalGeneration declares correct LoRA metadata.""" + + def test_omni_bagel_supports_lora(self): + from vllm_omni.model_executor.models.bagel.bagel import ( + OmniBagelForConditionalGeneration, + ) + + assert getattr(OmniBagelForConditionalGeneration, "supports_lora", False) is True + + def test_omni_bagel_packed_modules_mapping_complete(self): + from vllm_omni.model_executor.models.bagel.bagel import ( + OmniBagelForConditionalGeneration, + ) + + mapping = OmniBagelForConditionalGeneration.packed_modules_mapping + # Standard Qwen2 projections + assert mapping["qkv_proj"] == ["q_proj", "k_proj", "v_proj"] + assert mapping["gate_up_proj"] == ["gate_proj", "up_proj"] + # MoE generation-mode projections + assert mapping["qkv_proj_moe_gen"] == [ + "q_proj_moe_gen", + "k_proj_moe_gen", + "v_proj_moe_gen", + ] + assert mapping["mlp_moe_gen.gate_up_proj"] == [ + "mlp_moe_gen.gate_proj", + "mlp_moe_gen.up_proj", + ] + + +# --------------------------------------------------------------------------- +# Stage 1 (DiT / Diffusion) -- DiffusionLoRAManager with bagel component +# --------------------------------------------------------------------------- + + +class TestStage1DiTLoRA: + """Validate DiffusionLoRAManager discovers BAGEL's packed modules.""" + + def test_diffusion_lora_manager_discovers_bagel_packed_modules(self): + """Manager should derive packed→sublayer mapping from stacked_params_mapping.""" + pipeline = torch.nn.Module() + pipeline.bagel = torch.nn.Module() + + # Simulate a submodule that exposes stacked_params_mapping + # (as Bagel does after load_weights()) + language_model = torch.nn.Module() + language_model.stacked_params_mapping = [ + (".qkv_proj_moe_gen", ".q_proj_moe_gen", "q"), + (".qkv_proj_moe_gen", ".k_proj_moe_gen", "k"), + (".qkv_proj_moe_gen", ".v_proj_moe_gen", "v"), + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + pipeline.bagel.language_model = language_model + + manager = DiffusionLoRAManager( + pipeline=pipeline, + device=torch.device("cpu"), + dtype=torch.bfloat16, + max_cached_adapters=1, + ) + + mapping = manager._packed_modules_mapping + assert mapping["qkv_proj"] == ["q_proj", "k_proj", "v_proj"] + assert mapping["qkv_proj_moe_gen"] == [ + "q_proj_moe_gen", + "k_proj_moe_gen", + "v_proj_moe_gen", + ] + assert mapping["gate_up_proj"] == ["gate_proj", "up_proj"] + + def test_diffusion_lora_manager_replaces_bagel_packed_layer_via_sublayer_target(self, monkeypatch): + """Targeting sublayer 'q_proj' should replace the fused 'qkv_proj' under bagel.""" + import vllm_omni.diffusion.lora.manager as manager_mod + + monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", DummyBaseLayerWithLoRA) + + def _fake_from_layer_diffusion(*, layer, **_kwargs): + return DummyBaseLayerWithLoRA(layer) + + replace_calls: list[str] = [] + + monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion) + monkeypatch.setattr( + manager_mod, + "replace_submodule", + lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls), + ) + + # Build pipeline with bagel component + pipeline = torch.nn.Module() + pipeline.bagel = torch.nn.Module() + lm = torch.nn.Module() + lm.stacked_params_mapping = [ + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + ] + lm.attn = torch.nn.Module() + lm.attn.qkv_proj = _FakeLinearBase() + pipeline.bagel.language_model = lm + + manager = DiffusionLoRAManager( + pipeline=pipeline, + device=torch.device("cpu"), + dtype=torch.bfloat16, + max_cached_adapters=1, + ) + + # Treat qkv_proj as 3-slice packed layer + monkeypatch.setattr(manager, "_get_packed_modules_list", lambda _module: ["q", "k", "v"]) + + # Target sublayer "q_proj" -- manager should replace the packed "qkv_proj" + peft_helper = type("_PH", (), {"r": 1, "target_modules": ["q_proj"]})() + manager._replace_layers_with_lora(peft_helper) + + assert "language_model.attn.qkv_proj" in replace_calls + assert "bagel.language_model.attn.qkv_proj" in manager._lora_modules + # Verify the module was actually replaced in the tree (not just recorded) + assert isinstance(pipeline.bagel.language_model.attn.qkv_proj, DummyBaseLayerWithLoRA) + + +# --------------------------------------------------------------------------- +# Round-trip: synthetic checkpoint → set_active_adapter → verify weights +# --------------------------------------------------------------------------- + + +def _write_synthetic_lora( + adapter_dir: Path, + module_name: str, + rank: int, + in_dim: int, + out_dim: int, +) -> str: + """Write a minimal LoRA adapter (safetensors + config) to *adapter_dir*.""" + adapter_dir.mkdir(parents=True, exist_ok=True) + lora_a = torch.ones((rank, in_dim), dtype=torch.float32) + lora_b = torch.ones((out_dim, rank), dtype=torch.float32) * 2.0 + save_file( + { + f"base_model.model.{module_name}.lora_A.weight": lora_a, + f"base_model.model.{module_name}.lora_B.weight": lora_b, + }, + str(adapter_dir / "adapter_model.safetensors"), + ) + (adapter_dir / "adapter_config.json").write_text( + json.dumps({"r": rank, "lora_alpha": rank, "target_modules": [module_name]}), + encoding="utf-8", + ) + return str(adapter_dir) + + +class TestBagelLoRARoundTrip: + """End-to-end: synthetic checkpoint → load → activate → verify weights in fused layer.""" + + def test_set_active_adapter_loads_and_activates_bagel_lora(self, tmp_path, monkeypatch): + """Full round-trip through set_active_adapter for a bagel component module.""" + import vllm_omni.diffusion.lora.manager as manager_mod + + monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", DummyBaseLayerWithLoRA) + + # Build pipeline with bagel.language_model.foo (simple non-packed layer) + pipeline = torch.nn.Module() + pipeline.bagel = torch.nn.Module() + lm = torch.nn.Module() + lm.foo = _FakeLinearBase() + pipeline.bagel.language_model = lm + + def _fake_from_layer(*, layer, **_kwargs): + if isinstance(layer, FakeLinearBase): + return DummyBaseLayerWithLoRA(layer) + return layer + + monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer) + monkeypatch.setattr( + manager_mod, + "replace_submodule", + lambda root, name, sub: fake_replace_submodule(root, name, sub), + ) + + manager = DiffusionLoRAManager( + pipeline=pipeline, + device=torch.device("cpu"), + dtype=torch.bfloat16, + max_cached_adapters=1, + ) + + # Write synthetic adapter targeting bagel.language_model.foo + module_name = "bagel.language_model.foo" + rank = 2 + in_dim = 4 + out_dim = 4 + lora_dir = _write_synthetic_lora(tmp_path / "lora", module_name, rank, in_dim, out_dim) + + lora_request = LoRARequest( + lora_name="test_bagel", + lora_int_id=42, + lora_path=lora_dir, + ) + + # Full round-trip: load from disk → replace layer → activate weights + manager.set_active_adapter(lora_request, lora_scale=0.5) + + # Verify the layer was replaced and weights were set + replaced_layer = pipeline.bagel.language_model.foo + assert isinstance(replaced_layer, DummyBaseLayerWithLoRA), "Layer should be wrapped with LoRA" + assert len(replaced_layer.set_calls) == 1, "set_lora should have been called once" + + lora_a, lora_b = replaced_layer.set_calls[0] + # A weights should be ones (as written) + assert torch.all(lora_a == 1.0), f"lora_a should be all ones, got {lora_a}" + # B weights should be 2.0 * scale(0.5) = 1.0 + assert torch.allclose(lora_b, torch.ones_like(lora_b)), f"lora_b should be 2.0 * 0.5 = 1.0, got {lora_b}" diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 5f75e26cb..63e8d9a96 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -366,13 +366,17 @@ class DiffusionLoRAManager: fully_sharded_loras=False, ) - for component_name in ("transformer", "transformer_2", "dit"): + for component_name in ("transformer", "transformer_2", "dit", "bagel"): if not hasattr(self.pipeline, component_name): continue component = getattr(self.pipeline, component_name) if not isinstance(component, nn.Module): continue + # Collect replacements first to avoid mutating the module tree + # while iterating over named_modules(). + pending_replacements: list[tuple[str, str, nn.Module, list[str]]] = [] + for module_name, module in component.named_modules(remove_duplicate=False): # Don't recurse into already-replaced LoRA wrappers. Their # original LinearBase lives under "base_layer", and replacing @@ -401,6 +405,9 @@ class DiffusionLoRAManager: if not should_replace: continue + pending_replacements.append((module_name, full_module_name, module, packed_modules_list)) + + for module_name, full_module_name, module, packed_modules_list in pending_replacements: lora_layer = from_layer_diffusion( layer=module, max_loras=1, diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 685d14729..bbcd09dd5 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -929,17 +929,31 @@ class Qwen2MoTForCausalLM(Qwen2PreTrainedModel): (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), + # MLP gate/up projections — the DiT uses separate + # ColumnParallelLinear layers (no fused gate_up_proj), but + # these entries are needed so that DiffusionLoRAManager can + # derive the packed→sublayer mapping for LoRA checkpoints + # that store weights under fused gate_up_proj keys. + # The weight loader gracefully falls through to the + # non-stacked path when the fused parameter doesn't exist. + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), ] + self.stacked_params_mapping = stacked_params_mapping params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: + original_name = name for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) param = params_dict.get(name) if param is None: + # Fused param doesn't exist (e.g. gate_up_proj on DiT); + # restore original name and fall through to non-stacked path. + name = original_name break weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight, shard_id) diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index 3b4acae51..acbbc28b4 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -407,6 +407,22 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration): the DiT's denoising loop. """ + # LoRA packed→sublayer mapping for both standard Qwen2 projections + # and the MoE generation-mode projections added by _install_mot_modules(). + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + "qkv_proj_moe_gen": [ + "q_proj_moe_gen", + "k_proj_moe_gen", + "v_proj_moe_gen", + ], + "mlp_moe_gen.gate_up_proj": [ + "mlp_moe_gen.gate_proj", + "mlp_moe_gen.up_proj", + ], + } + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__(vllm_config=vllm_config, prefix=prefix) config = vllm_config.model_config.hf_config