A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 18:30 UTC
[Misc] Embedding model support LoRA (#14935)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Mar 18, 2025 at 20:07 UTC
db7c8ca910475d6cc4490d6a14dddc36c04b979b
1 file changed
+30
-2
vllm/lora/models.py
+30
-2
@@ -30,6 +30,7 @@ from vllm.lora.utils import (from_layer, from_layer_logits_processor,
30
is_regex_target_modules,
31
parse_fine_tuned_lora_name, replace_submodule)
32
from vllm.model_executor.models import SupportsLoRA, supports_multimodal
33
+from vllm.model_executor.models.interfaces import is_pooling_model
34
from vllm.model_executor.models.module_mapping import MultiModelKeys
35
from vllm.model_executor.models.utils import PPMissingLayer, WeightsMapper
36
from vllm.utils import is_pin_memory_available
@@ -104,6 +105,9 @@ class LoRAModel(AdapterModel):
105
"""Get LoRA for a given module by name"""
106
return self.loras.get(module_name, None)
107
108
+ def check_lora_name(self, lora_name: str) -> bool:
109
+ return lora_name in self.loras
110
+
111
# (yard1): TODO see if we can derive target_embedding_padding automatically
112
@classmethod
113
def from_lora_tensors(
@@ -335,6 +339,7 @@ class LoRAModelManager(AdapterModelManager):
339
# Used for long context lora.
340
self.scaling_factor_to_offset: Dict[float, int] = {}
341
super().__init__(model)
342
+
343
self.supported_lora_modules = get_supported_lora_modules(self.model)
344
assert self.supported_lora_modules, "No supported LoRA modules found in"
345
f"{self.model.__class__.__name__}."
@@ -350,6 +355,7 @@ class LoRAModelManager(AdapterModelManager):
355
# In case the model only supports LoRA for
356
# text modules (e.g. ChatGLM)
357
and hasattr(self.model, "get_mm_mapping"))
358
+ self.is_pooling_model = is_pooling_model(self.model)
359
self.packed_modules: Dict[str, List[str]] = {}
360
self.modules: Dict[str, BaseLayerWithLoRA] = {}
361
# Dict instead of a Set for compatibility with LRUCache.
@@ -389,7 +395,7 @@ class LoRAModelManager(AdapterModelManager):
395
lora_model.id, index)
396
self.lora_index_to_id[index] = lora_model.id
397
for module_name, module in self.modules.items():
392
- module_lora = lora_model.get_lora(module_name)
398
+ module_lora = self._get_lora_layer_weights(lora_model, module_name)
399
if module_lora:
400
module_lora.optimize()
401
# Bias is not explicitly enabled with the flag enable_lora_bias.
@@ -626,7 +632,7 @@ class LoRAModelManager(AdapterModelManager):
632
replaced_module: Set[str] = set()
633
has_replacement = False
634
for r in new_module_names:
629
- lora = lora_model.get_lora(r)
635
+ lora = self._get_lora_layer_weights(lora_model, r)
636
replacement_loras.append(lora)
637
if lora:
638
has_replacement = True
@@ -637,12 +643,34 @@ class LoRAModelManager(AdapterModelManager):
643
if replacement_loras[i]:
644
continue
645
replacement_loras[i] = None
646
+ # HACK Temporary solution for the pool model.
647
+ if self.is_pooling_model and not lora_model.check_lora_name(
648
+ module_name):
649
+ replaced_module_name = module_name.replace("model.", "")
650
+ if lora_model.check_lora_name(module_name):
651
+ module_name = replaced_module_name
652
lora_model.loras[module_name] = PackedLoRALayerWeights.pack(
653
replacement_loras)
654
# Remove the modules that have been replaced.
655
for module in replaced_module:
656
lora_model.loras.pop(module, None)
657
658
+ def _get_lora_layer_weights(
659
+ self, lora_model: LoRAModel,
660
+ module_name: str) -> Optional[LoRALayerWeights]:
661
+ org_module_name = module_name
662
+ if self.is_pooling_model and not lora_model.check_lora_name(
663
+ module_name):
664
+ # If it's a pool model, and the layer name is not found,
665
+ # remove the prefix 'model.' and search again.
666
+ module_name = module_name.replace("model.", "")
667
+ if lora_model.check_lora_name(module_name):
668
+ org_module_name = module_name
669
+ logger.info_once(
670
+ "For the pool model, successfully loaded the LoRA weights "
671
+ "after removing the prefix 'model.'.")
672
+ return lora_model.get_lora(org_module_name)
673
+
674
def deactivate_adapter(self, adapter_id: int) -> bool:
675
return deactivate_adapter(adapter_id, self._active_adapters,
676
self._deactivate_adapter)