ternary-qat-jax-tpu-demo

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)