bad_words_ids no longer slow on mps

[DWarez]

What does this PR do?

Using the bad_words_ids on mps is slowing down a lot the text generation, this PR tries to address that.

Fixes #39512

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?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@ArthurZucker

[xenova]

Thanks for identifying the issue and opening a PR! This does speed up things quite a lot, but it's still around 6 times slower than without setting bad_words_ids 🤔

Pretty much the same as just doing:

            eos_token_id_list = eos_token_id.tolist() # convert to python list before
            bad_words_ids = list(
                filter(lambda bad_token_seq: all(bad_token_seq != [i] for i in eos_token_id_list), bad_words_ids)
            )

[xenova]

Though... I guess, since this is just done once (in the init function), 2 seconds vs 36 seconds is still much better!

[ArthurZucker]

Already much better, forward pass could then be slow because of torch ops from SequenceBiasLogitsProcessor happy to merge as is, but if you have time nice to investigate in the forward see if we can squeeze more perfs!

[DWarez]

Hi @ArthurZucker, the problem was in the _prepare_bias_variables method, in particular here

transformers/src/transformers/generation/logits_process.py

Lines 1187 to 1190 in 3bc726b

	self.length_1_bias = torch.zeros((vocabulary_size,), dtype=torch.float, device=scores.device)
	for sequence_ids, bias in self.sequence_bias.items():
	if len(sequence_ids) == 1:
	self.length_1_bias[sequence_ids[-1]] = bias

In the last commit I made, I modified it with a vectorized access and it's much faster, almost reaching the perfs when not using the bad_words_ids parameter.

[ArthurZucker]

Nice let's go! 🤗

[ArthurZucker]

Could you check the failing test:

=========================== short test summary info ============================
FAILED tests/generation/test_logits_process.py::LogitsProcessorTest::test_no_bad_words_dist_processor - AssertionError: Lists differ: [[True, True, False, True, True], [True, True, True, False, True]] != [[True, True, False, True, False], [True, True, True, False, False]]

First differing element 0:
[True, True, False, True, True]
[True, True, False, True, False]

- [[True, True, False, True, True], [True, True, True, False, True]]
?                            ^^^                              ^^^

+ [[True, True, False, True, False], [True, True, True, False, False]]
?                            ^^^^                              ^^^^
============= 1 failed, 43 passed, 4 warnings in 107.08s (0:01:47) =============

[DWarez]

Oops, my bad, the previous fix had a small bug, which I fixed using @xenova's suggested code 🤗
We lose 0.1s on average with respect to the previous version, but at least it's correct.

[ArthurZucker]

@bot /style

[github-actions]

Style bot fixed some files and pushed the changes.

[ArthurZucker]

Thanks!