[BugFix] Fix Max Rank Handling in LoRA (#1397)

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

tests/diffusion/lora/test_lora_manager.py

@@ -508,3 +508,50 @@ def test_scale_keys_are_rounded():
     # Currently we round keys to 3 decimal places
     manager._update_adapter_scale(adapter_id, 0.0031)
     assert manager._adapter_scales[adapter_id] == 0.003
+
+
+def test_lora_manager_uses_valid_max_rank(monkeypatch):
+    """Ensure that the LoRA manager uses a valid max rank for vLLM."""
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+    )
+
+    # Ensure that the rank is correctly adjusted to the smallest valid max rank
+    supported_max_rank = 64
+    unsupported_max_rank = 63
+    assert supported_max_rank in DiffusionLoRAManager._VALID_MAX_RANKS
+    assert unsupported_max_rank not in DiffusionLoRAManager._VALID_MAX_RANKS
+
+    def _fake_load(_req: LoRARequest):
+        lora_model = type("LM", (), {"id": _req.lora_int_id})()
+        peft_helper = type("PH", (), {"r": unsupported_max_rank})()
+        return lora_model, peft_helper
+
+    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
+    req1 = _dummy_lora_request(1)
+    manager.add_adapter(req1)
+    assert manager._max_lora_rank == supported_max_rank
+
+
+@pytest.mark.parametrize("rank", [-1, 0, DiffusionLoRAManager._VALID_MAX_RANKS[-1] + 1])
+def test_lora_manager_max_rank_validation(monkeypatch, rank):
+    """Check that invalid max ranks are handled correctly."""
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+    )
+
+    lora_rank = rank
+
+    def _fake_load(_req: LoRARequest):
+        lora_model = type("LM", (), {"id": _req.lora_int_id})()
+        peft_helper = type("PH", (), {"r": lora_rank})()
+        return lora_model, peft_helper
+
+    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
+    req1 = _dummy_lora_request(1)
+    with pytest.raises(ValueError):
+        manager.add_adapter(req1)

vllm_omni/diffusion/lora/manager.py

@@ -4,9 +4,11 @@
 import math
 import time
 from collections import OrderedDict
+from typing import get_args
 
 import torch
 import torch.nn as nn
+from vllm.config.lora import MaxLoRARanks
 from vllm.logger import init_logger
 from vllm.lora.layers import BaseLayerWithLoRA
 from vllm.lora.lora_model import LoRAModel
@@ -38,6 +40,9 @@ class DiffusionLoRAManager:
     Uses LRU cache management similar to LRUCacheLoRAModelManager.
     """
 
+    # Valid max allowed ranks for LoRA in vLLM
+    _VALID_MAX_RANKS: list[int] = sorted(get_args(MaxLoRARanks))
+
     def __init__(
         self,
         pipeline: nn.Module,
@@ -413,11 +418,10 @@ class DiffusionLoRAManager:
         if min_rank <= self._max_lora_rank:
             return
 
-        if min_rank <= 0:
-            raise ValueError(f"Invalid LoRA rank: {min_rank}")
+        valid_max_rank = self._get_smallest_valid_max_rank(min_rank)
 
-        logger.info("Increasing max LoRA rank: %d -> %d", self._max_lora_rank, min_rank)
-        self._max_lora_rank = min_rank
+        logger.info("Increasing max LoRA rank: %d -> %d", self._max_lora_rank, valid_max_rank)
+        self._max_lora_rank = valid_max_rank
 
         if not self._lora_modules:
             return
@@ -441,6 +445,18 @@ class DiffusionLoRAManager:
             self._active_adapter_id = None
             self._activate_adapter(active_id, active_scale)
 
+    @classmethod
+    def _get_smallest_valid_max_rank(cls, min_rank: int) -> int:
+        """Given a LoRA rank, get the smallest max rank that can support it."""
+        if min_rank <= 0:
+            raise ValueError(f"Invalid LoRA rank: {min_rank}")
+
+        allowed_ranks = [rank for rank in cls._VALID_MAX_RANKS if rank >= min_rank]
+        if not allowed_ranks:
+            raise ValueError(f"LoRA rank of {min_rank} exceeds max allowed rank of {max(cls._VALID_MAX_RANKS)}")
+
+        return min(allowed_ranks)
+
     def _get_lora_weights(
         self,
         lora_model: LoRAModel,