MaskFormer, Mask2Former - replace einsum for tracing
Author: amyeroberts
Pull request: yes

# What does this PR do?

Maskformer cannot currently be traced because of einsum operations. This PR replaces the einsum operations with standard matmuls.  

With this PR, the following now runs: 

```python
import torch
from transformers import Mask2FormerForUniversalSegmentation

device = torch.device("cuda")
model = Mask2FormerForUniversalSegmentation.from_pretrained(
    "facebook/mask2former-swin-tiny-coco-instance",
    torchscript=True
).eval().to(device)
dummy_input = torch.randn((1,3,640,640)).to(device)
traced_model = torch.jit.trace(model, dummy_input)

with torch.no_grad():
    out = traced_model(torch.randn((2,3,640,640)).to(device))
    out = traced_model(torch.randn((2,3,640,640)).to(device))
```

Partially resolves #25261  - enables tracing but does not resolve the issue of different results between traced and non-traced model on GPU


## 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#start-contributing-pull-requests),
      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?


Comments:
_The documentation is not available anymore as the PR was closed or merged._

Sorry, let me check something first before merge 🙏 .

Hi @amyeroberts 

So the issue only happens at the 2nd `out = traced_model(torch.randn((2,3,640,640)).to(device))`. Before that (last) line, everything is fine. I am a dumb on this topic, but is it from the enisum? That's really weird (to me).

Still good for me, I am just going to run the torchscript tests and let you know --> ✅ .



@ydshieh Yes, it's bizarre, although it does work on torch nightly. If you run the script above, it complains about the batch dimension in the einsum operation, and removing einsum resolved it, so I'm pretty sure that's the cause. I just don't know why 🤷‍♀️ 

For setting `fx_compatible = True`, I didn't add initially because the tests would fail on the torchscipt runs even if I added it to one of the [compatible models](https://github.com/huggingface/transformers/blob/a6e6b1c622d8d08e2510a82cb6266d7b654f1cbf/src/transformers/utils/fx.py#L113). 

Essentially, the model can be traced using `torch.jit.trace` if you accept many warnings, but fails when we use the `HFTracer` to trace. The tracer warnings and failings when using `HFTracer` are the same: Mask2Former tries to iterate over tensors which breaks things e.g. [here](https://github.com/huggingface/transformers/blob/a6e6b1c622d8d08e2510a82cb6266d7b654f1cbf/src/transformers/models/mask2former/modeling_mask2former.py#L815). Resolving these is a bigger piece of work - I'm happy to add them to this PR or in a follow up. It's possible these are related to the [difference between traced and non-traced models on GPU](https://github.com/huggingface/transformers/issues/25261#issuecomment-1663526119).

Were you able to run the tests successfully? There could be something funny in my env too. 

Hi @amyeroberts Thanks for the info.

Forget about my comment about `fx_compatible = True`. The passed CI I mentioned is just the `test_torchscript_` tests (not the `test_torch_fx` - I didn't try as I saw it's not the raw `torch.jit.trace`)