rope positions need higher precision (#7)

Signed-off-by: Yu Chin Fabian Lim <[email protected]>

src/instructlab/dolomite/hf_models/modeling_utils/position_embedding/rope.py

@@ -55,9 +55,7 @@ def _set_cos_sin_cache(
         self, seq_len: int, device: torch.device, dtype: torch.dtype
     ) -> None:
         self.max_seq_len_cached = seq_len
-        t = torch.arange(
-            self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
-        )
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)
 
         freqs = torch.outer(t, self.inv_freq)
         # Different from paper, but it uses a different permutation in order to obtain the same calculation