Fix precision errors from casting rotary parameters to FP16 with AMP

[kevinhu]

What does this PR do?

When training with AMP, using einsum to multiply t and self.inv_freq will introduce precision errors because it casts the result to FP16. This can be avoided by using torch.outer instead, as originally mentioned here: https://github.com/Dao-AILab/flash-attention/blob/2c3baba4a63c4007c8a132c5380edc9430f88a22/flash_attn/layers/rotary.py#L396C1-L398C45

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]

Hey! Thanks for opening this PR, it seems to me that the issue lies with AMP no?
My only concern would have been performances, outer might be a little bit slower but it seems to be negligible so LGTM.
Let's make sure that the failing test is fixed!

[ArthurZucker]

I ran the following script for benchmarking:

import torch
from torch.utils import benchmark

results = []
for b in [10, 10000, 2000000]:
    for n in [10, 100, 10000, 1000000]:
        if b * n >= 1000000000:
            continue

        description = f'[{b}, {n}]'

        x = torch.rand(b, device='mps')
        y = torch.rand(n, device='mps')

        results.append(benchmark.Timer(
            stmt='torch.outer(x,y)',
            globals={'x': x, 'y': y},
            description=description,
        ).blocked_autorange())

        results.append(benchmark.Timer(
            stmt='torch.einsum("i,j->ij",x,y)',
            globals={'x': x, 'y': y},
            description=description,
        ).blocked_autorange())

compare = benchmark.Compare(results)
compare.trim_significant_figures()
compare.colorize()
compare.print()

Got the following:

So looks good to me 😉

[ArthurZucker]

failing test is unrelated to the PR i'll fix it on main

[HuggingFaceDocBuilderDev]

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

[ArthurZucker]

FYI @gante and @Rocketknight1 if we see anything failing. I ran slow tests locally and it was all good