RoPE loses precision for Llama / Gemma + Gemma logits.float()

src/transformers/models/gemma/modeling_gemma.py

@@ -110,7 +110,11 @@ def forward(self, x, position_ids, seq_len=None):
 
         inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
         position_ids_expanded = position_ids[:, None, :].float()
-        freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
+
+        # Force float32 since bfloat16 loses precision on long contexts
+        with torch.autocast(device_type=position_ids_expanded.device.type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+
         emb = torch.cat((freqs, freqs), dim=-1)
         return emb.cos().to(dtype=x.dtype), emb.sin().to(dtype=x.dtype)
 
@@ -1079,7 +1083,8 @@ def forward(
 
         hidden_states = outputs[0]
         logits = self.lm_head(hidden_states)
-
+        logits = logits.float()
+
         loss = None
         if labels is not None:
             # Shift so that tokens < n predict n

src/transformers/models/llama/modeling_llama.py

@@ -123,7 +123,9 @@ def forward(self, x, position_ids, seq_len=None):
         # x: [bs, num_attention_heads, seq_len, head_size]
         inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
         position_ids_expanded = position_ids[:, None, :].float()
-        freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
+        # Force float32 since bfloat16 loses precision on long contexts
+        with torch.autocast(device_type=position_ids_expanded.device.type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
         emb = torch.cat((freqs, freqs), dim=-1)
         cos = emb.cos().to(dtype=x.dtype)
         sin = emb.sin().to(dtype=x.dtype)