diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index e5f1c4700..ea967a7b0 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/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 diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 7fad1b9e7..3cccc732e 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/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 ) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 6a8494991..a20a26cac 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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() diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py index 072daee30..faa9cc968 100644 --- a/vllm_omni/entrypoints/async_omni_diffusion.py +++ b/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