Repository navigation
Qwen Image 2.1 with torch.compile triggers a recompile on each run with varying image size #15027
Description
Describe the bug
When running Qwen Image 2.1 as described in the documentation with QwenImage21FlexAttnProcessor, compiling using compile_repeated_blocks and generating images of different sizes with different condition image sizes, a recompilation is triggered on each generation. If running with fullgraph=True, a torch._dynamo.exc.FailOnRecompileLimitHit: Hard failure due to fullgraph=True error is thrown after a few generations.
Reproduction
import random
from PIL import Image
import torch
from diffusers import QwenImage21Pipeline
from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21FlexAttnProcessor
torch._logging.set_logs(
recompiles=True
)
pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1-Turbo", dtype=torch.bfloat16).to("cuda")
pipe.transformer.set_attn_processor(QwenImage21FlexAttnProcessor())
pipe.transformer.compile_repeated_blocks(fullgraph=True)
for i in range(12):
edit_image = Image.new(
"RGBA",
(640 + 32 * random.randint(0, 10), 640 + 32 * random.randint(0, 10)),
color=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 255)
)
image = pipe(
"A capybara wearing a wizard hat, oil painting",
image = edit_image,
height = 640 + 32 * random.randint(0, 10),
width = 640 + 32 * random.randint(0, 10)
).images[0]Logs
(Only including the last few lines)
12%|████████▍ | 1/8 [00:01<00:12, 1.78s/it]V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] Recompiling function forward in /workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py:615
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] triggered by the following guard failure(s):
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - 0/6: kv_cache_mode == 'extract' # if kv_cache_mode == "extract" and cache_write_slice is not None: # diffusers/models/transformers/transformer_qwenimage21.py:352 in _qwenimage21_prepare_qkv
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - User stack trace:
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 631, in forward
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - attn_output = self.attn(
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 580, in forward
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - return self.processor(self, hidden_states, **kwargs)
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 404, in __call__
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - query, key, value, seq_len_q = _qwenimage21_prepare_qkv(
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 352, in _qwenimage21_prepare_qkv
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - if kv_cache_mode == "extract" and cache_write_slice is not None:
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - 0/5: tensor 'layer_cache.k' size mismatch at index 1. expected 4121, actual 4169
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - 0/4: tensor 'layer_cache.k' size mismatch at index 1. expected 4121, actual 4169
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - 0/3: kv_cache_mode == 'extract' # if kv_cache_mode == "extract" and cache_write_slice is not None: # diffusers/models/transformers/transformer_qwenimage21.py:352 in _qwenimage21_prepare_qkv
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - User stack trace:
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 631, in forward
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - attn_output = self.attn(
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 580, in forward
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - return self.processor(self, hidden_states, **kwargs)
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 404, in __call__
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - query, key, value, seq_len_q = _qwenimage21_prepare_qkv(
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 352, in _qwenimage21_prepare_qkv
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - if kv_cache_mode == "extract" and cache_write_slice is not None:
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - 0/2: tensor 'layer_cache.k' size mismatch at index 1. expected 4121, actual 4169
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - 0/1: tensor 'layer_cache.k' size mismatch at index 1. expected 4121, actual 4169
V1014 08:27:17.745000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/7] [__recompiles] - 0/0: tensor 'rotary_emb' size mismatch at index 0. expected 6721, actual 2300
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] Recompiling function forward in /workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py:615
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] triggered by the following guard failure(s):
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - 0/7: 'NoneType' object has no attribute 'size'
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - 0/6: kv_cache_mode == 'extract' # if kv_cache_mode == "extract" and cache_write_slice is not None: # diffusers/models/transformers/transformer_qwenimage21.py:352 in _qwenimage21_prepare_qkv
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - User stack trace:
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 631, in forward
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - attn_output = self.attn(
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 580, in forward
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - return self.processor(self, hidden_states, **kwargs)
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 404, in __call__
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - query, key, value, seq_len_q = _qwenimage21_prepare_qkv(
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 352, in _qwenimage21_prepare_qkv
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - if kv_cache_mode == "extract" and cache_write_slice is not None:
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - 0/5: tensor 'layer_cache.k' size mismatch at index 1. expected 4121, actual 4169
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - 0/4: tensor 'layer_cache.k' size mismatch at index 1. expected 4121, actual 4169
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - 0/3: kv_cache_mode == 'extract' # if kv_cache_mode == "extract" and cache_write_slice is not None: # diffusers/models/transformers/transformer_qwenimage21.py:352 in _qwenimage21_prepare_qkv
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - User stack trace:
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 631, in forward
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - attn_output = self.attn(
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 580, in forward
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - return self.processor(self, hidden_states, **kwargs)
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 404, in __call__
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - query, key, value, seq_len_q = _qwenimage21_prepare_qkv(
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 352, in _qwenimage21_prepare_qkv
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - if kv_cache_mode == "extract" and cache_write_slice is not None:
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - 0/2: tensor 'layer_cache.k' size mismatch at index 1. expected 4121, actual 4169
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - 0/1: tensor 'layer_cache.k' size mismatch at index 1. expected 4121, actual 4169
V1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/guards.py:5621] [0/8] [__recompiles] - 0/0: tensor 'rotary_emb' size mismatch at index 0. expected 6721, actual 2300
W1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py:2008] [0/8] torch._dynamo hit config.recompile_limit (8)
W1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py:2008] [0/8] function: 'forward' (/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py:615)
W1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py:2008] [0/8] last reason: 0/7: 'NoneType' object has no attribute 'size'
W1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py:2008] [0/8] To log all recompilation reasons, use TORCH_LOGS="recompiles".
W1014 08:27:19.229000 28568 .venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py:2008] [0/8] To diagnose recompilation issues, see https://docs.pytorch.org/docs/main/user_guide/torch_compiler/compile/programming_model.recompilation.html
12%|████████▍ | 1/8 [00:03<00:25, 3.58s/it]
Traceback (most recent call last):
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 2038, in _compile
raise_unimplemented_cache_limit_exceeded()
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 2022, in raise_unimplemented_cache_limit_exceeded
unimplemented(
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/torch/_dynamo/exc.py", line 775, in unimplemented
raise Unsupported(msg, gb_type, skip_frame)
torch._dynamo.exc.Unsupported: Dynamo recompile limit exceeded
Explanation: Dynamo attempted to recompile the code object too many times, exceeding the recompile_limit cache size limit (currently set to 8). Excessive recompilations can degrade performance due to the compilation overhead of each recompilation.
Hint: To monitor recompilations, enable TORCH_LOGS=recompiles. If recompilations are expected, consider increasing torch._dynamo.config.recompile_limit to an appropriate value.
Hint: See https://docs.pytorch.org/docs/main/user_guide/torch_compiler/compile/programming_model.recompilation.html for tips on dealing with recompilations.
Developer debug context: Limit type: recompile_limit
For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0039.html
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/workspace/code/gpu_worker/python_worker/test.py", line 24, in <module>
image = pipe(
^^^^^
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py", line 784, in __call__
noise_pred = self.transformer(
^^^^^^^^^^^^^^^^^
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1783, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1794, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/utils/peft_utils.py", line 323, in wrapper
result = forward_fn(self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/diffusers/models/transformers/transformer_qwenimage21.py", line 999, in forward
joint_hidden_states = block(
^^^^^^
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1781, in _wrapped_call_impl
return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 1261, in compile_wrapper
result = fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1794, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 2649, in __call__
result = self._torchdynamo_orig_backend(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 778, in __call__
result = _compile(
^^^^^^^^^
File "/workspace/code/gpu_worker/python_worker/.venv/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 2045, in _compile
raise FailOnRecompileLimitHit(
torch._dynamo.exc.FailOnRecompileLimitHit: Hard failure due to fullgraph=TrueSystem Info
- 🤗 Diffusers version: 0.41.0.dev0
- Platform: Linux-7.0.0-1014-gcp-x86_64-with-glibc2.39
- Running on Google Colab?: No
- Python version: 3.11.17
- PyTorch version (GPU?): 2.14.1+cu130 (True)
- Huggingface_hub version: 1.33.0
- Transformers version: 5.14.1
- Accelerate version: 1.11.0
- PEFT version: 0.21.2
- bitsandbytes version: 0.48.1
- optimum-quanto version: 0.2.7
- torchao version: 0.17.0+cu130
- Safetensors version: 0.8.0
- xFormers version: 0.0.35
- Accelerator: NVIDIA RTX PRO 6000 Blackwell Server Edition, 97887 MiB
Activity
[-]Qwen Image 2.1 doesn't support Torch.compile with varying image size[/-][+]Qwen Image 2.1 with torch.compile trigger a recompile on each run with varying image size[/+][-]Qwen Image 2.1 with torch.compile trigger a recompile on each run with varying image size[/-][+]Qwen Image 2.1 with torch.compile triggers a recompile on each run with varying image size[/+]