[Bug] RuntimeError with torch.compile (reduce-overhead) and StaticCache on Qwen2.5
Author: lotato
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.57.3
- Platform: Linux-5.15.0-78-generic-x86_64-with-glibc2.35
- Python version: 3.12.12
- Huggingface_hub version: 0.36.0
- Safetensors version: 0.7.0
- Accelerate version: 1.12.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.9.1+cu128 (CUDA)
- 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?: No
- Using GPU in script?: Yes
- GPU type: NVIDIA RTX PRO 6000 Blackwell Server Edition

### Who can help?

@ArthurZucker @Cyrilvallez

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

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

### Reproduction

I am trying to use `StaticCache` combined with `torch.compile(mode="reduce-overhead")` as described in the official technical reports/guides. However, it crashes with a `RuntimeError` regarding CUDAGraphs overwriting tensors during the generation process.

The error specifically occurs when `torch.compile` is applied to `model.forward`.

Here is the reproduction script:

```python
import torch
import os
from transformers import AutoTokenizer, AutoModelForCausalLM, StaticCache

# Environment setup
torch.set_float32_matmul_precision('high')
os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Load model (Using Qwen2.5-1.5B-Instruct as an example)
model_id = "Qwen/Qwen2.5-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto", device_map="auto")

# Compile the forward pass
# NOTE: The crash happens due to this compilation line combined with StaticCache
model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)

# Prepare inputs
input_text = "The theory of special relativity states "
input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
prompt_length = input_ids.input_ids.shape[1]
model.generation_config.max_new_tokens = 16

# Initialize StaticCache
past_key_values = StaticCache(
    config=model.config,
    max_cache_len=prompt_length + (model.generation_config.max_new_tokens * 2),
)

print("--- First Generation ---")
# The error usually triggers here or during the second generation depending on when CUDAGraphs captures the execution
outputs = model.generate(**input_ids, past_key_values=past_key_values)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))

print("--- Second Generation (Continue) ---")
new_input_ids = outputs
outputs = model.generate(new_input_ids, past_key_values=past_key_values)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))

### Expected behavior

The model should generate text using the compiled graph and StaticCache without runtime errors.

The code raises the following RuntimeError (usually during the forward pass inside generate):
skipping cudagraphs due to mutated inputs (56 instances). Found from : 
   File ".../site-packages/transformers/utils/generic.py", line 918, in wrapper
    output = func(self, *args, **kwargs)
  File ".../site-packages/transformers/models/qwen2/modeling_qwen2.py", line 449, in forward
    outputs: BaseModelOutputWithPast = self.model(
  File ".../site-packages/transformers/utils/generic.py", line 1072, in wrapper
    outputs = func(self, *args, **kwargs)
  File ".../site-packages/transformers/models/qwen2/modeling_qwen2.py", line 384, in forward
    hidden_states = decoder_layer(
  File ".../site-packages/transformers/modeling_layers.py", line 94, in __call__
    return super().__call__(*args, **kwargs)
  File ".../site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
    return func(*args, **kwargs)
  File ".../site-packages/transformers/models/qwen2/modeling_qwen2.py", line 234, in forward
    hidden_states, _ = self.self_attn(
  File ".../site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
    return func(*args, **kwargs)
  File ".../site-packages/transformers/models/qwen2/modeling_qwen2.py", line 163, in forward
    key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
  File ".../site-packages/transformers/cache_utils.py", line 776, in update
    keys, values = self.layers[layer_idx].update(key_states, value_states, cache_kwargs)
  File ".../site-packages/transformers/cache_utils.py", line 332, in update
    self.keys.index_copy_(2, cache_position, key_states)

Traceback (most recent call last):
  File "test.py", line 21, in <module>
    outputs = model.generate(**input_ids, past_key_values=past_key_values)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/transformers/generation/utils.py", line 2564, in generate
    result = decoding_method(
             ^^^^^^^^^^^^^^^^
  File ".../site-packages/transformers/generation/utils.py", line 2787, in _sample
    outputs = model_forward(**model_inputs, return_dict=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_dynamo/eval_frame.py", line 832, in compile_wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_dynamo/eval_frame.py", line 832, in compile_wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/transformers/utils/generic.py", line 912, in wrapper
    @wraps(func)
  File ".../site-packages/torch/_dynamo/eval_frame.py", line 1044, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_functorch/aot_autograd.py", line 1130, in forward
    return compiled_fn(full_args)
           ^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 353, in runtime_wrapper
    all_outs = call_func_at_runtime_with_args(
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_functorch/_aot_autograd/utils.py", line 129, in call_func_at_runtime_with_args
    out = normalize_as_list(f(args))
                            ^^^^^^^
  File ".../site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 724, in inner_fn
    outs = compiled_fn(args)
           ^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 526, in wrapper
    return compiled_fn(runtime_args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_inductor/output_code.py", line 613, in __call__
    return self.current_callable(inputs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/torchinductor_root/f3/cf3i22fh2a32yvul4vu2ykz4gjsz35pysjkhpygo6m3u5dbfxcwk.py", line 4377, in call
    (buf663,) = self.partitions[0](partition0_args)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_inductor/compile_fx.py", line 1772, in run
    return compiled_fn(new_inputs)  # type: ignore[arg-type]
            ^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_inductor/cudagraph_trees.py", line 404, in deferred_cudagraphify
    fn, out = cudagraphify(model, inputs, new_static_input_idxs, *args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_inductor/cudagraph_trees.py", line 463, in cudagraphify
    return manager.add_function(
            ^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_inductor/cudagraph_trees.py", line 2316, in add_function
    return fn, fn(inputs)
                ^^^^^^^^^^
  File ".../site-packages/torch/_inductor/cudagraph_trees.py", line 2012, in run
    out = self._run(new_inputs, function_id)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_inductor/cudagraph_trees.py", line 2094, in _run
    return self.ids_to_funcs[function_id].model(new_inputs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/torchinductor_root/f3/cf3i22fh2a32yvul4vu2ykz4gjsz35pysjkhpygo6m3u5dbfxcwk.py", line 2161, in partition_0
    triton_poi_fused__to_copy_add_bmm_cat_cos_index_copy_mul_neg_sin_slice_transpose_unsqueeze_view_2.run(arg2_1, buf6, buf5, arg12_1, 256, stream=stream0)
  File ".../site-packages/torch/_inductor/runtime/triton_heuristics.py", line 1272, in run
    self.autotune_to_one_config(*args, **kwargs)
  File ".../site-packages/torch/_inductor/runtime/triton_heuristics.py", line 1048, in autotune_to_one_config
    timings = self.benchmark_all_configs(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_inductor/runtime/triton_heuristics.py", line 1023, in benchmark_all_configs
    launcher: self.bench(launcher, *args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_inductor/runtime/triton_heuristics.py", line 891, in bench
    return benchmarker.benchmark_gpu(kernel_call, rep=40)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_inductor/runtime/benchmarking.py", line 39, in wrapper
    return fn(self, *args, **kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_inductor/runtime/benchmarking.py", line 250, in benchmark_gpu
    _callable()
  File ".../site-packages/torch/_inductor/runtime/triton_heuristics.py", line 856, in kernel_call
    cloned_args, cloned_kwargs = self.maybe_clone_args(
                                 ^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_inductor/runtime/triton_heuristics.py", line 995, in maybe_clone_args
    prepare_arg(name, arg)
  File ".../site-packages/torch/_inductor/runtime/triton_heuristics.py", line 990, in prepare_arg
    return clone_preserve_strides(arg)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../site-packages/torch/_inductor/utils.py", line 3037, in clone_preserve_strides
    buffer = torch.as_strided(x, (needed_size,), (1,)).clone()
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run. Stack trace: File ".../site-packages/transformers/utils/generic.py", line 918, in wrapper
    output = func(self, *args, **kwargs)
  File ".../site-packages/transformers/models/qwen2/modeling_qwen2.py", line 449, in forward
    outputs: BaseModelOutputWithPast = self.model(
  File ".../site-packages/transformers/utils/generic.py", line 1072, in wrapper
    outputs = func(self, *args, **kwargs)
  File ".../site-packages/transformers/models/qwen2/modeling_qwen2.py", line 384, in forward
    hidden_states = decoder_layer(
  File ".../site-packages/transformers/modeling_layers.py", line 94, in __call__
    return super().__call__(*args, **kwargs)
  File ".../site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
    return func(*args, **kwargs)
  File ".../site-packages/transformers/models/qwen2/modeling_qwen2.py", line 234, in forward
    hidden_states, _ = self.self_attn(
  File ".../site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
    return func(*args, **kwargs)
  File ".../site-packages/transformers/models/qwen2/modeling_qwen2.py", line 163, in forward
    key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
  File ".../site-packages/transformers/cache_utils.py", line 776, in update
    keys, values = self.layers[layer_idx].update(key_states, value_states, cache_kwargs)
  File ".../site-packages/transformers/cache_utils.py", line 332, in update
    self.keys.index_copy_(2, cache_position, key_states). To prevent overwriting, clone the tensor outside of torch.compile() or call torch.compiler.cudagraph_mark_step_begin() before each model invocation.

Comments:
This is expected, you should not compile your model externally as then the prefill is compiled as well! Just use a StaticCache, and we compile internally for you (but not the prefill)!
Closing the issue, but feel free to reopen if you think something is still missing!

@Cyrilvallez 
Thanks for the clarification! That explains why the manual compilation on `forward` caused conflicts with the prefill stage.

However, I am working on a custom implementation of **Speculative Decoding** (using a 14B model as target and 1.5B as draft), which leads to two follow-up questions regarding your advice:

1.  **Multi-GPU behavior:** My actual deployment environment uses 2x RTX 4090s. The 14B model will likely need to be split across GPUs (using `device_map="auto"` or similar). Does the internal compilation logic (triggered by `StaticCache`) support model parallelism/pipeline parallelism robustly, or are there known caveats when using `torch.compile` + `StaticCache` with `device_map`?

2.  **Manual Forward Calls (Custom Loop):**
    Since I am modifying the speculative decoding logic, I often bypass `model.generate()` and manually invoke `model()` (or `model.forward`) in a loop to verify tokens, similar to:
    ```python
    outputs = self._model(
        new_tokens,
        past_key_values=self._past_key_values,
        use_cache=True,
        cache_position=cache_position
    )
    ```
    If I am **not** using `model.generate()`, I assume the "automatic internal compilation" you mentioned won't trigger. 
    In this case, how should I correctly apply `torch.compile` to strictly optimize the **decode** step without breaking the **prefill** step? Should I wrap the decode logic in a separate function and compile that instead of the whole model?

Thanks to our lazy initialization of the layers in the Cache, it will wait until it sees its computation device before instantiating the fixed size tensors, so you should have no problem with a device_map! It will naturally put everything on the correct device according to the device_map.
For your second point, you could simply declare a new function `compile_forward = torch.compile(forward,...)`, and use this compiled forward for decoding, and the original for prefill - this is more or less what we do internally. [Those lines](https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py#L2845-L2864) should give you a clearer idea of what we do!