Fix ONNX export for causal LM sequence classifiers by removing reverse indexing
Author: dwyatte
Pull request: yes

# What does this PR do?

Follow-up to https://github.com/huggingface/transformers/pull/27450 and another step to fixing https://github.com/huggingface/optimum/issues/1527. ONNX implements indexing using a combination of its own operators and when using reverse indexing (e.g., -1 to indicate 1 element from the right side of an array), it can produce incorrect results (see [PyTorch's ONNX export code](https://github.com/pytorch/pytorch/blob/71bedc3a69e3203fd8f76a68ecf2bd7c58d2e13e/torch/onnx/symbolic_opset9.py#L5859-L5865)). In practice, this can cause the batch dimension to get shuffled

Causal LM sequence were previously using `-1` for the last token. Adding `sequence_lengths = torch.where(sequence_lengths >= 0, sequence_lengths, input_ids.shape[-1] - 1)` effectively removes reverse indexing

While this could be fixed in https://github.com/huggingface/optimum by forcing the inputs used to trace the graph to contain a pad token and avoiding reverse indexing, it seems better to fix in `transformers` with the added benefit of bringing the code in line with TensorFlow implementations of the same code (e.g., https://github.com/huggingface/transformers/pull/25085/files#diff-7c6fdd54ac4b8ce0c09bb17da15f176d3e5827df39dd8234fd802631e99ef38dR801-R804)

## 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?


## Who can review?

@ArthurZucker, @amyeroberts, @younesbelkada (CC @fxmarty)


Comments:
Thanks @ArthurZucker @amyeroberts. I've unified the Bloom/Falcon/MPT implementations, but doing so triggered what looks to be an unrelated CI failure. Can someone take a look and fix/disable that test if it is indeed unrelated?

```
FAILED tests/models/seamless_m4t/test_modeling_seamless_m4t.py::SeamlessM4TModelWithTextInputTest::test_retain_grad_hidden_states_attentions - AttributeError: 'NoneType' object has no attribute 'retain_grad'
```

@dwyatte Yep, that's a flaky test. A patch to skip it in the testing suite was recently merged into main to prevent it affecting unrelated PRs like this one :) Could you rebase to include recent updates and trigger a new CI run? 

@amyeroberts Hm, [b0db02c](https://github.com/huggingface/transformers/commit/b0db02c395cee1e8b2ea73077c798617d14be2b4) contains the latest commit on `main` (224ab70969d1ac6c549f0beb3a8a71e2222e50f7), so I think `tests/models/seamless_m4t/test_modeling_seamless_m4t.py::SeamlessM4TModelWithTextInputTest::test_retain_grad_hidden_states_attentions` is still broken/flaking there

@dwyatte hm, that's odd. The test shouldn't even be running as [it's explicitly skipped](https://github.com/huggingface/transformers/blob/224ab70969d1ac6c549f0beb3a8a71e2222e50f7/tests/models/seamless_m4t/test_modeling_seamless_m4t.py#L616). In your local, on this branch, do you see this skip condition in `test_modeling_seamless_m4t.py`?

@amyeroberts I see what's going on -- the failure is on `SeamlessM4TModelWithTextInputTest` but the explicit skip exists on `SeamlessM4TModelWithSpeechInputTest`. Let me know if I should add the same skip to `SeamlessM4TModelWithTextInputTest` on my branch or if you prefer a different fix/PR

@dwyatte Ah! Gotcha. Yes please, could you open another separate PR to skip the retain grad tests for all the SeamlessMT4 models? 

Ok @amyeroberts @ArthurZucker, after rebasing on the above, this is ready for merging. Thanks both!

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_28144). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Why not just have a shared util for this, instead of repeating the code all over the place