Sandroeth

/

cali-0.1B

like 5

Loading Sandroeth/cali-0.1B with trust_remote_code=True in bfloat16 and running a forward/training step crashes inside the model's own attention:

RuntimeError: expected scalar type Float but found BFloat16

Where: modeling_cali.py, in GroupedQueryAttention.forward at out = torch.matmul(att, v).

Root cause: build_rope_cache builds cos/sin in float32 (it calls .float() on the frequency/time tensors). apply_rope(x, cos, sin) returns x * cos + torch.cat([-x2, x1], dim=-1) * sin without casting cos/sin back to x.dtype, so bf16 q/k are type-promoted to float32. The softmax line att = F.softmax(att.float(), dim=-1).to(q.dtype) then keeps att in float32 (because q.dtype is now float32), while v — which never passes through apply_rope — remains bfloat16. torch.matmul(att, v) therefore mixes float32 and bfloat16 operands and raises. Under torch.compile/inductor the same divergence appears as a meta_bmm "expected scalar type torch.float32 but found torch.bfloat16" on shapes (1, H, T, T) fp32 x (1, H, T, head_dim) bf16.

Minimal reproduction (single GPU, public):

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Sandroeth/cali-0.1B"
tok = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id, trust_remote_code=True, dtype=torch.bfloat16
).cuda()

ids = tok("The quick brown fox", return_tensors="pt").input_ids.cuda()
out = model(ids)          # RuntimeError: expected scalar type Float but found BFloat16

Expected behavior: a bf16 forward/training step completes without a dtype error.

Actual behavior: crashes at torch.matmul(att, v) with RuntimeError: expected scalar type Float but found BFloat16 (eager); equivalent fp32-vs-bf16 meta_bmm error under torch.compile.

Environment (one representative config; reproduces across GPUs):

• Model revision: 5a9cca33d2abd466235bd2fd730ec4734e41aa71 (current main; modeling_cali.py last changed 2026-05-22)
• GPU: NVIDIA B200 (also reproduces on H100 and GB200), driver 580.82.07, 1 GPU
• OS: Ubuntu 24.04, Python 3.12.3
• CUDA 13.4.1, cuDNN 9.25.0
• torch 2.14.0a0 (nightly), transformers 5.14.1, accelerate 1.9.0, peft 0.19.1
• Precision: bfloat16; both eager and torch.compile paths affected

Redacted stack trace:

File ".../modeling_cali.py", line 279, in forward
File ".../modeling_cali.py", line 211, in forward
File ".../modeling_cali.py", line 119, in forward
    attn_out, present = self.attn(...)
File ".../modeling_cali.py", line 94, in forward
    out = torch.matmul(att, v).transpose(1, 2).contiguous().view(B, T, self.num_heads * self.head_dim)
RuntimeError: expected scalar type Float but found BFloat16

Suggested fix (model-side, one line): keep RoPE in the tensor's own dtype so the whole attention path stays consistent. Either:

• in apply_rope, cast the cache to the input dtype:

  def apply_rope(x, cos, sin):
      cos, sin = cos.to(x.dtype), sin.to(x.dtype)
      half = x.shape[-1] // 2
      x1, x2 = x[..., :half], x[..., half:]
      return x * cos + torch.cat([-x2, x1], dim=-1) * sin
  

  (preferred — keeps q/k/att in the model dtype), or
• cast att to v's dtype right before the matmul: out = torch.matmul(att.to(v.dtype), v).

Both were validated to remove the crash in bf16 on B200, H100, and GB200 (eager and compiled).

This issue was drafted with assistance from the opus AI model.