[Bugfix] Fix LoRA Scaling on Active Adapters (#1421)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/diffusion/lora/test_lora_manager.py

@@ -38,21 +38,77 @@ class _FakeLinearBase(LinearBase):
         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
+
+
+class _DummyPipeline(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.transformer = torch.nn.Module()
+        self.transformer.foo = _FakeLinearBase()
+
+
+class _DummyLM(torch.nn.Module):
+    """LoRA enabled wrapper for _DummyPipeline."""
+
+    def __init__(self, rank: int):
+        super().__init__()
+        self.transformer = torch.nn.Module()
+        self.transformer.foo = _DummyBaseLayerWithLoRA(_FakeLinearBase())
+        self.rank = rank
+        self.loras = self.get_lora_modules()
+
+    def get_lora_modules(self):
+        return {"transformer.foo": self._get_initial_lora(self.rank)}
+
+    def get_lora(self, k: str) -> LoRALayerWeights:
+        """Get the unscaled LoRA weights for transformer.foo"""
+        return self.loras[k]
+
+    def _get_initial_lora(self, rank: int) -> LoRALayerWeights:
+        """Initializes a dummy LoRA for the current rank."""
+        A = torch.ones((rank, 4))
+        B = torch.ones((4, rank))
+        return LoRALayerWeights(
+            module_name="foo",
+            rank=rank,
+            lora_alpha=rank,
+            lora_a=A,
+            lora_b=B,
+        )
+
+
 def test_lora_manager_supported_modules_are_stable_with_wrapped_layers(monkeypatch):
     # Simulate a pipeline that already contains LoRA wrappers where the original
     # LinearBase is nested under ".base_layer".
     import vllm_omni.diffusion.lora.manager as manager_mod
 
-    class _DummyBaseLayerWithLoRA(torch.nn.Module):
-        def __init__(self, base_layer: torch.nn.Module):
-            super().__init__()
-            self.base_layer = base_layer
-
     monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
 
-    pipeline = torch.nn.Module()
-    pipeline.transformer = torch.nn.Module()
-    pipeline.transformer.foo = _DummyBaseLayerWithLoRA(_FakeLinearBase())
+    pipeline = _DummyLM(rank=2)
 
     # vLLM helper would see only the nested LinearBase and yield "base_layer".
     assert get_supported_lora_modules(pipeline) == ["base_layer"]
@@ -71,11 +127,6 @@ def test_lora_manager_supported_modules_are_stable_with_wrapped_layers(monkeypat
 def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
     import vllm_omni.diffusion.lora.manager as manager_mod
 
-    class _DummyBaseLayerWithLoRA(torch.nn.Module):
-        def __init__(self, base_layer: torch.nn.Module):
-            super().__init__()
-            self.base_layer = base_layer
-
     monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
 
     def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
@@ -116,11 +167,6 @@ def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
 def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch):
     import vllm_omni.diffusion.lora.manager as manager_mod
 
-    class _DummyBaseLayerWithLoRA(torch.nn.Module):
-        def __init__(self, base_layer: torch.nn.Module):
-            super().__init__()
-            self.base_layer = base_layer
-
     monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
 
     def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
@@ -189,9 +235,8 @@ def test_lora_manager_activates_fused_lora_on_packed_layer():
             },
         )()
     }
-    manager._adapter_scales = {7: 0.5}
 
-    manager._activate_adapter(7)
+    manager._activate_adapter(7, 0.5)
 
     assert packed_layer.reset_calls == 0
     assert len(packed_layer.set_calls) == 1
@@ -234,9 +279,8 @@ def test_lora_manager_activates_packed_lora_from_sublayers():
     manager._registered_adapters = {
         1: type("LM", (), {"id": 1, "loras": loras, "get_lora": lambda self, k: self.loras.get(k)})()
     }
-    manager._adapter_scales = {1: 2.0}
 
-    manager._activate_adapter(1)
+    manager._activate_adapter(1, scale=2.0)
 
     assert packed_layer.reset_calls == 0
     assert len(packed_layer.set_calls) == 1
@@ -274,7 +318,7 @@ def test_lora_manager_evicts_lru_adapter_when_cache_full(monkeypatch):
 
     monkeypatch.setattr(manager, "_load_adapter", _fake_load)
     monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
-    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id: None)
+    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id, scale: None)
 
     req1 = _dummy_lora_request(1)
     req2 = _dummy_lora_request(2)
@@ -306,7 +350,7 @@ def test_lora_manager_does_not_evict_pinned_adapter(monkeypatch):
 
     monkeypatch.setattr(manager, "_load_adapter", _fake_load)
     monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
-    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id: None)
+    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id, scale: None)
 
     manager.set_active_adapter(_dummy_lora_request(1), lora_scale=1.0)
     assert manager.pin_adapter(1)
@@ -332,7 +376,7 @@ def test_lora_manager_warns_when_all_adapters_pinned(monkeypatch):
 
     monkeypatch.setattr(manager, "_load_adapter", _fake_load)
     monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
-    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id: None)
+    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id, scale: None)
 
     manager.set_active_adapter(_dummy_lora_request(1), lora_scale=1.0)
     manager.set_active_adapter(_dummy_lora_request(2), lora_scale=1.0)
@@ -341,6 +385,118 @@ def test_lora_manager_warns_when_all_adapters_pinned(monkeypatch):
     assert manager.pin_adapter(2)
 
     manager.max_cached_adapters = 1
-    manager._evict_if_needed()
+    manager._evict_for_new_adapter()
 
     assert set(manager.list_adapters()) == {1, 2}
+
+
+def test_lora_manager_applies_multiple_scales_correctly(monkeypatch):
+    """Ensure that the LoRA manager applies scales correctly when the
+    active adapter receives a different scale, i.e., the rank is unchanged.
+    """
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyLoRALayer)
+
+    rank = 2
+    adapter_id = 7
+    req1 = _dummy_lora_request(adapter_id)
+    scale_1 = 0.25
+    scale_2 = 0.5
+
+    lora_model = _DummyLM(rank=rank)
+    manager = DiffusionLoRAManager(
+        pipeline=_DummyPipeline(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+    )
+
+    def _fake_load(_req: LoRARequest):
+        peft_helper = type("PH", (), {"r": rank})()
+        return lora_model, peft_helper
+
+    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
+    manager._registered_adapters = {
+        adapter_id: lora_model,
+    }
+    manager._lora_modules = {"transformer.foo": lora_model.transformer.foo}
+
+    # After the first scale, all B values should go from 1 -> scale_1
+    manager.set_active_adapter(req1, lora_scale=scale_1)
+    assert len(lora_model.transformer.foo.set_calls) == 1
+    lora_a, lora_b = lora_model.transformer.foo.set_calls[0]
+    assert torch.all(lora_a == 1)
+    assert torch.all(lora_b == scale_1)
+
+    # After the second scale, all B values should go from 1 -> scale_2
+    manager.set_active_adapter(req1, lora_scale=scale_2)
+    assert len(lora_model.transformer.foo.set_calls) == 2
+
+    lora_a, lora_b = lora_model.transformer.foo.set_calls[1]
+    assert torch.all(lora_a == 1)
+    assert torch.all(lora_b == scale_2)
+
+
+def test_lora_manager_scales_correctly_with_rank_changes(monkeypatch):
+    """Ensure that the LoRA manager correctly handles scaling when the rank
+    is changed and the buffers are reset + we reactivate.
+    """
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    rank = 2
+    adapter_id = 7
+    req1 = _dummy_lora_request(adapter_id)
+    initial_scale = 0.5
+
+    lora_model = _DummyLM(rank=rank)
+    manager = DiffusionLoRAManager(
+        pipeline=_DummyPipeline(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+    )
+
+    def _fake_load(_req: LoRARequest):
+        peft_helper = type("PH", (), {"r": rank})()
+        return lora_model, peft_helper
+
+    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
+    manager._registered_adapters = {
+        adapter_id: lora_model,
+    }
+    manager._lora_modules = {"transformer.foo": lora_model.transformer.foo}
+
+    # Activate adapter with initial scale
+    manager.set_active_adapter(req1, lora_scale=initial_scale)
+    assert lora_model.transformer.foo.create_calls == 0
+    assert len(lora_model.transformer.foo.set_calls) == 1
+    lora_a, lora_b = lora_model.transformer.foo.set_calls[0]
+    assert torch.all(lora_a == 1)
+    assert torch.all(lora_b == initial_scale)
+
+    # Increase the rank; this resets the buffers, so the adapter is activated again
+    manager._ensure_max_lora_rank(8)
+
+    # Ensure we actually took the rank expansion path, which recreates
+    # and sets the weight buffets, but that the scale didn't change
+    assert lora_model.transformer.foo.create_calls == 1
+    assert len(lora_model.transformer.foo.set_calls) == 2
+    lora_a, lora_b = lora_model.transformer.foo.set_calls[1]
+    assert torch.all(lora_a == 1)
+    assert torch.all(lora_b == initial_scale)
+
+
+def test_scale_keys_are_rounded():
+    """Ensure that added adapter scales are rounded to avoid lookup
+    issues due to precision differences, e.g., computed scales.
+    """
+    manager = DiffusionLoRAManager(
+        pipeline=_DummyPipeline(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+    )
+    adapter_id = 1
+    # Currently we round keys to 3 decimal places
+    manager._update_adapter_scale(adapter_id, 0.0031)
+    assert manager._adapter_scales[adapter_id] == 0.003

vllm_omni/diffusion/lora/manager.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import math
 import time
 from collections import OrderedDict
 
@@ -189,6 +190,10 @@ class DiffusionLoRAManager:
             logger.debug("No lora_request provided, deactivating all LoRA adapters")
             self._deactivate_all_adapters()
             return
+        elif math.isclose(0.0, lora_scale):
+            logger.warning("Received a request with LoRA scale 0; deactivating all LoRA adapters")
+            self._deactivate_all_adapters()
+            return
 
         adapter_id = lora_request.lora_int_id
         logger.debug(
@@ -202,16 +207,33 @@ class DiffusionLoRAManager:
         )
         if adapter_id not in self._registered_adapters:
             logger.info("Loading new adapter: id=%d, name=%s", adapter_id, lora_request.lora_name)
-            self.add_adapter(lora_request, lora_scale)
+            # Add the adapter + add to the cache
+            self.add_adapter(lora_request)
         else:
-            logger.debug("Adapter %d already loaded, activating", adapter_id)
+            # Just touch the cache access order
+            self._touch_adapter_info(adapter_id)
 
-            # update access order
-            self._adapter_scales[adapter_id] = lora_scale
-            self._adapter_access_order[adapter_id] = time.time()
-            self._adapter_access_order.move_to_end(adapter_id)
+        self._activate_adapter(adapter_id, lora_scale)
 
-        self._activate_adapter(adapter_id)
+    def _touch_adapter_info(self, adapter_id):
+        """Update the current caching ordering info."""
+        self._adapter_access_order[adapter_id] = time.time()
+        self._adapter_access_order.move_to_end(adapter_id)
+
+    def _update_adapter_scale(self, adapter_id: int, lora_scale: float):
+        """Update the adapter scale for a given adapter ID. To avoid potential
+        issues with using Floats as keys, for now, we round float values to
+        3 decimal points.
+        """
+        scale = DiffusionLoRAManager._get_rounded_scale(lora_scale)
+        self._adapter_scales[adapter_id] = scale
+
+    @staticmethod
+    def _get_rounded_scale(lora_scale: float):
+        """Normalizes a lora scale for use as a key in the _adapter_scales
+        dict; for now we just round scales to 3 decimal places.
+        """
+        return round(lora_scale, 3)
 
     def _load_adapter(
         self,
@@ -389,8 +411,9 @@ class DiffusionLoRAManager:
         # Re-apply active adapter if needed (buffers were reset).
         if self._active_adapter_id is not None:
             active_id = self._active_adapter_id
+            active_scale = self._adapter_scales[active_id]
             self._active_adapter_id = None
-            self._activate_adapter(active_id)
+            self._activate_adapter(active_id, active_scale)
 
     def _get_lora_weights(
         self,
@@ -416,9 +439,16 @@ class DiffusionLoRAManager:
         module_suffix = full_module_name.split(".")[-1]
         return lora_model.get_lora(module_suffix)
 
-    def _activate_adapter(self, adapter_id: int) -> None:
-        if self._active_adapter_id == adapter_id:
-            logger.debug("Adapter %d already active, skipping", adapter_id)
+    def _is_active_at_scale(self, adapter_id: int, scale: float) -> bool:
+        """True if the adapter_id is active and the current scale matches."""
+        rounded_scale = DiffusionLoRAManager._get_rounded_scale(scale)
+        is_active = self._active_adapter_id == adapter_id
+        matches_scale = self._adapter_scales.get(adapter_id) == rounded_scale
+        return is_active and matches_scale
+
+    def _activate_adapter(self, adapter_id: int, scale: float) -> None:
+        if self._is_active_at_scale(adapter_id, scale):
+            logger.debug("Adapter %d already active at scale %.3f skipping", adapter_id, scale)
             return
 
         logger.info("Activating adapter: id=%d", adapter_id)
@@ -453,7 +483,6 @@ class DiffusionLoRAManager:
                         lora_layer.reset_lora(0)
                         continue
 
-                    scale = self._adapter_scales.get(adapter_id, 1.0)
                     lora_a_list: list[torch.Tensor | None] = []
                     lora_b_list: list[torch.Tensor | None] = []
                     for sub_lora in sub_loras:
@@ -475,8 +504,6 @@ class DiffusionLoRAManager:
                     lora_layer.reset_lora(0)
                 continue
 
-            scale = self._adapter_scales.get(adapter_id, 1.0)
-
             # Packed LoRA weights already provide per-slice tensors.
             if isinstance(lora_weights, PackedLoRALayerWeights):
                 lora_a_list = lora_weights.lora_a
@@ -533,6 +560,7 @@ class DiffusionLoRAManager:
             )
 
         self._active_adapter_id = adapter_id
+        self._update_adapter_scale(adapter_id, scale)
 
     def _deactivate_all_adapters(self) -> None:
         logger.info("Deactivating all adapters: %d layers", len(self._lora_modules))
@@ -541,8 +569,10 @@ class DiffusionLoRAManager:
         self._active_adapter_id = None
         logger.debug("All adapters deactivated")
 
-    def _evict_if_needed(self) -> None:
-        while len(self._registered_adapters) > self.max_cached_adapters:
+    def _evict_for_new_adapter(self) -> None:
+        """Evict unpinned registered adapters until we have room for a new
+        adapter to be loaded."""
+        while len(self._registered_adapters) > (self.max_cached_adapters - 1):
             # Pick LRU among non-pinned adapters
             evict_candidates = [aid for aid in self._adapter_access_order.keys() if aid not in self._pinned_adapters]
             if not evict_candidates:
@@ -562,7 +592,7 @@ class DiffusionLoRAManager:
             )
             self.remove_adapter(lru_adapter_id)
 
-    def add_adapter(self, lora_request: LoRARequest, lora_scale: float = 1.0) -> bool:
+    def add_adapter(self, lora_request: LoRARequest) -> bool:
         """
         Add a new adapter to the cache without activating it.
         """
@@ -573,18 +603,18 @@ class DiffusionLoRAManager:
             return False
 
         logger.info("Adding new adapter: id=%d, name=%s", adapter_id, lora_request.lora_name)
+
+        # evict if cache full before adding the new adapter
+        # so that we don't go over capacity on the new load
+        self._evict_for_new_adapter()
+
         lora_model, peft_helper = self._load_adapter(lora_request)
+        self._touch_adapter_info(adapter_id)
+
         self._registered_adapters[adapter_id] = lora_model
-        self._adapter_scales[adapter_id] = lora_scale
 
         self._replace_layers_with_lora(peft_helper)
 
-        self._adapter_access_order[adapter_id] = time.time()
-        self._adapter_access_order.move_to_end(adapter_id)
-
-        # evict if cache full
-        self._evict_if_needed()
-
         logger.debug(
             "Adapter %d added, cache size: %d/%d", adapter_id, len(self._registered_adapters), self.max_cached_adapters
         )

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -192,8 +192,10 @@ class DiffusionWorker:
     def remove_lora(self, adapter_id: int) -> bool:
         return self.lora_manager.remove_adapter(adapter_id)
 
-    def add_lora(self, lora_request: LoRARequest, lora_scale: float = 1.0) -> bool:
-        return self.lora_manager.add_adapter(lora_request, lora_scale)
+    def add_lora(self, lora_request: LoRARequest) -> bool:
+        # NOTE (Alex): We have not implemented the API routing
+        # for the frontend server yet.
+        return self.lora_manager.add_adapter(lora_request)
 
     def list_loras(self) -> list[int]:
         return self.lora_manager.list_adapters()

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -263,7 +263,7 @@ class AsyncOmniDiffusion:
         )
         return all(results) if isinstance(results, list) else results
 
-    async def add_lora(self, lora_request: LoRARequest, lora_scale: float = 1.0) -> bool:
+    async def add_lora(self, lora_request: LoRARequest) -> bool:
         """Add a LoRA adapter"""
         loop = asyncio.get_event_loop()
         results = await loop.run_in_executor(
@@ -272,7 +272,7 @@ class AsyncOmniDiffusion:
             "add_lora",
             None,
             (),
-            {"lora_request": lora_request, "lora_scale": lora_scale},
+            {"lora_request": lora_request},
             None,
         )
         return all(results) if isinstance(results, list) else results