156 lines
7.0 KiB
Markdown
156 lines
7.0 KiB
Markdown
# ternary_optim — Wiring it into a training loop
|
||
|
||
This is the operational guide: the `materialize → forward → backward → update`
|
||
cycle, how gradients must be provided, the batch-counter contract, and where the
|
||
recipe knobs live. For the mechanism see [HOW_IT_WORKS.md](HOW_IT_WORKS.md); for
|
||
exact signatures see [API.md](API.md). A runnable end-to-end demo is
|
||
[`../example.c`](../example.c) (`make run`).
|
||
|
||
---
|
||
|
||
## 1. The one-matrix cycle
|
||
|
||
The optimizer owns exactly one weight matrix's `q` state. Per step:
|
||
|
||
```c
|
||
ternary_optim_materialize(opt, W); // 1. q -> ternary W (caller buffer)
|
||
// 2. forward: use W (and optional scales/bias) to compute activations
|
||
// 3. backward: produce grad_sum — dL/dW SUMMED over the batch
|
||
ternary_optim_step(opt, grad_sum, batch_size, counter, target_step); // 4.
|
||
counter++; // 5. exactly one increment per step
|
||
```
|
||
|
||
1. **Materialize** turns `q` into the ternary weights `{-1,0,+1}` your forward
|
||
math consumes. Do it every step — `q` moves under the update.
|
||
2. **Forward** is entirely yours (dense matmul, BLAS, conv, whatever). The
|
||
optimizer never sees your activations.
|
||
3. **Backward** is yours too; it must produce `grad_sum`, the per-weight gradient
|
||
**summed over the batch samples** (see §2).
|
||
4. **Step** integrates `grad_sum` into `q`.
|
||
5. **Increment** the counter (see §3).
|
||
|
||
---
|
||
|
||
## 2. How gradients must be provided
|
||
|
||
`grad_sum` is `dL/dW` **summed over the batch**, in the **same row-major layout as
|
||
the weight matrix** (element `(r,c)` at `r*cols + c`), at least `numel` floats.
|
||
|
||
- **Summed, not averaged.** The optimizer divides by `batch_size` internally
|
||
(`g = grad_sum / batch_size`) and then renormalizes per-matrix
|
||
(`s = target_step / (mean|g| + 1e-12)`). Because of that per-matrix
|
||
normalization the update is invariant to the *overall* gradient scale — but the
|
||
`batch_size` you pass must match the number of samples actually summed, or the
|
||
mean is wrong (this matters for a partial final batch).
|
||
- **It is the gradient w.r.t. the ternary weights** (`W`), i.e. the same `W` you
|
||
got from `materialize`. If you fold per-column scales into the forward
|
||
(`Z = scale ⊙ (W·x) + bias`), the scale must be pushed back into the delta
|
||
before the weight-grad matmul so `grad_sum` reflects the scaled forward — see
|
||
§5.
|
||
- **No gradient of the ternarization itself.** SR4 is straight-through in spirit:
|
||
you compute the gradient as if `W` were the real weights; the accumulator `q`
|
||
absorbs it. There is no `dW/dq` term.
|
||
|
||
`grad_sum` is read-only to the optimizer and stays caller-owned; reuse the same
|
||
buffer every step (zero it before accumulating).
|
||
|
||
---
|
||
|
||
## 3. The batch-counter contract
|
||
|
||
`batch_counter` is a `uint64_t` **global update index**, starting at 0, that drives
|
||
the stateless dither. Rules:
|
||
|
||
- **Start at 0, increment by exactly 1 per logical update step.** Think *updates*,
|
||
not epochs — the counter is the fundamental clock of the optimizer.
|
||
- **All matrices in one step share the same counter value.** If your model has
|
||
several `ternary_optim` handles (one per weight matrix), call `step` on each with
|
||
the *same* `batch_counter`, then increment once:
|
||
|
||
```c
|
||
ternary_optim_step(W1_opt, g1, bs, counter, ts);
|
||
ternary_optim_step(W2_opt, g2, bs, counter, ts); // same counter
|
||
counter++; // one increment for the step
|
||
```
|
||
|
||
Give each handle a distinct `idx_base` (non-overlapping `[idx_base,
|
||
idx_base+numel)` ranges) so their stochastic draws are independent even at the
|
||
same counter. E.g. `W1` uses `idx_base = 0`, `W2` uses `idx_base = W1.numel`.
|
||
- **Reproducibility.** Same `seed`, same `idx_base`, same counter sequence, same
|
||
gradients ⇒ bit-identical `q`. The counter must therefore be deterministic
|
||
across a resume (part of what a future checkpoint must persist).
|
||
|
||
---
|
||
|
||
## 4. The recipe lives outside the optimizer
|
||
|
||
The optimizer is just the int4 update. The accuracy recipe lives where it belongs:
|
||
|
||
| Knob | Where it lives | How |
|
||
|---|---|---|
|
||
| **Batch size** | training loop | dominant accuracy lever; bigger batch = cleaner per-step gradient = tighter seed band. Pass the true `batch_size` to `step`. |
|
||
| **Cosine step anneal** | training loop → `target_step` arg | compute the annealed value each step and pass it as `target_step` (pass `<= 0` to use the config default). Reference: cosine from `target_step` to `target_step/10` over all updates. |
|
||
| **Label smoothing** | loss | fold into your CE targets; the softmax gradient stays `p - t`. Zero optimizer involvement. |
|
||
| **Annealed dropout** | training loop | apply your own mask to activations in forward and to `dL/dh` in backward with identical keys; nothing to store if you use a counter-hash mask. Reference recipe: linear anneal of `p` to 0 over the first ~25% of training. |
|
||
| **Per-column scales** | the layer | see §5. |
|
||
|
||
Matched-updates thinking: when you change batch size, hold the number of *updates*
|
||
(hence the counter range and the cosine schedule length) fixed, not the number of
|
||
epochs.
|
||
|
||
---
|
||
|
||
## 5. Optional per-column output scales
|
||
|
||
Ternary weights lose magnitude; a small FP32 per-output-column scale vector
|
||
restores dynamic range. The vector is owned by **your layer**, not the optimizer.
|
||
Init to all-ones (exact identity) and train by plain SGD at a dedicated small LR
|
||
(reference: `SCALE_LR = 0.001`, ~10× below the bias LR).
|
||
|
||
Forward fold — `Z = scale ⊙ Z_raw + bias`:
|
||
|
||
```c
|
||
ternary_optim_materialize(opt, W);
|
||
// Z_raw = X @ W (your matmul)
|
||
memcpy(Z_raw_saved, Z_raw, rows*cols*sizeof(float)); // keep for the scale grad
|
||
ternary_apply_col_scale(Z, scale, rows, cols); // Z *= scale (per column)
|
||
// add bias, activation, ...
|
||
```
|
||
|
||
Backward — take the bias/scale grads from the **unscaled** delta, then push the
|
||
scale back into the delta before the weight-grad matmul:
|
||
|
||
```c
|
||
// D = dL/dZ (unscaled upstream delta)
|
||
ternary_col_scale_grad(D, Z_raw_saved, scale_grad, rows, cols); // dL/dscale
|
||
// bias_grad = column-sum of D (your code)
|
||
ternary_apply_col_scale(D, scale, rows, cols); // D *= scale -> dL/dZ_raw
|
||
// grad_sum = X^T @ D (your weight-grad matmul, feeds ternary_optim_step)
|
||
|
||
// SGD updates (your code):
|
||
for (c) scale[c] -= SCALE_LR * scale_grad[c] / batch_size;
|
||
for (c) bias[c] -= BIAS_LR * bias_grad[c] / batch_size;
|
||
```
|
||
|
||
The optimizer never touches `scale`/`bias`; they are ordinary FP32 SGD parameters
|
||
you maintain alongside `q`.
|
||
|
||
---
|
||
|
||
## 6. Evaluation / deployment
|
||
|
||
At eval time there is no update and no dropout — just `materialize` the ternary
|
||
weights, run the forward with the scales/biases, and read predictions. The
|
||
deployed artifact is `ternarize(q)` (the ternary matrix) plus the FP32 scales and
|
||
biases; nothing else needs to persist.
|
||
|
||
---
|
||
|
||
## 7. Checkpoint / resume (currently stubbed)
|
||
|
||
`ternary_optim_save`/`_load` return `TERNARY_OPTIM_ENOTIMPL` today. Until they
|
||
land, to snapshot state yourself, persist: the packed `q`
|
||
(`ternary_optim_q(opt)`, `numel` bytes — 4 b/w of real information), the shape,
|
||
the config, your scale/bias vectors, and the `batch_counter`. The dither/masks
|
||
need not be saved — they regenerate from `(seed, counter, index)`.
|