Fix dtype mismatch in RoPE cos/sin for mixed precision training

Cast cos/sin buffers (float32) to input dtype in _apply_rotation
to prevent implicit upcasting when q/k are bf16/fp16.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

llm_lab/model/rope.py

@@ -93,8 +93,9 @@ class RotaryPositionalEmbedding(nn.Module):
         x_odd  = x[..., 1::2]  # odd indices
 
         # Adjust dimensions for broadcasting: (seq_len, dim/2) → (1, 1, seq_len, dim/2)
-        cos = cos.unsqueeze(0).unsqueeze(0)
-        sin = sin.unsqueeze(0).unsqueeze(0)
+        # Cast to input dtype to support bf16/fp16 mixed precision training
+        cos = cos.unsqueeze(0).unsqueeze(0).to(x.dtype)
+        sin = sin.unsqueeze(0).unsqueeze(0).to(x.dtype)
 
         # Apply rotation
         rotated_even = x_even * cos - x_odd * sin