[Feature] Enable LoRA adapter injection for BAGEL (#2490)

Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu>
Signed-off-by: Claude <noreply@anthropic.com>
Co-authored-by: Claude <noreply@anthropic.com>

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)

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)

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}"

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,

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)

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