Foundry/ternary_optim/README.md

6.6 KiB

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.

Quick start

#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.