[Bugfix] Resolve seq-cls `num_labels` from the top-level config for multimodal checkpoints (#50950)
Signed-off-by: Lazurite <43494437+Rapisurazurite@users.noreply.github.com>
Lazurite committed
Aug 4, 2026 at 19:39 UTC
6a9fdf0dc75844de3fff993c32e6f2a35f6fbf0d
2 files changed
+71
-2
tests/models/test_adapters.py
+47
-1
@@ -4,8 +4,12 @@
4
5
import pytest
6
import torch
7
+from transformers import Gemma3Config, PretrainedConfig, Qwen2Config
8
8
-from vllm.model_executor.models.adapters import _create_pooling_model_cls
9
+from vllm.model_executor.models.adapters import (
10
+ _create_pooling_model_cls,
11
+ _resolve_num_labels,
12
+)
13
from vllm.model_executor.models.utils import AutoWeightsLoader, StageMissingLayer
14
15
pytestmark = pytest.mark.cpu_test
@@ -146,3 +150,45 @@ def test_pooling_load_weights_clones_probed_weights():
150
expected = {n: p.data.clone() for n, p in ground_truth.named_parameters()}
151
152
_load_and_compare(_make_pooling_model(PackedWeightModel), ref, expected)
153
+
154
+
155
+def _composite_config(outer_labels=None, inner_labels=None):
156
+ """Build a multimodal-style config whose text_config is a separate object."""
157
+ config = Gemma3Config(text_config={"num_hidden_layers": 1})
158
+ if outer_labels is not None:
159
+ config.num_labels = outer_labels
160
+ if inner_labels is not None:
161
+ config.get_text_config().num_labels = inner_labels
162
+ return config
163
+
164
+
165
+def test_resolve_num_labels_text_only_config():
166
+ config = Qwen2Config(num_labels=7)
167
+ assert config.get_text_config() is config
168
+ assert _resolve_num_labels(config, config.get_text_config()) == 7
169
+
170
+
171
+def test_resolve_num_labels_defaults_when_undeclared():
172
+ config = _composite_config()
173
+ assert (
174
+ _resolve_num_labels(config, config.get_text_config())
175
+ == PretrainedConfig().num_labels
176
+ )
177
+
178
+
179
+def test_resolve_num_labels_declared_on_outer_config():
180
+ """Multimodal checkpoints keep id2label/problem_type on the top-level config."""
181
+ config = _composite_config(outer_labels=20)
182
+ assert config.get_text_config().num_labels == PretrainedConfig().num_labels
183
+ assert _resolve_num_labels(config, config.get_text_config()) == 20
184
+
185
+
186
+def test_resolve_num_labels_declared_on_text_config():
187
+ """Overrides written into text_config keep working."""
188
+ config = _composite_config(inner_labels=5)
189
+ assert _resolve_num_labels(config, config.get_text_config()) == 5
190
+
191
+
192
+def test_resolve_num_labels_outer_wins_when_both_declared():
193
+ config = _composite_config(outer_labels=20, inner_labels=5)
194
+ assert _resolve_num_labels(config, config.get_text_config()) == 20
vllm/model_executor/models/adapters.py
+24
-1
@@ -261,6 +261,29 @@ def as_embedding_model(cls: _T) -> _T:
261
return ModelForEmbedding # type: ignore
262
263
264
+def _resolve_num_labels(hf_config: Any, text_config: Any) -> int:
265
+ """Resolve the label count for a sequence classification head.
266
+
267
+ ``PretrainedConfig.num_labels`` is derived from ``id2label``, which always
268
+ carries a default of two entries. Composite configs (such as multimodal
269
+ checkpoints) declare their label space on the top-level config, so reading
270
+ ``num_labels`` from ``get_text_config()`` silently returns that default and
271
+ builds a score head of the wrong size.
272
+
273
+ Prefer the top-level config whenever it declares a label space of its own,
274
+ mirroring the ``classifier_from_token`` / ``method`` lookups in
275
+ ``as_seq_cls_model``.
276
+ """
277
+ if text_config is hf_config:
278
+ return hf_config.num_labels
279
+
280
+ from transformers import PretrainedConfig
281
+
282
+ if hf_config.num_labels != PretrainedConfig().num_labels:
283
+ return hf_config.num_labels
284
+ return text_config.num_labels
285
+
286
+
287
def as_seq_cls_model(cls: _T) -> _T:
288
"""
289
Subclass an existing vLLM model to support classify and score tasks.
@@ -320,7 +343,7 @@ def as_seq_cls_model(cls: _T) -> _T:
343
344
self.score = ReplicatedLinear(
345
model_config.get_hidden_size(),
323
- text_config.num_labels,
346
+ _resolve_num_labels(hf_config, text_config),
347
bias=False,
348
params_dtype=model_config.head_dtype,
349
quant_config=quant_config,