"""Proof: the update that vanished. In bf16, a weight near 1.0 has a
local step of about 0.0078; an update smaller than half that step
rounds away, and training silently stops for that weight. The fp32
master copy is the fix; stochastic rounding is the other fix."""
import torch

w = torch.tensor(1.0, dtype=torch.bfloat16)
u = 1e-3                                   # lr * grad, a normal size
print(f"bf16 next after 1.0 : {torch.nextafter(torch.tensor(1.0, dtype=torch.bfloat16), torch.tensor(2.0, dtype=torch.bfloat16)).item():.6f}")
print(f"w + {u} in bf16     : {(w + torch.tensor(u, dtype=torch.bfloat16)).item():.6f}  (unchanged: {(w + torch.tensor(u, dtype=torch.bfloat16)) == w})")

# 1000 such updates, three ways
steps = 1000
wb = torch.tensor(1.0, dtype=torch.bfloat16)
for _ in range(steps):
    wb = wb + torch.tensor(u, dtype=torch.bfloat16)
wm = torch.tensor(1.0, dtype=torch.float32)     # master copy
for _ in range(steps):
    wm = wm + u
g = torch.Generator().manual_seed(0)            # stochastic rounding, emulated
ws = torch.tensor(1.0, dtype=torch.float32)
for _ in range(steps):
    hi = torch.nextafter(ws.to(torch.bfloat16), torch.tensor(2.0, dtype=torch.bfloat16)).float()
    lo = ws.to(torch.bfloat16).float()
    x = ws + u
    p = ((x - lo) / (hi - lo)).clamp(0, 1)
    ws = hi if torch.rand((), generator=g) < p else lo
print(f"after {steps} updates of {u}:")
print(f"  bf16 accumulate    : {wb.item():.4f}   (true answer 2.0)")
print(f"  fp32 master        : {wm.item():.4f}")
print(f"  bf16 + stochastic  : {ws.item():.4f}")