[Bugfix] Fix ColQwen3.5 pooler projector initialization (#54847)
Signed-off-by: Divakar Verma <divakar.verma@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Divakar Verma committed
Sep 1, 2026 at 22:35 UTC
dbf1a044eacdd5a5a063bab3701e63549ac87a8e
1 file changed
+10
-17
vllm/model_executor/models/colqwen3_5.py
+10
-17
@@ -121,13 +121,13 @@ class ColQwen3_5Model(
121
linear projection layer for per-token embeddings. It supports:
122
- "token_embed" task: Per-token embeddings for late interaction scoring
123
124
- The model produces per-token embeddings by:
124
+ The model produces L2-normalized per-token embeddings by:
125
1. Running the Qwen3.5 backbone (vision + language) to get hidden states
126
- 2. Projecting hidden states through a linear layer (hidden_size -> embed_dim)
127
- 3. L2 normalization is handled by the pooler via PoolerNormalize
126
+ 2. Projection and L2-normalization via the pooler (TokenEmbeddingPoolerHead)
127
128
Attributes:
130
- custom_text_proj: Linear projection from hidden_size to embed_dim
129
+ custom_text_proj: Linear projection from hidden_size to embed_dim.
130
+ This is passed to the pooler as pooler.head.projector.
131
"""
132
133
# Mark this as a pooling model so vLLM routes to pooler path
@@ -187,7 +187,7 @@ class ColQwen3_5Model(
187
assert pooler_config is not None
188
self.pooler = pooler_for_token_embed(
189
pooler_config,
190
- projector=None,
190
+ projector=self.custom_text_proj,
191
)
192
193
def forward(
@@ -198,8 +198,8 @@ class ColQwen3_5Model(
198
inputs_embeds: torch.Tensor | None = None,
199
**kwargs: object,
200
) -> torch.Tensor:
201
- """Run forward pass producing per-token embeddings."""
202
- hidden_states = super().forward(
201
+ """Run forward pass returning hidden states for pooler."""
202
+ return super().forward(
203
input_ids=input_ids,
204
positions=positions,
205
intermediate_tensors=intermediate_tensors,
@@ -207,16 +207,6 @@ class ColQwen3_5Model(
207
**kwargs,
208
)
209
210
- if not isinstance(hidden_states, torch.Tensor):
211
- return hidden_states # type: ignore
212
-
213
- proj_dtype = self.custom_text_proj.weight.dtype
214
- if hidden_states.dtype != proj_dtype:
215
- hidden_states = hidden_states.to(proj_dtype)
216
-
217
- # Project to embedding dimension (normalization handled by pooler)
218
- return self.custom_text_proj(hidden_states)
219
-
210
# Names used for the projection layer across different ColQwen3.5 variants
211
_PROJ_LAYER_NAMES = {
212
"custom_text_proj", # ColPali naming
@@ -249,5 +239,8 @@ class ColQwen3_5Model(
239
weight = weight.to(device=param.device, dtype=param.dtype)
240
default_weight_loader(param, weight)
241
loaded.add(f"custom_text_proj.{param_name}")
242
+ # Also mark as loaded under pooler path since custom_text_proj
243
+ # is assigned to pooler.head.projector
244
+ loaded.add(f"pooler.head.projector.{param_name}")
245
246
return loaded