diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index d63fb567ea..93fb5f5b81 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -54,6 +54,7 @@ steps: - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_gptoss_tp.py - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s -x lora/test_qwen35_densemodel_lora.py - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s -x lora/test_gemma4_tp.py + - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp mirror: amd: label: ":amd: (MI300) LoRA TP (Distributed) %N" @@ -79,3 +80,4 @@ steps: - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_gptoss_tp.py - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s -x lora/test_qwen35_densemodel_lora.py - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s -x lora/test_gemma4_tp.py + - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp diff --git a/docs/features/lora.md b/docs/features/lora.md index 839624542e..78004722a9 100644 --- a/docs/features/lora.md +++ b/docs/features/lora.md @@ -422,6 +422,18 @@ vllm serve ibm-granite/granite-speech-3.3-2b \ Note: Default multimodal LoRAs are currently only available for `.generate` and chat completions. +## Sequence-Classification LoRA Adapters + +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`. + +See [classification_with_lora_offline.py](../../examples/pooling/classify/classification_with_lora_offline.py) for an offline classification example using a LoRA adapter. + +This support has the following limitations: + +- 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. +- A classification head stored as float32 is converted to the runtime head dtype when it is loaded. +- Token-classification adapters are not supported by this feature. + ## Using Tips ### Configuring `max_lora_rank` diff --git a/examples/pooling/classify/classification_with_lora_offline.py b/examples/pooling/classify/classification_with_lora_offline.py new file mode 100644 index 0000000000..a44066a653 --- /dev/null +++ b/examples/pooling/classify/classification_with_lora_offline.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Example of offline classification with a LoRA adapter.""" + +from vllm import LLM +from vllm.lora.request import LoRARequest +from vllm.transformers_utils.repo_utils import hf_api + + +def main(): + lora_path = hf_api().snapshot_download( + repo_id="AmirMohseni/skywork-qwen3-0.6b-reward-lora" + ) + llm = LLM( + model="Skywork/Skywork-Reward-V2-Qwen3-0.6B", + runner="pooling", + enable_lora=True, + max_lora_rank=8, + max_model_len=512, + ) + + prompt = "Which response is more helpful and correct?" + (output,) = llm.classify( + prompt, + lora_request=LoRARequest("reward-lora", 1, lora_path), + ) + print(f"Prompt: {prompt!r}") + print(f"Class Probabilities: {output.outputs.probs}") + + +if __name__ == "__main__": + main() diff --git a/tests/entrypoints/serve/lora/test_lora_adapters.py b/tests/entrypoints/serve/lora/test_lora_adapters.py index a22f0b3899..8a860dc625 100644 --- a/tests/entrypoints/serve/lora/test_lora_adapters.py +++ b/tests/entrypoints/serve/lora/test_lora_adapters.py @@ -26,7 +26,7 @@ BADREQUEST_CASES = [ ( "test_modules_to_save", {"modules_to_save": ["lm_head"]}, - "only supports modules_to_save being None", + "Unsupported modules_to_save", ), ] diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py index 2c64fb2245..8635779cff 100644 --- a/tests/lora/conftest.py +++ b/tests/lora/conftest.py @@ -270,6 +270,30 @@ def qwen3_woofing_lora_files(): return hf_api().snapshot_download(repo_id="Jackmin108/Qwen3-0.6B-Woof-LoRA") +@pytest.fixture(scope="session") +def qwen3_guard_star_trek_lora_files(): + return hf_api().snapshot_download( + repo_id="geoffmunn/Qwen3Guard-StarTrek-Classification-0.6B", + allow_patterns=["adapter_config.json", "adapter_model.safetensors"], + ) + + +@pytest.fixture(scope="session") +def qwen3_guard_new_zealand_lora_files(): + return hf_api().snapshot_download( + repo_id="geoffmunn/Qwen3Guard-NewZealand-Classification-0.6B", + allow_patterns=["adapter_config.json", "adapter_model.safetensors"], + ) + + +@pytest.fixture(scope="session") +def skywork_qwen3_reward_lora_files(): + return hf_api().snapshot_download( + repo_id="AmirMohseni/skywork-qwen3-0.6b-reward-lora", + allow_patterns=["adapter_config.json", "adapter_model.safetensors"], + ) + + @pytest.fixture(scope="session") def deepseekv2_lora_files(): return hf_api().snapshot_download(repo_id="wuchen01/DeepSeek-V2-Lite-Chat-All-LoRA") diff --git a/tests/lora/test_peft_helper.py b/tests/lora/test_peft_helper.py index 23a8aad1a5..9fb52ea727 100644 --- a/tests/lora/test_peft_helper.py +++ b/tests/lora/test_peft_helper.py @@ -20,7 +20,7 @@ ERROR_CASES = [ ( "test_modules_to_save", {"modules_to_save": ["lm_head"]}, - "only supports modules_to_save being None", + "Unsupported modules_to_save", ), ("test_rank_zero", {"r": 0}, "must be a positive integer"), ("test_rank_negative", {"r": -8}, "must be a positive integer"), diff --git a/tests/lora/test_sequence_classification.py b/tests/lora/test_sequence_classification.py new file mode 100644 index 0000000000..86721bfb6e --- /dev/null +++ b/tests/lora/test_sequence_classification.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from tests.utils import multi_gpu_test +from vllm import LLM, PoolingParams +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.lora.request import LoRARequest + +MODEL_NAME = "Qwen/Qwen3-0.6B" +NATIVE_MODEL_NAME = "Skywork/Skywork-Reward-V2-Qwen3-0.6B" +PROMPTS = { + "star_trek": "Does warp drive appear in Star Trek?", + "new_zealand": "Wellington is the capital of New Zealand.", +} +EXPECTED_OUT = { + "star_trek": [13.6640625, -13.6640625], + "new_zealand": [-4.80078125, 8.6171875], + "native": [-1.52734375], +} + + +def _classify_logits( + llm: LLM, + prompts: list[str], + requests: LoRARequest | list[LoRARequest], +) -> torch.Tensor: + params = [PoolingParams(task="classify", use_activation=False) for _ in prompts] + outputs = llm.classify( + prompts, + pooling_params=params, + lora_request=requests, + use_tqdm=False, + ) + return torch.tensor([output.outputs.probs for output in outputs]) + + +def test_converted_model_with_modules_to_save( + qwen3_guard_star_trek_lora_files: str, +) -> None: + prompt = PROMPTS["star_trek"] + llm = LLM( + model=MODEL_NAME, + runner="pooling", + convert="classify", + hf_overrides={"num_labels": 2}, + dtype="float16", + enable_lora=True, + max_lora_rank=16, + enforce_eager=True, + max_model_len=512, + gpu_memory_utilization=0.5, + ) + actual = _classify_logits( + llm, + [prompt], + LoRARequest("star-trek", 1, qwen3_guard_star_trek_lora_files), + )[0] + + torch.testing.assert_close( + actual, + torch.tensor(EXPECTED_OUT["star_trek"]), + atol=2e-2, + rtol=2e-2, + ) + + del llm + cleanup_dist_env_and_memory() + + +def test_native_classification_model_with_modules_to_save( + skywork_qwen3_reward_lora_files: str, +) -> None: + prompt = "Which response is more helpful and correct?" + llm = LLM( + model=NATIVE_MODEL_NAME, + runner="pooling", + dtype="float16", + enable_lora=True, + max_lora_rank=8, + enforce_eager=True, + max_model_len=512, + gpu_memory_utilization=0.5, + ) + actual = _classify_logits( + llm, + [prompt], + LoRARequest("native", 1, skywork_qwen3_reward_lora_files), + )[0] + + torch.testing.assert_close( + actual, + torch.tensor(EXPECTED_OUT["native"]), + atol=2e-2, + rtol=2e-2, + ) + + del llm + cleanup_dist_env_and_memory() + + +def _test_batched_loras( + qwen3_guard_star_trek_lora_files: str, + qwen3_guard_new_zealand_lora_files: str, + tp_size: int, +) -> None: + names = ["star_trek", "new_zealand"] + prompts = [PROMPTS[name] for name in names] + adapter_paths = [ + qwen3_guard_star_trek_lora_files, + qwen3_guard_new_zealand_lora_files, + ] + requests = [ + LoRARequest(name, index, adapter_path) + for index, (name, adapter_path) in enumerate(zip(names, adapter_paths), start=1) + ] + + llm = LLM( + model=MODEL_NAME, + runner="pooling", + convert="classify", + hf_overrides={"num_labels": 2}, + dtype="float16", + enable_lora=True, + max_loras=2, + max_cpu_loras=2, + max_lora_rank=16, + enforce_eager=True, + max_model_len=512, + gpu_memory_utilization=0.5, + tensor_parallel_size=tp_size, + ) + actual = _classify_logits(llm, prompts, requests) + + torch.testing.assert_close( + actual, + torch.tensor([EXPECTED_OUT[name] for name in names]), + atol=2e-2, + rtol=2e-2, + ) + + del llm + cleanup_dist_env_and_memory() + + +def test_batched_loras( + qwen3_guard_star_trek_lora_files: str, + qwen3_guard_new_zealand_lora_files: str, +) -> None: + _test_batched_loras( + qwen3_guard_star_trek_lora_files, + qwen3_guard_new_zealand_lora_files, + tp_size=1, + ) + + +@multi_gpu_test(num_gpus=2) +def test_batched_loras_tp( + qwen3_guard_star_trek_lora_files: str, + qwen3_guard_new_zealand_lora_files: str, +) -> None: + _test_batched_loras( + qwen3_guard_star_trek_lora_files, + qwen3_guard_new_zealand_lora_files, + tp_size=2, + ) diff --git a/vllm/lora/layers/__init__.py b/vllm/lora/layers/__init__.py index e0c5f202bd..2f88640e5f 100644 --- a/vllm/lora/layers/__init__.py +++ b/vllm/lora/layers/__init__.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from vllm.lora.layers.base import BaseLayerWithLoRA +from vllm.lora.layers.classifier import ClassificationHeadWithLoRA from vllm.lora.layers.column_parallel_linear import ( ColumnParallelLinearWithLoRA, ColumnParallelLinearWithShardedLoRA, @@ -42,4 +43,5 @@ __all__ = [ "LoRAMappingType", "FusedMoEWithLoRA", "FusedMoE3DWithLoRA", + "ClassificationHeadWithLoRA", ] diff --git a/vllm/lora/layers/classifier.py b/vllm/lora/layers/classifier.py new file mode 100644 index 0000000000..aa42023763 --- /dev/null +++ b/vllm/lora/layers/classifier.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch +from transformers import PretrainedConfig + +from vllm.config.lora import LoRAConfig +from vllm.platforms import current_platform + +from .replicated_linear import ReplicatedLinearWithLoRA + + +class ClassificationHeadWithLoRA(ReplicatedLinearWithLoRA): + def create_lora_weights( + self, + max_loras: int, + lora_config: LoRAConfig, + model_config: PretrainedConfig | None = None, + ) -> None: + # Preserve ordinary LoRA A/B support for classification heads. + super().create_lora_weights(max_loras, lora_config, model_config) + self.full_weight_stacked = torch.zeros( + max_loras, + 1, + self.output_size, + self.input_size, + dtype=lora_config.lora_dtype, + device=self.device, + ) + self.full_bias_stacked = torch.zeros( + max_loras, + self.output_size, + dtype=self.base_layer.params_dtype, + device=self.device, + ) + self.full_module_enabled = torch.zeros( + max_loras, dtype=torch.bool, device=self.device + ) + + def reset_module_to_save(self, index: int) -> None: + self.full_weight_stacked[index].zero_() + self.full_bias_stacked[index].zero_() + self.full_module_enabled[index] = False + + def set_module_to_save( + self, + index: int, + weight: torch.Tensor, + bias: torch.Tensor | None, + ) -> None: + self.reset_module_to_save(index) + self.full_weight_stacked[index, 0].copy_(weight, non_blocking=True) + if bias is not None: + self.full_bias_stacked[index].copy_(bias, non_blocking=True) + self.full_module_enabled[index] = True + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + # TODO base_result maybe don't need compute truly. + base_result = super().forward(input_) + output = base_result[0] if isinstance(base_result, tuple) else base_result + + full_output = self.punica_wrapper.apply_lora_full_linear( + output, + input_.to(self.full_weight_stacked.dtype), + self.full_weight_stacked, + self.full_bias_stacked, + self.full_module_enabled, + ) + + if full_output is not None: + output = full_output + return output + + def _apply_lora_to_output( + self, x: torch.Tensor, output: torch.Tensor + ) -> torch.Tensor: + original_shape = output.shape if output.ndim == 3 else None + if x.ndim == 3 and output.ndim == 3: + output = output.flatten(0, 1) + x = x.flatten(0, 1) + + # Classification outputs are request-level, so use prompt-based LoRA routing. + lora_output: torch.Tensor | None = self.punica_wrapper.add_lora_logits( + output, + x.to(self.lora_a_stacked[0].dtype), + self.lora_a_stacked[0], + self.lora_b_stacked[0], + 1.0, + ) + if not current_platform.can_update_inplace(): + output = lora_output + + if original_shape is not None: + output = output.reshape(original_shape) + + return output diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py index f094f11f09..609f14f85e 100644 --- a/vllm/lora/lora_model.py +++ b/vllm/lora/lora_model.py @@ -8,7 +8,7 @@ import safetensors import torch from vllm.logger import init_logger -from vllm.lora.lora_weights import LoRALayerWeights +from vllm.lora.lora_weights import LoRAFullModuleWeights, LoRALayerWeights from vllm.lora.peft_helper import PEFTHelper from vllm.lora.utils import ( get_lora_id, @@ -66,6 +66,8 @@ class LoRAModel: rank: int, loras: dict[str, LoRALayerWeights], is_3d_lora_weight: bool = False, + *, + modules_to_save: dict[str, LoRAFullModuleWeights] | None = None, ) -> None: """ Args: @@ -85,6 +87,7 @@ class LoRAModel: ) self.rank = rank self.loras: dict[str, LoRALayerWeights] = loras + self.modules_to_save = modules_to_save or {} self.is_3d_lora_weight = is_3d_lora_weight def clone(self, lora_model_id: int) -> "LoRAModel": @@ -96,12 +99,16 @@ class LoRAModel: rank=self.rank, loras=self.loras.copy(), is_3d_lora_weight=self.is_3d_lora_weight, + modules_to_save=self.modules_to_save.copy(), ) def get_lora(self, module_name: str) -> LoRALayerWeights | None: """Get LoRA for a given module by name""" return self.loras.get(module_name, None) + def get_module_to_save(self, module_name: str) -> LoRAFullModuleWeights | None: + return self.modules_to_save.get(module_name) + def check_lora_name(self, lora_name: str) -> bool: return lora_name in self.loras @@ -127,6 +134,9 @@ class LoRAModel: ) -> "LoRAModel": """Create a LoRAModel from a dictionary of tensors.""" pin_memory = str(device) == "cpu" and PIN_MEMORY + + modules_to_save_names = peft_helper.modules_to_save or [] + full_parameters: dict[str, dict[str, torch.Tensor]] = {} loras: dict[str, LoRALayerWeights] = {} for tensor_name, tensor in tensors.items(): if is_base_embedding_weights(tensor_name): @@ -137,6 +147,12 @@ class LoRAModel: module_name, is_lora_a = parse_fine_tuned_lora_name( tensor_name, weights_mapper ) + if module_name in modules_to_save_names: + full_parameters.setdefault(module_name, {})[ + tensor_name.split(".")[-1] + ] = tensor.to(device=device) + continue + if module_name not in loras: loras[module_name] = LoRALayerWeights.from_config( module_name, peft_helper @@ -161,7 +177,23 @@ class LoRAModel: if pin_memory: loras[module_name].lora_b = loras[module_name].lora_b.pin_memory() - return cls(lora_model_id, peft_helper.r, loras) + modules_to_save = {} + for module_name, parameters in full_parameters.items(): + weight = parameters.get("weight") + if weight is None: + raise ValueError(f"Full module {module_name!r} is missing weight.") + modules_to_save[module_name] = LoRAFullModuleWeights( + module_name=module_name, + weight=weight, + bias=parameters.get("bias"), + ) + if len(modules_to_save) > 1: + raise ValueError( + "Only one full classification module is supported per " + f"adapter, received {sorted(modules_to_save)}." + ) + + return cls(lora_model_id, peft_helper.r, loras, modules_to_save=modules_to_save) @classmethod def from_local_checkpoint( @@ -223,6 +255,7 @@ class LoRAModel: ): continue module_name, _ = parse_fine_tuned_lora_name(lora_module, weights_mapper) + base_name = module_name.rsplit(".", 1)[-1] # Case for expert lora weights if ".experts" in module_name: expert_idx = module_name.find(".experts") @@ -230,7 +263,9 @@ class LoRAModel: if expert_suffix not in expected_lora_modules: unexpected_modules.append(module_name) - elif module_name.rsplit(".", 1)[-1] not in expected_lora_modules: + elif base_name not in expected_lora_modules and base_name not in ( + peft_helper.modules_to_save or () + ): unexpected_modules.append(module_name) if unexpected_modules: diff --git a/vllm/lora/lora_weights.py b/vllm/lora/lora_weights.py index bb241db77e..d9b200f0bb 100644 --- a/vllm/lora/lora_weights.py +++ b/vllm/lora/lora_weights.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Sequence as GenericSequence +from dataclasses import dataclass import torch import torch.types @@ -281,3 +282,17 @@ class PackedLoRALayerWeights(LoRALayerWeights): @property def is_packed(self) -> bool: return True + + +@dataclass +class LoRAFullModuleWeights: + """LoRA weights for classification layers.""" + + module_name: str + weight: torch.Tensor + bias: torch.Tensor | None = None + + def pin_memory(self) -> None: + self.weight = self.weight.pin_memory() + if self.bias is not None: + self.bias = self.bias.pin_memory() diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 7791e1b762..2f3b069f07 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -13,16 +13,22 @@ from vllm.config.lora import LoRAConfig from vllm.logger import init_logger from vllm.lora.layers import ( BaseLayerWithLoRA, + ClassificationHeadWithLoRA, FusedMoE3DWithLoRA, FusedMoEWithLoRA, LoRAMapping, LoRAMappingType, ) from vllm.lora.lora_model import LoRAModel, MoEEPLoadSpec -from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights +from vllm.lora.lora_weights import ( + LoRAFullModuleWeights, + LoRALayerWeights, + PackedLoRALayerWeights, +) from vllm.lora.punica_wrapper import PunicaWrapperBase, get_punica_wrapper from vllm.lora.utils import ( from_layer, + from_layer_classification, from_layer_logits_processor, get_supported_lora_modules, is_in_target_modules, @@ -114,6 +120,13 @@ class LoRAModelManager: self.vocab_size = vocab_size self.is_pooling_model = is_pooling_model(self.model) + self._pooling_task = ( + vllm_config.model_config.get_pooling_task( + tuple(self.model.pooler.get_supported_tasks()) + ) + if self.is_pooling_model + else None + ) self.packed_modules: dict[str, list[str]] = {} self.modules: dict[str, BaseLayerWithLoRA] = {} self._last_mapping: LoRAMapping | None = None @@ -121,6 +134,13 @@ class LoRAModelManager: is_moe = is_moe_model(self.model) self._is_moe = is_moe + self.supported_modules_to_save = ( + {"score", "classifier"} + if self.is_pooling_model + and vllm_config.model_config.score_type == "cross-encoder" + else set() + ) + # When the engine is started with enable_mixed_moe_lora_format=True # we force the universal 2D wrapper (FusedMoEWithLoRA) regardless of # the model's 3D flag, so 2D and 3D adapters can coexist. @@ -147,6 +167,20 @@ class LoRAModelManager: self._init_punica_wrapper(max_num_batched_tokens, vllm_config) self._create_lora_modules() + self._classification_head: tuple[str, ClassificationHeadWithLoRA] | None = None + if self.is_pooling_model: + classification_heads = [ + (module_name, module) + for module_name, module in self.modules.items() + if isinstance(module, ClassificationHeadWithLoRA) + ] + if classification_heads: + assert len(classification_heads) == 1, ( + "Expected 1 classification head, but found " + f"{len(classification_heads)}." + ) + self._classification_head = classification_heads[0] + self.moe_ep_load_spec: MoEEPLoadSpec | None = self._build_moe_ep_load_spec() self.model.lora_manager = self @@ -338,6 +372,15 @@ class LoRAModelManager: self.lora_index_to_id[index] = lora_model.id num_lora_weights_applied = 0 for module_name, module in self.modules.items(): + if isinstance(module, ClassificationHeadWithLoRA): + full_module = self._get_module_to_save_weights(lora_model, module_name) + if full_module is None: + module.reset_module_to_save(index) + else: + module.set_module_to_save( + index, full_module.weight, full_module.bias + ) + module_lora = self._get_lora_layer_weights(lora_model, module_name) if not module_lora: module.reset_lora(index) @@ -427,6 +470,8 @@ class LoRAModelManager: if isinstance(module, PPMissingLayer): continue + is_classifier_head = module_name in self.supported_modules_to_save + if self.lora_config.target_modules is None: if not is_supported_lora_module( module_name, @@ -434,7 +479,8 @@ class LoRAModelManager: self.supported_lora_modules, ): continue - elif not self._match_target_modules(module_name): + elif not self._match_target_modules(module_name) and not is_classifier_head: + # Full classification heads bypass the A/B LoRA target filter. continue punica_wrapper = self._get_punica_wrapper(module_name) @@ -486,17 +532,25 @@ class LoRAModelManager: # LoRA weights of w1 and w3 have already been fused on disk. packed_moduled_lst = ["w13"] if self._is_3d_moe_model else ["w1", "w3"] - new_module = replace_submodule( - self.model, - module_name, - from_layer( + if is_classifier_head: + new_module = from_layer_classification( + module, + self.lora_slots, + self.lora_config, + self.model.config, + ) + else: + new_module = from_layer( module, self.lora_slots, self.lora_config, packed_moduled_lst, self.model.config, - ), - ) + ) + new_module = replace_submodule(self.model, module_name, new_module) + + if is_classifier_head: + self.model.pooler.replace_classifier(new_module) if isinstance(new_module, BaseLayerWithLoRA): wrapped_by_id[id(module)] = new_module wrapped_by_id[id(new_module)] = new_module @@ -746,6 +800,12 @@ class LoRAModelManager: if not self.supports_mm: return self.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY] + if module_name in self.supported_modules_to_save: + lm_wrapper = self.punica_wrapper_mapping.get( + self.mm_mapping.language_model[0] + ) + return lm_wrapper + # For multimodal model # NOTE Sort by prefix length (descending) to match the longest prefix first # e.g., 'visual.merger' should match 'visual.merger' instead of 'visual.' @@ -871,6 +931,10 @@ class LoRAModelManager: lora.lora_a = lora.lora_a.pin_memory() lora.lora_b = lora.lora_b.pin_memory() + if lora_model.modules_to_save: + for full_module in lora_model.modules_to_save.values(): + full_module.pin_memory() + def _stack_moe_lora_weights( self, lora_model: LoRAModel, module: FusedMoE3DWithLoRA, module_name: str ): @@ -1164,6 +1228,82 @@ class LoRAModelManager: ) return lora_model.get_lora(org_module_name) + def _get_module_to_save_weights( + self, lora_model: LoRAModel, module_name: str + ) -> LoRAFullModuleWeights | None: + weights = lora_model.get_module_to_save(module_name) + if weights is not None: + return weights + + unprefixed_module_name = module_name.removeprefix("model.") + weights = lora_model.get_module_to_save(unprefixed_module_name) + if weights is not None and unprefixed_module_name != module_name: + logger.info_once( + "For the pool model, successfully loaded the full module weights " + "after removing the prefix 'model.'." + ) + return weights + + def _validate_modules_to_save(self, lora_model: LoRAModel) -> None: + if not lora_model.modules_to_save: + return + saved_module_name = next(iter(lora_model.modules_to_save)) + if self._classification_head is None: + raise ValueError( + f"Cannot load full module {saved_module_name!r}: the model does " + "not expose a unique ClassificationHeadWithLoRA." + ) + + module_name, wrapper = self._classification_head + full_module = self._get_module_to_save_weights(lora_model, module_name) + if full_module is None: + raise ValueError( + f"Full module {saved_module_name!r} does not match the model's " + f"classification head {module_name!r}." + ) + + expected_weight_shape = (wrapper.output_size, wrapper.input_size) + received_weight_shape = tuple(full_module.weight.shape) + if received_weight_shape != expected_weight_shape: + raise ValueError( + f"Full module {saved_module_name!r} for {module_name!r} has " + "an incompatible weight shape: expected " + f"{expected_weight_shape}, received {received_weight_shape}." + ) + + if full_module.bias is None: + return + expected_bias_shape = (wrapper.output_size,) + received_bias_shape = tuple(full_module.bias.shape) + if received_bias_shape != expected_bias_shape: + raise ValueError( + f"Full module {saved_module_name!r} for {module_name!r} has " + "an incompatible bias shape: expected " + f"{expected_bias_shape}, received {received_bias_shape}." + ) + + def _validate_token_classification_lora(self, lora_model: LoRAModel) -> None: + if self._pooling_task != "token_classify": + return + if self._classification_head is None: + return + + module_name, _ = self._classification_head + has_full_module = ( + self._get_module_to_save_weights(lora_model, module_name) is not None + ) + has_lora_weights = ( + self._get_lora_layer_weights(lora_model, module_name) is not None + ) + if not has_full_module and not has_lora_weights: + return + + raise ValueError( + f"LoRA adapter {lora_model.id} contains weights for classification " + f"head {module_name!r}, but token_classify only supports LoRA on " + "the model backbone." + ) + def deactivate_adapter(self, adapter_id: int) -> bool: if adapter_id not in self._active_adapters: return False diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py index a0bb8bf3c2..8b40af0345 100644 --- a/vllm/lora/peft_helper.py +++ b/vllm/lora/peft_helper.py @@ -45,7 +45,15 @@ class PEFTHelper: """ error_msg = [] if self.modules_to_save: - error_msg.append("vLLM only supports modules_to_save being None.") + unsupported_modules = [ + m for m in self.modules_to_save if m not in ["classifier", "score"] + ] + if unsupported_modules: + error_msg.append( + "vLLM only supports modules_to_save being either None " + 'or ["classifier", "score"] for classification models. ' + f"Unsupported modules_to_save: {unsupported_modules}" + ) if self.use_dora: error_msg.append("vLLM does not yet support DoRA.") return error_msg diff --git a/vllm/lora/punica_wrapper/punica_base.py b/vllm/lora/punica_wrapper/punica_base.py index f5c7ca9cb7..7ab309515c 100644 --- a/vllm/lora/punica_wrapper/punica_base.py +++ b/vllm/lora/punica_wrapper/punica_base.py @@ -120,6 +120,18 @@ class PunicaWrapperABC(ABC): """ raise NotImplementedError + @abstractmethod + def apply_lora_full_linear( + self, + y: torch.Tensor, + x: torch.Tensor, + weight_stacked: torch.Tensor, + bias_stacked: torch.Tensor, + module_enabled: torch.Tensor, + ) -> torch.Tensor | None: + """Apply request-routed full linear weights to selected rows.""" + raise NotImplementedError + class PunicaWrapperBase(PunicaWrapperABC): """ @@ -217,6 +229,22 @@ class PunicaWrapperBase(PunicaWrapperABC): self.token_nums = token_nums self.no_lora = no_lora + def _select_full_linear_output( + self, + y: torch.Tensor, + adapter_y: torch.Tensor, + bias_stacked: torch.Tensor, + module_enabled: torch.Tensor, + ) -> torch.Tensor: + indices = self.sampler_indices + assert indices.size(0) == y.size(0), ( + "Full linear rows do not match LoRA request mapping" + ) + safe_indices = indices.clamp_min(0) + use_full = (indices >= 0) & module_enabled[safe_indices] + adapter_y = adapter_y.to(y.dtype) + bias_stacked[safe_indices].to(y.dtype) + return torch.where(use_full.unsqueeze(-1), adapter_y, y) + @property def prefill_metadata( self, diff --git a/vllm/lora/punica_wrapper/punica_cpu.py b/vllm/lora/punica_wrapper/punica_cpu.py index 1a700d9bf1..d462ce4476 100644 --- a/vllm/lora/punica_wrapper/punica_cpu.py +++ b/vllm/lora/punica_wrapper/punica_cpu.py @@ -349,3 +349,23 @@ class PunicaWrapperCPU(PunicaWrapperBase): bgmv_shrink(x, lora_a_stacked, buffer, self.sampler_indices, scale) bgmv_expand(buffer, lora_b_stacked, y, self.sampler_indices, add_inputs=True) y = y.view_as(y_org) + + def apply_lora_full_linear( + self, + y: torch.Tensor, + x: torch.Tensor, + weight_stacked: torch.Tensor, + bias_stacked: torch.Tensor, + module_enabled: torch.Tensor, + ) -> None: + indices = self.sampler_indices + adapter_y = torch.zeros( + (x.size(0), weight_stacked.size(-2)), + dtype=torch.float32, + device=x.device, + ) + bgmv_shrink(x, weight_stacked, adapter_y, indices, 1.0) + result = self._select_full_linear_output( + y, adapter_y, bias_stacked, module_enabled + ) + y.copy_(result) diff --git a/vllm/lora/punica_wrapper/punica_gpu.py b/vllm/lora/punica_wrapper/punica_gpu.py index 2f746f967e..7deb97e513 100644 --- a/vllm/lora/punica_wrapper/punica_gpu.py +++ b/vllm/lora/punica_wrapper/punica_gpu.py @@ -327,6 +327,33 @@ class PunicaWrapperGPU(PunicaWrapperBase): ) y = y.view_as(y_org) + def apply_lora_full_linear( + self, + y: torch.Tensor, + x: torch.Tensor, + weight_stacked: torch.Tensor, + bias_stacked: torch.Tensor, + module_enabled: torch.Tensor, + ) -> None: + adapter_y = torch.empty( + (x.size(0), weight_stacked.size(-2)), + dtype=torch.float32, + device=x.device, + ) + lora_shrink( + x, + [weight_stacked], + adapter_y.unsqueeze(0), + *self.prompt_mapping_meta.meta_args( + x.size(0), self.lora_config.specialize_active_lora + ), + 1.0, + ) + result = self._select_full_linear_output( + y, adapter_y, bias_stacked, module_enabled + ) + y.copy_(result) + def moe_lora_align_block_size( self, topk_ids: torch.Tensor, diff --git a/vllm/lora/punica_wrapper/punica_xpu.py b/vllm/lora/punica_wrapper/punica_xpu.py index 7f040732ee..3d1cf29683 100755 --- a/vllm/lora/punica_wrapper/punica_xpu.py +++ b/vllm/lora/punica_wrapper/punica_xpu.py @@ -333,6 +333,26 @@ class PunicaWrapperXPU(PunicaWrapperBase): bgmv_expand(buffer, lora_b_stacked, y, sampler_indices, add_inputs=True) y = y.view_as(y_org) + def apply_lora_full_linear( + self, + y: torch.Tensor, + x: torch.Tensor, + weight_stacked: torch.Tensor, + bias_stacked: torch.Tensor, + module_enabled: torch.Tensor, + ) -> None: + indices = self.sampler_indices + adapter_y = torch.zeros( + (x.size(0), weight_stacked.size(-2)), + dtype=torch.float32, + device=x.device, + ) + bgmv_shrink(x, weight_stacked, adapter_y, indices, 1.0) + result = self._select_full_linear_output( + y, adapter_y, bias_stacked, module_enabled + ) + y.copy_(result) + def moe_lora_align_block_size( self, topk_ids: torch.Tensor, diff --git a/vllm/lora/utils.py b/vllm/lora/utils.py index f2a9dc902f..fba78eddcf 100644 --- a/vllm/lora/utils.py +++ b/vllm/lora/utils.py @@ -15,6 +15,7 @@ from vllm.logger import init_logger # being imported for _all_lora_classes below from vllm.lora.layers import ( BaseLayerWithLoRA, + ClassificationHeadWithLoRA, ColumnParallelLinearWithLoRA, ColumnParallelLinearWithShardedLoRA, FusedMoE3DWithLoRA, @@ -143,6 +144,17 @@ def from_layer_logits_processor( return ret +def from_layer_classification( + layer: nn.Module, + max_loras: int, + lora_config: LoRAConfig, + model_config: PretrainedConfig | None = None, +) -> ClassificationHeadWithLoRA: + instance_layer = ClassificationHeadWithLoRA(layer) + instance_layer.create_lora_weights(max_loras, lora_config, model_config) + return instance_layer + + def replace_submodule( model: nn.Module, module_name: str, new_module: nn.Module ) -> nn.Module: @@ -193,13 +205,14 @@ def parse_fine_tuned_lora_name( start_index = 2 if name.startswith("base_model.model.") else 0 parts = name.split(".") - if ( - parts[-1] == "weight" - and len(parts) >= 2 - and (parts[-2] == "lora_A" or parts[-2] == "lora_B") - ): - new_name = ".".join(parts[start_index:-2]) - return new_name, parts[-2] == "lora_A" + if (parts[-1] == "weight" or parts[-1] == "bias") and len(parts) >= 2: + if parts[-2] in ["lora_A", "lora_B"]: + new_name = ".".join(parts[start_index:-2]) + return new_name, parts[-2] == "lora_A" + # For modules_to_save in classification. + elif parts[-2] in ["score", "classifier"]: + new_name = parts[-2] + return new_name, False if parts[-1] == "lora_embedding_A" or parts[-1] == "lora_embedding_B": new_name = ".".join(parts[start_index:-1]) diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index 4c00c6531b..93624523e2 100644 --- a/vllm/lora/worker_manager.py +++ b/vllm/lora/worker_manager.py @@ -157,6 +157,10 @@ class WorkerLoRAManager: # 3D->2D conversion when running under the universal 2D wrapper. lora.is_3d_lora_weight = lora_request.is_3d_lora_weight + # Validate classification-head weights. + self._adapter_manager._validate_modules_to_save(lora) + self._adapter_manager._validate_token_classification_lora(lora) + except FileNotFoundError as e: # FileNotFoundError should be raised if both # - No adapter found to download from huggingface (or in diff --git a/vllm/model_executor/layers/pooler/special.py b/vllm/model_executor/layers/pooler/special.py index 090991bedb..23354d66f3 100644 --- a/vllm/model_executor/layers/pooler/special.py +++ b/vllm/model_executor/layers/pooler/special.py @@ -74,6 +74,20 @@ class DispatchPooler(Pooler): def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: return self.poolers_by_task[task].get_pooling_updates(task) + def replace_classifier( + self, + new_classifier: ClassifierFn, + ) -> None: + """Replaces the classifier to the LoRA-wrapped version.""" + for task, pooler in self.poolers_by_task.items(): + if task not in {"classify"}: + # Now LoRA only supports classification tasks, + # so we only replace the classifier for "classify" task. + continue + head = getattr(pooler, "head", None) + if head is not None and hasattr(head, "classifier"): + head.classifier = new_classifier + def forward( self, hidden_states: torch.Tensor,