Embeddings differ between colpali engine & vLLM

#1
by Samoed - opened

Hi, @athrael-soju ! I tried to comapred a colpali engine implementation with vLLM and found a bit difference in embeddings:

import torch
from colpali_engine.models import ColQwen3_5, ColQwen3_5Processor

model = ColQwen3_5.from_pretrained(
    "vultr/VultronRetrieverFlash-Qwen3.5-0.8B",
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",   # required (see above)
    device_map="cuda:0",
).eval()

processor = ColQwen3_5Processor.from_pretrained(
    "vultr/VultronRetrieverFlash-Qwen3.5-0.8B",
    max_num_visual_tokens=1792,
)

queries = ["What was Q3 revenue?", "Summarize the safety findings."]

with torch.no_grad():
    qry_emb = model(**processor.process_queries(queries).to(model.device))

print(qry_emb)
tensor([[[-0.0000,  0.0000, -0.0000,  ...,  0.0000, -0.0000, -0.0000],
         [-0.0098, -0.0532,  0.0282,  ..., -0.0199,  0.0530,  0.0020],
         [ 0.0518,  0.0157, -0.1123,  ...,  0.0105,  0.0364,  0.0425],
         ...,
         [ 0.0449, -0.1016, -0.0310,  ..., -0.0859, -0.0454,  0.0139],
         [ 0.0322, -0.1006, -0.0312,  ..., -0.0898, -0.0479,  0.0141],
         [ 0.0216, -0.1001, -0.0312,  ..., -0.0938, -0.0530,  0.0135]],

        [[ 0.0417, -0.0349, -0.0435,  ..., -0.0396,  0.0576,  0.0282],
         [-0.0339, -0.0559, -0.0045,  ..., -0.0535,  0.0332,  0.0106],
         [-0.0479, -0.0354, -0.0009,  ..., -0.0315,  0.0640, -0.0120],
         ...,
         [-0.0410, -0.0093, -0.0227,  ..., -0.0280, -0.0267,  0.0165],
         [-0.0415, -0.0108, -0.0215,  ..., -0.0275, -0.0275,  0.0167],
         [-0.0415, -0.0115, -0.0216,  ..., -0.0269, -0.0275,  0.0167]]],
       device='cuda:0', dtype=torch.bfloat16)
import torch
from vllm import LLM


MODEL = "vultr/VultronRetrieverFlash-Qwen3.5-0.8B"
MAX_PIXELS = 1792 * 32 * 32   # max_num_visual_tokens * (patch_size 16 * merge_size 2)^2
llm = LLM(
    model=MODEL,
    runner="pooling",
    dtype="bfloat16",
    enable_prefix_caching=False,
    enable_chunked_prefill=False,
    mm_processor_kwargs={"min_pixels": 65536, "max_pixels": MAX_PIXELS},
)

def query_prompt(q):
    return q + "<|endoftext|>" * 10   # query augmentation

queries = ["What was Q3 revenue?", "Summarize the safety findings."]
qry_out = llm.encode([query_prompt(q) for q in queries], pooling_task="token_embed")
def mv(o):  # one [num_tokens, 320] multi-vector per item, L2-normalized per token
    t = torch.as_tensor(o.outputs.data, dtype=torch.float32)
    return torch.nn.functional.normalize(t, p=2, dim=-1)
qrys = [mv(o) for o in qry_out]
print(qrys)
[tensor([[-0.0134, -0.0545,  0.0315,  ..., -0.0169,  0.0523,  0.0004],
        [ 0.0535,  0.0159, -0.1114,  ...,  0.0084,  0.0366,  0.0438],
        [ 0.1022, -0.1022, -0.0065,  ..., -0.0775,  0.0665, -0.0115],
        ...,
        [ 0.0438, -0.1008, -0.0315,  ..., -0.0863, -0.0442,  0.0134],
        [ 0.0301, -0.1009, -0.0313,  ..., -0.0903, -0.0473,  0.0132],
        [ 0.0226, -0.0996, -0.0301,  ..., -0.0921, -0.0502,  0.0150]]), tensor([[ 0.0397, -0.0383, -0.0446,  ..., -0.0417,  0.0567,  0.0306],
        [-0.0342, -0.0560, -0.0058,  ..., -0.0524,  0.0342,  0.0105],
        [-0.0486, -0.0339,  0.0006,  ..., -0.0321,  0.0646, -0.0124],
        ...,
        [-0.0420, -0.0109, -0.0230,  ..., -0.0263, -0.0261,  0.0162],
        [-0.0416, -0.0124, -0.0224,  ..., -0.0264, -0.0273,  0.0168],
        [-0.0417, -0.0126, -0.0222,  ..., -0.0255, -0.0292,  0.0173]])]
qrys[-1] - qry_emb[-1], (qrys[-1] - qry_emb[-1]).mean()
tensor([[-2.0363e-03, -3.3458e-03, -1.1890e-03,  ..., -2.1827e-03,
         -8.9351e-04,  2.3923e-03],
        [-2.5149e-04, -5.2899e-05, -1.3261e-03,  ...,  1.1103e-03,
          9.4742e-04, -8.6430e-05],
        [-7.0890e-04,  1.4785e-03,  1.4692e-03,  ..., -6.3630e-04,
          6.6428e-04, -4.7570e-04],
        ...,
        [-1.0184e-03, -1.5462e-03, -3.0975e-04,  ...,  1.6311e-03,
          6.7541e-04, -2.8464e-04],
        [-1.2470e-04, -1.6343e-03, -8.7845e-04,  ...,  1.0165e-03,
          1.7914e-04,  2.8048e-05],
        [-2.1007e-04, -1.0888e-03, -6.0353e-04,  ...,  1.3782e-03,
         -1.7731e-03,  5.4824e-04]]), tensor(4.9784e-06)

Lib versions:

uv pip freeze | grep -E "vllm|colpali-engine|transformers"

colpali-engine==0.3.17
transformers==5.14.1
vllm==0.25.1+cu129

I think this is most probably due to different attentions in vllm and colpali implementations. In https://github.com/vllm-project/vllm/pull/46108 you've changed attention of ColQwen3.5 models to bidirectional, but in transformers code seems this is not possible to make attention bidirectional for qwen3.5 https://github.com/huggingface/transformers/blob/7ea2320c76117e6742364808a666ef6f2fb40a67/src/transformers/models/qwen3_5/modeling_qwen3_5.py#L660

Hey @Samoed nice catch. I’ll take a look at this.

Great! But I think there still will be problem for athrael-soju/colqwen3.5-4.5B-v3, because this model is expected to be bidirectional, but this is not possible with current transformers implementation which probably should be updated

Great! But I think there still will be problem for athrael-soju/colqwen3.5-4.5B-v3, because this model is expected to be bidirectional, but this is not possible with current transformers implementation which probably should be updated

That's right, #49372 fixes vLLM only. Transformers still hard-codes Qwen3.5 attention and will need a separate fix. Will look into it, If I get the chance.

Sign up or log in to comment