added unsqueeze_dim to apply_rotary_pos_emb

[ShashankMosaicML]

Making the unsqueeze dimension parameterized in the apply_rotary_pos_emb function in modeling_llama.py

This PR introduces a new parameter to the apply_rotary_pos_emb function which allows the specification of the dimension along which to unsqueeze to make the cosine and sine rotary tensors broadcastable to the query and key tensors. This will make the function compatible with codebases that have different shapes for the query and key tensors without needing any back-and-forth transposing.

Fixes #26948

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. Link to the issue
• 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?

@gante , @ArthurZucker

[gante]

Thank you for opening the PR @ShashankMosaicML! And nice touch, the docstring :)

CI is failing because of the # Copied from ... statement, which is used by our internal tools and CI to maintain one file per model policy while enabling the propagation of bugfixes and upgrades.

In this particular case, there is also a bug on our end -- recently we've made Llama the base implementation, from which others are copied from. As such, you'll have to:

1. Remove the # Copied from ... statement in L194
2. Run make fix-copies in your terminal
3. Push the changes

You may also need to run make fixup afterwards, which applies automatic code formatting.

[ShashankMosaicML]

Thank you @gante for reviewing and providing me the steps to fix the errors! However I still see some errors after following the steps, and on clicking the details for the error, I see the following. Could you please let me know how to fix this error? Thank you!

[amyeroberts]

Thanks for working on this!

Change to the function LGTM - there's just some addition changes in this diff to docstrings which need to be removed before merging.

[ShashankMosaicML]

Hi @amyeroberts , thanks for suggesting the changes. I have incorporated those, but some quality checks are still failing. Could you take a look?

[gante]

@ShashankMosaicML running make fixup then committing the changes doesn't fix it?

[ShashankMosaicML]

@gante , I think that worked! (I wasn't running make fixup properly earlier 😅 )

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[amyeroberts]

Thanks for iterating!