Foundry/ternary_optim/docs/OPERATING_GUIDE.md

98 lines
5.5 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# SR4 Operating Guide — how it behaves, and how to train with it well
> This is the **operational** knowledge about the SR4 optimizer — distilled from the
> SR4 MNIST research campaign (reached the FP32 ceiling, memory-neutral) and the first RWKV-x
> LM runs. `HOW_IT_WORKS.md` covers the *mechanism*; this covers *how to actually get good
> results with it and how to read its behavior.* If you are training with SR4, read this first.
## The one root fact everything follows from
**SR4 is stochastic on every update.** The int4 accumulator `q` is integrated by *stochastic
rounding* — so the per-update gradient **noise is amplified** into the weights, every step,
every weight. A deterministic optimizer (SGD/Adam) only touches randomness at init + data
order; SR4 threads it through all ~10⁵ updates. Almost every operating rule below is a
consequence of this.
## Rule 1 — Batch size is the DOMINANT lever. Never run tiny batches.
Bigger batch → lower per-update gradient variance → the stochastic rounding integrates a
**cleaner** signal. This is the single most important knob.
- On MNIST, going batch **128 → 2048** was the path from "SR4 works" to "SR4 *matches
full-precision-dense*." Every variance-reducer we found was subsumed by big batch at the top.
- **`batch=1` is pathological** — maximum per-update noise, and SR4 will **plateau** while
Adam (batch-insensitive due to per-coordinate adaptivity) descends right past it. If you see
that pattern, the fix is **batch, not learning rate.**
- Diagnostic signature of "SR4 starved of batch": **noisy/spiky train loss + an early
plateau**, while a reference optimizer keeps improving.
## Rule 2 — SR4 has its own learning rate, and it is NOT Adam's
The step size is `TARGET_STEP` (the mean proposed step in LSBs per update). It must be tuned
**per optimizer and per scale** — do a small sweep at each new rung. Observed optima:
MNIST-MLP `0.001`, LM memorization `0.003`, LM real-data `~0.001`. **Never share an LR with
Adam** in an A/B — that rigs the comparison. Compare each optimizer *at its own best LR*, at
matched **tokens**, and report wall-clock separately.
## Rule 3 — Report seed BANDS, never single runs
Because SR4 is stochastic every step, final quality varies by seed. Run **≥3 seeds** and
compare **band means with a noise margin** (MNIST band ≈ 0.25pp; a gain must *clear* the band
to count). A single lucky/unlucky seed is not a result.
## Rule 4 — Variance-reducers compound (super-additively)
Anything that lowers the noise floor lets small mean-gains become visible and stackable:
- **Low-discrepancy (Weyl) dither** — built into the core; variance-reduced *unbiased*
stochastic rounding (free, always on).
- **Batch** — the big one (Rule 1).
- **Dropout** — a *tunable reliability dial*: more dropout → tighter seed band (its cost is
mean, so anneal/gentle-decay it; regime-dependent — see Rule 6).
Small sub-band mechanisms that look like noise individually start *stacking* once the floor
drops.
## Rule 5 — The recipe (core vs knobs)
The **core** is small: int4 accumulator + stochastic-rounded integer update + ternary forward.
Everything else is **recipe**, and it lives in different places:
- **Batch size** → training loop (Rule 1, dominant).
- **Co-trained per-column FP32 scales** → per layer; restore the dynamic range ternary weights
throw away (`Z = scale ⊙ (ternaryW·x)`). ~0.04 b/w. Basically required for quality.
- **Label smoothing** → loss side (biggest single regularizer on MNIST).
- **Dropout / annealed dropout** → the reliability dial (Rule 4).
The core got MNIST to ~97; the recipe got it to the FP32 ceiling (~98.3). *Don't judge SR4
without its recipe.*
## Rule 6 — Regime matters
A lever's value depends on its base. Examples that actually flipped: early-release dropout was
the **worst** schedule at batch 1024 but the **best** at batch 2048; a higher step helped at
moderate batch and was null at large batch. **Tune levers in the target regime**; do not assume
a lever that helped at one scale/batch transfers.
## Rule 7 — Which params are ternary (the split)
SR4 ternarizes the **large weight matrices** (the projections — the bulk of params, the memory
win). Keep **small/precision-sensitive params continuous** (float, plain optimizer): embeddings,
the LM head, LayerNorm γ/β, LoRA/low-rank gates, biases, and the per-column scales themselves.
This mirrors the BitNet/ternary-LLM convention. Each ternary matrix gets its own `idx_base` so
the SR dither lanes stay independent.
## Rule 8 — Think in updates/tokens, not epochs
The meaningful invariant is the number of weight **updates** (and, at LM scale, **tokens**) —
not passes over the data. When you change batch size, hold updates constant (scale steps
accordingly), or you are changing two things at once.
## Quick reference — reading an SR4 run
| You see | It usually means | Do |
|---|---|---|
| Spiky train loss + early plateau; Adam descends past it | per-update noise too high | **raise batch** (Rule 1) |
| Whole seed band wide | noise floor high | add a variance-reducer (batch / dither / dropout) |
| Loss won't drop at all | LR wrong or scales missing | check `TARGET_STEP` sweep + that per-column scales are on |
| Matches Adam only after the recipe | expected | that's the point — SR4 core + recipe ≈ full precision, at ~24× less optimizer memory |
---
*Provenance: the SR4 MNIST research campaign and the RWKV-x SR4-vs-AdamW scale test (both
unpublished). Update this as behavior at LM scale is learned.*