ONNX export failure for models invoking SDPA attention
Author: BowenBao
Pull request: no

> ValueError: Attention using SDPA can not be traced with torch.jit.trace when no attention_mask is provided. To solve this issue, please either load your model with the argument `attn_implementation="eager"` or pass an attention_mask input when tracing the model. 

There has been some discussion about its possible resolutions in the ONNX exporter team. I'd like to post an issue here as well to seek advice and preferences.

1. Check `torch.jit.is_tracing()` and fallback to eager attn implementation if needed.
2. Create `attention_mask` before passing to SDPA if it is None.
3. Support SDPA tracing w/o attention_mask (not sure how feasible this is).

Comments:
cc @fxmarty 

Thank you for the ping, thank you @BowenBao. cc @drisspg and linking relevant issues as well: https://github.com/pytorch/pytorch/issues/110681 & https://github.com/pytorch/pytorch/issues/108108

Solution 3. SDPA tracing without attention_mask is I think not possible due to the data-dependent controlflow here: https://github.com/huggingface/transformers/blob/e201864bcb147ce3c2374605fac9f0e5043f7e43/src/transformers/models/llama/modeling_llama.py#L735 `q_len > 1`. The reason for this controlflow is that SDPA attention_mask from `is_causal` is top-left aligned.

The same issue exists when tracing SDPA with symbolic_trace or with dynamo + fullgraph=True (https://pytorch.slack.com/archives/C033H6DJSJU/p1702029349053049?thread_ts=1694790001.945579&cid=C033H6DJSJU).

Solution 1. is what the error suggests. I don't think it would be easy to implement (would need to have `torch.jit.is_tracing()` controlflow that does magic on the model).

Solution 2. is probably the most doable (we would need to look at pad tokens). Currently we try as much as possible to pass a `attn_mask=None` since SDPA is able to dispatch to some mem-efficient attention & flash attention path only in the case. We already avoid setting the `attention_mask` to None in case we are tracing: https://github.com/huggingface/transformers/blob/e201864bcb147ce3c2374605fac9f0e5043f7e43/src/transformers/modeling_attn_mask_utils.py#L371-L381



> Thank you for the ping, thank you @BowenBao. cc @drisspg and linking relevant issues as well: [pytorch/pytorch#110681](https://github.com/pytorch/pytorch/issues/110681) & [pytorch/pytorch#108108](https://github.com/pytorch/pytorch/issues/108108)
> 
> Solution 3. SDPA tracing without attention_mask is I think not possible due to the data-dependent controlflow here:
> 
> https://github.com/huggingface/transformers/blob/e201864bcb147ce3c2374605fac9f0e5043f7e43/src/transformers/models/llama/modeling_llama.py#L735
> 
> `q_len > 1`. The reason for this controlflow is that SDPA attention_mask from `is_causal` is top-left aligned.
> The same issue exists when tracing SDPA with symbolic_trace or with dynamo + fullgraph=True (https://pytorch.slack.com/archives/C033H6DJSJU/p1702029349053049?thread_ts=1694790001.945579&cid=C033H6DJSJU).
> 
> Solution 1. is what the error suggests. I don't think it would be easy to implement (would need to have `torch.jit.is_tracing()` controlflow that does magic on the model).
> 
> Solution 2. is probably the most doable (we would need to look at pad tokens). Currently we try as much as possible to pass a `attn_mask=None` since SDPA is able to dispatch to some mem-efficient attention & flash attention path only in the case. We already avoid setting the `attention_mask` to None in case we are tracing:
> 
> https://github.com/huggingface/transformers/blob/e201864bcb147ce3c2374605fac9f0e5043f7e43/src/transformers/modeling_attn_mask_utils.py#L371-L381

Hi @fxmarty why Solution 1 is not easy to implement? I was thining something like

```python
if torch.jit.is_tracing():
    attn_mask = old_attention()
else:
    attn_mask = new_sdpa_attention()
```

Thanks for your reply and context @fxmarty 

I have a local fix using solution 1, will put up a PR to unblock exporter in the short term, while waiting on https://github.com/pytorch/pytorch/issues/108108 .

@fxmarty I might be wrong but, installing from the latest source, I still have the same issue for BART based model export without attention_mask. Is is something planned to be supported?

@LoicDagnas Yes, it is expected. If you want to trace the model without an `attention_mask` input, you should load your model with the argument `attn_implementation="eager"` passed to `from_pretrained`, as suggested in the error that should be raised:

https://github.com/huggingface/transformers/blob/ece1b62b93cde70233f235f6a4c84e37bfc8eba0/src/transformers/modeling_attn_mask_utils.py#L391-L393

Note: this is due to the following controlflow https://github.com/huggingface/transformers/blob/ece1b62b93cde70233f235f6a4c84e37bfc8eba0/src/transformers/models/bart/modeling_bart.py#L597

See for reference https://github.com/pytorch/pytorch/issues/110681 & https://github.com/pytorch/pytorch/issues/108108