Upload Ivme-Coder-v1 (Otter 1): safetensors + custom modeling code

model/model.py

@@ -64,8 +64,14 @@ def precompute_rope(dim, max_len, theta=10000.0, device="cpu"):
 
 def apply_rope(x, cos, sin):
     T = x.size(2)
-    cos = cos[:T].unsqueeze(0).unsqueeze(0)
-    sin = sin[:T].unsqueeze(0).unsqueeze(0)
+    # Cast cos/sin to match x's dtype. precompute_rope always builds these tables in
+    # fp32 for precision, but if q/k arrive in bf16 (as they do when the model is
+    # loaded with dtype=torch.bfloat16), multiplying against fp32 cos/sin silently
+    # promotes the result back to fp32 - which then doesn't match `v` (which never
+    # passes through this function and stays in bf16), and scaled_dot_product_attention
+    # requires q, k, v to share one dtype.
+    cos = cos[:T].unsqueeze(0).unsqueeze(0).to(x.dtype)
+    sin = sin[:T].unsqueeze(0).unsqueeze(0).to(x.dtype)
     x1, x2 = x[..., 0::2], x[..., 1::2]
     rot1 = x1 * cos - x2 * sin
     rot2 = x1 * sin + x2 * cos

modeling_ivme_coder.py

@@ -64,8 +64,14 @@ def precompute_rope(dim, max_len, theta=10000.0, device="cpu"):
 
 def apply_rope(x, cos, sin):
     T = x.size(2)
-    cos = cos[:T].unsqueeze(0).unsqueeze(0)
-    sin = sin[:T].unsqueeze(0).unsqueeze(0)
+    # Cast cos/sin to match x's dtype. precompute_rope always builds these tables in
+    # fp32 for precision, but if q/k arrive in bf16 (as they do when the model is
+    # loaded with dtype=torch.bfloat16), multiplying against fp32 cos/sin silently
+    # promotes the result back to fp32 - which then doesn't match `v` (which never
+    # passes through this function and stays in bf16), and scaled_dot_product_attention
+    # requires q, k, v to share one dtype.
+    cos = cos[:T].unsqueeze(0).unsqueeze(0).to(x.dtype)
+    sin = sin[:T].unsqueeze(0).unsqueeze(0).to(x.dtype)
     x1, x2 = x[..., 0::2], x[..., 1::2]
     rot1 = x1 * cos - x2 * sin
     rot2 = x1 * sin + x2 * cos