[LoRA] Support variable num_labels for sequence classification (#57766)

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

docs/features/lora.md

@@ -428,9 +428,16 @@ vLLM supports PEFT sequence-classification adapters that save a complete, single
 
 See [classification_with_lora_offline.py](../../examples/pooling/classify/classification_with_lora_offline.py) for an offline classification example using a LoRA adapter.
 
+To batch adapters with different `num_labels`, set the maximum number of labels:
+
+```bash
+vllm serve model --enable-lora --max-lora-cls-labels 8
+```
+
+The equivalent `LLM` argument is `max_lora_cls_labels`. It defaults to the base model's `num_labels`, and each request returns its adapter's number of labels.
+
 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.
 

tests/lora/test_sequence_classification.py

@@ -120,12 +120,13 @@ def _test_batched_loras(
         model=MODEL_NAME,
         runner="pooling",
         convert="classify",
-        hf_overrides={"num_labels": 2},
+        hf_overrides={"num_labels": 1},
         dtype="float16",
         enable_lora=True,
         max_loras=2,
         max_cpu_loras=2,
         max_lora_rank=16,
+        max_lora_cls_labels=2,
         enforce_eager=True,
         max_model_len=512,
         gpu_memory_utilization=0.5,

vllm/config/lora.py

@@ -84,6 +84,9 @@ class LoRAConfig:
     experts (stored once with expert-dim 1) instead of per-expert. The shared
     factors are broadcast to the expert count at kernel time. Only meaningful for
     MoE models whose adapters use this layout; ignored otherwise."""
+    max_lora_cls_labels: int | None = Field(default=None, ge=1)
+    """Maximum output size for LoRA classification heads. Defaults to the
+    base classification head size."""
 
     def compute_hash(self) -> str:
         """WARNING: Whenever a new field is added to this config,
@@ -104,6 +107,7 @@ class LoRAConfig:
         factors.append(self.enable_tower_connector_lora)
         factors.append(self.enable_mixed_moe_lora_format)
         factors.append(self.enable_moe_shared_loras)
+        factors.append(self.max_lora_cls_labels)
         # target_modules affects which modules get LoRA applied
         factors.append(
             tuple(sorted(self.target_modules)) if self.target_modules else None

vllm/engine/arg_utils.py

@@ -676,6 +676,7 @@ class EngineArgs:
     specialize_active_lora: bool = LoRAConfig.specialize_active_lora
     enable_mixed_moe_lora_format: bool = LoRAConfig.enable_mixed_moe_lora_format
     enable_moe_shared_loras: bool = LoRAConfig.enable_moe_shared_loras
+    max_lora_cls_labels: int | None = LoRAConfig.max_lora_cls_labels
 
     ray_workers_use_nsight: bool = ParallelConfig.ray_workers_use_nsight
     num_gpu_blocks_override: int | None = CacheConfig.num_gpu_blocks_override
@@ -1571,6 +1572,10 @@ class EngineArgs:
             "--enable-moe-shared-loras",
             **lora_kwargs["enable_moe_shared_loras"],
         )
+        lora_group.add_argument(
+            "--max-lora-cls-labels",
+            **lora_kwargs["max_lora_cls_labels"],
+        )
 
         # Logging arguments
         logging_group = parser.add_argument_group(
@@ -2587,6 +2592,7 @@ class EngineArgs:
                 specialize_active_lora=self.specialize_active_lora,
                 enable_mixed_moe_lora_format=self.enable_mixed_moe_lora_format,
                 enable_moe_shared_loras=self.enable_moe_shared_loras,
+                max_lora_cls_labels=self.max_lora_cls_labels,
                 max_cpu_loras=self.max_cpu_loras
                 if self.max_cpu_loras and self.max_cpu_loras > 0
                 else None,

vllm/lora/layers/classifier.py

@@ -20,28 +20,38 @@ class ClassificationHeadWithLoRA(ReplicatedLinearWithLoRA):
     ) -> None:
         # Preserve ordinary LoRA A/B support for classification heads.
         super().create_lora_weights(max_loras, lora_config, model_config)
+
+        self.max_lora_cls_labels = lora_config.max_lora_cls_labels or self.output_size
+        self.padded_num_labels = max(self.output_size, self.max_lora_cls_labels)
         self.full_weight_stacked = torch.zeros(
             max_loras,
             1,
-            self.output_size,
+            self.padded_num_labels,
             self.input_size,
             dtype=lora_config.lora_dtype,
             device=self.device,
         )
         self.full_bias_stacked = torch.zeros(
             max_loras,
-            self.output_size,
+            self.padded_num_labels,
             dtype=self.base_layer.params_dtype,
             device=self.device,
         )
         self.full_module_enabled = torch.zeros(
             max_loras, dtype=torch.bool, device=self.device
         )
+        # Zero marks an unused slot; positive values are unpadded label counts.
+        self.full_module_num_labels = [0] * max_loras
+        self._output_lora_indices: tuple[int, ...] = ()
+
+    def set_output_mapping(self, slot_indices: tuple[int, ...]) -> None:
+        self._output_lora_indices = slot_indices
 
     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
+        self.full_module_num_labels[index] = 0
 
     def set_module_to_save(
         self,
@@ -50,17 +60,33 @@ class ClassificationHeadWithLoRA(ReplicatedLinearWithLoRA):
         bias: torch.Tensor | None,
     ) -> None:
         self.reset_module_to_save(index)
-        self.full_weight_stacked[index, 0].copy_(weight, non_blocking=True)
+        num_labels = weight.size(0)
+        self.full_weight_stacked[index, 0, :num_labels].copy_(weight, non_blocking=True)
         if bias is not None:
-            self.full_bias_stacked[index].copy_(bias, non_blocking=True)
+            self.full_bias_stacked[index, :num_labels].copy_(bias, non_blocking=True)
+        self.full_module_num_labels[index] = num_labels
         self.full_module_enabled[index] = True
 
-    def forward(self, input_: torch.Tensor) -> torch.Tensor:
+    def forward(self, input_: torch.Tensor) -> torch.Tensor | list[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
+        base_result = base_result[0] if isinstance(base_result, tuple) else base_result
 
-        full_output = self.punica_wrapper.apply_lora_full_linear(
+        slot_indices = self._output_lora_indices
+        assert len(slot_indices) == base_result.size(0), (
+            "Classification rows do not match LoRA request mapping"
+        )
+        num_labels = [
+            self.full_module_num_labels[index] if index >= 0 else 0
+            for index in slot_indices
+        ]
+
+        if not any(num_labels):
+            return base_result
+
+        output = base_result.new_zeros(base_result.size(0), self.padded_num_labels)
+        output[:, : self.output_size].copy_(base_result)
+        self.punica_wrapper.apply_lora_full_linear(
             output,
             input_.to(self.full_weight_stacked.dtype),
             self.full_weight_stacked,
@@ -68,8 +94,12 @@ class ClassificationHeadWithLoRA(ReplicatedLinearWithLoRA):
             self.full_module_enabled,
         )
 
-        if full_output is not None:
-            output = full_output
+        # Rows without a full head keep the base classifier output size.
+        num_labels = [size or self.output_size for size in num_labels]
+        if len(set(num_labels)) == 1:
+            return output[:, : num_labels[0]]
+
+        output = [output[row, :size] for row, size in enumerate(num_labels)]
         return output
 
     def _apply_lora_to_output(

vllm/lora/model_manager.py

@@ -451,6 +451,16 @@ class LoRAModelManager:
             self.vocab_size,
         )
 
+        if self._classification_head is not None:
+            _, classification_head = self._classification_head
+            if classification_head.punica_wrapper is punica_wrapper:
+                classification_head.set_output_mapping(
+                    tuple(
+                        self.lora_index_to_id.index(lora_id) if lora_id > 0 else -1
+                        for lora_id in mapping.prompt_mapping
+                    )
+                )
+
     def remove_all_adapters(self):
         """Remove all LoRAModels from the manager."""
         self._registered_adapters.clear()
@@ -1262,18 +1272,23 @@ class LoRAModelManager:
                 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:
+        if (
+            full_module.weight.ndim != 2
+            or full_module.weight.size(0) < 1
+            or full_module.weight.size(0) > wrapper.max_lora_cls_labels
+            or full_module.weight.size(1) != wrapper.input_size
+        ):
             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}."
+                f"(1..{wrapper.max_lora_cls_labels}, {wrapper.input_size}), "
+                f"received {received_weight_shape}."
             )
 
         if full_module.bias is None:
             return
-        expected_bias_shape = (wrapper.output_size,)
+        expected_bias_shape = (full_module.weight.size(0),)
         received_bias_shape = tuple(full_module.bias.shape)
         if received_bias_shape != expected_bias_shape:
             raise ValueError(

vllm/model_executor/layers/pooler/seqwise/heads.py

@@ -171,16 +171,24 @@ class ClassifierPoolerHead(SequencePoolerHead):
             pooled_data = pooled_data.to(self.head_dtype)
 
         if self.classifier is not None:
-            logits = self.classifier(pooled_data)
+            logits: SequencePoolerHeadOutput = self.classifier(pooled_data)
         else:
             logits = pooled_data
 
         # logits shape: [batchsize, num_labels]
         # Affine score calibration: activation((logit - mean) / sigma)
         if self.logit_mean is not None:
-            logits = logits - self.logit_mean
+            logits = (
+                [value - self.logit_mean for value in logits]
+                if isinstance(logits, list)
+                else logits - self.logit_mean
+            )
         if self.logit_sigma is not None:
-            logits = logits / self.logit_sigma
+            logits = (
+                [value / self.logit_sigma for value in logits]
+                if isinstance(logits, list)
+                else logits / self.logit_sigma
+            )
 
         if self.activation is not None:
             flags = [p.use_activation for p in pooling_params]

vllm/model_executor/layers/pooler/seqwise/poolers.py

@@ -130,7 +130,7 @@ def pooler_for_classify(
         logit_mean=model_config.pooler_config.logit_mean,
         logit_sigma=model_config.pooler_config.logit_sigma,
         activation=resolve_classifier_act_fn(
-            model_config, static_num_labels=True, act_fn=act_fn
+            model_config, static_num_labels=False, act_fn=act_fn
         ),
     )