models.py 33.9 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
import math
import os
6
7
from collections.abc import Sequence
from typing import Any, Callable, Optional, Union
8

9
import regex as re
10
11
12
13
import safetensors.torch
import torch
from torch import nn

14
15
16
17
18
from vllm.adapter_commons.models import (AdapterLRUCache, AdapterModel,
                                         AdapterModelManager)
from vllm.adapter_commons.utils import (add_adapter, deactivate_adapter,
                                        get_adapter, list_adapters,
                                        remove_adapter, set_adapter_mapping)
19
from vllm.config import LoRAConfig
20
from vllm.logger import init_logger
21
from vllm.lora.layers import BaseLayerWithLoRA, LoRAMapping
22
from vllm.lora.lora import LoRALayerWeights, PackedLoRALayerWeights
23
from vllm.lora.peft_helper import PEFTHelper
24
from vllm.lora.punica_wrapper import get_punica_wrapper
25
from vllm.lora.utils import (from_layer, from_layer_logits_processor,
26
                             get_supported_lora_modules,
27
                             is_regex_target_modules,
28
                             parse_fine_tuned_lora_name, replace_submodule)
29
from vllm.model_executor.layers.fused_moe import FusedMoE
30
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
31
from vllm.model_executor.models import SupportsLoRA, supports_multimodal
32
from vllm.model_executor.models.interfaces import is_pooling_model
33
from vllm.model_executor.models.module_mapping import MultiModelKeys
34
from vllm.model_executor.models.utils import PPMissingLayer, WeightsMapper
35
from vllm.model_executor.utils import get_packed_modules_mapping
36
from vllm.utils import is_pin_memory_available
37

38
logger = init_logger(__name__)
39
40
41
42
43
44
45
46
47
48

_GLOBAL_LORA_ID = 0


def get_lora_id():
    global _GLOBAL_LORA_ID
    _GLOBAL_LORA_ID += 1
    return _GLOBAL_LORA_ID


49
50
51
52
53
54
55
56
57
58
59
def is_moe_model(model: nn.Module) -> bool:
    """Checks if the model contains FusedMoE layers and warns the user."""
    if any(isinstance(module, FusedMoE) for module in model.modules()):
        logger.warning_once(
            "For MoE models, vLLM currently does not support fused MoE LoRA "
            "inference. Please ensure that the loaded LoRA model does not "
            "contain expert weights.")
        return True
    return False


60
class LoRAModel(AdapterModel):
61
62
63
64
65
66
    """A LoRA fine-tuned model."""

    def __init__(
        self,
        lora_model_id: int,
        rank: int,
67
        loras: dict[str, LoRALayerWeights],
68
    ) -> None:
69
70
71
72
73
        """
        Args:
            lora_model_id: The integer id for the lora model.
            rank: lora rank.
            loras: module name -> weights for lora-replaced layers.
74

75
        """
76
        self.id = lora_model_id
77

78
79
80
        assert (
            lora_model_id
            > 0), f"a valid lora id should be greater than 0, got {self.id}"
81
        self.rank = rank
82
        self.loras: dict[str, LoRALayerWeights] = loras
83

84
85
86
87
88
89
90
91
92
93
    def clone(self, lora_model_id: int) -> "LoRAModel":
        """Return a copy of the object with different ids.

        Will share the underlying tensors."""
        return self.__class__(
            lora_model_id,
            rank=self.rank,
            loras=self.loras.copy(),
        )

94
95
96
97
98
99
100
101
102
    @property
    def extra_vocab_size(self) -> int:
        return max(lora.extra_vocab_size
                   for lora in self.loras.values()) if self.loras else 0

    def get_lora(self, module_name: str) -> Optional[LoRALayerWeights]:
        """Get LoRA for a given module by name"""
        return self.loras.get(module_name, None)

103
104
105
    def check_lora_name(self, lora_name: str) -> bool:
        return lora_name in self.loras

106
107
108
109
110
    # (yard1): TODO see if we can derive target_embedding_padding automatically
    @classmethod
    def from_lora_tensors(
        cls,
        lora_model_id: int,
111
        tensors: dict[str, torch.Tensor],
112
        peft_helper: PEFTHelper,
113
114
        device: str = "cuda",
        dtype: Optional[torch.dtype] = None,
115
        embeddings: Optional[dict[str, torch.Tensor]] = None,
116
        target_embedding_padding: Optional[int] = None,
117
118
        embedding_modules: Optional[dict[str, str]] = None,
        embedding_padding_modules: Optional[list[str]] = None,
119
        weights_mapper: Optional[WeightsMapper] = None,
120
121
    ) -> "LoRAModel":
        """Create a LoRAModel from a dictionary of tensors."""
122
        pin_memory = str(device) == "cpu" and is_pin_memory_available()
123
        loras: dict[str, LoRALayerWeights] = {}
124
        for tensor_name, tensor in tensors.items():
125
            module_name, is_lora_a, is_bias = parse_fine_tuned_lora_name(
126
                tensor_name, weights_mapper)
127
128
129
            if module_name not in loras:
                lora_embeddings_tensor = None
                if embeddings:
130
                    assert embedding_modules is not None
131
                    embeddings_module = next(
Terry's avatar
Terry committed
132
                        (k for k in embedding_modules if k in module_name),
133
134
135
                        None)
                    if embeddings_module:
                        lora_embeddings_tensor = embeddings[
Terry's avatar
Terry committed
136
                            embedding_modules[embeddings_module]].to(
137
138
139
140
                                device=device, dtype=dtype)
                        if pin_memory:
                            lora_embeddings_tensor = (
                                lora_embeddings_tensor.pin_memory())
141
142
143
                loras[module_name] = LoRALayerWeights.from_config(
                    module_name, peft_helper, lora_embeddings_tensor)

144
145
146
147
148
149
150
151
            if is_bias:
                loras[module_name].bias = tensor.to(device=device,
                                                    dtype=dtype).t()
                bias = tensor.to(device=device, dtype=dtype).t()
                if pin_memory:
                    bias = bias.pin_memory()
                loras[module_name].bias = bias
            elif is_lora_a:
152
153
154
155
156
157
158
159
                loras[module_name].lora_a = tensor.to(device=device,
                                                      dtype=dtype).t()
                if pin_memory:
                    loras[module_name].lora_a = loras[
                        module_name].lora_a.pin_memory()
            else:
                loras[module_name].lora_b = tensor.to(device=device,
                                                      dtype=dtype).t()
160
                assert embedding_padding_modules is not None
161
                if any(name in module_name
Terry's avatar
Terry committed
162
                       for name in embedding_padding_modules
163
164
165
166
167
168
169
170
171
172
173
174
                       ) and target_embedding_padding is not None:
                    lora_b = loras[module_name].lora_b
                    assert target_embedding_padding >= lora_b.shape[1]
                    addition = target_embedding_padding - lora_b.shape[1]
                    loras[module_name].lora_b = torch.nn.functional.pad(
                        lora_b, (0, addition))
                if pin_memory:
                    loras[module_name].lora_b = loras[
                        module_name].lora_b.pin_memory()

        for lora in loras.values():
            lora.optimize()
175

176
        return cls(lora_model_id, peft_helper.r, loras)
177
178
179

    @classmethod
    def from_local_checkpoint(
180
181
182
183
184
185
186
187
188
189
190
191
192
            cls,
            lora_dir: str,
            expected_lora_modules: list[str],
            peft_helper: PEFTHelper,
            *,
            lora_model_id: Optional[int] = None,
            device: str = "cuda",
            dtype: Optional[torch.dtype] = None,
            target_embedding_padding: Optional[int] = None,
            embedding_modules: Optional[dict[str, str]] = None,
            embedding_padding_modules: Optional[list[str]] = None,
            weights_mapper: Optional[WeightsMapper] = None,
            tensorizer_config_dict: Optional[dict] = None) -> "LoRAModel":
193
        """Create a LoRAModel from a local checkpoint.
194

195
196
197
198
        Args:
            lora_dir: The local path that has lora data.
            expected_lora_modules: Name of modules that are expected to be
                replaced by lora.
199
            peft_helper: Loaded lora configuration information.
200
            lora_model_id: LoRA model id. If not given, automatically set by
201
202
203
204
205
206
207
                a global counter.
            device: Device where the lora model is loaded.
            dtype: dtype of the lora model weights.

        Returns:
            Loaded LoRA Model.
        """
208
209
        lora_tensor_path = os.path.join(lora_dir, "adapter_model.safetensors")
        lora_bin_file_path = os.path.join(lora_dir, "adapter_model.bin")
210
        lora_pt_file_path = os.path.join(lora_dir, "adapter_model.pt")
211
212
213
214
        new_embeddings_tensor_path = os.path.join(
            lora_dir, "new_embeddings.safetensors")
        new_embeddings_bin_file_path = os.path.join(lora_dir,
                                                    "new_embeddings.bin")
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
        tensors: dict[str, torch.Tensor] = {}
        unexpected_modules: list[Union[list[str], str]] = []

        def check_unexpected_modules(modules: dict):
            for lora_module in modules.keys():  # noqa
                module_name, _, _ = parse_fine_tuned_lora_name(
                    lora_module, weights_mapper)
                part_name = module_name.split(".")[-1]
                if part_name not in expected_lora_modules:
                    unexpected_modules.append(module_name)
            if unexpected_modules:
                raise ValueError(
                    f"While loading {lora_dir}, expected"
                    f" target modules in {expected_lora_modules}"
                    f" but received {unexpected_modules}."
                    f" Please verify that the loaded LoRA module is correct")

        if tensorizer_config_dict:
            from tensorizer import TensorDeserializer
234

235
236
237
238
            tensorizer_config = TensorizerConfig(**tensorizer_config_dict)
            lora_tensor_path = os.path.join(tensorizer_config.tensorizer_dir,
                                            "adapter_model.tensors")
            tensorizer_args = tensorizer_config._construct_tensorizer_args()
239
240
241
242
            tensors = TensorDeserializer(
                lora_tensor_path,
                dtype=tensorizer_config.dtype,
                **tensorizer_args.deserialization_kwargs)
243
            check_unexpected_modules(tensors)
244

245
        elif os.path.isfile(lora_tensor_path):
246
247
248
249
250
251
252
253
254
255
            # Find unexpected modules.
            # Use safetensor key as a source of truth to find expected modules.
            # in peft if you have target_modules A, B, C and C does not exist
            # in the model it won’t error and model will be trained with A, B
            # loraified. C won’t exist in the safetensor but it will exist in
            # the target_modules of the adapter_config.json.
            unexpected_modules = []
            with safetensors.safe_open(lora_tensor_path,
                                       framework="pt") as f:  # type: ignore
                # Load tensors if there are only expected modules.
256
                check_unexpected_modules(f)
257
258
                for module in f.keys():  # noqa
                    tensors[module] = f.get_tensor(module)
259
260
261
262
        elif os.path.isfile(lora_bin_file_path) or os.path.isfile(
                lora_pt_file_path):
            # When a bin/pt file is provided, we rely on config to find
            # unexpected modules.
263
            unexpected_modules = []
264
            target_modules = peft_helper.target_modules
265
266
            if not isinstance(target_modules, list):
                target_modules = [target_modules]
267
268
269
270
271
272
273
274
275
276
            for module in target_modules:
                # Compatible with more modules,
                # such as:layers.11.self_attn.k_proj
                part_name = module.split(".")[-1]
                if part_name not in expected_lora_modules:
                    unexpected_modules.append(module)
            # loaded lora's target modules must be a subset of
            # expected_lora_modules. It is not reliable. See
            # https://github.com/vllm-project/vllm/pull/5909. But there's no
            # other better mechanism.
277
            if unexpected_modules and not is_regex_target_modules(
278
                    peft_helper.target_modules, expected_lora_modules):
279
280
281
282
283
                raise ValueError(
                    f"While loading {lora_dir}, expected"
                    f" target modules in {expected_lora_modules}"
                    f" but received {unexpected_modules}."
                    f" Please verify that the loaded LoRA module is correct")
284
285
286
287
            lora_file_path = (lora_bin_file_path
                              if os.path.isfile(lora_bin_file_path) else
                              lora_pt_file_path)
            tensors = torch.load(lora_file_path,
cyyever's avatar
cyyever committed
288
289
                                 map_location=device,
                                 weights_only=True)
290
291
292
293
294
295
296
297
        else:
            raise ValueError(f"{lora_dir} doesn't contain tensors")

        embeddings = None
        if os.path.isfile(new_embeddings_tensor_path):
            embeddings = safetensors.torch.load_file(
                new_embeddings_tensor_path)
        elif os.path.isfile(new_embeddings_bin_file_path):
298
            embeddings = torch.load(new_embeddings_bin_file_path,
299
300
                                    map_location=device,
                                    weights_only=True)
301
302
303
304
305

        return cls.from_lora_tensors(
            lora_model_id=get_lora_id()
            if lora_model_id is None else lora_model_id,
            tensors=tensors,
306
            peft_helper=peft_helper,
307
308
309
310
            device=device,
            dtype=dtype,
            embeddings=embeddings,
            target_embedding_padding=target_embedding_padding,
Terry's avatar
Terry committed
311
            embedding_modules=embedding_modules,
312
313
            embedding_padding_modules=embedding_padding_modules,
            weights_mapper=weights_mapper)
314
315


316
class LoRAModelManager(AdapterModelManager):
317
318
319
320
    """A manager that manages multiple LoRA-fine-tuned models."""

    def __init__(
        self,
321
        model: SupportsLoRA,
322
323
324
325
        max_num_seqs: int,
        max_num_batched_tokens: int,
        vocab_size: int,
        lora_config: LoRAConfig,
326
        device: torch.device,
327
328
329
330
331
332
333
334
335
336
337
338
339
    ):
        """Create a LoRAModelManager and adapter for a given model.

        Args:
            model: the model to be adapted.
            max_num_seqs: the maximum number of sequences model can run in a
                single batch.
            max_num_batched_tokens: the maximum number of tokens model can run
                in a single batch.
            vocab_size: the vocab size of the model.
            lora_config: the LoRA configuration.
        """
        self.lora_config = lora_config
340
        self.device = device
341
342
343
        self.max_num_seqs = max_num_seqs
        assert self.capacity >= self.lora_slots
        self.max_num_batched_tokens = math.ceil(max_num_batched_tokens / 8) * 8
344
        self.lora_index_to_id: list[Optional[int]] = [None] * self.lora_slots
345
        self.vocab_size = vocab_size
346
347
348
349
350
        self.punica_wrapper = get_punica_wrapper(
            max_num_batched_tokens,
            max_batches=self.max_num_seqs,
            device=self.device,
            max_loras=self.lora_config.max_loras)
351

352
        super().__init__(model)
353
354
        self.supported_lora_modules = get_supported_lora_modules(self.model)
        assert self.supported_lora_modules, "No supported LoRA modules found in"
355
        f" {self.model.__class__.__name__}."
356
        
357
358
        if lora_config.lora_target_modules is not None:
            self.supported_lora_modules = lora_config.lora_target_modules
359
360

        self.packed_modules_mapping = get_packed_modules_mapping(self.model)
361
        # Used to indicate whether the model is a multimodal model
362
363
364
365
366
        self.supports_mm: bool = (
            supports_multimodal(self.model)
            # In case the model only supports LoRA for
            # text modules (e.g. ChatGLM)
            and hasattr(self.model, "get_mm_mapping"))
367
        self.is_pooling_model = is_pooling_model(self.model)
368
        self.is_moe_model = is_moe_model(self.model)
369
370
371
        self.packed_modules: dict[str, list[str]] = {}
        self.modules: dict[str, BaseLayerWithLoRA] = {}
        # Dict instead of a set for compatibility with LRUCache.
372
        self._last_mapping: Optional[LoRAMapping] = None
373
        self._create_lora_modules()
374
        self.model.lora_manager = self
375
        self.adapter_type = 'LoRA'
376
377
378
379
380
381
382
383
384

    @property
    def capacity(self) -> int:
        return self.lora_config.max_cpu_loras

    @property
    def lora_slots(self) -> int:
        return self.lora_config.max_loras

385
386
387
    @property
    def adapter_slots(self) -> int:
        return self.lora_slots
388

389
    def activate_adapter(
390
391
392
393
        self,
        lora_id: int,
    ) -> bool:
        """Move LoRA into a GPU buffer to be used in the forward pass."""
394
        if lora_id in self._active_adapters:
395
396
397
398
399
400
401
            return False
        first_free_slot = next(
            ((i, lora_id) for i, lora_id in enumerate(self.lora_index_to_id)
             if lora_id is None), None)
        if first_free_slot is None:
            raise ValueError("No free lora slots")
        index, _ = first_free_slot
402
403
        self._active_adapters[lora_id] = None
        lora_model = self._registered_adapters[lora_id]
404
405
        logger.debug("Activating LoRA. int id: %d, slot index: %d",
                     lora_model.id, index)
406
407
        self.lora_index_to_id[index] = lora_model.id
        for module_name, module in self.modules.items():
408
            module_lora = self._get_lora_layer_weights(lora_model, module_name)
409
410
            if module_lora:
                module_lora.optimize()
411
412
413
414
415
416
417
418
419
420
                # Bias is not explicitly enabled with the flag enable_lora_bias.
                bias = module_lora.bias
                if ((torch.is_tensor(bias) or
                     (isinstance(bias, Sequence) and any(b is not None
                                                         for b in bias)))
                        and not self.lora_config.bias_enabled):
                    module_lora.bias = None
                    raise ValueError(
                        f"Adapter bias cannot be used for {module_name}"
                        " without --enable-lora-bias.")
421
                module.set_lora(index, module_lora.lora_a, module_lora.lora_b,
422
423
                                module_lora.embeddings_tensor,
                                module_lora.bias)
424
425
426
427
            else:
                module.reset_lora(index)
        return True

428
    def _deactivate_adapter(self, lora_id: int):
429
430
431
432
433
434
        try:
            index = self.lora_index_to_id.index(lora_id)
            self.lora_index_to_id[index] = None
        except ValueError:
            pass

435
    def _add_adapter(self, lora: LoRAModel):
436
        self._create_merged_loras_inplace(lora)
437
        self._registered_adapters[lora.id] = lora
438

439
    def pin_adapter(self, lora_id: int) -> bool:
440
441
        """Pin a LoRAModel in the manager cache."""
        raise NotImplementedError(
442
            "Pinning is not supported in LoRAModelManager. "
443
444
            "Use LRUCacheLoRAModelManager for pinning")  # type: ignore

445
    def _set_adapter_mapping(self, mapping: LoRAMapping) -> None:
446
447
448
449
450
451
452
453
        # update lora states
        self.punica_wrapper.update_metadata(
            mapping,
            self.lora_index_to_id,
            self.lora_slots + 1,
            self.vocab_size,
            self.lora_config.lora_extra_vocab_size,
        )
454

455
    def remove_all_adapters(self):
456
        """Remove all LoRAModels from the manager."""
457
        self._registered_adapters.clear()
458
        self.lora_index_to_id = [None] * self.lora_slots
459
        self._active_adapters.clear()
460
461

    def _create_lora_modules(self):
462
463
464
465
466
467
468
469

        def _parent_module(module_name: str) -> str:
            # module name is a dot separated name.
            # for example:
            #  - given an input 'x.y.z' return 'x.y'
            #  - given an input 'x' return ''
            return module_name.rpartition('.')[0]

470
471
        for module_name, module in self.model.named_modules(
                remove_duplicate=False):
472
473
            if isinstance(module, PPMissingLayer):
                continue
474
475
            if not self._match_target_modules(module_name):
                continue
476
477
478
479
480
481
482
483
484
            # A temporary approach for multimodal models to support LoRA
            # TODO: Remove this restriction
            if self._filter_unsupported_mm_module(module_name):
                logger.warning(
                    "Regarding multimodal models, vLLM currently only supports "
                    "adding LoRA to language model, %s will be ignored.",
                    module_name,
                )
                continue
485
486
            parts = module_name.split(".")[-1]
            packed_moduled_lst = self.packed_modules_mapping.get(parts, [])
487
488
489
            new_module = replace_submodule(
                self.model, module_name,
                from_layer(module, self.lora_slots, self.lora_config,
490
                           packed_moduled_lst, self.model.config))
491

492
493
            # (yard1): TODO make this more robust
            if "lm_head" in module_name:
494
495
496
497
498
499
                logits_processor_module_name = 'logits_processor'
                parent_module = _parent_module(module_name)
                if parent_module:
                    logits_processor_module_name = (
                        f"{parent_module}.{logits_processor_module_name}")

500
                logits_processor_module = self.model.get_submodule(
501
502
                    logits_processor_module_name)

503
                new_module = replace_submodule(
504
                    self.model, logits_processor_module_name,
505
506
507
508
                    from_layer_logits_processor(logits_processor_module,
                                                module, self.lora_slots,
                                                self.lora_config,
                                                self.model.config))
509
510
511
512
513
514
515
516
517

            # In some models, especially multimodal ones, layers with the same
            # name may have different types, such as nn.Linear and
            # ReplicatedLinear. The nn.Linear layers cannot be replaced with
            # LoRA layers, leading to assertion error. The following check
            # aims to prevent this error
            if self.supports_mm and not isinstance(new_module,
                                                   BaseLayerWithLoRA):
                continue
518
519
            self.register_module(module_name, new_module)
            self._register_packed_modules(module_name)
520
521
            # All lora layers share the same punica_wrapper based on reference.
            new_module.set_mapping(self.punica_wrapper)
522
523
524
525
526

    def register_module(self, module_name: str, module: "BaseLayerWithLoRA"):
        assert isinstance(module, BaseLayerWithLoRA)
        self.modules[module_name] = module

Terry's avatar
Terry committed
527
528
529
530
    def create_dummy_lora(
            self,
            lora_id: int,
            rank: int,
531
            embedding_modules: Optional[dict[str, str]] = None) -> LoRAModel:
532
        """Create zero-initialized LoRAModel for warmup."""
533
        model = LoRAModel(lora_id, rank, {})
534
        for module_name, module in self.model.named_modules():
535
            bias_enabled = self.lora_config.bias_enabled
536
537
538
            if (not self._match_target_modules(module_name)
                    or not isinstance(module, BaseLayerWithLoRA)
                    or self._filter_unsupported_mm_module(module_name)):
539
540
541
                continue
            parts = module_name.split(".")
            if module_name not in self.packed_modules:
542
                assert embedding_modules is not None
Terry's avatar
Terry committed
543
                if parts[-1] in embedding_modules:
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
                    input_dim = (module.base_layer.org_vocab_size +
                                 self.lora_config.lora_extra_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]
                    embeddings_tensor_dim = (module.base_layer.embedding_dim if
                                             hasattr(module.base_layer,
                                                     "embedding_dim") else
                                             module.base_layer.weight.shape[1])
                    lora = LoRALayerWeights.create_dummy_lora_weights(
                        module_name,
                        input_dim,
                        output_dim,
                        rank,
560
                        module.lora_a_stacked[0].dtype,
561
                        "cpu",
562
563
                        embeddings_tensor_dim=embeddings_tensor_dim,
                        bias_enabled=bias_enabled)
564
565
566
                else:
                    lora = LoRALayerWeights.create_dummy_lora_weights(
                        module_name,
567
568
                        module.lora_a_stacked[0].shape[-1],
                        module.lora_b_stacked[0].shape[-2],
569
                        rank,
570
                        module.lora_a_stacked[0].dtype,
571
                        "cpu",
572
                        bias_enabled=bias_enabled,
573
574
575
576
577
                    )
                lora.optimize()
            else:
                parts = module_name.split(".")
                replacements = self.packed_modules_mapping[parts[-1]]
578
                subloras: list[Optional[LoRALayerWeights]] = []
579
580
581
582
583
584
585
586
                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",
587
                        bias_enabled=bias_enabled,
588
589
590
591
592
593
594
595
596
597
598
599
                    )
                    lora.optimize()
                    subloras.append(lora)
                lora = PackedLoRALayerWeights.pack(subloras)
            model.loras[module_name] = lora
        return model

    def _match_target_modules(self, module_name: str):
        return any(
            re.match(
                r".*\.{target_module}$".format(target_module=target_module),
                module_name) or target_module == module_name
Terry's avatar
Terry committed
600
            for target_module in self.supported_lora_modules)
601

602
603
604
    def _filter_unsupported_mm_module(self, module_name: str) -> bool:
        """
        Regarding multimodal models, vLLM currently only supports adding LoRA to
605
        language model. LoRA for other modules, such as the vision tower, will
606
607
608
609
        be filtered out.
        """
        if self.supports_mm:
            module_mapping: MultiModelKeys = self.model.get_mm_mapping()
610
611
612
            prefix_lst = module_mapping.connector + module_mapping.tower_model
            return any(
                [module_name.startswith(prefix) for prefix in prefix_lst])
613
614
        return False

615
616
617
    def _register_packed_modules(self, module_full_name: str) -> None:
        parts = module_full_name.split(".")
        module_name = parts[-1]
618
619
620
621
        replacements = self.packed_modules_mapping.get(module_name, [])
        # When replacements is less than or equal to 1, it indicates that this
        # module is not a packed module.
        if len(replacements) <= 1:
622
623
624
625
626
627
628
629
            return
        prefix = ".".join(parts[:-1])
        self.packed_modules[module_full_name] = [
            prefix + "." + r if prefix else r for r in replacements
        ]

    def _create_merged_loras_inplace(self, lora_model: LoRAModel) -> None:
        for module_name, new_module_names in self.packed_modules.items():
630
631
            replacement_loras: list[Optional[LoRALayerWeights]] = []
            replaced_module: set[str] = set()
632
633
            has_replacement = False
            for r in new_module_names:
634
                lora = self._get_lora_layer_weights(lora_model, r)
635
636
637
                replacement_loras.append(lora)
                if lora:
                    has_replacement = True
638
                    replaced_module.add(r)
639
640
641
642
643
644
            if not has_replacement:
                continue
            for i in range(len(replacement_loras)):
                if replacement_loras[i]:
                    continue
                replacement_loras[i] = None
645
646
647
648
649
650
            # HACK Temporary solution for the pool model.
            if self.is_pooling_model and not lora_model.check_lora_name(
                    module_name):
                replaced_module_name = module_name.replace("model.", "")
                if lora_model.check_lora_name(module_name):
                    module_name = replaced_module_name
651
652
            lora_model.loras[module_name] = PackedLoRALayerWeights.pack(
                replacement_loras)
653
654
655
            # Remove the modules that have been replaced.
            for module in replaced_module:
                lora_model.loras.pop(module, None)
656

657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
    def _get_lora_layer_weights(
            self, lora_model: LoRAModel,
            module_name: str) -> Optional[LoRALayerWeights]:
        org_module_name = module_name
        if self.is_pooling_model and not lora_model.check_lora_name(
                module_name):
            # If it's a pool model, and the layer name is not found,
            # remove the prefix 'model.' and search again.
            module_name = module_name.replace("model.", "")
            if lora_model.check_lora_name(module_name):
                org_module_name = module_name
                logger.info_once(
                    "For the pool model, successfully loaded the LoRA weights "
                    "after removing the prefix 'model.'.")
        return lora_model.get_lora(org_module_name)

673
674
675
676
677
    def deactivate_adapter(self, adapter_id: int) -> bool:
        return deactivate_adapter(adapter_id, self._active_adapters,
                                  self._deactivate_adapter)

    def add_adapter(self, adapter: LoRAModel) -> bool:
678
679
        logger.debug("Adding lora. Model id: %d, "
                     "int id: %d", adapter.id, adapter.id)
680
681
        return add_adapter(adapter, self._registered_adapters, self.capacity,
                           self._add_adapter)
682

683
684
685
686
687
688
689
690
    def set_adapter_mapping(self, mapping: LoRAMapping) -> None:
        self._last_mapping = set_adapter_mapping(mapping, self._last_mapping,
                                                 self._set_adapter_mapping)

    def remove_adapter(self, adapter_id: int) -> bool:
        return remove_adapter(adapter_id, self._registered_adapters,
                              self.deactivate_adapter)

691
    def list_adapters(self) -> dict[int, Any]:
692
693
694
695
696
697
698
        return list_adapters(self._registered_adapters)

    def get_adapter(self, adapter_id: int) -> Optional[Any]:
        return get_adapter(adapter_id, self._registered_adapters)


class LoRALRUCache(AdapterLRUCache[LoRAModel]):
699

700
701
    def __init__(self, capacity: int, deactivate_lora_fn: Callable[[int],
                                                                   bool]):
702
        super().__init__(capacity, deactivate_lora_fn)
703
704
705
706
707


class LRUCacheLoRAModelManager(LoRAModelManager):
    """A model manager that manages multiple LoRAs with LRU cache."""

708
709
710
    def __init__(self, model: nn.Module, max_num_seqs: int,
                 max_num_batched_tokens: int, vocab_size: int,
                 lora_config: LoRAConfig, device: torch.device):
711
        super().__init__(model, max_num_seqs, max_num_batched_tokens,
712
                         vocab_size, lora_config, device)
713
714
715
716
        self._registered_adapters: LoRALRUCache = LoRALRUCache(
            self.capacity, self.deactivate_adapter)
        self._active_adapters: LoRALRUCache = LoRALRUCache(
            self.lora_slots, self._deactivate_adapter)
717

718
    def list_adapters(self) -> dict[int, LoRAModel]:
719
        """List all registered LoRAModels."""
720
        return dict(self._registered_adapters.cache)
721

722
    def add_adapter(self, lora: LoRAModel) -> bool:
723
        """Add a LoRAModel to the manager."""
724
725
        logger.debug("Adding lora. Model id: %d, "
                     "int id: %d", lora.id, lora.id)
726
727
        if lora.id not in self._registered_adapters:
            self._add_adapter(lora)
728
729
730
            was_added = True
        else:
            # We always touch to update the LRU cache order
731
            self._registered_adapters.touch(lora.id)
732
733
734
            was_added = False
        return was_added

735
    def activate_adapter(
736
737
738
        self,
        lora_id: int,
    ) -> bool:
739
740
741
742
        if lora_id not in self._active_adapters and len(
                self._active_adapters) >= self.lora_slots:
            self._active_adapters.remove_oldest()
        result = super().activate_adapter(lora_id)
743
        # We always touch to update the LRU cache order
744
        self._active_adapters.touch(lora_id)
745
746
        return result

747
748
749
    def remove_oldest_adapter(self) -> bool:
        if len(self._registered_adapters) > 0:
            self._registered_adapters.remove_oldest()
750
751
752
            return True
        return False

753
    def pin_adapter(self, lora_id: int) -> bool:
754
755
756
757
758
759
760
        """Pin a LoRAModel in the manager cache."""
        self._pin_lora_in_cpu_cache(lora_id)
        self._pin_lora_in_gpu_cache(lora_id)
        return True

    def _pin_lora_in_cpu_cache(self, lora_id: int):
        try:
761
            self._registered_adapters.pin(lora_id)
762
763
764
765
766
        except ValueError as err:
            raise ValueError("Pinning failed. "
                             f"LoRA {lora_id} is not registered.") from err

    def _pin_lora_in_gpu_cache(self, lora_id: int):
767
        if lora_id not in self._active_adapters:
768
            # move lora to gpu if not already active
769
            self.activate_adapter(lora_id)
770

771
        self._active_adapters.pin(lora_id)
772

773
774
775
776
777
778
779

def create_lora_manager(
        model: nn.Module,
        max_num_seqs: int,
        max_num_batched_tokens: int,
        vocab_size: int,
        lora_config: LoRAConfig,
780
        device: torch.device,
781
        lora_manager_cls: type[LoRAModelManager] = LoRAModelManager,
782
783
        **kwargs) -> LoRAModelManager:
    """Create a LoRA adapter for a given model."""
784
    if not isinstance(model, SupportsLoRA):
785
786
787
788
789
790
791
        raise ValueError(f"Model {type(model)} is not supported for LoRA.")
    lora_manager = lora_manager_cls(
        model=model,
        max_num_seqs=max_num_seqs,
        max_num_batched_tokens=max_num_batched_tokens,
        vocab_size=vocab_size,
        lora_config=lora_config,
792
        device=device,
793
794
        **kwargs)
    return lora_manager