Bugfix/inv freq

[cyr0930]

What does this PR do?

Fixes # (issue)

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

[Rocketknight1]

Hi, I don't see the bug in the original code! self.inv_freq.to(x.device) and inv_freq = self.inv_freq.to(x.device) are both equivalent - in either case, the original tensor is moved.

[vasqu]

@Rocketknight1 It is not due to any missing assignment, you can verify this with the following script:

import torch     # v2.3.1

x = torch.randn(size=(2, 4, 8)).to(device='cpu')
print(x.device)  # cpu

x.to(device='cuda')
print(x.device)  # cpu

Hence, you need an assignment - be it with a new var or directly on self.inv_freq.

Edit: Maybe for more clarification. The in-place movement for a device only happens when we call it on a nn.Module for example, not on a tensor / parameter. You can can also verify this with the following script:

import torch


class TestModel(torch.nn.Module):
    def __init__(self):
        super().__init__()

        inv_freq = torch.tensor(1.0)
        self.register_buffer("inv_freq", tensor=inv_freq, persistent=False)

    def forward(self, x):
        self.inv_freq.to('cuda')
        print(self.inv_freq.device)


model = TestModel() 

# you will see 2x cpu
print(model.inv_freq.device)
model('test')

# in-place movement 
model.to('cuda')

# you will see 2x cuda:x
print(model.inv_freq.device)
model('test')

[Rocketknight1]

You're totally right, my bad! That's an embarrassing mistake to make, and my only defence is that I had a lot of notifications that day. The new code is absolutely more correct!

cc @LysandreJik for core maintainer review

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

[vasqu]

No worries! Happens to all of us :) and tbh, it is some niche behaviour on torch's side on how device movement is handled.

[LysandreJik]

I think cc @ArthurZucker would like to review this

[ArthurZucker]

Hey! This was introduced by #30775, but TBH I am not convinced that we need to even put self.inv_freq to device. See Llama model as it does not have this and there are no reports about failures in this case!