A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 11:45 UTC
Refactor: Pass num_labels explicitly to PoolerClassify instead of reading from global config (#42851)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
May 17, 2026 at 09:40 UTC
1c8e9c0399f6a6a98f406dce5947a2ad318e195a
2 files changed
+34
-40
tests/model_executor/layers/test_pooler_activations.py
+26
-29
@@ -8,7 +8,6 @@ import pytest
8
import torch
9
import torch.nn as nn
10
11
-from vllm.config import VllmConfig, set_current_vllm_config
11
from vllm.model_executor.layers.pooler.activations import (
12
LambdaPoolerActivation,
13
PoolerClassify,
@@ -20,15 +19,6 @@ from vllm.model_executor.layers.pooler.activations import (
19
)
20
21
23
-@pytest.fixture
24
-def vllm_config():
25
- """VllmConfig with a minimal model_config stub for PoolerClassify."""
26
- config = VllmConfig()
27
- config.model_config = SimpleNamespace(hf_config=SimpleNamespace(num_labels=0))
28
- with set_current_vllm_config(config):
29
- yield config
30
-
31
-
22
# ---------------------------------------------------------------------------
23
# PoolerIdentity
24
# ---------------------------------------------------------------------------
@@ -102,36 +92,39 @@ class TestPoolerMultiLabelClassify:
92
# PoolerClassify
93
# ---------------------------------------------------------------------------
94
class TestPoolerClassify:
105
- def test_dynamic_softmax_when_num_labels_ge_2(self):
106
- pooler = PoolerClassify(static_num_labels=False)
95
+ def test_infers_from_shape_when_num_labels_none(self):
96
+ pooler = PoolerClassify(num_labels=None)
97
assert pooler.num_labels is None
98
x = torch.randn(2, 5)
99
out = pooler(x)
100
sums = out.sum(dim=-1)
101
assert torch.allclose(sums, torch.ones(2), atol=1e-5)
102
113
- def test_dynamic_sigmoid_when_num_labels_lt_2(self):
114
- pooler = PoolerClassify(static_num_labels=False)
103
+ def test_sigmoid_when_num_labels_lt_2(self):
104
+ pooler = PoolerClassify(num_labels=1)
105
x = torch.zeros(1, 1)
106
out = pooler(x)
107
assert torch.allclose(out, torch.tensor([[0.5]]), atol=1e-5)
108
119
- def test_static_default_num_labels_zero_uses_sigmoid(self, vllm_config):
120
- pooler = PoolerClassify(static_num_labels=True)
109
+ def test_num_labels_zero_uses_sigmoid(self):
110
+ pooler = PoolerClassify(num_labels=0)
111
assert pooler.num_labels == 0
112
x = torch.zeros(1, 3)
113
out = pooler(x)
114
assert torch.allclose(out, torch.full((1, 3), 0.5), atol=1e-5)
115
126
- def test_static_num_labels_ge_2_uses_softmax(self, vllm_config):
127
- vllm_config.model_config.hf_config.num_labels = 4
128
- pooler = PoolerClassify(static_num_labels=True)
116
+ def test_num_labels_ge_2_uses_softmax(self):
117
+ pooler = PoolerClassify(num_labels=4)
118
assert pooler.num_labels == 4
119
x = torch.randn(2, 4)
120
out = pooler(x)
121
sums = out.sum(dim=-1)
122
assert torch.allclose(sums, torch.ones(2), atol=1e-5)
123
124
+ def test_default_num_labels_is_none(self):
125
+ pooler = PoolerClassify()
126
+ assert pooler.num_labels is None
127
+
128
129
# ---------------------------------------------------------------------------
130
# LambdaPoolerActivation
@@ -160,24 +153,25 @@ class TestGetActFn:
153
def _make_config(**kwargs):
154
return SimpleNamespace(**kwargs)
155
163
- def test_regression(self, vllm_config):
156
+ def test_regression(self):
157
cfg = self._make_config(problem_type="regression")
158
result = get_act_fn(cfg)
159
assert isinstance(result, PoolerIdentity)
160
168
- def test_single_label_classification(self, vllm_config):
161
+ def test_single_label_classification(self):
162
cfg = self._make_config(
163
problem_type="single_label_classification", num_labels=3
164
)
165
result = get_act_fn(cfg)
166
assert isinstance(result, PoolerClassify)
167
+ assert result.num_labels == 3
168
175
- def test_multi_label_classification(self, vllm_config):
169
+ def test_multi_label_classification(self):
170
cfg = self._make_config(problem_type="multi_label_classification")
171
result = get_act_fn(cfg)
172
assert isinstance(result, PoolerMultiLabelClassify)
173
180
- def test_sentence_transformers_activation(self, vllm_config):
174
+ def test_sentence_transformers_activation(self):
175
cfg = self._make_config(
176
problem_type="",
177
sentence_transformers={
@@ -187,7 +181,7 @@ class TestGetActFn:
181
result = get_act_fn(cfg)
182
assert isinstance(result, PoolerClassify)
183
190
- def test_sbert_activation(self, vllm_config):
184
+ def test_sbert_activation(self):
185
cfg = self._make_config(
186
problem_type="",
187
sbert_ce_default_activation_function=(
@@ -197,12 +191,12 @@ class TestGetActFn:
191
result = get_act_fn(cfg)
192
assert isinstance(result, PoolerClassify)
193
200
- def test_default_fallback(self, vllm_config):
194
+ def test_default_fallback(self):
195
cfg = self._make_config(problem_type="")
196
result = get_act_fn(cfg)
197
assert isinstance(result, PoolerClassify)
198
205
- def test_sentence_transformers_takes_priority(self, vllm_config):
199
+ def test_sentence_transformers_takes_priority(self):
200
cfg = self._make_config(
201
problem_type="",
202
sentence_transformers={"activation_fn": "torch.nn.modules.linear.Identity"},
@@ -213,7 +207,7 @@ class TestGetActFn:
207
result = get_act_fn(cfg)
208
assert isinstance(result, PoolerIdentity)
209
216
- def test_rejects_non_torch_activation(self, vllm_config):
210
+ def test_rejects_non_torch_activation(self):
211
cfg = self._make_config(
212
problem_type="",
213
sentence_transformers={"activation_fn": "os.system"},
@@ -226,10 +220,13 @@ class TestGetActFn:
220
# resolve_classifier_act_fn
221
# ---------------------------------------------------------------------------
222
class TestResolveClassifierActFn:
229
- def test_delegates_to_get_act_fn_when_none(self, vllm_config):
230
- model_config = vllm_config.model_config
223
+ def test_delegates_to_get_act_fn_when_none(self):
224
+ model_config = SimpleNamespace(
225
+ hf_config=SimpleNamespace(num_labels=3, problem_type="")
226
+ )
227
result = resolve_classifier_act_fn(model_config, act_fn=None)
228
assert isinstance(result, PoolerClassify)
229
+ assert result.num_labels == 3
230
231
def test_passes_through_provided_act_fn(self):
232
custom = PoolerIdentity()
vllm/model_executor/layers/pooler/activations.py
+8
-11
@@ -9,7 +9,7 @@ import torch.nn as nn
9
import torch.nn.functional as F
10
from transformers import PretrainedConfig
11
12
-from vllm.config import ModelConfig, get_current_vllm_config
12
+from vllm.config import ModelConfig
13
from vllm.logger import init_logger
14
from vllm.utils.import_utils import resolve_obj_by_qualname
15
@@ -23,11 +23,15 @@ def get_act_fn(
23
# get classification act_fn
24
# Implement alignment with transformers ForSequenceClassificationLoss
25
# https://github.com/huggingface/transformers/blob/57bb6db6ee4cfaccc45b8d474dfad5a17811ca60/src/transformers/loss/loss_utils.py#L92
26
+ num_labels: int | None = None
27
+ if static_num_labels:
28
+ num_labels = getattr(config, "num_labels", 0)
29
+
30
problem_type = getattr(config, "problem_type", "")
31
if problem_type == "regression":
32
return PoolerIdentity()
33
if problem_type == "single_label_classification":
30
- return PoolerClassify(static_num_labels=static_num_labels)
34
+ return PoolerClassify(num_labels=num_labels)
35
if problem_type == "multi_label_classification":
36
return PoolerMultiLabelClassify()
37
@@ -52,7 +56,7 @@ def get_act_fn(
56
fn = resolve_obj_by_qualname(function_name)()
57
return PoolerActivation.wraps(fn)
58
55
- return PoolerClassify(static_num_labels=static_num_labels)
59
+ return PoolerClassify(num_labels=num_labels)
60
61
62
def resolve_classifier_act_fn(
@@ -110,16 +114,9 @@ class PoolerMultiLabelClassify(PoolerActivation):
114
115
116
class PoolerClassify(PoolerActivation):
113
- def __init__(self, *, static_num_labels: bool = True) -> None:
117
+ def __init__(self, *, num_labels: int | None = None) -> None:
118
super().__init__()
119
116
- if static_num_labels:
117
- vllm_config = get_current_vllm_config()
118
- model_config = vllm_config.model_config
119
- num_labels = getattr(model_config.hf_config, "num_labels", 0)
120
- else:
121
- num_labels = None
122
-
120
if num_labels == 0:
121
logger.warning(
122
"num_labels should be > 0 for classification "