283 lines
13 KiB
C
283 lines
13 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file ternary_optim.h
|
|
* @brief SR4 — 4-bit Stochastic-Rounded Latent Accumulator (ternary optimizer).
|
|
*
|
|
* SR4 trains a network whose DEPLOYED weights are ternary ({-1, 0, +1}) while
|
|
* carrying an ENTIRE persistent per-weight training state of a single 4-bit
|
|
* signed integer accumulator `q` — no FP32 master weight, no momentum, no Adam
|
|
* moments. Stochastic rounding (driven by a stateless counter-hash) supplies the
|
|
* sub-integer resolution that a low-bit state would otherwise lack.
|
|
*
|
|
* This header is the CLEAN, MNIST-decoupled public API of the optimizer core,
|
|
* generalized to an arbitrary weight matrix of shape (rows x cols). It knows
|
|
* nothing about MNIST, the 784-256-10 MLP, data loading, or the training loop —
|
|
* those live in the caller.
|
|
*
|
|
* The numerics (stochastic rounding, low-discrepancy dither, per-matrix step
|
|
* normalization, saturating add, the SplitMix64 counter-hash) are lifted
|
|
* faithfully and bit-for-bit from the SR4 research prototype's `sr_*` core.
|
|
*
|
|
* SCOPE (honest): this is a first-pass CPU REFERENCE library of the core
|
|
* mechanism, validated on MNIST-MLP. See README.md for what is stubbed/deferred
|
|
* (state serialization, GPU kernels, foundry tensor/optim integration,
|
|
* scale-generalization testing).
|
|
*
|
|
* Threading / determinism: every stochastic draw is a pure function of
|
|
* (seed, batch_counter, global weight index) via an integer counter-hash — no
|
|
* stored RNG stream. Results are reproducible per seed and bit-identical across
|
|
* compiler optimization / fast-math flags.
|
|
*/
|
|
|
|
#ifndef TERNARY_OPTIM_H
|
|
#define TERNARY_OPTIM_H
|
|
|
|
#include <stdint.h>
|
|
#include <stddef.h>
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* Return codes
|
|
* ------------------------------------------------------------------------- */
|
|
enum {
|
|
TERNARY_OPTIM_OK = 0, /**< success */
|
|
TERNARY_OPTIM_EINVAL = -1, /**< invalid argument */
|
|
TERNARY_OPTIM_ENOMEM = -2, /**< allocation failure */
|
|
TERNARY_OPTIM_ENOTIMPL = -3 /**< not implemented (stubbed feature) */
|
|
};
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* Configuration
|
|
* ------------------------------------------------------------------------- */
|
|
|
|
/**
|
|
* @brief Immutable configuration for one optimizer instance.
|
|
*
|
|
* Populate with ternary_optim_config_default() and then override fields as
|
|
* needed. All fields are copied into the optimizer at create time.
|
|
*/
|
|
typedef struct {
|
|
/** Average proposed step, in LSB, for a weight whose |gradient| equals the
|
|
* matrix mean |gradient|. Small (~0.001): `q` INTEGRATES gradient
|
|
* direction over many updates rather than jumping. The training loop may
|
|
* additionally anneal an effective per-step value (e.g. cosine) and pass
|
|
* it to ternary_optim_step(); this default is used when a step call passes
|
|
* a non-positive target_step. */
|
|
float target_step;
|
|
|
|
/** Symmetric saturation of the accumulator: q is clamped to [-q_clamp,
|
|
* +q_clamp]. Reference value 7 (15 states -> a 4-bit nibble). */
|
|
int q_clamp;
|
|
|
|
/** Dead-zone threshold: the deployed weight is 0 while |q| < theta, else
|
|
* sign(q). Reference value 2 (a weight must build conviction |q|>=2 before
|
|
* it fires; hysteresis prevents chatter). */
|
|
int theta;
|
|
|
|
/** 1 = stratified Weyl low-discrepancy dither (variance-reduced, the
|
|
* reference/default); 0 = plain i.i.d. counter-hash dither (unbiased but
|
|
* noisier). Both are unbiased; LD tightens the seed-to-seed band. */
|
|
int use_dither;
|
|
|
|
/** Seed threaded through EVERY weight on EVERY update via the counter-hash.
|
|
* Also seeds the initial q distribution. Two instances that must diverge
|
|
* should differ in either `seed` or `idx_base`. */
|
|
uint64_t seed;
|
|
|
|
/** Global weight-index base for this matrix in the counter-hash lane. When
|
|
* a model has several matrices, give each a distinct, non-overlapping
|
|
* [idx_base, idx_base+numel) range so their stochastic draws are
|
|
* independent. A single-matrix caller may leave this 0. */
|
|
int64_t idx_base;
|
|
} ternary_optim_config;
|
|
|
|
/**
|
|
* @brief Fill @p cfg with the reference defaults.
|
|
*
|
|
* target_step=0.001, q_clamp=7, theta=2, use_dither=1, seed=42, idx_base=0.
|
|
*
|
|
* @param cfg out: configuration to initialize (must be non-NULL).
|
|
*/
|
|
void ternary_optim_config_default(ternary_optim_config *cfg);
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* Optimizer handle
|
|
* ------------------------------------------------------------------------- */
|
|
|
|
/** Opaque optimizer / param-state handle. One instance owns the `q` state for
|
|
* exactly one (rows x cols) weight matrix. */
|
|
typedef struct ternary_optim ternary_optim;
|
|
|
|
/**
|
|
* @brief Create an optimizer for a (rows x cols) row-major weight matrix.
|
|
*
|
|
* Allocates the int4 accumulator state (one signed byte per weight; the
|
|
* INFORMATION content is 4 bits/weight) and initializes it to {-3, 0, +3} at
|
|
* {33.5%, 33%, 33.5%} — ~67% of weights active, signs uniform, magnitude 3 =
|
|
* theta+1 giving every active weight one LSB of dead-zone margin. The initial
|
|
* distribution is drawn from an internal xorshift RNG seeded from
|
|
* (cfg->seed ^ cfg->idx_base), so distinct matrices get distinct inits.
|
|
*
|
|
* @param rows matrix rows (> 0).
|
|
* @param cols matrix columns (> 0).
|
|
* @param cfg configuration; copied. If NULL, defaults are used.
|
|
* @return owning handle, or NULL on invalid args / allocation failure.
|
|
* Free with ternary_optim_free().
|
|
*/
|
|
ternary_optim *ternary_optim_create(int rows, int cols,
|
|
const ternary_optim_config *cfg);
|
|
|
|
/**
|
|
* @brief Release an optimizer and its state. Safe to call with NULL.
|
|
* @param opt handle to free (may be NULL).
|
|
*/
|
|
void ternary_optim_free(ternary_optim *opt);
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* Core operations
|
|
* ------------------------------------------------------------------------- */
|
|
|
|
/**
|
|
* @brief Materialize the ternary deployed weights from `q` into a caller buffer.
|
|
*
|
|
* out[i] = 0 if |q[i]| < theta
|
|
* +1 if q[i] >= +theta
|
|
* -1 if q[i] <= -theta
|
|
*
|
|
* Call this before each forward pass; the result is the actual weight matrix
|
|
* the forward math consumes.
|
|
*
|
|
* @param opt optimizer (non-NULL).
|
|
* @param out_weights out: caller-owned buffer of at least numel floats,
|
|
* filled row-major with values in {-1, 0, +1}.
|
|
*/
|
|
void ternary_optim_materialize(const ternary_optim *opt, float *out_weights);
|
|
|
|
/**
|
|
* @brief One SR4 update step: integrate the batch gradient into `q`.
|
|
*
|
|
* Given the per-weight gradient SUMMED over the batch (same row-major layout as
|
|
* the weight matrix) and the batch size, this:
|
|
* 1. normalizes the step per-matrix: s = target_step / (mean|g| + 1e-12),
|
|
* where g = grad_sum / batch_size — so a typical weight proposes ~
|
|
* target_step LSB regardless of raw gradient scale;
|
|
* 2. snaps the proposed step to 8.8 fixed point and splits it into an integer
|
|
* part (applied deterministically) and a fractional part (applied
|
|
* stochastically as an occasional +1), so E[increment] is exact (unbiased);
|
|
* 3. draws the stochastic bit from the counter-hash — either the stratified
|
|
* Weyl low-discrepancy dither (use_dither=1) or plain i.i.d. (use_dither=0);
|
|
* 4. saturating-adds into q, clamped to [-q_clamp, +q_clamp].
|
|
*
|
|
* The only side effect is the integer add into the internal q state. Nothing is
|
|
* stored beyond q: the dither is regenerated from (seed, batch_counter, index).
|
|
*
|
|
* @param opt optimizer (non-NULL).
|
|
* @param grad_sum per-weight gradient summed over the batch (>= numel
|
|
* floats, row-major, matching the weight layout).
|
|
* @param batch_size number of samples summed into grad_sum (> 0). Used for
|
|
* the /batch_size mean; may vary (e.g. a partial final
|
|
* batch).
|
|
* @param batch_counter monotonically increasing global update index, starting
|
|
* at 0. All matrices updated within the SAME step use the
|
|
* SAME counter value; the caller increments it by 1 AFTER
|
|
* stepping every matrix. Drives the stateless dither; see
|
|
* docs/USAGE.md.
|
|
* @param target_step effective per-step target (e.g. a cosine-annealed
|
|
* value). If <= 0, the config default is used.
|
|
*/
|
|
void ternary_optim_step(ternary_optim *opt, const float *grad_sum,
|
|
int batch_size, uint64_t batch_counter,
|
|
float target_step);
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* Accessors (diagnostics)
|
|
* ------------------------------------------------------------------------- */
|
|
|
|
/** @return number of weights (rows*cols), or 0 if opt is NULL. */
|
|
int ternary_optim_numel(const ternary_optim *opt);
|
|
/** @return matrix rows, or 0 if opt is NULL. */
|
|
int ternary_optim_rows(const ternary_optim *opt);
|
|
/** @return matrix cols, or 0 if opt is NULL. */
|
|
int ternary_optim_cols(const ternary_optim *opt);
|
|
/**
|
|
* @brief Read-only view of the raw int4 accumulator (one signed byte/weight,
|
|
* clamped to [-q_clamp, +q_clamp]). For diagnostics/fingerprints only —
|
|
* the update rule never reads it back. Valid until the optimizer is
|
|
* freed.
|
|
* @return pointer to numel int8 values, or NULL if opt is NULL.
|
|
*/
|
|
const int8_t *ternary_optim_q(const ternary_optim *opt);
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* Optional per-column output scales (recipe adjunct, decoupled)
|
|
*
|
|
* Ternary weights throw away magnitude. These free helpers restore per-output-
|
|
* column dynamic range via an FP32 scale vector folded into the forward pass:
|
|
* Z = scale (*) Z_raw + bias.
|
|
* The scale vector is NOT owned by the optimizer (it belongs to the layer);
|
|
* these operate on caller buffers over an arbitrary (rows x cols) matrix. The
|
|
* caller allocates the scale vector (length cols, init 1.0) and trains it by
|
|
* plain SGD at a small dedicated learning rate. See docs/USAGE.md.
|
|
* ------------------------------------------------------------------------- */
|
|
|
|
/**
|
|
* @brief In place, multiply each column c of a row-major (rows x cols) matrix
|
|
* by scale[c]. Exact identity when scale == all-ones.
|
|
* @param Z in/out: row-major matrix (rows*cols floats).
|
|
* @param scale length-cols per-column scale vector.
|
|
* @param rows matrix rows.
|
|
* @param cols matrix cols.
|
|
*/
|
|
void ternary_apply_col_scale(float *Z, const float *scale, int rows, int cols);
|
|
|
|
/**
|
|
* @brief Per-column scale gradient: out[c] = sum_r D[r,c] * Z_raw[r,c].
|
|
*
|
|
* D is the UNSCALED upstream delta (dL/dZ before folding the scale back in) and
|
|
* Z_raw the pre-scale, pre-bias matmul output. `out` is OVERWRITTEN (not
|
|
* accumulated). Apply to the scale vector by plain SGD at a dedicated small LR.
|
|
*
|
|
* @param D upstream delta, row-major (rows*cols floats).
|
|
* @param Z_raw pre-scale/pre-bias matmul output, row-major (rows*cols floats).
|
|
* @param out out: length-cols gradient (overwritten).
|
|
* @param rows matrix rows.
|
|
* @param cols matrix cols.
|
|
*/
|
|
void ternary_col_scale_grad(const float *D, const float *Z_raw, float *out,
|
|
int rows, int cols);
|
|
|
|
/* ---------------------------------------------------------------------------
|
|
* State serialization (STUBBED — TODO: checkpoint / resume)
|
|
*
|
|
* The thing to persist is the packed `q` array (4 bits/weight) plus the shape
|
|
* and config; the dither and any masks are regenerated from the counter-hash so
|
|
* they never need saving. These are intentionally NOT implemented in this
|
|
* first-pass reference library: both return TERNARY_OPTIM_ENOTIMPL. See
|
|
* README.md "Deferred / stubbed".
|
|
* ------------------------------------------------------------------------- */
|
|
|
|
/**
|
|
* @brief TODO/STUB: save optimizer state (packed q + shape + config) to @p path.
|
|
* @return TERNARY_OPTIM_ENOTIMPL always (not yet implemented).
|
|
*/
|
|
int ternary_optim_save(const ternary_optim *opt, const char *path);
|
|
|
|
/**
|
|
* @brief TODO/STUB: load optimizer state previously written by
|
|
* ternary_optim_save().
|
|
* @param path path to a saved state file.
|
|
* @param out out: receives a newly created handle on success.
|
|
* @return TERNARY_OPTIM_ENOTIMPL always (not yet implemented); *out set to NULL.
|
|
*/
|
|
int ternary_optim_load(const char *path, ternary_optim **out);
|
|
|
|
#ifdef __cplusplus
|
|
} /* extern "C" */
|
|
#endif
|
|
|
|
#endif /* TERNARY_OPTIM_H */
|