Ternary QAT in JAX on TPU — infrastructure demo / experiment
This is a public experiment and infrastructure test, not a usable model. It exists to demonstrate one specific thing: BitNet b1.58-style ternary quantization-aware training (QAT) with a straight-through estimator converges if and only if master weights are kept in fp32 — and to show that working end-to-end in JAX/Flax on TPU.
The failure this disproves
In a prior PyTorch run of the same QAT scheme with bf16 master weights, training silently no-opped: Adam updates (~5e-5) fall below the bf16 ULP at typical weight magnitudes (8 mantissa bits -> spacing ~1e-4 near |w| ~ 0.03), so every update rounds away. Across 936 steps, 0 of 318M sampled weights changed ternary value.
The result here (fp32 masters)
- Ternary flip check: 288,655 / 3,145,728 BitLinear weights changed bucket (9.18%)
- Flip structure: q/k projections flip ~2-4x more than v/o; layer 0 flips most and rates decay with depth; ff_out flips ~2x ff_in. Full per-tensor table in the notebook.
- Post-training ternary distribution: -1: 35.2%, 0: 29.6%, +1: 35.1% (all three levels in use)
- Final training loss: 3.818 (from ~11.1 at init), 800 steps — see
loss_curve.png - Hardware: 1x TPU v6 lite, JAX 0.10.2
- Fully deterministic: seeded init + deterministic data stream; reruns reproduce the flip count exactly.
What it is
~16M-param decoder-only transformer (4 layers, d_model 256,
8 heads, seq len 256); all attention and FFN projections are
BitLinear (weights ternarized to {-1, 0, +1} x mean(|W|) per-tensor scale on every
forward pass, STE backward: w + stop_gradient(quantize(w) - w)); embeddings, norms
and the tied output head are ordinary fp32. ~3.15M of 16.08M params are ternary.
Trained on streamed TinyStories with Adam (peak lr 0.0003, 50 warmup steps,
cosine decay), data-parallel via jax.jit + NamedSharding — runs unmodified on 1
(Colab v6e-1) or 8 (Kaggle v5e-8) TPU devices.
Sample output
After 800 steps this produces TinyStories-flavored babble, exactly as expected at this scale and budget:
Once upon a time, there lived a small measure the little boy named white car. We "Wow, always took it was a toy.
The man began to only a big owl and used that she saw a ribbon. Tom. The man.
Lily and arrived, Lily and worked it back to play off it in the lovely everywhere on, Lily and said, the sleepy.
Reproduce
Run ternary_qat_jax_tpu.ipynb (in this repo) top-to-bottom on any Colab or Kaggle
TPU runtime. Restore the checkpoint:
import orbax.checkpoint as ocp
with ocp.StandardCheckpointer() as ckptr:
params = ckptr.restore("checkpoint/params", nnx.state(model, nnx.Param))
nnx.update(model, params)