Qwen2-VL Doesn't Execute on TPUs

System Info

transformers version: 4.45.0.dev0
Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.31
Python version: 3.10.14
Huggingface_hub version: 0.24.6
Safetensors version: 0.4.4
Accelerate version: 0.33.0
Accelerate config: not found
PyTorch version (GPU?): 2.5.0.dev20240830+cpu (False)
Tensorflow version (GPU?): not installed (NA)
Flax version (CPU?/GPU?/TPU?): not installed (NA)
Jax version: not installed
JaxLib version: not installed
Using distributed or parallel set-up in script?:

Who can help?

No response

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

#Following this Qwen2-VL guide => https://huggingface.co/Qwen/Qwen2-VL-2B-Instruct#quickstart

Script

from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from qwen_vl_utils import process_vision_info
import numpy as np
import torch
import torch_xla as xla
import torch_xla.core.xla_model as xm
import torch_xla.distributed.spmd as xs

from torch.distributed._tensor import DeviceMesh, distribute_module
from torch_xla.distributed.spmd import auto_policy

from torch_xla import runtime as xr
from torch_xla.experimental.spmd_fully_sharded_data_parallel import (
    _prepare_spmd_partition_spec,
    SpmdFullyShardedDataParallel as FSDPv2,
)

import time

start = time.time()

device = xm.xla_device()

# default: Load the model on the available device(s)
model = Qwen2VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2-VL-2B-Instruct",
    torch_dtype=torch.bfloat16,
    attn_implementation="eager",
).to(device)


print(model.device)

# default processer
processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct-GPTQ-Int4")


message = [
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "image": "image1.jpg",
            },
            {"type": "text", "text": "Describe this image in detail."},
        ],
    }
]

all_messages = [[message] for _ in range(1)]
for messages in all_messages:

    # Preparation for inference
    texts = [
        processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
        for msg in messages
    ]

    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=texts,
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )
    inputs = inputs.to(device)

    # Inference: Generation of the output
    generated_ids = model.generate(**inputs, max_new_tokens=512)
    generated_ids_trimmed = [
        out_ids[len(in_ids) :]
        for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )
    for i, text in enumerate(output_text):
        print(f"Output {i}: {text}")

print(f"Time taken: {time.time() - start}")

Output Logs

kojoe@t1v-n-cb70f560-w-0:~/EasyAnimate/easyanimate/image_caption$ python caption.py
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
Loading checkpoint shards: 100%|██████████████████████████████████████████████████| 2/2 [00:00<00:00,  4.39it/s]
xla:0

Expected behavior

The model works fine when chaging device to "cpu", but stuck executing on TPUs. The model should run on TPUs

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

@tengomucho

Hey! Don't think we officially test nor support TPU for this model 🤗 I can't really reproduce 😢
@tengomucho might have an idea

@radna0

@radna0 transformers does not support officially TPU, but I think things might work if:

you use the raw model forward API, not generate
you use the static KV cache, that works best on TPUs with Pytorch XLA.
I haven't tried this model myself, but you can try, using inspiration from this script. If it does not work, feel free to open a PR on optimum-tpu and we will try to help you.

@ArthurZucker

Hey @ArthurZucker following up -- do you guys generally support Qwen-series models on TPU? Saw the HF trainer API was made compatible with XLA so curious if the status on this has changed.

We support xla so in general it shouuld all be supported.

@ArthurZucker

@ArthurZucker Unfortunately, the source code for Qwen2.5VL in the HuggingFace library (that defines the forward logic) is not XLA friendly, as far as I can tell. It utilizes a lot of tensor materialization, and therefore there are many recompilations within the forward pass.

We can take the example code from the huggingface model card (shown below). Note: I had to replace the .generate() function with just the forward() pass, because .generate() hangs indefinitely so you can't get a good time estimate from it. I made one modification in order to batch the same input multiple times, so we can see how the latency builds as batch size increases. But otherwise its exactly the same code as shown on the model card

Regardless, for a batch size of 1, the Qwen2.5VL-7B running on a v4-8 TPU gives the following forward pass times:
Forward pass 1: 13.58 seconds
Forward pass 2: 8.43 seconds
Forward pass 3: 8.13 seconds
Forward pass 4: 8.03 seconds
Forward pass 5: 7.96 seconds

For a batch size of 4, the same setup gives the following forward pass times:
Forward pass 1: 55.02 seconds
Forward pass 2: 34.99 seconds
Forward pass 3: 34.05 seconds
Forward pass 4: 33.80 seconds
Forward pass 5: 34.13 seconds

Again, I want to reiterate, this is for a relatively short input with relatively small batch sizes. You can see the exact code I use at the bottom of this comment.

Why is this happening? We can inspect where unexpected compilations are happening by setting the following environment variable: os.environ["PT_XLA_DEBUG_LEVEL"] = '2'

I won't go through the logs because they are enormous (thousands of lines of recompilations). However, I can, from them, point to a few functions that cause unexpected compilations (all in src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py):

def get_rope_index() -- line 957 in modeling_qwen2_5_vl.py
def get_window_index() -- line 365 in modeling_qwen2_5_vl.py
def rot_pos_emb() -- line 336 in modeling_qwen2_5_vl.py

There are more unexpected compilations in the forward pass of the ViT, as well as the overall model. Those are harder for me to pinpoint, as while this debugging strategy will return the function name where a compilation happened (and then I can track down the line number), generic "forward" functions are harder to pinpoint because the line number given doesn't exactly match up with meaningful line numbers in the actual file modeling_qwen2_5_vl.

For my use case, I was going to rewrite some of the forward logic to be XLA friendly and not cause recompilations. The task was too time consuming and out of the scope of my research, I had to move on to using other hardware. I suspect this issue would hold across other VLMs as well.

Here is the code I'm running:

from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from qwen_vl_utils import process_vision_info
import time
import os

# default: Load the model on the available device(s)
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2.5-VL-7B-Instruct", torch_dtype="auto", device_map="auto"
)
model = model.to(device="xla")
print(model.device)


# default processer
processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")

BATCH_SIZE = 1
messages = [
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
            },
            {"type": "text", "text": "Describe this image."},
        ],
    }
]
# batch the messages
messages = [messages]*BATCH_SIZE
# Preparation for inference
text = processor.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
)
image_inputs, video_inputs = process_vision_info(messages)
inputs = processor(
    text=text,
    images=image_inputs,
    videos=video_inputs,
    padding=True,
    return_tensors="pt",
)
inputs = inputs.to("xla")

# Inference: Generation of the output
num_iters = 1
os.environ["PT_XLA_DEBUG_LEVEL"] = '2'
for i in range(num_iters):
    start_time = time.time()
    output = model(**inputs)
    end_time = time.time()
    print(f"Forward pass: {end_time - start_time}")

Environment:
torch: 2.8.0
torch-xla: 2.8.0
torchvision: 0.23.0
transformers: 4.56.1
qwen-vl-utils: 0.0.11
accelerate: 1.10.1
python-3.11.13
TPU VM: v4-8 node

It makes a lot of sense, and I am sorry that you had to go through this!
If someone from the community is able to make the code a bit more TPU (but also probably just compile) friendly / more efficient, that would be awesome!

From our team I am not sure we'll have the bandwidth for now but happy to review PRs

As discussed, the Qwen2-VL/Qwen2.5-VL vision path materializes data-dependent shapes (e.g. grid_thw-driven cu_seqlens, window indices, repeat_interleave), which triggers XLA recompilations. Contributions making these static-shape friendly in modeling_qwen2_vl.py and modular_qwen2_5_vl.py (then make fix-repo) are welcome; they would also help torch.compile. See https://huggingface.co/docs/transformers/main/en/google_tpu.

@ArthurZucker

Hi @ArthurZucker — I'd like to pick this up, following your Oct 20 note inviting contributions that make the vision path static-shape friendly.

What I reproduced (TPU v5e-1, single chip, torch 2.8.0 + torch_xla 2.8.0, Qwen2-VL-2B-Instruct, transformers 5.20.1):

224×224 image: forward 41.9s / 0.2s / 0.2s; generate(20) did not finish within a 5-minute window.

What I plan to do (patch drafted and CPU-verified, covering both files you named):
vision_utils.py already documents that grid_thw-derived tensors should be precomputed on the host before tracing and offers the precomputed-kwargs injection channel — but its own TODO admits this is currently only a "power feature" and the default forward path doesn't use it. The patch wires up the default path: in Qwen2VLModel.forward (modeling_qwen2_vl.py) and Qwen2_5_VLModel (modular_qwen2_5_vl.py + regenerated modeling_qwen2_5_vl.py), compute cu_seqlens / position ids / max_seqlen (plus window_index / cu_window_seqlens / max_window_seqlen for 2.5-VL) on CPU from grid_thw right before calling the vision encoder, injected via the existing image_*/video_* prefixed kwargs. Caller-provided precomputed tensors are never overridden, and no consumer-side code changes.

Verification so far: pre/post-patch logits max diff 0.0 on both models, repo test subsets 4+3 passed, ruff clean. I'll run make fix-repo before submitting.

If this direction looks OK, I'll open a PR.

fix(qwen_vl): precompute vision metadata on CPU to eliminate XLA reco…

Navigation Menu

Repository navigation

Qwen2-VL Doesn't Execute on TPUs #33289

System Info

Who can help?

Information

Tasks

Reproduction

Expected behavior

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Qwen2-VL Doesn't Execute on TPUs #33289

Description

System Info

Who can help?

Information

Tasks

Reproduction

Expected behavior

Activity

github-actions commented on Oct 4, 2024

ArthurZucker commented on Oct 5, 2024

tengomucho commented on Oct 7, 2024

jereminuer commented on Sep 6, 2025

ArthurZucker commented on Sep 11, 2025

jereminuer commented on Sep 13, 2025

ArthurZucker commented on Sep 18, 2025

ArthurZucker commented on Oct 20, 2026

notpad commented on Oct 22, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions