diff --git a/vllm/lora/layers/logits_processor.py b/vllm/lora/layers/logits_processor.py index 237a61eace1e..194438c877d7 100644 --- a/vllm/lora/layers/logits_processor.py +++ b/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: diff --git a/vllm/lora/layers/utils.py b/vllm/lora/layers/utils.py index c19b097586f5..7571787c57e4 100644 --- a/vllm/lora/layers/utils.py +++ b/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): diff --git a/vllm/lora/layers/vocal_parallel_embedding.py b/vllm/lora/layers/vocal_parallel_embedding.py index 05e7cfa06c85..3f1206964a02 100644 --- a/vllm/lora/layers/vocal_parallel_embedding.py +++ b/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,6 +49,14 @@ 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, @@ -53,7 +64,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_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 ) diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 9d3772560433..f36cccab75fa 100644 --- a/vllm/lora/model_manager.py +++ b/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], diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index bea6d015e0a6..b000be62368b 100644 --- a/vllm/lora/worker_manager.py +++ b/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 ) diff --git a/vllm/v1/worker/lora_model_runner_mixin.py b/vllm/v1/worker/lora_model_runner_mixin.py index 53873d156f88..f5a4c4501713 100644 --- a/vllm/v1/worker/lora_model_runner_mixin.py +++ b/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