From 953f90d25fea2240965878a1aca66cec519685b0 Mon Sep 17 00:00:00 2001 From: linitra24 Date: Mon, 28 Sep 2026 22:58:35 +0800 Subject: [PATCH] [LoRA] Support variable num_labels for sequence classification (#57766) Signed-off-by: linitra24 Co-authored-by: Jee Jee Li --- docs/features/lora.md | 9 +++- tests/lora/test_sequence_classification.py | 3 +- vllm/config/lora.py | 4 ++ vllm/engine/arg_utils.py | 6 +++ vllm/lora/layers/classifier.py | 48 +++++++++++++++---- vllm/lora/model_manager.py | 23 +++++++-- .../layers/pooler/seqwise/heads.py | 14 ++++-- .../layers/pooler/seqwise/poolers.py | 2 +- 8 files changed, 90 insertions(+), 19 deletions(-) diff --git a/docs/features/lora.md b/docs/features/lora.md index 78004722a9..4c95bde8cd 100644 --- a/docs/features/lora.md +++ b/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. diff --git a/tests/lora/test_sequence_classification.py b/tests/lora/test_sequence_classification.py index 86721bfb6e..773e5448c8 100644 --- a/tests/lora/test_sequence_classification.py +++ b/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, diff --git a/vllm/config/lora.py b/vllm/config/lora.py index 24e42a42fc..e5c1d14f6d 100644 --- a/vllm/config/lora.py +++ b/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 diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 3e84f50997..d641b1fbab 100644 --- a/vllm/engine/arg_utils.py +++ b/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, diff --git a/vllm/lora/layers/classifier.py b/vllm/lora/layers/classifier.py index 58160d1931..55dd4e2e7d 100644 --- a/vllm/lora/layers/classifier.py +++ b/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( diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index b47c22d6aa..adc2353b3e 100644 --- a/vllm/lora/model_manager.py +++ b/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( diff --git a/vllm/model_executor/layers/pooler/seqwise/heads.py b/vllm/model_executor/layers/pooler/seqwise/heads.py index c72532cc10..40624cc17c 100644 --- a/vllm/model_executor/layers/pooler/seqwise/heads.py +++ b/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] diff --git a/vllm/model_executor/layers/pooler/seqwise/poolers.py b/vllm/model_executor/layers/pooler/seqwise/poolers.py index 2115e5184e..c2ac967a15 100644 --- a/vllm/model_executor/layers/pooler/seqwise/poolers.py +++ b/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 ), )