Regression in SmolVLM results in different vision embeddings

[yfw]

System Info

This is an issue since v4.55.1

Who can help?

@zucchini-nlp

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Please run the following reproducer script: uv run reproducer.py

#!/usr/bin/env -S uv run --script
# /// script
# dependencies = [
#   "num2words",
#   "pillow",
#   "torch",
#   "torchvision",
#   "transformers==4.54.1",
# ]
# ///

import transformers
print(transformers.__version__)

from transformers import AutoProcessor, AutoModelForImageTextToText
import torch

model_path = "HuggingFaceTB/SmolVLM2-2.2B-Instruct"
processor = AutoProcessor.from_pretrained(model_path)
model = AutoModelForImageTextToText.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
).to("cuda")

messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg"},
            {"type": "text", "text": "Can you describe this image?"},
        ]
    },
]

inputs = processor.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device, dtype=torch.bfloat16)

patch_size = model.model.vision_model.patch_size
pixel_values = inputs['pixel_values']
batch_size, num_images, num_channels, height, width = pixel_values.shape
pixel_values = pixel_values.view(batch_size * num_images, *pixel_values.shape[2:])

patch_attention_mask = torch.ones(
    (
        batch_size,
        pixel_values.size(2) // patch_size,
        pixel_values.size(3) // patch_size,
    )
)
patch_attention_mask = patch_attention_mask.to(dtype=torch.bool, device=pixel_values.device)

embeddings_model = model.model.vision_model.embeddings
embeddings = embeddings_model(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask)
print(embeddings[0][-1][:10])

For transformers==4.54.1, this will produce:

tensor([-0.5938, -0.5117, -0.7305, -1.1797, -0.5977, -0.7305, -0.7070, -0.6484,
        -0.5547, -0.6758], device='cuda:0', dtype=torch.bfloat16,
       grad_fn=<SliceBackward0>)

Changing the dependencies to "transformers==4.55.1" and rerunning the script will produce:

tensor([-0.5977, -0.5742, -0.6875, -1.0625, -0.5977, -0.8398, -0.7031, -0.6406,
        -0.5078, -0.6445], device='cuda:0', dtype=torch.bfloat16,
       grad_fn=<SliceBackward0>)

The issue is the logic to calculate position_ids changed slightly from 4.54.1 (https://github.com/huggingface/transformers/blob/4.54.1/src/transformers/models/smolvlm/modeling_smolvlm.py#L131-L156) to 4.55.1 (https://github.com/huggingface/transformers/blob/v4.55.1/src/transformers/models/smolvlm/modeling_smolvlm.py#L131-L162). This results in different position_embeddings which affect the final embeddings.

Expected behavior

We expect the embeddings to be the same between versions. Furthermore, the current version differs from vllm's implementation (which is aligned with 4.54.1). This causes issues in RL when we expect the inference and training implementations to be aligned.

[zucchini-nlp]

AFAIR that was changed to accommodate for export and the ops look similar visually. Ideally we'd love to keep exportability and BC. Would you like to submit a PR keeping both?

[yfw]

AFAIR that was changed to accommodate for export and the ops look similar visually. Ideally we'd love to keep exportability and BC. Would you like to submit a PR keeping both?

Can you explain what is needed for exportability? Seems the problem is with some numerical differences in fractional_coords due to a slightly different implementation. The difference is small but becomes larger when we look at the final embeddings.

4.54.1:

fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h)
fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / nb_patches_w)

results in:

tensor([0.00000000, 0.03703704, 0.07407407, 0.11111111, 0.14814815, 0.18518519,
        0.22222222, 0.25925925, 0.29629630, 0.33333334, 0.37037039, 0.40740740,
        0.44444445, 0.48148149, 0.51851851, 0.55555558, 0.59259260, 0.62962961,
        0.66666669, 0.70370370, 0.74074078, 0.77777779, 0.81481481, 0.85185188,
        0.88888890, 0.92592591, 0.96296299])

4.55.1:

h_indices = torch.arange(nb_patches_h, device=position_ids.device, dtype=position_ids.dtype)
w_indices = torch.arange(nb_patches_w, device=position_ids.device, dtype=position_ids.dtype)
fractional_coords_h = h_indices / nb_patches_h * (1 - 1e-6)
fractional_coords_w = w_indices / nb_patches_w * (1 - 1e-6)

results in:

tensor([0.00000000, 0.03703700, 0.07407400, 0.11111100, 0.14814800, 0.18518499,
        0.22222200, 0.25925899, 0.29629600, 0.33333302, 0.37036997, 0.40740699,
        0.44444400, 0.48148102, 0.51851797, 0.55555505, 0.59259200, 0.62962896,
        0.66666603, 0.70370299, 0.74073994, 0.77777702, 0.81481397, 0.85185105,
        0.88888800, 0.92592496, 0.96296203], device='cuda:0')

[zucchini-nlp]

For testing we can run the tests from #39614 to check if exportability is still working. I didn't work on it, visually looks like dynamic control flow issues

[i3hz]

Hello, I think I have a viable approach for this issue and I just wanted to confirm it's correct before submitting a PR .
We basically have to replicate the old implementation.

h_indices = torch.arange(nb_patches_h, device=position_ids.device, dtype=position_ids.dtype)
w_indices = torch.arange(nb_patches_w, device=position_ids.device, dtype=position_ids.dtype)

fractional_coords_h = h_indices / nb_patches_h * (1 - 1e-6)
fractional_coords_w = w_indices / nb_patches_w * (1 - 1e-6)

by using an integer arange which should be friendly to ONNX exporter .
So my approach is something like this where we pre calculate the step.

# calculate the float step size, like the old arange logic.
step_h = 1.0 / nb_patches_h
step_w = 1.0 / nb_patches_w

# create integer indices and multiply
h_indices = torch.arange(nb_patches_h, device=position_ids.device, dtype=torch.float32)
w_indices = torch.arange(nb_patches_w, device=position_ids.device, dtype=torch.float32)

fractional_coords_h = h_indices * step_h
fractional_coords_w = w_indices * step_w


fractional_coords_h = torch.clamp(fractional_coords_h, max=(1.0 - 1e-6))
fractional_coords_w = torch.clamp(fractional_coords_w, max=(1.0 - 1e-6))

# cast back to the original dtype 
fractional_coords_h = fractional_coords_h.to(position_ids.dtype)
fractional_coords_w = fractional_coords_w.to(position_ids.dtype)

I also created a short script to test the outputs

import torch

# We use the same parameters  to ensure a valid comparison.
nb_patches_h = 27
nb_patches_w = 27
device = "cuda" 

print(f"Using device: {device}")
print(f"Testing with nb_patches_h = {nb_patches_h}")
print("-" * 40)


print("1. Running the v4.54.1 logic")
original_coords = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h, device=device).to(torch.float32)
print(f"   Output: {original_coords[:10]}")
print(f"   Shape: {original_coords.shape}")


print("\n2. Running the v4.55.1 logic")
h_indices_buggy = torch.arange(nb_patches_h, device=device)
buggy_coords = (h_indices_buggy / nb_patches_h * (1 - 1e-6)).to(torch.float32)
print(f"   Output: {buggy_coords[:10]}")
print(f"   Shape: {buggy_coords.shape}")

print("Running the Proposed Fix ")
step_h = 1.0 / nb_patches_h
h_indices_fix = torch.arange(nb_patches_h, device=device, dtype=torch.float32)
fixed_coords = h_indices_fix * step_h
fixed_coords = torch.clamp(fixed_coords, max=(1.0 - 1e-6))
print(f"   Output: {fixed_coords[:10]}")
print(f"   Shape: {fixed_coords.shape}")

print("\n" + "-" * 40)
print("--- VERIFICATION ---")

is_buggy_correct = torch.equal(original_coords, buggy_coords)
is_fix_correct = torch.equal(original_coords, fixed_coords)

print(f"\nIs the buggy version numerically IDENTICAL to the original? \n>>> {is_buggy_correct}")
print(f"\nIs the proposed fix numerically IDENTICAL to the original? \n>>> {is_fix_correct}")

print("\n" + "-" * 40)
if is_fix_correct:
    print(" The proposed fix is numerically equivalent to the original code.")
else:
    print("The proposed fix is NOT numerically equivalent.")

which does confirm that the fix works .

@zucchini-nlp Is this the correct implementation? If so, I'll move forward with creating the PR.

[zucchini-nlp]

@i3hz if the computed positions are passing the check for equivalence, it looks good to me. Can you open a PR and run slow tests in SmolVLM (RUN_SLOW=1 pytest tests/models/smolvlm/test_modeling_smolvlm.py) to make sure it is indeed equivalent and the torch export is still supported?

[i3hz]

I can work on that . I had a question , in the src/transformers/models/smolvlm/modeling_smolvlm.py file there is a comment stating .

so after looking into modular_smolvlm.py I found out that the vision embedding logic is just inherited from idefics3 model . So I wanted to make sure that is where I'm supposed to make the change (src/transformers/models/idefics3/modeling_idefics3.py)? Thanks for the help @zucchini-nlp

[zucchini-nlp]

Yep, we need to fix the root and then run make fix-copies :)

[i3hz]

Thanks again .
I have created a PR that hopefully fixed this issue XD
#41374
Please look into and let me know if there's anything missing @zucchini-nlp