Make multilora_inference work

vllm/lora/layers/logits_processor.py

@@ -98,7 +98,7 @@ def create_lora_weights(
                 self.hidden_size,
             ),
             dtype=lora_config.lora_dtype,
-            device=self.device,
+            device="cpu",  #self.device,
         )
         self.lora_b_stacked = torch.zeros(
             (
@@ -108,7 +108,7 @@ def create_lora_weights(
                 lora_config.max_lora_rank,
             ),
             dtype=lora_config.lora_dtype,
-            device=self.device,
+            device="cpu", # self.device,
         )
 
         if self.sharded_to_full_mapping is not None:

vllm/lora/layers/utils.py

@@ -7,9 +7,11 @@
 import torch
 import torch.nn as nn
 
+from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.fused_moe import try_get_optimal_moe_config
 from vllm.utils.math_utils import next_power_of_2
 
+logger = init_logger(__name__)
 
 class LoRAMappingType(Enum):
     LANGUAGE = 1
@@ -32,26 +34,40 @@ def __post_init__(self):
 def _get_lora_device(base_layer: nn.Module) -> torch.device:
     # code borrowed from https://github.com/fmmoret/vllm/blob/fm-support-lora-on-quantized-models/vllm/lora/layers.py#L34
     """Returns the device for where to place the LoRA tensors."""
-    # unquantizedLinear
-    if hasattr(base_layer, "weight"):
-        return base_layer.weight.device
-    # Compressed Tensor
-    elif hasattr(base_layer, "weight_packed"):
-        return base_layer.weight_packed.device
-    # GPTQ/AWQ
-    elif hasattr(base_layer, "qweight"):
-        return base_layer.qweight.device
-    # MoE layer
-    elif hasattr(base_layer, "w2_weight"):
-        return base_layer.w2_weight.device
-    # MoE Compressed Tensor
-    elif hasattr(base_layer, "w2_weight_packed"):
-        return base_layer.w2_weight_packed.device
-    # MoE GPTQ/AWQ/GGUF
-    elif hasattr(base_layer, "w2_qweight"):
-        return base_layer.w2_qweight.device
-    else:
-        raise ValueError(f"Unsupported base layer: {base_layer}")
+
+    # In case some module wrap the Tensor in ParameterList
+    def get_dev(obj):
+        dev = None
+        if obj is not None:
+            if hasattr(obj, "device"):
+                dev = obj.device
+                logger.debug(f"get_dev type of obj = {type(obj)} dev = {dev}")
+            elif isinstance(obj, (nn.ParameterList, list, tuple)) and len(obj) > 0:
+                if hasattr(obj[0], "device"):
+                    dev = obj[0].device
+                    logger.debug(f"get_dev type of obj[0] = {type(obj[0])} dev = {dev}")
+        logger.debug(f"get_dev final return dev = {dev}")
+        return dev
+
+    attr_names = ["weight", # unquantizedLinear
+                  "weight_packed", # Compressed Tensor
+                  "qweight", # GPTQ/AWQ
+                  "w2_weight", # MoE layer
+                  "w2_weight_packed", # MoE Compressed Tensor
+                  "w2_qweight", # MoE GPTQ/AWQ/GGUF
+                 ]
+    for attr in attr_names:
+        logger.debug(f"lora base_layer = {base_layer} attr_name = {attr}")
+        target = getattr(base_layer, attr, None)
+        dev = get_dev(target)
+        if dev is not None:
+            return dev
+
+    try:
+        return next(base_layer.parameters()).device
+    except StopIteration:
+        logger.debug("lora base_layer = {base_layer} in except StopInteration return cpu")
+        return torch.device("cpu")
 
 
 def _not_fully_sharded_can_replace(can_replace):

vllm/lora/layers/vocal_parallel_embedding.py

@@ -4,15 +4,18 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import traceback
 from transformers import PretrainedConfig
 
 from vllm.config.lora import LoRAConfig
+from vllm.logger import init_logger
 from vllm.model_executor.custom_op import maybe_get_oot_by_class
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from vllm.platforms import current_platform
 
 from .base import BaseLayerWithLoRA
 
+logger = init_logger(__name__)
 
 class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
     def __init__(self, base_layer: VocabParallelEmbedding) -> None:
@@ -46,14 +49,22 @@ def create_lora_weights(
             self.embeddings_slice = None
             self.embeddings_weights = None
 
+        logger.info(f"Before lora_a_stacked created device = {self.base_layer.weight.device} ")
+        def nested_function():
+            # Capture current call stack
+            stack_list = traceback.format_stack()
+            stack_str = "".join(stack_list)
+            logger.info(f"Current Stack Trace:\n{stack_str}")
+        nested_function()
+
         self.lora_a_stacked = torch.zeros(
             (
                 max_loras,
                 self.base_layer.org_vocab_size,
                 lora_config.max_lora_rank,
             ),
             dtype=lora_config.lora_dtype,
-            device=self.base_layer.weight.device,
+            device="cpu", #self.base_layer.weight.device,
         )
         self.lora_b_stacked = torch.zeros(
             (
@@ -63,7 +74,7 @@ def create_lora_weights(
                 lora_config.max_lora_rank,
             ),
             dtype=lora_config.lora_dtype,
-            device=self.base_layer.weight.device,
+            device="cpu", #self.base_layer.weight.device,
         )
         self.lora_a_stacked_2d = self.lora_a_stacked.view(
             self.lora_a_stacked.shape[0] * self.lora_a_stacked.shape[1],
@@ -85,7 +96,8 @@ def set_lora(
         self.reset_lora(index)
         # NOTE self.lora_a_stacked is row-major, and lora_a is col-major,
         # so we need transpose here
-
+        logger.info(f"Before transpose self.lora_a_stacked type {type(self.lora_a_stacked)} ")
+        logger.info(f"Before transpose lora_a type {type(lora_a)} ")
         self.lora_a_stacked[index, : lora_a.shape[1], : lora_a.shape[0]].copy_(
             lora_a.T, non_blocking=True
         )

vllm/lora/model_manager.py

@@ -509,6 +509,7 @@ def create_dummy_lora(
                             if hasattr(module.base_layer, "embedding_dim")
                             else module.base_layer.weight.shape[0]
                         )
+                    logger.info(f"Before create_dummy_lora_weights {type(module.lora_a_stacked)}")
                     lora = LoRALayerWeights.create_dummy_lora_weights(
                         module_name,
                         input_dim,
@@ -521,6 +522,7 @@ def create_dummy_lora(
                 elif module.__class__.__name__ == "FusedMoE3DWithLoRA":
                     # Case for 3D moe model
                     # w2
+                    logger.info(f"Before create_dummy_lora_weights FusedMoE3DWithLoRA w2 {type(module.w2_lora_a_stacked)}")
                     lora = LoRALayerWeights.create_dummy_lora_weights(
                         module_name,
                         module.w2_input_size,
@@ -531,6 +533,7 @@ def create_dummy_lora(
                     )
                     model.loras[module_name] = lora
                     # w13
+                    logger.info(f"Before create_dummy_lora_weights FusedMoE3DWithLoRA w13 {type(module.w13_lora_a_stacked)}")
                     lora = LoRALayerWeights.create_dummy_lora_weights(
                         module_name,
                         module.w13_input_size,
@@ -542,6 +545,7 @@ def create_dummy_lora(
                     )
                     model.loras[module_name + ".base_layer"] = lora
                 else:
+                    logger.info(f"Before create_dummy_lora_weights else lora_a_stacked type {type(module.lora_a_stacked)} ")
                     lora = LoRALayerWeights.create_dummy_lora_weights(
                         module_name,
                         module.lora_a_stacked[0].shape[-1],

vllm/lora/worker_manager.py

@@ -186,10 +186,13 @@ def _load_adapter(self, lora_request: LoRARequest) -> LoRAModel:
 
     def add_dummy_lora(self, lora_request: LoRARequest, rank: int) -> bool:
         if lora_request.lora_int_id in self.list_adapters():
+            logger.info("Before return False")
             return False
         if isinstance(self._cached_dummy_lora, LoRAModel):
+            logger.info("Before _cached_dummy_lora.clone")
             dummy_lora = self._cached_dummy_lora.clone(lora_request.lora_int_id)
         else:
+            logger.info("Before create_dummy_lora")
             dummy_lora = self._adapter_manager.create_dummy_lora(
                 lora_request.lora_int_id, rank, self.embedding_modules
             )

vllm/v1/worker/lora_model_runner_mixin.py

@@ -113,11 +113,12 @@ def maybe_setup_dummy_loras(
                 )
                 for lora_id in range(1, num_loras + 1)
             }
-
+            logger.info("Before self.lora_manager.dummy_lora_cache")
             with self.lora_manager.dummy_lora_cache():
                 # Add the dummy LoRAs here so _set_active_loras doesn't try to
                 # load from disk.
                 for lr in lora_requests:
+                    logger.info("Before self.lora_manager.add_dummy_lora")
                     self.lora_manager.add_dummy_lora(lr, rank=lora_warmup_rank)
 
                 yield