Fix ONNX export for causal LM sequence classifiers by removing reverse indexing

[dwyatte]

What does this PR do?

Follow-up to #27450 and another step to fixing huggingface/optimum#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). 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).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? 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, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@ArthurZucker, @amyeroberts, @younesbelkada (CC @fxmarty)

[amyeroberts]

Thanks for fixing this!

General comment on the indexing assumptions made here

[amyeroberts]

• Let's split across lines - it'll use fewer lines overall
• Can we instead use modulo to convert to the equivalent negative index? The logic at the moment assumes we only ever want to take the final index and I think it'll be faster than torch.where

Suggested change

	sequence_lengths = torch.where(sequence_lengths >= 0, sequence_lengths, input_ids.shape[-1] - 1).to(
	logits.device
	)
	sequence_lengths = sequence_lengths % input_ids.shape[-1]
	sequence_lengths = sequence_lengths.to(logits.device)

[dwyatte]

That should work. Updated all changes

[amyeroberts]

LGTM - thanks for iterating!

I'd like to have another approval from either @younesbelkada or @ArthurZucker before merging as they know the causal lm models well

[ArthurZucker]

Thanks both, looks good yep!
More models that can benefit from this: and that use sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(dim=-1) - 1).to(logits.device)
(not sure if argmax is faster?)

• bloom
• falcon
• mpt

let's try to unify this!

[dwyatte]

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'

[amyeroberts]

@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?

[dwyatte]

@amyeroberts Hm, b0db02c contains the latest commit on main (224ab70), so I think tests/models/seamless_m4t/test_modeling_seamless_m4t.py::SeamlessM4TModelWithTextInputTest::test_retain_grad_hidden_states_attentions is still broken/flaking there

[amyeroberts]

@dwyatte hm, that's odd. The test shouldn't even be running as it's explicitly skipped. In your local, on this branch, do you see this skip condition in test_modeling_seamless_m4t.py?

[dwyatte]

@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

[amyeroberts]

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

[dwyatte]

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

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sentialx]

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