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

4
from typing import TYPE_CHECKING, Any, Literal
5
6

import torch
7
from pydantic import ConfigDict, Field, model_validator
8
from pydantic.dataclasses import dataclass
9
from typing_extensions import Self
10
11
12

from vllm.config.utils import config
from vllm.logger import init_logger
13
from vllm.utils.hashing import safe_hash
14
15
16
17
18
19
20
21
22
23
24

if TYPE_CHECKING:
    from vllm.config import ModelConfig
    from vllm.config.cache import CacheConfig
else:
    ModelConfig = Any
    CacheConfig = Any

logger = init_logger(__name__)

LoRADType = Literal["auto", "float16", "bfloat16"]
25
26
MaxLoRARanks = Literal[1, 8, 16, 32, 64, 128, 256, 320, 512]
LoRAExtraVocabSize = Literal[256, 512]
27
28
29
30
31
32
33


@config
@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
class LoRAConfig:
    """Configuration for LoRA."""

34
    max_lora_rank: MaxLoRARanks = 16
35
    """Max LoRA rank."""
36
    max_loras: int = Field(default=1, ge=1)
37
38
39
40
41
42
    """Max number of LoRAs in a single batch."""
    fully_sharded_loras: bool = False
    """By default, only half of the LoRA computation is sharded with tensor
    parallelism. Enabling this will use the fully sharded layers. At high
    sequence length, max rank or tensor parallel size, this is likely faster.
    """
43
    max_cpu_loras: int | None = None
44
45
    """Maximum number of LoRAs to store in CPU memory. Must be >= than
    `max_loras`."""
46
    lora_dtype: torch.dtype | LoRADType = "auto"
47
    """Data type for LoRA. If auto, will default to base model dtype."""
48
    default_mm_loras: dict[str, str] | None = None
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
    """Dictionary mapping specific modalities to LoRA model paths; this field
    is only applicable to multimodal models and should be leveraged when a
    model always expects a LoRA to be active when a given modality is present.
    Note that currently, if a request provides multiple additional
    modalities, each of which have their own LoRA, we do NOT apply
    default_mm_loras because we currently only support one lora adapter
    per prompt. When run in offline mode, the lora IDs for n modalities
    will be automatically assigned to 1-n with the names of the modalities
    in alphabetic order."""

    def compute_hash(self) -> str:
        """
        WARNING: Whenever a new field is added to this config,
        ensure that it is included in the factors list if
        it affects the computation graph.

        Provide a hash that uniquely identifies all the configs
        that affect the structure of the computation
        graph from input ids/embeddings to the final hidden states,
        excluding anything before input ids/embeddings and after
        the final hidden states.
        """
        factors: list[Any] = []
        factors.append(self.max_lora_rank)
        factors.append(self.max_loras)
        factors.append(self.fully_sharded_loras)
        factors.append(self.lora_dtype)
76

77
        hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
78
79
        return hash_str

80
81
    @model_validator(mode="after")
    def _validate_lora_config(self) -> Self:
82
83
84
85
86
        if self.max_cpu_loras is None:
            self.max_cpu_loras = self.max_loras
        elif self.max_cpu_loras < self.max_loras:
            raise ValueError(
                f"max_cpu_loras ({self.max_cpu_loras}) must be >= "
87
88
                f"max_loras ({self.max_loras})"
            )
89

90
91
        return self

92
93
94
95
96
    def verify_with_model_config(self, model_config: ModelConfig):
        if self.lora_dtype in (None, "auto"):
            self.lora_dtype = model_config.dtype
        elif isinstance(self.lora_dtype, str):
            self.lora_dtype = getattr(torch, self.lora_dtype)