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, 13:00 UTC
Add unit tests for pooler activation functions (#42824)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
May 16, 2026 at 14:58 UTC
787bc0d0313840c16e403dfa2d135781d41d3614
2 files changed
+238
-1
tests/model_executor/layers/test_pooler_activations.py
new
+237
@@ -0,0 +1,237 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for vllm.model_executor.layers.pooler.activations."""
4
+
5
+from types import SimpleNamespace
6
+
7
+import pytest
8
+import torch
9
+import torch.nn as nn
10
+
11
+from vllm.config import VllmConfig, set_current_vllm_config
12
+from vllm.model_executor.layers.pooler.activations import (
13
+ LambdaPoolerActivation,
14
+ PoolerClassify,
15
+ PoolerIdentity,
16
+ PoolerMultiLabelClassify,
17
+ PoolerNormalize,
18
+ get_act_fn,
19
+ resolve_classifier_act_fn,
20
+)
21
+
22
+
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
+
32
+# ---------------------------------------------------------------------------
33
+# PoolerIdentity
34
+# ---------------------------------------------------------------------------
35
+class TestPoolerIdentity:
36
+ def test_returns_input_unchanged(self):
37
+ pooler = PoolerIdentity()
38
+ x = torch.randn(4, 128)
39
+ out = pooler(x)
40
+ assert torch.equal(out, x)
41
+
42
+ def test_forward_list(self):
43
+ pooler = PoolerIdentity()
44
+ tensors = [torch.randn(128), torch.randn(256)]
45
+ out = pooler(tensors)
46
+ assert len(out) == 2
47
+ for orig, result in zip(tensors, out):
48
+ assert torch.equal(orig, result)
49
+
50
+
51
+# ---------------------------------------------------------------------------
52
+# PoolerNormalize
53
+# ---------------------------------------------------------------------------
54
+class TestPoolerNormalize:
55
+ def test_output_has_unit_norm(self):
56
+ pooler = PoolerNormalize()
57
+ x = torch.randn(4, 128)
58
+ out = pooler(x)
59
+ norms = torch.linalg.norm(out, dim=-1)
60
+ assert torch.allclose(norms, torch.ones(4), atol=1e-5)
61
+
62
+ def test_single_vector(self):
63
+ pooler = PoolerNormalize()
64
+ x = torch.randn(1, 64)
65
+ out = pooler(x)
66
+ norm = torch.linalg.norm(out, dim=-1)
67
+ assert torch.allclose(norm, torch.ones(1), atol=1e-5)
68
+
69
+ def test_forward_list(self):
70
+ pooler = PoolerNormalize()
71
+ tensors = [torch.randn(1, 64), torch.randn(1, 128)]
72
+ out = pooler(tensors)
73
+ for t in out:
74
+ norm = torch.linalg.norm(t, dim=-1)
75
+ assert torch.allclose(norm, torch.ones(1), atol=1e-5)
76
+
77
+
78
+# ---------------------------------------------------------------------------
79
+# PoolerMultiLabelClassify
80
+# ---------------------------------------------------------------------------
81
+class TestPoolerMultiLabelClassify:
82
+ def test_output_in_zero_one(self):
83
+ pooler = PoolerMultiLabelClassify()
84
+ x = torch.randn(4, 10)
85
+ out = pooler(x)
86
+ assert (out >= 0).all() and (out <= 1).all()
87
+
88
+ def test_large_positive_maps_near_one(self):
89
+ pooler = PoolerMultiLabelClassify()
90
+ x = torch.full((1, 3), 100.0)
91
+ out = pooler(x)
92
+ assert torch.allclose(out, torch.ones(1, 3), atol=1e-4)
93
+
94
+ def test_large_negative_maps_near_zero(self):
95
+ pooler = PoolerMultiLabelClassify()
96
+ x = torch.full((1, 3), -100.0)
97
+ out = pooler(x)
98
+ assert torch.allclose(out, torch.zeros(1, 3), atol=1e-4)
99
+
100
+
101
+# ---------------------------------------------------------------------------
102
+# PoolerClassify
103
+# ---------------------------------------------------------------------------
104
+class TestPoolerClassify:
105
+ def test_dynamic_softmax_when_num_labels_ge_2(self):
106
+ pooler = PoolerClassify(static_num_labels=False)
107
+ assert pooler.num_labels is None
108
+ x = torch.randn(2, 5)
109
+ out = pooler(x)
110
+ sums = out.sum(dim=-1)
111
+ assert torch.allclose(sums, torch.ones(2), atol=1e-5)
112
+
113
+ def test_dynamic_sigmoid_when_num_labels_lt_2(self):
114
+ pooler = PoolerClassify(static_num_labels=False)
115
+ x = torch.zeros(1, 1)
116
+ out = pooler(x)
117
+ assert torch.allclose(out, torch.tensor([[0.5]]), atol=1e-5)
118
+
119
+ def test_static_default_num_labels_zero_uses_sigmoid(self, vllm_config):
120
+ pooler = PoolerClassify(static_num_labels=True)
121
+ assert pooler.num_labels == 0
122
+ x = torch.zeros(1, 3)
123
+ out = pooler(x)
124
+ assert torch.allclose(out, torch.full((1, 3), 0.5), atol=1e-5)
125
+
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)
129
+ assert pooler.num_labels == 4
130
+ x = torch.randn(2, 4)
131
+ out = pooler(x)
132
+ sums = out.sum(dim=-1)
133
+ assert torch.allclose(sums, torch.ones(2), atol=1e-5)
134
+
135
+
136
+# ---------------------------------------------------------------------------
137
+# LambdaPoolerActivation
138
+# ---------------------------------------------------------------------------
139
+class TestLambdaPoolerActivation:
140
+ def test_applies_custom_fn(self):
141
+ pooler = LambdaPoolerActivation(nn.ReLU())
142
+ x = torch.tensor([[-1.0, 2.0, -3.0]])
143
+ out = pooler(x)
144
+ expected = torch.tensor([[0.0, 2.0, 0.0]])
145
+ assert torch.equal(out, expected)
146
+
147
+ def test_forward_list(self):
148
+ pooler = LambdaPoolerActivation(nn.ReLU())
149
+ tensors = [torch.tensor([-1.0, 2.0]), torch.tensor([3.0, -4.0])]
150
+ out = pooler(tensors)
151
+ assert torch.equal(out[0], torch.tensor([0.0, 2.0]))
152
+ assert torch.equal(out[1], torch.tensor([3.0, 0.0]))
153
+
154
+
155
+# ---------------------------------------------------------------------------
156
+# get_act_fn factory
157
+# ---------------------------------------------------------------------------
158
+class TestGetActFn:
159
+ @staticmethod
160
+ def _make_config(**kwargs):
161
+ return SimpleNamespace(**kwargs)
162
+
163
+ def test_regression(self, vllm_config):
164
+ cfg = self._make_config(problem_type="regression")
165
+ result = get_act_fn(cfg)
166
+ assert isinstance(result, PoolerIdentity)
167
+
168
+ def test_single_label_classification(self, vllm_config):
169
+ cfg = self._make_config(
170
+ problem_type="single_label_classification", num_labels=3
171
+ )
172
+ result = get_act_fn(cfg)
173
+ assert isinstance(result, PoolerClassify)
174
+
175
+ def test_multi_label_classification(self, vllm_config):
176
+ cfg = self._make_config(problem_type="multi_label_classification")
177
+ result = get_act_fn(cfg)
178
+ assert isinstance(result, PoolerMultiLabelClassify)
179
+
180
+ def test_sentence_transformers_activation(self, vllm_config):
181
+ cfg = self._make_config(
182
+ problem_type="",
183
+ sentence_transformers={
184
+ "activation_fn": "torch.nn.modules.activation.Sigmoid"
185
+ },
186
+ )
187
+ result = get_act_fn(cfg)
188
+ assert isinstance(result, PoolerClassify)
189
+
190
+ def test_sbert_activation(self, vllm_config):
191
+ cfg = self._make_config(
192
+ problem_type="",
193
+ sbert_ce_default_activation_function=(
194
+ "torch.nn.modules.activation.Sigmoid"
195
+ ),
196
+ )
197
+ result = get_act_fn(cfg)
198
+ assert isinstance(result, PoolerClassify)
199
+
200
+ def test_default_fallback(self, vllm_config):
201
+ cfg = self._make_config(problem_type="")
202
+ result = get_act_fn(cfg)
203
+ assert isinstance(result, PoolerClassify)
204
+
205
+ def test_sentence_transformers_takes_priority(self, vllm_config):
206
+ cfg = self._make_config(
207
+ problem_type="",
208
+ sentence_transformers={"activation_fn": "torch.nn.modules.linear.Identity"},
209
+ sbert_ce_default_activation_function=(
210
+ "torch.nn.modules.activation.Sigmoid"
211
+ ),
212
+ )
213
+ result = get_act_fn(cfg)
214
+ assert isinstance(result, PoolerIdentity)
215
+
216
+ def test_rejects_non_torch_activation(self, vllm_config):
217
+ cfg = self._make_config(
218
+ problem_type="",
219
+ sentence_transformers={"activation_fn": "os.system"},
220
+ )
221
+ with pytest.raises(AssertionError, match="restricted"):
222
+ get_act_fn(cfg)
223
+
224
+
225
+# ---------------------------------------------------------------------------
226
+# resolve_classifier_act_fn
227
+# ---------------------------------------------------------------------------
228
+class TestResolveClassifierActFn:
229
+ def test_delegates_to_get_act_fn_when_none(self, vllm_config):
230
+ model_config = vllm_config.model_config
231
+ result = resolve_classifier_act_fn(model_config, act_fn=None)
232
+ assert isinstance(result, PoolerClassify)
233
+
234
+ def test_passes_through_provided_act_fn(self):
235
+ custom = PoolerIdentity()
236
+ result = resolve_classifier_act_fn(None, act_fn=custom)
237
+ assert result is custom
vllm/model_executor/layers/pooler/activations.py
+1
-1
@@ -123,7 +123,7 @@ class PoolerClassify(PoolerActivation):
123
if num_labels == 0:
124
logger.warning(
125
"num_labels should be > 0 for classification "
126
- "models, falling back to softmax. "
126
+ "models, falling back to sigmoid. "
127
"Please check if the configuration is correct."
128
)
129