A TPU precision trap that silently changes fp32 matmuls

JFPuget · x · 2026-07-20

The post highlights a TPU-specific pitfall: although users may request fp32 matmuls, TPUs can still execute them in bf16 by default.

A training flag called --lm-head-precision highest may look like it protects quality, but the author says it does nothing in practice because of this TPU behavior. The screenshot shows a JAX example where an fp32 input still ends up using the default precision path on TPU, and the point is that this is documented behavior — not a hidden bug.

The author says even Google researchers have been bitten by this issue repeatedly. The takeaway is that precision assumptions on TPU can be misleading, so practitioners need to understand the platform’s default execution semantics rather than relying on flags that appear reassuring.

Original post →

More from Infra

Infra channel →