sexpgpu/optim

The standard optimizer module, compiled into the binary: (require "sexpgpu/optim" name ...). Every optimizer accepts a number or a schedule for each numeric hyperparameter and calls update-diagnostics.

Optimizers

definition	keywords	states
`sgd`	`(lr 0.01) (weight-decay 0.0)`	none
`momentum-sgd`	`(lr 0.01) (momentum 0.9) (nesterov false)`	`velocity`
`adamw`	`(lr 1e-3) (betas [0.9 0.95]) (eps 1e-8) (weight-decay 0.0)`	`m`, `v`
`muon`	`(lr 0.02) (weight-decay 0.0) (momentum 0.95) (nesterov true) (steps 12)`	`momentum-buffer`

sgd: param - lr * (grad + weight-decay * param).
adamw: bias-corrected with t = step + 1; weight decay is added to the normalized step, then scaled by lr.
muon follows modded-nanogpt: heavy-ball momentum, then the orthogonalization of that momentum by newton-schulz, scaled by sqrt(max(1, out/in)) in PyTorch's layout, with decoupled weight decay. Matrices only: (assert (= 2 (rank param))) fails otherwise. Give the rest to adamw with a group; see optimizer groups.

Helpers

definition	signature and meaning
`schedule-value`	`(schedule-value value ctx)`: `(value ctx)` when it is a function, otherwise `value`; see writing optimizers
`update-diagnostics`	`(update-diagnostics param grad next)`: the diagnostics `optimizer/grad-rms` and `optimizer/update-ratio`, both `:max`
`newton-schulz`	`(newton-schulz x &key (steps 12))`: the orthogonal factor of a matrix by Newton-Schulz with coefficients (2, -1.5, 0.5) from `x / ‖x‖`, in `bf16`; a tall matrix is transposed in and back out, decided at compile time

Schedules

definition	signature and meaning
`stable-then-decay`	`(stable-then-decay peak &key (decay-fraction 0.7))`: `peak` until `1 - decay-fraction` of the run, then linear to zero
`warmup-stable-decay`	`(warmup-stable-decay peak &key (warmup-steps 0) (decay-fraction 0.7))`: the same times a linear warmup `min(1, (step + 1) / warmup-steps)`, so the first update already moves

(adamw :lr (warmup-stable-decay 0.0015 :warmup-steps 25) :weight-decay 0.1)

Layout

Weights are stored [in, out], the transpose of PyTorch's [out, in]. newton-schulz commutes with transposition, so the muon update is exactly the transpose of the modded-nanogpt update; only the scale is read in PyTorch's layout: out/in is (dim param 1) over (dim param 0).

Diagnostics

name	reducer	what it is
`optimizer/grad-rms`	max	the gradient's root mean square, the group's largest
`optimizer/update-ratio`	max	the step's root mean square over the parameter's: how far the update moves the parameter against its own size

See optimizer diagnostics for how they fold and what makes one noisy.