148 lines
6.6 KiB
Markdown
148 lines
6.6 KiB
Markdown
# ternary_optim — SR4 4-bit Stochastic-Rounded Latent Accumulator
|
|
|
|
A small, standalone **C reference library** for the SR4 ternary optimizer: a
|
|
from-scratch training optimizer whose **entire persistent per-weight state is a
|
|
single 4-bit integer** — no FP32 master weight, no momentum, no Adam moments —
|
|
which trains a network whose **deployed weights are ternary** (`{-1, 0, +1}`).
|
|
|
|
This library is the **extraction step**: the SR4 core lifted out of the SR4 research
|
|
prototype (an MNIST training harness) into a reusable, MNIST-decoupled API
|
|
generalized to arbitrary `(rows x cols)` weight matrices. The numerics
|
|
(stochastic rounding, low-discrepancy dither, per-matrix step normalization,
|
|
saturating add, counter-hash) are bit-for-bit faithful to the research
|
|
prototype.
|
|
|
|
- `ternary_optim.h` / `.c` — the library.
|
|
- `example.c` — a self-contained demo (recover a ternary teacher via regression).
|
|
- **`docs/OPERATING_GUIDE.md` — how SR4 *behaves* and how to train with it well (batch is
|
|
the dominant lever, seed bands, the recipe, regime-matters). READ THIS BEFORE TRAINING.**
|
|
- `docs/HOW_IT_WORKS.md` — the mechanism, step by step, with the math.
|
|
- `docs/API.md` — every public function: signature, semantics, ownership.
|
|
- `docs/USAGE.md` — how to wire it into a training loop.
|
|
|
|
## Why it exists — the mission
|
|
|
|
Train large models on commodity hardware by making the *training state* almost
|
|
free. Standard optimizers carry a lot of state per weight:
|
|
|
|
| Optimizer | Persistent training state / weight | @ 3.6B params |
|
|
|---|---|---|
|
|
| Adam / AdamW (mixed precision) | FP32 master (4B) + m (4B) + v (4B) = **12 B** | ~43 GB |
|
|
| 8-bit Adam | master (4B) + m (1B) + v (1B) = **6 B** | ~21 GB |
|
|
| SGD + momentum | master (4B) + momentum (4B) = **8 B** | ~29 GB |
|
|
| **SR4** | one int4 accumulator = **0.5 B (4 bits)** | **~1.8 GB** |
|
|
|
|
**~24x less training-state memory than Adam**, and the weights come out ternary
|
|
(cheap inference too).
|
|
|
|
## The core idea
|
|
|
|
Separate two things people usually conflate:
|
|
|
|
- **Deployed-weight precision** — how precise the weights you *ship* are: here
|
|
**ternary** (`{-1, 0, +1}`, ~1.58 bits).
|
|
- **Training-state precision** — how precise the state you must *hold during
|
|
training* is: here **4 bits** — a tiny signed integer accumulator `q` per
|
|
weight.
|
|
|
|
Classic straight-through training keeps an FP32 shadow weight (32 bits of
|
|
training state) and quantizes it on the forward pass. SR4 keeps the same *idea*
|
|
of a latent accumulator, but at **4 bits instead of 32**. Stochastic rounding
|
|
supplies the sub-integer resolution the missing 28 bits used to provide. No FP32
|
|
shadow weights anywhere.
|
|
|
|
## How it works at a glance
|
|
|
|
Each weight has one integer accumulator `q ∈ [-7, +7]` (15 states → a nibble):
|
|
|
|
- **Forward (ternarize):** `w = 0 if |q| < 2 else sign(q)`. A weight must build
|
|
*conviction* (`|q| ≥ 2`) before it fires — a built-in hysteresis / dead-zone.
|
|
- **Update (integrate):** each batch, nudge `q` toward the negative gradient by a
|
|
tiny amount, using **stochastic rounding** so fractional-LSB steps are realized
|
|
as an occasional ±1. `q` is a slow integer *integrator* of gradient direction.
|
|
|
|
The stochastic rounding is driven by a **stateless counter-hash** (SplitMix64
|
|
keyed on `(seed, counter, weight-index)`) — nothing is stored, reproducible,
|
|
parallelizable, bit-identical across compiler flags. See
|
|
[`docs/HOW_IT_WORKS.md`](docs/HOW_IT_WORKS.md).
|
|
|
|
## Quick start
|
|
|
|
```c
|
|
#include "ternary_optim.h"
|
|
|
|
ternary_optim_config cfg;
|
|
ternary_optim_config_default(&cfg); /* target_step=0.001, q_clamp=7,
|
|
theta=2, dither on, seed=42 */
|
|
ternary_optim *opt = ternary_optim_create(rows, cols, &cfg);
|
|
|
|
float *W = malloc(rows*cols*sizeof(float)); /* caller-owned scratch */
|
|
float *grad_sum = malloc(rows*cols*sizeof(float));
|
|
|
|
uint64_t counter = 0;
|
|
for (int step = 0; step < n_steps; step++) {
|
|
ternary_optim_materialize(opt, W); /* q -> ternary weights */
|
|
/* ... forward with W, backward -> fill grad_sum (summed over the batch) ...*/
|
|
ternary_optim_step(opt, grad_sum, batch_size, counter, /*target_step*/ -1.0f);
|
|
counter++; /* one increment per update */
|
|
}
|
|
ternary_optim_free(opt);
|
|
```
|
|
|
|
Build and run the demo:
|
|
|
|
```
|
|
make run
|
|
```
|
|
|
|
## What's validated (and what isn't)
|
|
|
|
Validated on MNIST (784-256-10 MLP), memory-neutral at 4 b/w:
|
|
|
|
- Ternary SR4 reaches **~98.3**, essentially the **FP32 non-ternary ceiling
|
|
(98.36)** of the same architecture.
|
|
- Beat the prior V1 gold standard (INT16 accumulator, 16 b/w) at a quarter of the
|
|
memory.
|
|
|
|
**Not yet validated:** LM / transformer scale. Adam earns its state most on hard,
|
|
large-scale loss surfaces; whether SR4 holds parity there is the open question.
|
|
|
|
## The recipe (optimizer core vs. the knobs)
|
|
|
|
The **core** — the int4 accumulator + stochastic-rounded integer update +
|
|
ternary forward + low-discrepancy dither — is what this library ships. Everything
|
|
else is a **recipe** that lives *outside* the optimizer, in the training loop or
|
|
the loss:
|
|
|
|
- **Batch size** — the dominant accuracy lever. Training loop.
|
|
- **Co-trained per-column output scales** — restore ternary dynamic range
|
|
(~0.042 b/w). Provided here as *optional* decoupled helpers
|
|
(`ternary_apply_col_scale`, `ternary_col_scale_grad`); the scale vector belongs
|
|
to the layer, not the optimizer.
|
|
- **Label smoothing** — loss-side regularizer. Loss.
|
|
- **Annealed dropout** — variance/reliability control. Training loop.
|
|
- **Cosine step-size anneal** — schedule; pass the annealed value as the
|
|
`target_step` argument to `ternary_optim_step()`.
|
|
|
|
## Deferred / stubbed (honest scope)
|
|
|
|
This is a **first-pass CPU reference** of the core mechanism. Explicitly NOT done
|
|
here:
|
|
|
|
- **State serialization** — `ternary_optim_save`/`ternary_optim_load` are
|
|
**stubs** returning `TERNARY_OPTIM_ENOTIMPL`. The thing to persist is the
|
|
packed `q` (4 b/w) + shape + config; the format is a TODO.
|
|
- **GPU kernels** — CPU scalar reference only.
|
|
- **foundry integration** — not wired into dcherrera/foundry's tensor + optim
|
|
framework; that is the follow-up.
|
|
- **Scale-generalization testing** — the per-column scale helpers are generalized
|
|
to `(rows, cols)` but have only been validated in the MNIST MLP shape.
|
|
- **Multi-layer convenience** — one handle owns one matrix; a model wires several
|
|
handles (give each a distinct `idx_base`). No container type yet.
|
|
|
|
## Provenance
|
|
|
|
Mechanism source of truth: the `sr_*` core of the SR4 research prototype (not
|
|
published) and its design notes, which `docs/HOW_IT_WORKS.md` adapts. The update
|
|
arithmetic here is a faithful, decoupled lift of that code.
|