Improve greedy search memory usage

[regisss]

What does this PR do?

When doing greedy search, inputs go through _expand_inputs_for_generation here where they are expanded (see here). As the expand size is always 1 in the case of greedy search, torch.repeat_interleave do not modify the inputs. However, it does increase the memory usage as the input to torch.repeat_interleave is cloned.

Here is a code snippet to check this behaviour:

import torch

a = torch.ones(1000, 1000, 1000, device="cuda")
print(torch.cuda.max_memory_allocated())

expand_size = 1
a = a.repeat_interleave(expand_size, dim=0)
print(torch.cuda.max_memory_allocated())

which returns

4000000000
8000000000

Thus, if the expand size is 1, we can return the model inputs before calling torch.repeat_interleave. That's the change introduced in this PR.

More context in this Slack thread: https://huggingface.slack.com/archives/C01N44FJDHT/p1723827436938589

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.

[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.

[amyeroberts]

cc @gante @zucchini-nlp

[regisss]

CI failed but it doesn't seem to be related to this PR

[zucchini-nlp]

Wow, interesting finding! Thanks for handling

[gante]

Thank you for the fix 🙏

(CI is failing for reasons of our knowledge, will take care of rebasing and merging when the root cause is fixed cc @amyeroberts )

[amyeroberts]

Thanks to adding this improvement!