Rope Positions Require Higher Precision

[fabianlim]

Previously the rope positions were created like this

t = torch.arange(
     self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
)

Problem

It could happen that inv_freq can end up having a low precision dtype:

• even if inv_freq was created in float32, it could happen that the model's parameters ended up being downcasted.
• if so, then it is not reliable to infer the dtype of inv_freq since that may change from creation time.
• If inv_freq was downcasted to bfloat16, then it has a limited mantissa, and cannot represent values even up to 3000+, of which self.max_seq_len_cached can exceed.

Fix
The fix here is to create the t vector in float32 everytime, and not infer the dtype of inv_freq.

Loss Comparisons

[mayank31398]

LGTM
this is not an issue since FSDP does downcasting itself after reset_parameters is called.
but its better to merge this

[fabianlim]

@mayank31398 i see, i did observe the downcasting for DeepSpeed, but I didnt debug it very carefully to ascertain when the downcasting happens.

[mayank31398]

I think its fine, this still helps with inference.

[Maxusmusti]

LGTM