[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>

.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

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`

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()

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",
     ),
 ]
 

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")

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"),

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,
+    )

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",
 ]

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

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:

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()

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

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

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,

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)

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,

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,

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])

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

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,