@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c81eb3e69c

[LoRA] Support modules_to_save for sequence classification (#53555)

Signed-off-by: linitra24 <Joy25810@foxmail.com> Signed-off-by: linitra24 <renshuang.zhou@daocloud.io> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

linitra24 committed Sep 16, 2026 at 22:15 UTC c81eb3e69c40ab17cb3c4cd8350e4a60b1be70d4
20 files changed +681 -21
.buildkite/test_areas/lora.yaml
+2
@@ -54,6 +54,7 @@ steps:
54 - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_gptoss_tp.py
55 - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s -x lora/test_qwen35_densemodel_lora.py
56 - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s -x lora/test_gemma4_tp.py
57 + - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp
58 mirror:
59 amd:
60 label: ":amd: (MI300) LoRA TP (Distributed) %N"
@@ -79,3 +80,4 @@ steps:
80 - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_gptoss_tp.py
81 - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s -x lora/test_qwen35_densemodel_lora.py
82 - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s -x lora/test_gemma4_tp.py
83 + - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp
docs/features/lora.md
+12
@@ -422,6 +422,18 @@ vllm serve ibm-granite/granite-speech-3.3-2b \
422
423 Note: Default multimodal LoRAs are currently only available for `.generate` and chat completions.
424
425 +## Sequence-Classification LoRA Adapters
426 +
427 +vLLM supports PEFT sequence-classification adapters that save a complete, single-layer linear classification head through `modules_to_save`. The saved module must be named `score` or `classifier`.
428 +
429 +See [classification_with_lora_offline.py](../../examples/pooling/classify/classification_with_lora_offline.py) for an offline classification example using a LoRA adapter.
430 +
431 +This support has the following limitations:
432 +
433 +- All adapters in one engine must have the same `num_labels` and hidden size as the base classification head. An incompatible adapter is rejected when it is loaded.
434 +- A classification head stored as float32 is converted to the runtime head dtype when it is loaded.
435 +- Token-classification adapters are not supported by this feature.
436 +
437 ## Using Tips
438
439 ### Configuring `max_lora_rank`
examples/pooling/classify/classification_with_lora_offline.py new
+32
@@ -0,0 +1,32 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Example of offline classification with a LoRA adapter."""
4 +
5 +from vllm import LLM
6 +from vllm.lora.request import LoRARequest
7 +from vllm.transformers_utils.repo_utils import hf_api
8 +
9 +
10 +def main():
11 + lora_path = hf_api().snapshot_download(
12 + repo_id="AmirMohseni/skywork-qwen3-0.6b-reward-lora"
13 + )
14 + llm = LLM(
15 + model="Skywork/Skywork-Reward-V2-Qwen3-0.6B",
16 + runner="pooling",
17 + enable_lora=True,
18 + max_lora_rank=8,
19 + max_model_len=512,
20 + )
21 +
22 + prompt = "Which response is more helpful and correct?"
23 + (output,) = llm.classify(
24 + prompt,
25 + lora_request=LoRARequest("reward-lora", 1, lora_path),
26 + )
27 + print(f"Prompt: {prompt!r}")
28 + print(f"Class Probabilities: {output.outputs.probs}")
29 +
30 +
31 +if __name__ == "__main__":
32 + main()
tests/entrypoints/serve/lora/test_lora_adapters.py
+1 -1
@@ -26,7 +26,7 @@ BADREQUEST_CASES = [
26 (
27 "test_modules_to_save",
28 {"modules_to_save": ["lm_head"]},
29 - "only supports modules_to_save being None",
29 + "Unsupported modules_to_save",
30 ),
31 ]
32
tests/lora/conftest.py
+24
@@ -270,6 +270,30 @@ def qwen3_woofing_lora_files():
270 return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA")
271
272
273 +@pytest.fixture(scope="session")
274 +def qwen3_guard_star_trek_lora_files():
275 + return hf_api().snapshot_download(
276 + repo_id="geoffmunn/Qwen3Guard-StarTrek-Classification-0.6B",
277 + allow_patterns=["adapter_config.json", "adapter_model.safetensors"],
278 + )
279 +
280 +
281 +@pytest.fixture(scope="session")
282 +def qwen3_guard_new_zealand_lora_files():
283 + return hf_api().snapshot_download(
284 + repo_id="geoffmunn/Qwen3Guard-NewZealand-Classification-0.6B",
285 + allow_patterns=["adapter_config.json", "adapter_model.safetensors"],
286 + )
287 +
288 +
289 +@pytest.fixture(scope="session")
290 +def skywork_qwen3_reward_lora_files():
291 + return hf_api().snapshot_download(
292 + repo_id="AmirMohseni/skywork-qwen3-0.6b-reward-lora",
293 + allow_patterns=["adapter_config.json", "adapter_model.safetensors"],
294 + )
295 +
296 +
297 @pytest.fixture(scope="session")
298 def deepseekv2_lora_files():
299 return hf_api().snapshot_download(repo_id="wuchen01/DeepSeek-V2-Lite-Chat-All-LoRA")
tests/lora/test_peft_helper.py
+1 -1
@@ -20,7 +20,7 @@ ERROR_CASES = [
20 (
21 "test_modules_to_save",
22 {"modules_to_save": ["lm_head"]},
23 - "only supports modules_to_save being None",
23 + "Unsupported modules_to_save",
24 ),
25 ("test_rank_zero", {"r": 0}, "must be a positive integer"),
26 ("test_rank_negative", {"r": -8}, "must be a positive integer"),
tests/lora/test_sequence_classification.py new
+167
@@ -0,0 +1,167 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from tests.utils import multi_gpu_test
7 +from vllm import LLM, PoolingParams
8 +from vllm.distributed import cleanup_dist_env_and_memory
9 +from vllm.lora.request import LoRARequest
10 +
11 +MODEL_NAME = "Qwen/Qwen3-0.6B"
12 +NATIVE_MODEL_NAME = "Skywork/Skywork-Reward-V2-Qwen3-0.6B"
13 +PROMPTS = {
14 + "star_trek": "Does warp drive appear in Star Trek?",
15 + "new_zealand": "Wellington is the capital of New Zealand.",
16 +}
17 +EXPECTED_OUT = {
18 + "star_trek": [13.6640625, -13.6640625],
19 + "new_zealand": [-4.80078125, 8.6171875],
20 + "native": [-1.52734375],
21 +}
22 +
23 +
24 +def _classify_logits(
25 + llm: LLM,
26 + prompts: list[str],
27 + requests: LoRARequest | list[LoRARequest],
28 +) -> torch.Tensor:
29 + params = [PoolingParams(task="classify", use_activation=False) for _ in prompts]
30 + outputs = llm.classify(
31 + prompts,
32 + pooling_params=params,
33 + lora_request=requests,
34 + use_tqdm=False,
35 + )
36 + return torch.tensor([output.outputs.probs for output in outputs])
37 +
38 +
39 +def test_converted_model_with_modules_to_save(
40 + qwen3_guard_star_trek_lora_files: str,
41 +) -> None:
42 + prompt = PROMPTS["star_trek"]
43 + llm = LLM(
44 + model=MODEL_NAME,
45 + runner="pooling",
46 + convert="classify",
47 + hf_overrides={"num_labels": 2},
48 + dtype="float16",
49 + enable_lora=True,
50 + max_lora_rank=16,
51 + enforce_eager=True,
52 + max_model_len=512,
53 + gpu_memory_utilization=0.5,
54 + )
55 + actual = _classify_logits(
56 + llm,
57 + [prompt],
58 + LoRARequest("star-trek", 1, qwen3_guard_star_trek_lora_files),
59 + )[0]
60 +
61 + torch.testing.assert_close(
62 + actual,
63 + torch.tensor(EXPECTED_OUT["star_trek"]),
64 + atol=2e-2,
65 + rtol=2e-2,
66 + )
67 +
68 + del llm
69 + cleanup_dist_env_and_memory()
70 +
71 +
72 +def test_native_classification_model_with_modules_to_save(
73 + skywork_qwen3_reward_lora_files: str,
74 +) -> None:
75 + prompt = "Which response is more helpful and correct?"
76 + llm = LLM(
77 + model=NATIVE_MODEL_NAME,
78 + runner="pooling",
79 + dtype="float16",
80 + enable_lora=True,
81 + max_lora_rank=8,
82 + enforce_eager=True,
83 + max_model_len=512,
84 + gpu_memory_utilization=0.5,
85 + )
86 + actual = _classify_logits(
87 + llm,
88 + [prompt],
89 + LoRARequest("native", 1, skywork_qwen3_reward_lora_files),
90 + )[0]
91 +
92 + torch.testing.assert_close(
93 + actual,
94 + torch.tensor(EXPECTED_OUT["native"]),
95 + atol=2e-2,
96 + rtol=2e-2,
97 + )
98 +
99 + del llm
100 + cleanup_dist_env_and_memory()
101 +
102 +
103 +def _test_batched_loras(
104 + qwen3_guard_star_trek_lora_files: str,
105 + qwen3_guard_new_zealand_lora_files: str,
106 + tp_size: int,
107 +) -> None:
108 + names = ["star_trek", "new_zealand"]
109 + prompts = [PROMPTS[name] for name in names]
110 + adapter_paths = [
111 + qwen3_guard_star_trek_lora_files,
112 + qwen3_guard_new_zealand_lora_files,
113 + ]
114 + requests = [
115 + LoRARequest(name, index, adapter_path)
116 + for index, (name, adapter_path) in enumerate(zip(names, adapter_paths), start=1)
117 + ]
118 +
119 + llm = LLM(
120 + model=MODEL_NAME,
121 + runner="pooling",
122 + convert="classify",
123 + hf_overrides={"num_labels": 2},
124 + dtype="float16",
125 + enable_lora=True,
126 + max_loras=2,
127 + max_cpu_loras=2,
128 + max_lora_rank=16,
129 + enforce_eager=True,
130 + max_model_len=512,
131 + gpu_memory_utilization=0.5,
132 + tensor_parallel_size=tp_size,
133 + )
134 + actual = _classify_logits(llm, prompts, requests)
135 +
136 + torch.testing.assert_close(
137 + actual,
138 + torch.tensor([EXPECTED_OUT[name] for name in names]),
139 + atol=2e-2,
140 + rtol=2e-2,
141 + )
142 +
143 + del llm
144 + cleanup_dist_env_and_memory()
145 +
146 +
147 +def test_batched_loras(
148 + qwen3_guard_star_trek_lora_files: str,
149 + qwen3_guard_new_zealand_lora_files: str,
150 +) -> None:
151 + _test_batched_loras(
152 + qwen3_guard_star_trek_lora_files,
153 + qwen3_guard_new_zealand_lora_files,
154 + tp_size=1,
155 + )
156 +
157 +
158 +@multi_gpu_test(num_gpus=2)
159 +def test_batched_loras_tp(
160 + qwen3_guard_star_trek_lora_files: str,
161 + qwen3_guard_new_zealand_lora_files: str,
162 +) -> None:
163 + _test_batched_loras(
164 + qwen3_guard_star_trek_lora_files,
165 + qwen3_guard_new_zealand_lora_files,
166 + tp_size=2,
167 + )
vllm/lora/layers/__init__.py
+2
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from vllm.lora.layers.base import BaseLayerWithLoRA
4 +from vllm.lora.layers.classifier import ClassificationHeadWithLoRA
5 from vllm.lora.layers.column_parallel_linear import (
6 ColumnParallelLinearWithLoRA,
7 ColumnParallelLinearWithShardedLoRA,
@@ -42,4 +43,5 @@ __all__ = [
43 "LoRAMappingType",
44 "FusedMoEWithLoRA",
45 "FusedMoE3DWithLoRA",
46 + "ClassificationHeadWithLoRA",
47 ]
vllm/lora/layers/classifier.py new
+97
@@ -0,0 +1,97 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import torch
6 +from transformers import PretrainedConfig
7 +
8 +from vllm.config.lora import LoRAConfig
9 +from vllm.platforms import current_platform
10 +
11 +from .replicated_linear import ReplicatedLinearWithLoRA
12 +
13 +
14 +class ClassificationHeadWithLoRA(ReplicatedLinearWithLoRA):
15 + def create_lora_weights(
16 + self,
17 + max_loras: int,
18 + lora_config: LoRAConfig,
19 + model_config: PretrainedConfig | None = None,
20 + ) -> None:
21 + # Preserve ordinary LoRA A/B support for classification heads.
22 + super().create_lora_weights(max_loras, lora_config, model_config)
23 + self.full_weight_stacked = torch.zeros(
24 + max_loras,
25 + 1,
26 + self.output_size,
27 + self.input_size,
28 + dtype=lora_config.lora_dtype,
29 + device=self.device,
30 + )
31 + self.full_bias_stacked = torch.zeros(
32 + max_loras,
33 + self.output_size,
34 + dtype=self.base_layer.params_dtype,
35 + device=self.device,
36 + )
37 + self.full_module_enabled = torch.zeros(
38 + max_loras, dtype=torch.bool, device=self.device
39 + )
40 +
41 + def reset_module_to_save(self, index: int) -> None:
42 + self.full_weight_stacked[index].zero_()
43 + self.full_bias_stacked[index].zero_()
44 + self.full_module_enabled[index] = False
45 +
46 + def set_module_to_save(
47 + self,
48 + index: int,
49 + weight: torch.Tensor,
50 + bias: torch.Tensor | None,
51 + ) -> None:
52 + self.reset_module_to_save(index)
53 + self.full_weight_stacked[index, 0].copy_(weight, non_blocking=True)
54 + if bias is not None:
55 + self.full_bias_stacked[index].copy_(bias, non_blocking=True)
56 + self.full_module_enabled[index] = True
57 +
58 + def forward(self, input_: torch.Tensor) -> torch.Tensor:
59 + # TODO base_result maybe don't need compute truly.
60 + base_result = super().forward(input_)
61 + output = base_result[0] if isinstance(base_result, tuple) else base_result
62 +
63 + full_output = self.punica_wrapper.apply_lora_full_linear(
64 + output,
65 + input_.to(self.full_weight_stacked.dtype),
66 + self.full_weight_stacked,
67 + self.full_bias_stacked,
68 + self.full_module_enabled,
69 + )
70 +
71 + if full_output is not None:
72 + output = full_output
73 + return output
74 +
75 + def _apply_lora_to_output(
76 + self, x: torch.Tensor, output: torch.Tensor
77 + ) -> torch.Tensor:
78 + original_shape = output.shape if output.ndim == 3 else None
79 + if x.ndim == 3 and output.ndim == 3:
80 + output = output.flatten(0, 1)
81 + x = x.flatten(0, 1)
82 +
83 + # Classification outputs are request-level, so use prompt-based LoRA routing.
84 + lora_output: torch.Tensor | None = self.punica_wrapper.add_lora_logits(
85 + output,
86 + x.to(self.lora_a_stacked[0].dtype),
87 + self.lora_a_stacked[0],
88 + self.lora_b_stacked[0],
89 + 1.0,
90 + )
91 + if not current_platform.can_update_inplace():
92 + output = lora_output
93 +
94 + if original_shape is not None:
95 + output = output.reshape(original_shape)
96 +
97 + return output
vllm/lora/lora_model.py
+38 -3
@@ -8,7 +8,7 @@ import safetensors
8 import torch
9
10 from vllm.logger import init_logger
11 -from vllm.lora.lora_weights import LoRALayerWeights
11 +from vllm.lora.lora_weights import LoRAFullModuleWeights, LoRALayerWeights
12 from vllm.lora.peft_helper import PEFTHelper
13 from vllm.lora.utils import (
14 get_lora_id,
@@ -66,6 +66,8 @@ class LoRAModel:
66 rank: int,
67 loras: dict[str, LoRALayerWeights],
68 is_3d_lora_weight: bool = False,
69 + *,
70 + modules_to_save: dict[str, LoRAFullModuleWeights] | None = None,
71 ) -> None:
72 """
73 Args:
@@ -85,6 +87,7 @@ class LoRAModel:
87 )
88 self.rank = rank
89 self.loras: dict[str, LoRALayerWeights] = loras
90 + self.modules_to_save = modules_to_save or {}
91 self.is_3d_lora_weight = is_3d_lora_weight
92
93 def clone(self, lora_model_id: int) -> "LoRAModel":
@@ -96,12 +99,16 @@ class LoRAModel:
99 rank=self.rank,
100 loras=self.loras.copy(),
101 is_3d_lora_weight=self.is_3d_lora_weight,
102 + modules_to_save=self.modules_to_save.copy(),
103 )
104
105 def get_lora(self, module_name: str) -> LoRALayerWeights | None:
106 """Get LoRA for a given module by name"""
107 return self.loras.get(module_name, None)
108
109 + def get_module_to_save(self, module_name: str) -> LoRAFullModuleWeights | None:
110 + return self.modules_to_save.get(module_name)
111 +
112 def check_lora_name(self, lora_name: str) -> bool:
113 return lora_name in self.loras
114
@@ -127,6 +134,9 @@ class LoRAModel:
134 ) -> "LoRAModel":
135 """Create a LoRAModel from a dictionary of tensors."""
136 pin_memory = str(device) == "cpu" and PIN_MEMORY
137 +
138 + modules_to_save_names = peft_helper.modules_to_save or []
139 + full_parameters: dict[str, dict[str, torch.Tensor]] = {}
140 loras: dict[str, LoRALayerWeights] = {}
141 for tensor_name, tensor in tensors.items():
142 if is_base_embedding_weights(tensor_name):
@@ -137,6 +147,12 @@ class LoRAModel:
147 module_name, is_lora_a = parse_fine_tuned_lora_name(
148 tensor_name, weights_mapper
149 )
150 + if module_name in modules_to_save_names:
151 + full_parameters.setdefault(module_name, {})[
152 + tensor_name.split(".")[-1]
153 + ] = tensor.to(device=device)
154 + continue
155 +
156 if module_name not in loras:
157 loras[module_name] = LoRALayerWeights.from_config(
158 module_name, peft_helper
@@ -161,7 +177,23 @@ class LoRAModel:
177 if pin_memory:
178 loras[module_name].lora_b = loras[module_name].lora_b.pin_memory()
179
164 - return cls(lora_model_id, peft_helper.r, loras)
180 + modules_to_save = {}
181 + for module_name, parameters in full_parameters.items():
182 + weight = parameters.get("weight")
183 + if weight is None:
184 + raise ValueError(f"Full module {module_name!r} is missing weight.")
185 + modules_to_save[module_name] = LoRAFullModuleWeights(
186 + module_name=module_name,
187 + weight=weight,
188 + bias=parameters.get("bias"),
189 + )
190 + if len(modules_to_save) > 1:
191 + raise ValueError(
192 + "Only one full classification module is supported per "
193 + f"adapter, received {sorted(modules_to_save)}."
194 + )
195 +
196 + return cls(lora_model_id, peft_helper.r, loras, modules_to_save=modules_to_save)
197
198 @classmethod
199 def from_local_checkpoint(
@@ -223,6 +255,7 @@ class LoRAModel:
255 ):
256 continue
257 module_name, _ = parse_fine_tuned_lora_name(lora_module, weights_mapper)
258 + base_name = module_name.rsplit(".", 1)[-1]
259 # Case for expert lora weights
260 if ".experts" in module_name:
261 expert_idx = module_name.find(".experts")
@@ -230,7 +263,9 @@ class LoRAModel:
263 if expert_suffix not in expected_lora_modules:
264 unexpected_modules.append(module_name)
265
233 - elif module_name.rsplit(".", 1)[-1] not in expected_lora_modules:
266 + elif base_name not in expected_lora_modules and base_name not in (
267 + peft_helper.modules_to_save or ()
268 + ):
269 unexpected_modules.append(module_name)
270
271 if unexpected_modules:
vllm/lora/lora_weights.py
+15
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Sequence as GenericSequence
5 +from dataclasses import dataclass
6
7 import torch
8 import torch.types
@@ -281,3 +282,17 @@ class PackedLoRALayerWeights(LoRALayerWeights):
282 @property
283 def is_packed(self) -> bool:
284 return True
285 +
286 +
287 +@dataclass
288 +class LoRAFullModuleWeights:
289 + """LoRA weights for classification layers."""
290 +
291 + module_name: str
292 + weight: torch.Tensor
293 + bias: torch.Tensor | None = None
294 +
295 + def pin_memory(self) -> None:
296 + self.weight = self.weight.pin_memory()
297 + if self.bias is not None:
298 + self.bias = self.bias.pin_memory()
vllm/lora/model_manager.py
+148 -8
@@ -13,16 +13,22 @@ from vllm.config.lora import LoRAConfig
13 from vllm.logger import init_logger
14 from vllm.lora.layers import (
15 BaseLayerWithLoRA,
16 + ClassificationHeadWithLoRA,
17 FusedMoE3DWithLoRA,
18 FusedMoEWithLoRA,
19 LoRAMapping,
20 LoRAMappingType,
21 )
22 from vllm.lora.lora_model import LoRAModel, MoEEPLoadSpec
22 -from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
23 +from vllm.lora.lora_weights import (
24 + LoRAFullModuleWeights,
25 + LoRALayerWeights,
26 + PackedLoRALayerWeights,
27 +)
28 from vllm.lora.punica_wrapper import PunicaWrapperBase, get_punica_wrapper
29 from vllm.lora.utils import (
30 from_layer,
31 + from_layer_classification,
32 from_layer_logits_processor,
33 get_supported_lora_modules,
34 is_in_target_modules,
@@ -114,6 +120,13 @@ class LoRAModelManager:
120 self.vocab_size = vocab_size
121
122 self.is_pooling_model = is_pooling_model(self.model)
123 + self._pooling_task = (
124 + vllm_config.model_config.get_pooling_task(
125 + tuple(self.model.pooler.get_supported_tasks())
126 + )
127 + if self.is_pooling_model
128 + else None
129 + )
130 self.packed_modules: dict[str, list[str]] = {}
131 self.modules: dict[str, BaseLayerWithLoRA] = {}
132 self._last_mapping: LoRAMapping | None = None
@@ -121,6 +134,13 @@ class LoRAModelManager:
134 is_moe = is_moe_model(self.model)
135 self._is_moe = is_moe
136
137 + self.supported_modules_to_save = (
138 + {"score", "classifier"}
139 + if self.is_pooling_model
140 + and vllm_config.model_config.score_type == "cross-encoder"
141 + else set()
142 + )
143 +
144 # When the engine is started with enable_mixed_moe_lora_format=True
145 # we force the universal 2D wrapper (FusedMoEWithLoRA) regardless of
146 # the model's 3D flag, so 2D and 3D adapters can coexist.
@@ -147,6 +167,20 @@ class LoRAModelManager:
167 self._init_punica_wrapper(max_num_batched_tokens, vllm_config)
168 self._create_lora_modules()
169
170 + self._classification_head: tuple[str, ClassificationHeadWithLoRA] | None = None
171 + if self.is_pooling_model:
172 + classification_heads = [
173 + (module_name, module)
174 + for module_name, module in self.modules.items()
175 + if isinstance(module, ClassificationHeadWithLoRA)
176 + ]
177 + if classification_heads:
178 + assert len(classification_heads) == 1, (
179 + "Expected 1 classification head, but found "
180 + f"{len(classification_heads)}."
181 + )
182 + self._classification_head = classification_heads[0]
183 +
184 self.moe_ep_load_spec: MoEEPLoadSpec | None = self._build_moe_ep_load_spec()
185
186 self.model.lora_manager = self
@@ -338,6 +372,15 @@ class LoRAModelManager:
372 self.lora_index_to_id[index] = lora_model.id
373 num_lora_weights_applied = 0
374 for module_name, module in self.modules.items():
375 + if isinstance(module, ClassificationHeadWithLoRA):
376 + full_module = self._get_module_to_save_weights(lora_model, module_name)
377 + if full_module is None:
378 + module.reset_module_to_save(index)
379 + else:
380 + module.set_module_to_save(
381 + index, full_module.weight, full_module.bias
382 + )
383 +
384 module_lora = self._get_lora_layer_weights(lora_model, module_name)
385 if not module_lora:
386 module.reset_lora(index)
@@ -427,6 +470,8 @@ class LoRAModelManager:
470 if isinstance(module, PPMissingLayer):
471 continue
472
473 + is_classifier_head = module_name in self.supported_modules_to_save
474 +
475 if self.lora_config.target_modules is None:
476 if not is_supported_lora_module(
477 module_name,
@@ -434,7 +479,8 @@ class LoRAModelManager:
479 self.supported_lora_modules,
480 ):
481 continue
437 - elif not self._match_target_modules(module_name):
482 + elif not self._match_target_modules(module_name) and not is_classifier_head:
483 + # Full classification heads bypass the A/B LoRA target filter.
484 continue
485
486 punica_wrapper = self._get_punica_wrapper(module_name)
@@ -486,17 +532,25 @@ class LoRAModelManager:
532 # LoRA weights of w1 and w3 have already been fused on disk.
533
534 packed_moduled_lst = ["w13"] if self._is_3d_moe_model else ["w1", "w3"]
489 - new_module = replace_submodule(
490 - self.model,
491 - module_name,
492 - from_layer(
535 + if is_classifier_head:
536 + new_module = from_layer_classification(
537 + module,
538 + self.lora_slots,
539 + self.lora_config,
540 + self.model.config,
541 + )
542 + else:
543 + new_module = from_layer(
544 module,
545 self.lora_slots,
546 self.lora_config,
547 packed_moduled_lst,
548 self.model.config,
498 - ),
499 - )
549 + )
550 + new_module = replace_submodule(self.model, module_name, new_module)
551 +
552 + if is_classifier_head:
553 + self.model.pooler.replace_classifier(new_module)
554 if isinstance(new_module, BaseLayerWithLoRA):
555 wrapped_by_id[id(module)] = new_module
556 wrapped_by_id[id(new_module)] = new_module
@@ -746,6 +800,12 @@ class LoRAModelManager:
800 if not self.supports_mm:
801 return self.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY]
802
803 + if module_name in self.supported_modules_to_save:
804 + lm_wrapper = self.punica_wrapper_mapping.get(
805 + self.mm_mapping.language_model[0]
806 + )
807 + return lm_wrapper
808 +
809 # For multimodal model
810 # NOTE Sort by prefix length (descending) to match the longest prefix first
811 # e.g., 'visual.merger' should match 'visual.merger' instead of 'visual.'
@@ -871,6 +931,10 @@ class LoRAModelManager:
931 lora.lora_a = lora.lora_a.pin_memory()
932 lora.lora_b = lora.lora_b.pin_memory()
933
934 + if lora_model.modules_to_save:
935 + for full_module in lora_model.modules_to_save.values():
936 + full_module.pin_memory()
937 +
938 def _stack_moe_lora_weights(
939 self, lora_model: LoRAModel, module: FusedMoE3DWithLoRA, module_name: str
940 ):
@@ -1164,6 +1228,82 @@ class LoRAModelManager:
1228 )
1229 return lora_model.get_lora(org_module_name)
1230
1231 + def _get_module_to_save_weights(
1232 + self, lora_model: LoRAModel, module_name: str
1233 + ) -> LoRAFullModuleWeights | None:
1234 + weights = lora_model.get_module_to_save(module_name)
1235 + if weights is not None:
1236 + return weights
1237 +
1238 + unprefixed_module_name = module_name.removeprefix("model.")
1239 + weights = lora_model.get_module_to_save(unprefixed_module_name)
1240 + if weights is not None and unprefixed_module_name != module_name:
1241 + logger.info_once(
1242 + "For the pool model, successfully loaded the full module weights "
1243 + "after removing the prefix 'model.'."
1244 + )
1245 + return weights
1246 +
1247 + def _validate_modules_to_save(self, lora_model: LoRAModel) -> None:
1248 + if not lora_model.modules_to_save:
1249 + return
1250 + saved_module_name = next(iter(lora_model.modules_to_save))
1251 + if self._classification_head is None:
1252 + raise ValueError(
1253 + f"Cannot load full module {saved_module_name!r}: the model does "
1254 + "not expose a unique ClassificationHeadWithLoRA."
1255 + )
1256 +
1257 + module_name, wrapper = self._classification_head
1258 + full_module = self._get_module_to_save_weights(lora_model, module_name)
1259 + if full_module is None:
1260 + raise ValueError(
1261 + f"Full module {saved_module_name!r} does not match the model's "
1262 + f"classification head {module_name!r}."
1263 + )
1264 +
1265 + expected_weight_shape = (wrapper.output_size, wrapper.input_size)
1266 + received_weight_shape = tuple(full_module.weight.shape)
1267 + if received_weight_shape != expected_weight_shape:
1268 + raise ValueError(
1269 + f"Full module {saved_module_name!r} for {module_name!r} has "
1270 + "an incompatible weight shape: expected "
1271 + f"{expected_weight_shape}, received {received_weight_shape}."
1272 + )
1273 +
1274 + if full_module.bias is None:
1275 + return
1276 + expected_bias_shape = (wrapper.output_size,)
1277 + received_bias_shape = tuple(full_module.bias.shape)
1278 + if received_bias_shape != expected_bias_shape:
1279 + raise ValueError(
1280 + f"Full module {saved_module_name!r} for {module_name!r} has "
1281 + "an incompatible bias shape: expected "
1282 + f"{expected_bias_shape}, received {received_bias_shape}."
1283 + )
1284 +
1285 + def _validate_token_classification_lora(self, lora_model: LoRAModel) -> None:
1286 + if self._pooling_task != "token_classify":
1287 + return
1288 + if self._classification_head is None:
1289 + return
1290 +
1291 + module_name, _ = self._classification_head
1292 + has_full_module = (
1293 + self._get_module_to_save_weights(lora_model, module_name) is not None
1294 + )
1295 + has_lora_weights = (
1296 + self._get_lora_layer_weights(lora_model, module_name) is not None
1297 + )
1298 + if not has_full_module and not has_lora_weights:
1299 + return
1300 +
1301 + raise ValueError(
1302 + f"LoRA adapter {lora_model.id} contains weights for classification "
1303 + f"head {module_name!r}, but token_classify only supports LoRA on "
1304 + "the model backbone."
1305 + )
1306 +
1307 def deactivate_adapter(self, adapter_id: int) -> bool:
1308 if adapter_id not in self._active_adapters:
1309 return False
vllm/lora/peft_helper.py
+9 -1
@@ -45,7 +45,15 @@ class PEFTHelper:
45 """
46 error_msg = []
47 if self.modules_to_save:
48 - error_msg.append("vLLM only supports modules_to_save being None.")
48 + unsupported_modules = [
49 + m for m in self.modules_to_save if m not in ["classifier", "score"]
50 + ]
51 + if unsupported_modules:
52 + error_msg.append(
53 + "vLLM only supports modules_to_save being either None "
54 + 'or ["classifier", "score"] for classification models. '
55 + f"Unsupported modules_to_save: {unsupported_modules}"
56 + )
57 if self.use_dora:
58 error_msg.append("vLLM does not yet support DoRA.")
59 return error_msg
vllm/lora/punica_wrapper/punica_base.py
+28
@@ -120,6 +120,18 @@ class PunicaWrapperABC(ABC):
120 """
121 raise NotImplementedError
122
123 + @abstractmethod
124 + def apply_lora_full_linear(
125 + self,
126 + y: torch.Tensor,
127 + x: torch.Tensor,
128 + weight_stacked: torch.Tensor,
129 + bias_stacked: torch.Tensor,
130 + module_enabled: torch.Tensor,
131 + ) -> torch.Tensor | None:
132 + """Apply request-routed full linear weights to selected rows."""
133 + raise NotImplementedError
134 +
135
136 class PunicaWrapperBase(PunicaWrapperABC):
137 """
@@ -217,6 +229,22 @@ class PunicaWrapperBase(PunicaWrapperABC):
229 self.token_nums = token_nums
230 self.no_lora = no_lora
231
232 + def _select_full_linear_output(
233 + self,
234 + y: torch.Tensor,
235 + adapter_y: torch.Tensor,
236 + bias_stacked: torch.Tensor,
237 + module_enabled: torch.Tensor,
238 + ) -> torch.Tensor:
239 + indices = self.sampler_indices
240 + assert indices.size(0) == y.size(0), (
241 + "Full linear rows do not match LoRA request mapping"
242 + )
243 + safe_indices = indices.clamp_min(0)
244 + use_full = (indices >= 0) & module_enabled[safe_indices]
245 + adapter_y = adapter_y.to(y.dtype) + bias_stacked[safe_indices].to(y.dtype)
246 + return torch.where(use_full.unsqueeze(-1), adapter_y, y)
247 +
248 @property
249 def prefill_metadata(
250 self,
vllm/lora/punica_wrapper/punica_cpu.py
+20
@@ -349,3 +349,23 @@ class PunicaWrapperCPU(PunicaWrapperBase):
349 bgmv_shrink(x, lora_a_stacked, buffer, self.sampler_indices, scale)
350 bgmv_expand(buffer, lora_b_stacked, y, self.sampler_indices, add_inputs=True)
351 y = y.view_as(y_org)
352 +
353 + def apply_lora_full_linear(
354 + self,
355 + y: torch.Tensor,
356 + x: torch.Tensor,
357 + weight_stacked: torch.Tensor,
358 + bias_stacked: torch.Tensor,
359 + module_enabled: torch.Tensor,
360 + ) -> None:
361 + indices = self.sampler_indices
362 + adapter_y = torch.zeros(
363 + (x.size(0), weight_stacked.size(-2)),
364 + dtype=torch.float32,
365 + device=x.device,
366 + )
367 + bgmv_shrink(x, weight_stacked, adapter_y, indices, 1.0)
368 + result = self._select_full_linear_output(
369 + y, adapter_y, bias_stacked, module_enabled
370 + )
371 + y.copy_(result)
vllm/lora/punica_wrapper/punica_gpu.py
+27
@@ -327,6 +327,33 @@ class PunicaWrapperGPU(PunicaWrapperBase):
327 )
328 y = y.view_as(y_org)
329
330 + def apply_lora_full_linear(
331 + self,
332 + y: torch.Tensor,
333 + x: torch.Tensor,
334 + weight_stacked: torch.Tensor,
335 + bias_stacked: torch.Tensor,
336 + module_enabled: torch.Tensor,
337 + ) -> None:
338 + adapter_y = torch.empty(
339 + (x.size(0), weight_stacked.size(-2)),
340 + dtype=torch.float32,
341 + device=x.device,
342 + )
343 + lora_shrink(
344 + x,
345 + [weight_stacked],
346 + adapter_y.unsqueeze(0),
347 + *self.prompt_mapping_meta.meta_args(
348 + x.size(0), self.lora_config.specialize_active_lora
349 + ),
350 + 1.0,
351 + )
352 + result = self._select_full_linear_output(
353 + y, adapter_y, bias_stacked, module_enabled
354 + )
355 + y.copy_(result)
356 +
357 def moe_lora_align_block_size(
358 self,
359 topk_ids: torch.Tensor,
vllm/lora/punica_wrapper/punica_xpu.py
+20
@@ -333,6 +333,26 @@ class PunicaWrapperXPU(PunicaWrapperBase):
333 bgmv_expand(buffer, lora_b_stacked, y, sampler_indices, add_inputs=True)
334 y = y.view_as(y_org)
335
336 + def apply_lora_full_linear(
337 + self,
338 + y: torch.Tensor,
339 + x: torch.Tensor,
340 + weight_stacked: torch.Tensor,
341 + bias_stacked: torch.Tensor,
342 + module_enabled: torch.Tensor,
343 + ) -> None:
344 + indices = self.sampler_indices
345 + adapter_y = torch.zeros(
346 + (x.size(0), weight_stacked.size(-2)),
347 + dtype=torch.float32,
348 + device=x.device,
349 + )
350 + bgmv_shrink(x, weight_stacked, adapter_y, indices, 1.0)
351 + result = self._select_full_linear_output(
352 + y, adapter_y, bias_stacked, module_enabled
353 + )
354 + y.copy_(result)
355 +
356 def moe_lora_align_block_size(
357 self,
358 topk_ids: torch.Tensor,
vllm/lora/utils.py
+20 -7
@@ -15,6 +15,7 @@ from vllm.logger import init_logger
15 # being imported for _all_lora_classes below
16 from vllm.lora.layers import (
17 BaseLayerWithLoRA,
18 + ClassificationHeadWithLoRA,
19 ColumnParallelLinearWithLoRA,
20 ColumnParallelLinearWithShardedLoRA,
21 FusedMoE3DWithLoRA,
@@ -143,6 +144,17 @@ def from_layer_logits_processor(
144 return ret
145
146
147 +def from_layer_classification(
148 + layer: nn.Module,
149 + max_loras: int,
150 + lora_config: LoRAConfig,
151 + model_config: PretrainedConfig | None = None,
152 +) -> ClassificationHeadWithLoRA:
153 + instance_layer = ClassificationHeadWithLoRA(layer)
154 + instance_layer.create_lora_weights(max_loras, lora_config, model_config)
155 + return instance_layer
156 +
157 +
158 def replace_submodule(
159 model: nn.Module, module_name: str, new_module: nn.Module
160 ) -> nn.Module:
@@ -193,13 +205,14 @@ def parse_fine_tuned_lora_name(
205 start_index = 2 if name.startswith("base_model.model.") else 0
206
207 parts = name.split(".")
196 - if (
197 - parts[-1] == "weight"
198 - and len(parts) >= 2
199 - and (parts[-2] == "lora_A" or parts[-2] == "lora_B")
200 - ):
201 - new_name = ".".join(parts[start_index:-2])
202 - return new_name, parts[-2] == "lora_A"
208 + if (parts[-1] == "weight" or parts[-1] == "bias") and len(parts) >= 2:
209 + if parts[-2] in ["lora_A", "lora_B"]:
210 + new_name = ".".join(parts[start_index:-2])
211 + return new_name, parts[-2] == "lora_A"
212 + # For modules_to_save in classification.
213 + elif parts[-2] in ["score", "classifier"]:
214 + new_name = parts[-2]
215 + return new_name, False
216
217 if parts[-1] == "lora_embedding_A" or parts[-1] == "lora_embedding_B":
218 new_name = ".".join(parts[start_index:-1])
vllm/lora/worker_manager.py
+4
@@ -157,6 +157,10 @@ class WorkerLoRAManager:
157 # 3D->2D conversion when running under the universal 2D wrapper.
158 lora.is_3d_lora_weight = lora_request.is_3d_lora_weight
159
160 + # Validate classification-head weights.
161 + self._adapter_manager._validate_modules_to_save(lora)
162 + self._adapter_manager._validate_token_classification_lora(lora)
163 +
164 except FileNotFoundError as e:
165 # FileNotFoundError should be raised if both
166 # - No adapter found to download from huggingface (or in
vllm/model_executor/layers/pooler/special.py
+14
@@ -74,6 +74,20 @@ class DispatchPooler(Pooler):
74 def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
75 return self.poolers_by_task[task].get_pooling_updates(task)
76
77 + def replace_classifier(
78 + self,
79 + new_classifier: ClassifierFn,
80 + ) -> None:
81 + """Replaces the classifier to the LoRA-wrapped version."""
82 + for task, pooler in self.poolers_by_task.items():
83 + if task not in {"classify"}:
84 + # Now LoRA only supports classification tasks,
85 + # so we only replace the classifier for "classify" task.
86 + continue
87 + head = getattr(pooler, "head", None)
88 + if head is not None and hasattr(head, "classifier"):
89 + head.classifier = new_classifier
90 +
91 def forward(
92 self,
93 hidden_states: torch.Tensor,