@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 953f90d25f

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

linitra24 committed Sep 28, 2026 at 22:58 UTC 953f90d25fea2240965878a1aca66cec519685b0

8 files changed +90 -19

docs/features/lora.md

+8 -1

@@ -428,9 +428,16 @@ vLLM supports PEFT sequence-classification adapters that save a complete, single

428

429 See [classification_with_lora_offline.py](../../examples/pooling/classify/classification_with_lora_offline.py) for an offline classification example using a LoRA adapter.

430

431 +To batch adapters with different `num_labels`, set the maximum number of labels:

432 +

433 +```bash

434 +vllm serve model --enable-lora --max-lora-cls-labels 8

435 +```

436 +

437 +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.

438 +

439 This support has the following limitations:

440

433 -- All adapters in one engine must have the same `num_labels` and hidden size as the base classification head. An incompatible adapter is rejected when it is loaded.

441 - A classification head stored as float32 is converted to the runtime head dtype when it is loaded.

442 - Token-classification adapters are not supported by this feature.

443

tests/lora/test_sequence_classification.py

+2 -1

@@ -120,12 +120,13 @@ def _test_batched_loras(

120 model=MODEL_NAME,

121 runner="pooling",

122 convert="classify",

123 - hf_overrides={"num_labels": 2},

123 + hf_overrides={"num_labels": 1},

124 dtype="float16",

125 enable_lora=True,

126 max_loras=2,

127 max_cpu_loras=2,

128 max_lora_rank=16,

129 + max_lora_cls_labels=2,

130 enforce_eager=True,

131 max_model_len=512,

132 gpu_memory_utilization=0.5,

vllm/config/lora.py

@@ -84,6 +84,9 @@ class LoRAConfig:

84 experts (stored once with expert-dim 1) instead of per-expert. The shared

85 factors are broadcast to the expert count at kernel time. Only meaningful for

86 MoE models whose adapters use this layout; ignored otherwise."""

87 + max_lora_cls_labels: int | None = Field(default=None, ge=1)

88 + """Maximum output size for LoRA classification heads. Defaults to the

89 + base classification head size."""

91 def compute_hash(self) -> str:

92 """WARNING: Whenever a new field is added to this config,

@@ -104,6 +107,7 @@ class LoRAConfig:

107 factors.append(self.enable_tower_connector_lora)

108 factors.append(self.enable_mixed_moe_lora_format)

109 factors.append(self.enable_moe_shared_loras)

110 + factors.append(self.max_lora_cls_labels)

111 # target_modules affects which modules get LoRA applied

112 factors.append(

113 tuple(sorted(self.target_modules)) if self.target_modules else None

vllm/engine/arg_utils.py

@@ -676,6 +676,7 @@ class EngineArgs:

676 specialize_active_lora: bool = LoRAConfig.specialize_active_lora

677 enable_mixed_moe_lora_format: bool = LoRAConfig.enable_mixed_moe_lora_format

678 enable_moe_shared_loras: bool = LoRAConfig.enable_moe_shared_loras

679 + max_lora_cls_labels: int | None = LoRAConfig.max_lora_cls_labels

680

681 ray_workers_use_nsight: bool = ParallelConfig.ray_workers_use_nsight

682 num_gpu_blocks_override: int | None = CacheConfig.num_gpu_blocks_override

@@ -1571,6 +1572,10 @@ class EngineArgs:

1572 "--enable-moe-shared-loras",

1573 **lora_kwargs["enable_moe_shared_loras"],

1574 )

1575 + lora_group.add_argument(

1576 + "--max-lora-cls-labels",

1577 + **lora_kwargs["max_lora_cls_labels"],

1578 + )

1579

1580 # Logging arguments

1581 logging_group = parser.add_argument_group(

@@ -2587,6 +2592,7 @@ class EngineArgs:

2592 specialize_active_lora=self.specialize_active_lora,

2593 enable_mixed_moe_lora_format=self.enable_mixed_moe_lora_format,

2594 enable_moe_shared_loras=self.enable_moe_shared_loras,

2595 + max_lora_cls_labels=self.max_lora_cls_labels,

2596 max_cpu_loras=self.max_cpu_loras

2597 if self.max_cpu_loras and self.max_cpu_loras > 0

2598 else None,

vllm/lora/layers/classifier.py

+39 -9

@@ -20,28 +20,38 @@ class ClassificationHeadWithLoRA(ReplicatedLinearWithLoRA):

20 ) -> None:

21 # Preserve ordinary LoRA A/B support for classification heads.

22 super().create_lora_weights(max_loras, lora_config, model_config)

23 +

24 + self.max_lora_cls_labels = lora_config.max_lora_cls_labels or self.output_size

25 + self.padded_num_labels = max(self.output_size, self.max_lora_cls_labels)

26 self.full_weight_stacked = torch.zeros(

27 max_loras,

28 1,

26 - self.output_size,

29 + self.padded_num_labels,

30 self.input_size,

31 dtype=lora_config.lora_dtype,

32 device=self.device,

33 )

34 self.full_bias_stacked = torch.zeros(

35 max_loras,

33 - self.output_size,

36 + self.padded_num_labels,

37 dtype=self.base_layer.params_dtype,

38 device=self.device,

39 )

40 self.full_module_enabled = torch.zeros(

41 max_loras, dtype=torch.bool, device=self.device

42 )

43 + # Zero marks an unused slot; positive values are unpadded label counts.

44 + self.full_module_num_labels = [0] * max_loras

45 + self._output_lora_indices: tuple[int, ...] = ()

46 +

47 + def set_output_mapping(self, slot_indices: tuple[int, ...]) -> None:

48 + self._output_lora_indices = slot_indices

50 def reset_module_to_save(self, index: int) -> None:

51 self.full_weight_stacked[index].zero_()

52 self.full_bias_stacked[index].zero_()

53 self.full_module_enabled[index] = False

54 + self.full_module_num_labels[index] = 0

56 def set_module_to_save(

57 self,

@@ -50,17 +60,33 @@ class ClassificationHeadWithLoRA(ReplicatedLinearWithLoRA):

60 bias: torch.Tensor | None,

61 ) -> None:

62 self.reset_module_to_save(index)

53 - self.full_weight_stacked[index, 0].copy_(weight, non_blocking=True)

63 + num_labels = weight.size(0)

64 + self.full_weight_stacked[index, 0, :num_labels].copy_(weight, non_blocking=True)

65 if bias is not None:

55 - self.full_bias_stacked[index].copy_(bias, non_blocking=True)

66 + self.full_bias_stacked[index, :num_labels].copy_(bias, non_blocking=True)

67 + self.full_module_num_labels[index] = num_labels

68 self.full_module_enabled[index] = True

58 - def forward(self, input_: torch.Tensor) -> torch.Tensor:

70 + def forward(self, input_: torch.Tensor) -> torch.Tensor | list[torch.Tensor]:

71 # TODO base_result maybe don't need compute truly.

72 base_result = super().forward(input_)

61 - output = base_result[0] if isinstance(base_result, tuple) else base_result

73 + base_result = base_result[0] if isinstance(base_result, tuple) else base_result

74 +

75 + slot_indices = self._output_lora_indices

76 + assert len(slot_indices) == base_result.size(0), (

77 + "Classification rows do not match LoRA request mapping"

78 + )

79 + num_labels = [

80 + self.full_module_num_labels[index] if index >= 0 else 0

81 + for index in slot_indices

82 + ]

83 +

84 + if not any(num_labels):

85 + return base_result

63 - full_output = self.punica_wrapper.apply_lora_full_linear(

87 + output = base_result.new_zeros(base_result.size(0), self.padded_num_labels)

88 + output[:, : self.output_size].copy_(base_result)

89 + self.punica_wrapper.apply_lora_full_linear(

90 output,

91 input_.to(self.full_weight_stacked.dtype),

92 self.full_weight_stacked,

@@ -68,8 +94,12 @@ class ClassificationHeadWithLoRA(ReplicatedLinearWithLoRA):

94 self.full_module_enabled,

95 )

71 - if full_output is not None:

72 - output = full_output

97 + # Rows without a full head keep the base classifier output size.

98 + num_labels = [size or self.output_size for size in num_labels]

99 + if len(set(num_labels)) == 1:

100 + return output[:, : num_labels[0]]

101 +

102 + output = [output[row, :size] for row, size in enumerate(num_labels)]

103 return output

104

105 def _apply_lora_to_output(

vllm/lora/model_manager.py

+19 -4

@@ -451,6 +451,16 @@ class LoRAModelManager:

451 self.vocab_size,

452 )

453

454 + if self._classification_head is not None:

455 + _, classification_head = self._classification_head

456 + if classification_head.punica_wrapper is punica_wrapper:

457 + classification_head.set_output_mapping(

458 + tuple(

459 + self.lora_index_to_id.index(lora_id) if lora_id > 0 else -1

460 + for lora_id in mapping.prompt_mapping

461 + )

462 + )

463 +

464 def remove_all_adapters(self):

465 """Remove all LoRAModels from the manager."""

466 self._registered_adapters.clear()

@@ -1262,18 +1272,23 @@ class LoRAModelManager:

1272 f"classification head {module_name!r}."

1273 )

1274

1265 - expected_weight_shape = (wrapper.output_size, wrapper.input_size)

1275 received_weight_shape = tuple(full_module.weight.shape)

1267 - if received_weight_shape != expected_weight_shape:

1276 + if (

1277 + full_module.weight.ndim != 2

1278 + or full_module.weight.size(0) < 1

1279 + or full_module.weight.size(0) > wrapper.max_lora_cls_labels

1280 + or full_module.weight.size(1) != wrapper.input_size

1281 + ):

1282 raise ValueError(

1283 f"Full module {saved_module_name!r} for {module_name!r} has "

1284 "an incompatible weight shape: expected "

1271 - f"{expected_weight_shape}, received {received_weight_shape}."

1285 + f"(1..{wrapper.max_lora_cls_labels}, {wrapper.input_size}), "

1286 + f"received {received_weight_shape}."

1287 )

1288

1289 if full_module.bias is None:

1290 return

1276 - expected_bias_shape = (wrapper.output_size,)

1291 + expected_bias_shape = (full_module.weight.size(0),)

1292 received_bias_shape = tuple(full_module.bias.shape)

1293 if received_bias_shape != expected_bias_shape:

1294 raise ValueError(

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

+11 -3

@@ -171,16 +171,24 @@ class ClassifierPoolerHead(SequencePoolerHead):

171 pooled_data = pooled_data.to(self.head_dtype)

172

173 if self.classifier is not None:

174 - logits = self.classifier(pooled_data)

174 + logits: SequencePoolerHeadOutput = self.classifier(pooled_data)

175 else:

176 logits = pooled_data

177

178 # logits shape: [batchsize, num_labels]

179 # Affine score calibration: activation((logit - mean) / sigma)

180 if self.logit_mean is not None:

181 - logits = logits - self.logit_mean

181 + logits = (

182 + [value - self.logit_mean for value in logits]

183 + if isinstance(logits, list)

184 + else logits - self.logit_mean

185 + )

186 if self.logit_sigma is not None:

183 - logits = logits / self.logit_sigma

187 + logits = (

188 + [value / self.logit_sigma for value in logits]

189 + if isinstance(logits, list)

190 + else logits / self.logit_sigma

191 + )

192

193 if self.activation is not None:

194 flags = [p.use_activation for p in pooling_params]

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

+1 -1

@@ -130,7 +130,7 @@ def pooler_for_classify(

130 logit_mean=model_config.pooler_config.logit_mean,

131 logit_sigma=model_config.pooler_config.logit_sigma,

132 activation=resolve_classifier_act_fn(

133 - model_config, static_num_labels=True, act_fn=act_fn

133 + model_config, static_num_labels=False, act_fn=act_fn

134 ),

135 )

136