@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 136 lines 5.14 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Adapted from: https://github.com/huggingface/peft/blob/main/src/peft/tuners/lora/config.py
5
6 import json
7 import math
8 import os
9 from dataclasses import MISSING, dataclass, field, fields
10 from typing import Literal
11
12 from vllm.config.lora import LoRAConfig
13 from vllm.logger import init_logger
14 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
15
16 logger = init_logger(__name__)
17
18
19 @dataclass
20 class PEFTHelper:
21 """A helper class for PEFT configurations, specifically designed for LoRA.
22 This class handles configuration validation, compatibility checks for
23 various LoRA implementations.
24 """
25
26 # Required fields
27 r: int
28 lora_alpha: int
29 target_modules: list[str] | str
30
31 bias: Literal["none"] = field(default="none")
32 modules_to_save: list[str] | None = field(default=None)
33 # True to use Rank-Stabilized LoRA (rsLoRA, see: https://arxiv.org/abs/2312.03732)
34 use_rslora: bool = field(default=False)
35 # True to use Weight-Decomposed Low-Rank Adaptation (DoRA, see: https://arxiv.org/abs/2402.09353)
36 use_dora: bool = field(default=False)
37 # Extra vllm field, start with 'vllm_' to avoid conflict
38 vllm_lora_scaling_factor: float = field(default=1.0)
39 vllm_max_position_embeddings: int | None = field(default=False)
40
41 def _validate_features(self) -> list[str]:
42 """Check if there are any unsupported LoRA features."""
43 error_msg = []
44 if self.modules_to_save:
45 unsupported_modules = [
46 m for m in self.modules_to_save if m not in ["classifier", "score"]
47 ]
48 if unsupported_modules:
49 error_msg.append(
50 "vLLM only supports modules_to_save being either None "
51 'or ["classifier", "score"] for classification models. '
52 f"Unsupported modules_to_save: {unsupported_modules}"
53 )
54 if self.use_dora:
55 error_msg.append("vLLM does not yet support DoRA.")
56 return error_msg
57
58 def __post_init__(self):
59 if self.r <= 0:
60 raise ValueError(f"LoRA rank `r` must be a positive integer, got {self.r}.")
61 if self.use_rslora:
62 logger.info_once("Loading LoRA weights trained with rsLoRA.")
63 self.vllm_lora_scaling_factor = self.lora_alpha / math.sqrt(self.r)
64 else:
65 self.vllm_lora_scaling_factor = self.lora_alpha / self.r
66
67 @classmethod
68 def from_dict(cls, config_dict: dict) -> "PEFTHelper":
69 # Get all field information from the class
70 class_fields = {f.name: f for f in fields(cls)}
71 # Check for required fields
72 required_fields = {
73 name
74 for name, f in class_fields.items()
75 if f.default is MISSING and f.default_factory is MISSING
76 }
77
78 # Identify any missing required fields
79 missing_fields = required_fields - set(config_dict.keys())
80 if missing_fields:
81 raise ValueError(f"Missing required configuration fields: {missing_fields}")
82
83 # Filter out fields that aren't defined in the class
84 filtered_dict = {k: v for k, v in config_dict.items() if k in class_fields}
85 return cls(**filtered_dict)
86
87 @classmethod
88 def from_local_dir(
89 cls,
90 lora_path: str,
91 max_position_embeddings: int | None,
92 tensorizer_config_dict: dict | None = None,
93 ) -> "PEFTHelper":
94 lora_config_path = os.path.join(lora_path, "adapter_config.json")
95
96 if tensorizer_config_dict:
97 tensorizer_config = TensorizerConfig(**tensorizer_config_dict)
98 tensorizer_args = tensorizer_config._construct_tensorizer_args()
99 from tensorizer.stream_io import open_stream
100
101 tensorizer_dir = tensorizer_config.tensorizer_dir
102 if tensorizer_dir is None:
103 raise ValueError("tensorizer_dir must be set in tensorizer config.")
104
105 lora_config_path = os.path.join(tensorizer_dir, "adapter_config.json")
106 with open_stream(
107 lora_config_path, mode="rb", **tensorizer_args.stream_kwargs
108 ) as f:
109 config = json.load(f)
110
111 logger.info(
112 "Successfully deserialized LoRA config from %s",
113 tensorizer_config.tensorizer_dir,
114 )
115
116 else:
117 with open(lora_config_path) as f:
118 config = json.load(f)
119
120 config["vllm_max_position_embeddings"] = max_position_embeddings
121 return cls.from_dict(config)
122
123 def validate_legal(self, lora_config: LoRAConfig) -> None:
124 """Validates the LoRA configuration settings against application
125 constraints and requirements.
126 """
127 error_msg = self._validate_features()
128 if self.r > lora_config.max_lora_rank:
129 error_msg.append(
130 f"LoRA rank {self.r} is greater than max_lora_rank"
131 f" {lora_config.max_lora_rank}."
132 )
133 if self.bias != "none":
134 error_msg.append("Adapter bias is not supported.")
135 if error_msg:
136 raise ValueError(f"{' '.join(error_msg)}")