diff --git a/vllm/lora/layers/utils.py b/vllm/lora/layers/utils.py index c19b097586f5..d942228ddd1c 100644 --- a/vllm/lora/layers/utils.py +++ b/vllm/lora/layers/utils.py @@ -32,26 +32,27 @@ 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}") + def get_dev(obj): + if obj is None: + return None + if isinstance(obj, (nn.ParameterList, list, tuple)): + return obj[0].device if len(obj) > 0 else None + if hasattr(obj, "device"): + return obj.device + return None + + attr_names = ["weight", "weight_packed", "qweight", + "w2_weight", "w2_weight_packed", "w2_qweight"] + for attr in attr_names: + 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: + return torch.device("cpu") def _not_fully_sharded_can_replace(can_replace): diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 9d3772560433..97fc6fdaa7b9 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -6,6 +6,7 @@ from typing import TypeVar import torch +import torchax from torch import nn from vllm.config import VllmConfig @@ -481,101 +482,102 @@ def create_dummy_lora( ) -> LoRAModel: """Create zero-initialized LoRAModel for warmup.""" model = LoRAModel(lora_id, rank, {}) - for module_name, module in self.model.named_modules(): - if ( - not self._match_target_modules(module_name) - or not isinstance(module, BaseLayerWithLoRA) - or self._get_punica_wrapper(module_name) is None - ): - continue - parts = module_name.split(".") - if module_name not in self.packed_modules: - assert embedding_modules is not None - if parts[-1] in embedding_modules: - # Special-case lm_head: wrapped by LogitsProcessorWithLoRA. - # LoRA input dim is hidden_size, output dim is vocab size. - # LogitsProcessorWithLoRA handles extra vocab size directly. - if parts[-1] == "lm_head": - input_dim = module.lora_a_stacked[0].shape[-1] - output_dim = module.lora_b_stacked[0].shape[-2] - else: - input_dim = ( - module.base_layer.org_vocab_size - if hasattr(module.base_layer, "org_vocab_size") - else module.base_layer.weight.shape[1] + with torchax.default_env(): + for module_name, module in self.model.named_modules(): + if ( + not self._match_target_modules(module_name) + or not isinstance(module, BaseLayerWithLoRA) + or self._get_punica_wrapper(module_name) is None + ): + continue + parts = module_name.split(".") + if module_name not in self.packed_modules: + assert embedding_modules is not None + if parts[-1] in embedding_modules: + # Special-case lm_head: wrapped by LogitsProcessorWithLoRA. + # LoRA input dim is hidden_size, output dim is vocab size. + # LogitsProcessorWithLoRA handles extra vocab size directly. + if parts[-1] == "lm_head": + input_dim = module.lora_a_stacked[0].shape[-1] + output_dim = module.lora_b_stacked[0].shape[-2] + else: + input_dim = ( + module.base_layer.org_vocab_size + if hasattr(module.base_layer, "org_vocab_size") + else module.base_layer.weight.shape[1] + ) + output_dim = ( + module.base_layer.embedding_dim + if hasattr(module.base_layer, "embedding_dim") + else module.base_layer.weight.shape[0] + ) + lora = LoRALayerWeights.create_dummy_lora_weights( + module_name, + input_dim, + output_dim, + rank, + module.lora_a_stacked[0].dtype, + "cpu", ) - output_dim = ( - module.base_layer.embedding_dim - if hasattr(module.base_layer, "embedding_dim") - else module.base_layer.weight.shape[0] + model.loras[module_name] = lora + elif module.__class__.__name__ == "FusedMoE3DWithLoRA": + # Case for 3D moe model + # w2 + lora = LoRALayerWeights.create_dummy_lora_weights( + module_name, + module.w2_input_size, + module.w2_output_size, + rank * module.w2_lora_a_stacked[0].shape[1], # rank*num_experts + module.w2_lora_a_stacked[0].dtype, + "cpu", ) - lora = LoRALayerWeights.create_dummy_lora_weights( - module_name, - input_dim, - output_dim, - rank, - module.lora_a_stacked[0].dtype, - "cpu", - ) - model.loras[module_name] = lora - elif module.__class__.__name__ == "FusedMoE3DWithLoRA": - # Case for 3D moe model - # w2 - lora = LoRALayerWeights.create_dummy_lora_weights( - module_name, - module.w2_input_size, - module.w2_output_size, - rank * module.w2_lora_a_stacked[0].shape[1], # rank*num_experts - module.w2_lora_a_stacked[0].dtype, - "cpu", - ) - model.loras[module_name] = lora - # w13 - lora = LoRALayerWeights.create_dummy_lora_weights( - module_name, - module.w13_input_size, - module.w13_output_size, - rank - * module.w13_lora_a_stacked[0].shape[1], # rank*num_experts - module.w13_lora_a_stacked[0].dtype, - "cpu", - ) - model.loras[module_name + ".base_layer"] = lora + model.loras[module_name] = lora + # w13 + lora = LoRALayerWeights.create_dummy_lora_weights( + module_name, + module.w13_input_size, + module.w13_output_size, + rank + * module.w13_lora_a_stacked[0].shape[1], # rank*num_experts + module.w13_lora_a_stacked[0].dtype, + "cpu", + ) + model.loras[module_name + ".base_layer"] = lora + else: + lora = LoRALayerWeights.create_dummy_lora_weights( + module_name, + module.lora_a_stacked[0].shape[-1], + module.lora_b_stacked[0].shape[-2], + rank, + module.lora_a_stacked[0].dtype, + "cpu", + ) + model.loras[module_name] = lora else: - lora = LoRALayerWeights.create_dummy_lora_weights( - module_name, - module.lora_a_stacked[0].shape[-1], - module.lora_b_stacked[0].shape[-2], - rank, - module.lora_a_stacked[0].dtype, - "cpu", - ) + parts = module_name.split(".") + replacements = self.packed_modules_mapping[parts[-1]] + subloras: list[LoRALayerWeights | None] = [] + for i, r in enumerate(replacements): + lora = LoRALayerWeights.create_dummy_lora_weights( + module_name + "." + r, + module.lora_a_stacked[i].shape[-1], + module.lora_b_stacked[i].shape[-2], + rank, + module.lora_a_stacked[i].dtype, + "cpu", + ) + subloras.append(lora) + if module.__class__.__name__ == "FusedMoEWithLoRA": + # For non-gated MoE, pad subloras to 3 elements per expert + # to match pack_moe expectations (w1, w2, None for w3) + if self._is_non_gated_moe and len(subloras) > 0: + subloras = self._pad_lora_pairs_to_triplets(subloras) + lora = PackedLoRALayerWeights.pack_moe( + subloras, module_name, is_non_gated_moe=self._is_non_gated_moe + ) + else: + lora = PackedLoRALayerWeights.pack(subloras) model.loras[module_name] = lora - else: - parts = module_name.split(".") - replacements = self.packed_modules_mapping[parts[-1]] - subloras: list[LoRALayerWeights | None] = [] - for i, r in enumerate(replacements): - lora = LoRALayerWeights.create_dummy_lora_weights( - module_name + "." + r, - module.lora_a_stacked[i].shape[-1], - module.lora_b_stacked[i].shape[-2], - rank, - module.lora_a_stacked[i].dtype, - "cpu", - ) - subloras.append(lora) - if module.__class__.__name__ == "FusedMoEWithLoRA": - # For non-gated MoE, pad subloras to 3 elements per expert - # to match pack_moe expectations (w1, w2, None for w3) - if self._is_non_gated_moe and len(subloras) > 0: - subloras = self._pad_lora_pairs_to_triplets(subloras) - lora = PackedLoRALayerWeights.pack_moe( - subloras, module_name, is_non_gated_moe=self._is_non_gated_moe - ) - else: - lora = PackedLoRALayerWeights.pack(subloras) - model.loras[module_name] = lora return model def _match_target_modules(self, module_name: str) -> bool: