Attention refactor in #35235 adds a `__getitem__` into the forward pass, which causes errors with torch dynamo.
Author: i-colbert
Pull request: no
Labels: name: bug

name: ExecuTorch

### System Info

- `transformers` version: 4.51.3
- Platform: Linux-5.15.0-122-generic-x86_64-with-glibc2.31
- Python version: 3.12.8
- Huggingface_hub version: 0.31.1
- Safetensors version: 0.5.2
- Accelerate version: 1.3.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.4.0+rocm6.1 (True)
- 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?: N/A
- Using GPU in script?: NA
- GPU type: AMD Instinct MI210

### Who can help?

@ArthurZucker, error seems to be introduced with #35235 

### 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

You can reproduce the error with:

```python
import torch
from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", torch_dtype=torch.bfloat16)
model.eval()

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")

data = tokenizer("Once upon a time, ", return_tensors='pt')
inputs = {
    'input_ids': data.input_ids,
    'attention_mask': torch.ones_like(data.input_ids)}
model.config.use_cache = False
with torch.no_grad():
    model(**inputs)  # verify forward pass works
    model, guards = torch._dynamo.export(model)(**inputs)
    model(**inputs)  # verify forward pass works
```

and the error received looks like
```shell
Traceback (most recent call last):
  File "/home/repro.py", line 17, in <module>
    model, guards = torch._dynamo.export(model)(**inputs)
  # ... removed for brevity...
torch._dynamo.exc.Unsupported: call_method GetAttrVariable(UserDefinedObjectVariable(AttentionInterface), _global_mapping) __getitem__ (ConstantVariable(),) {}

from user code:
  # ...removed for brevity... 
  File "/site-packages/transformers/models/llama/modeling_llama.py", line 274, in forward
    attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
  File "/site-packages/transformers/modeling_utils.py", line 5841, in __getitem__
    return self._global_mapping[key]
```

### Expected behavior

Expected behavior is no errors during the dynamo export.

A possible fix is

```python
def forward(self, ...):
  # ...ignored for brevity...
  if self.attention_interface is None:
      self.attention_interface: Callable = eager_attention_forward
      if self.config._attn_implementation != "eager":
          if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
              logger.warning_once(
                  "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "
                  'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
              )
          else:
              self.attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

  attn_output, attn_weights = self.attention_interface(
      self,
      query_states,
      key_states,
      value_states,
      attention_mask,
      dropout=0.0 if not self.training else self.attention_dropout,
      scaling=self.scaling,
      **kwargs,
  )
```

where `self.attention_interface = None` in the init

Comments:
cc @Cyrilvallez we need to either register the op, or we clarify `torch.export` vs `torch._dynamo.export` but not super super important IMO @i-colbert as `torch._dynamo` is private

Thanks for the reply. This is infrastructure that we leverage to get modular fx graphs within Brevitas. I'll loop in @Giuseppe5 and @nickfraser for further comment as it was suggested by the PyTorch folks to use this.

super happy to have a PR for a fix btw and sorry we broke this was not expecting get_item to not be compatible 

cc @vasqu so that you can follow

It may be an issue only on older `torch` versions. Could you check again on more recent torch @i-colbert please? 🤗

I am running the example on https://huggingface.co/google/gemma-3-1b-it, and I get the same error on torch 2.2.2, cuda12.1, transformers 4.53.0, and get the same error. Any idea which `torch` version solves the problem?  

Recent version (2.6, 2.7.1) should solve it. Also note that `export` may need tweaking to work correctly, which is why we have primitives to do it in `executorch.py` (see the `TorchExportableModuleForDecoderOnlyLM` class)

@farrokhsiar we should be able to just disable the context derocator when using export, can you confirm this works for torch >= 2.3? 

@ArthurZucker: could you elaborate what do you mean by disabling the context decorator?  

The export path is here: https://github.com/huggingface/transformers/blob/fix-modular/src/transformers/integrations/executorch.py#L71-L71, the decorator which takes core of `output_attentions` can be easily disable if we pass something like `is_exporting_to_executorch` to the forward https://github.com/huggingface/transformers/blob/fix-modular/src/transformers/utils/generic.py#L996-L996 

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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.