Fix executorch export with dynamic shapes
Author: justinchuby
Pull request: yes

# What does this PR do?

This PR addd shape compatibility check for attention mask to ensure torch.export can reason about the logic without failing _when exporting with dynamic shapes_.

It additionally simplified the sdpa forward function for export only usage.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@CyrilVallez @jackzhxng @guangy10


Comments:
Is fx tracing still supported? It doesn’t seem to be compatible with torch._check here?

Hey @justinchuby! Sorry but I don't get why this is needed. `torch._check` are very internal dynamo checks, and I don't see why we would have this here. Moreover, `torch.export` is only supported with a mask through the `executorch` integration, as otherwise the mask creation breaks due to using `vmap`

@Cyrilvallez thanks for the suggestion - I can test with the executorch integration. For context, we are looking to enable onnx export via torch.export. 

Could you share guidance on a potential integration « onnx » directory next to the current executorch directory, for correct usages of onnx export calls? I assume we can follow a similar model with executorch where we have an integration and a separate optimum-onnx project. 

torch.export recently resolved the vmap issues, and the torch._check was useful for the dynamic shapes engine of the tracer to understand equivalence. I agree that it should be part of the torch.export integration. 

I updated the PR to fix the executorch integration. Please take another look. Thanks!

>torch.export recently resolved the vmap issues

Ouuhhh that's very very nice, I wasn't aware of it!

BTW, we will probably add a non-vmap path to the sdpa mask creation due to vmap being much slower, and the impact being noticeable for small models (https://github.com/huggingface/transformers/issues/41639). This PR is still welcome though, as its focus is on the sdpa attention, not sdpa mask, but I thought I'd mention it just in case!

>Is fx tracing still supported? It doesn’t seem to be compatible with torch._check here?

Concerning this, the answer is no: https://github.com/huggingface/transformers/pull/41683!

I've tried this PR and found it didn't solve the performance regression in #41639 . Will it be fixed in the next PR? 

cc @justinchuby  @Cyrilvallez 

> I've tried this PR and found it didn't solve the performance regression in #41639 . Will it be fixed in the next PR?
> 
> cc @justinchuby @Cyrilvallez

I believe this PR is unrelated to the issue you linked

I think we can merge https://github.com/huggingface/transformers/pull/41900, then I will create a separate PR to guard torch.export on dynamic shapes.