Open JerrickLiu opened 1 month ago
BF16 has lower precision than FP16. Normally I use 0.001 as absolute tolerance for float16, while 0.01 for bfloat16.
>>> import torch
>>> x = torch.rand(1000, 1000)
>>> ((x + 1e-6).to(torch.bfloat16) - x.to(torch.bfloat16)).abs().max()
tensor(0.0039, dtype=torch.bfloat16)
>>> ((x + 1e-6).to(torch.float16) - x.to(torch.float16)).abs().max()
tensor(0.0005, dtype=torch.float16)
Right, still hitting some mismatches even when setting atol=1e-2
Please compare (flashattn in bf16 - reference attn in fp32) vs (reference attn in bf16 - reference attn in fp32)
Repro
This prints out infinite SNR and both
out
andout_flash
are allclose.However, when changing the dtype to
torch.bfloat16
, there is like a SNR of 45 and the tensors do not match.Running on 80 GB SXM A100