Training: modules, losses, optimizers, schedules, checkpoints

Normative for package RecursiveAI.OneModel.Training. It is written on the tensor layer and its reverse-mode autodiff (autograd.md) and uses nothing below it. References: Kingma and Ba 2015 (Adam); Loshchilov and Hutter 2019 (AdamW); Reddi, Kale and Kumar 2018 (AMSGrad); Sutskever et al. 2013 (Nesterov momentum); Loshchilov and Hutter 2017 (cosine annealing); Smith and Topin 2018 (one cycle); Vaswani et al. 2017 (transformer encoder); Steele, Lea and Flood 2014 (SplitMix64). PyTorch's public documentation is the reference for the conventions it shares (reduction, label smoothing, parameter shapes); it is also the test oracle (section 8).

1. Modules

NnModule owns parameters (contiguous float32 leaves that require gradients), buffers (saved, not trained: batch-norm running statistics, the dropout draw counter) and named children. Names contain no dot. The state order is fixed: a module's parameters in registration order, then its buffers in registration order, then each child, recursively, with dotted paths (layer0.attn.q.weight). NamedParameters() is the parameter subsequence of it, StateDict() the whole, and optimizers and checkpoints use these orders. Train(bool) / Eval() set IsTraining on the module and every child. LoadStateDict copies by name into the existing tensors (identity, gradients and optimizer references stay valid) and rejects any missing name, extra name, type or shape difference before copying anything.

1.1 Initialisation

Weights come from Generator, SplitMix64: state += 0x9E3779B97F4A7C15, then z = (z ^ z>>30) * 0xBF58476D1CE4E5B9, z = (z ^ z>>27) * 0x94D049BB133111EB, output z ^ z>>31. NextSingle is (output >> 40) / 2^24; Uniform(a, b) = a + (b−a)·NextSingle(). Linear and Conv2d draw weight then bias from Uniform(−1/√fanIn, +1/√fanIn) in row-major order (fanIn = in, resp. in/groups·k·k; the square root is IEEE). Embedding draws from Uniform(−√3, √3) (unit variance). Children draw in construction order (MultiHeadAttention: q, k, v, o; TransformerEncoderLayer: attention, ff1, ff2). The streams are not PyTorch's, so a run that must match PyTorch loads its initial weights (section 8).

1.2 Dropout

Training mode with p > 0: let d be the value of draws (then draws += 1), stream = Mix(seed + d·0x9E3779B97F4A7C15) (all unsigned 64-bit, wrapping) and for element i (row-major, from 0) bits = Mix(stream + (i+1)·0x9E3779B97F4A7C15), u = (bits >> 40) / 2^24. The element is kept, multiplied by 1/(1−p) (float32), when u >= p, else zeroed. The mask depends on nothing else, so it is the same on every platform and thread count, and because draws is a saved buffer a resumed run continues the stream. Evaluation mode, or p = 0, returns the input itself. The mask multiplies through Ops.Mul, so backward is the mask.

2. Losses (Losses)

All return a rank-0 float32 tensor and are compositions of differentiable ops. mean (default) divides by the number of counted elements; otherwise the sum is returned.

• CrossEntropy(logits[N,C], targets[N], labelSmoothing ε, ignoreIndex = −100): NLL(LogSoftmax(logits)); with ε > 0: (1−ε)·NLL + ε·S, S = Σ_rows (−mean_c logp[n,c]) / counted over rows whose target is not the ignore index. No counted row: NaN mean.

• NllLoss: Ops.NllLoss.

• Mse: mean((p−t)²). L1: mean(|p−t|), with |d| built as relu(d) + relu(−d) (gradient 0 at d = 0, as in PyTorch).

• BceWithLogits(x, t): per element softplus(x) − x·t, softplus(x) = −LogSoftmax([x, 0])[1], stable for any |x| (no overflow, no cancellation). No pos_weight, no class weights.

• Ctc(logProbs[T,B,C], targets, inputLengths, targetLengths, blank, reduction, zeroInfinity): Ops.CtcLoss.

3. Optimizers

One parameter group per optimizer: one learning rate, one set of hyper-parameters. Parameters are the optimizer's constructor list; those without a gradient at a step are skipped (their state does not change). A step applies, parameter by parameter and element by element in index order, plain float32 + − × ÷ sqrt with no fused multiply-add (the JIT does not contract float expressions, and the code must not use Fma), so the update is bit-identical in fast and exact mode, on every platform and thread count. State is float32, shaped like the parameter, zero at creation, and listed by StateTensors (momentum_buffer.i, exp_avg.i, exp_avg_sq.i, max_exp_avg_sq.i, and an Int32 initialized flag per parameter for SGD). StepCount counts completed steps.

SGD (lr, μ momentum, d dampening, wd, nesterov), per element with gradient g: g += wd·p; if μ ≠ 0: b = first ? g : μ·b + (1−d)·g (first = the first step in which this parameter has a gradient), g = nesterov ? g + μ·b : b; then p −= lr·g. Nesterov needs μ > 0 and d = 0.

Adam / AdamW (lr, β₁, β₂, ε, wd, amsgrad), t = StepCount + 1: bc1 = 1 − β₁ᵗ, bc2 = 1 − β₂ᵗ (powers by binary exponentiation in double, round to nearest at each product), stepSize = float(lr / bc1), sqrtBc2 = float(√bc2) (double sqrt). Per element: Adam with wd adds wd·p to g; AdamW first sets p *= 1 − lr·wd; then m = β₁·m + (1−β₁)·g, v = β₂·v + (1−β₂)·g·g (evaluated left to right: ((1−β₂)·g)·g), vd = amsgrad ? max(vmax, v) : v (and vmax = vd), denom = √vd / sqrtBc2 + ε, p −= stepSize·m / denom. AdamW's default wd is 0.01.

Differences from PyTorch's implementation, all within float rounding of one step (measured, section 8): the first moment uses β₁m + (1−β₁)g where PyTorch uses lerp; the step number is global, not per parameter; the learning-rate argument is read at each step (a schedule sets it). Not provided: parameter groups, maximize, foreach/fused variants, Adamax, RMSprop and the rest.

4. Gradient clipping (GradClip)

ClipNorm(params, maxNorm, infinityNorm): norm = √Σg² over every element of every gradient in order (squares added in double), or the largest |g|; coef = maxNorm / (norm + 1e-6) in double, rounded to float; if coef < 1 every gradient element is multiplied by it. Returns float(norm). A non-finite norm is returned without scaling anything. ClipValue clamps each element to ±c (NaN stays).

5. Learning-rate schedules

A schedule is a pure function LrAt(stepIndex) of the number of optimizer steps already taken (0 for the first step); Apply(optimizer) sets the rate for optimizer.StepCount. There is no schedule state, so a checkpoint needs only the optimizer's step count. Arithmetic is IEEE double (+ − × ÷, no Math function), rounded to float at the end:

• StepLr(base, size, γ): base·γ^⌊s/size⌋ (integer power by squaring).
• CosineLr(base, T, min): min + (base−min)(1 + cos(π·s/T))/2, s/T clamped to 1 (the rate stays at min after T). PyTorch's recursive form turns back up after T_max; this one does not.
• WarmupLr(after, W, f): for s < W, r₀·(f + (1−f)·s/W) with r₀ = after.LrAt(0); for s ≥ W, after.LrAt(s − W).
• OneCycleLr(max, T, pct 0.3, div 25, finalDiv 1e4): with initial = max/div, min = initial/finalDiv, up = pct·T − 1, end = T − 1: half-cosine from initial to max over steps 0..up, then from max to min over up..end, anneal(a,b,x) = b + (a−b)/2·(cos(πx)+1). Momentum is not cycled.

LrSchedule.Cos(t) = cos(π t), t ∈ [0,1]: t ≤ 0 gives 1, t ≥ 1 gives −1; otherwise y = t or 1−t (sign −1 for t > 0.5), a = π·y with π the double 3.14159265358979323846, a² = a·a, and Horner's scheme r = 1; for k = 13 down to 1: r = 1 − a²/((2k−1)(2k))·r; the result is sign·r. This is the Maclaurin series to degree 26 (truncation below 1e-21). Measured error against Math.Cos over 10 001 points: 3.9e-16. It is used instead of Math.Cos because the C library's cosine may differ between platforms in the last bit, which would break exact-mode checkpoints and hashes.

6. Checkpoints (Checkpoint)

File layout, little endian, the same frame as the model file (format.md) with a different magic:

offset 0                64-byte header
offset 64               manifest, UTF-8 JSON, manifest_length bytes
then zero padding to the next multiple of 64 (= data_offset)
data_offset             data region, data_length bytes; end of file = data_offset + data_length

Header: 8 bytes ASCII ONEMCKPT; u32 version = 1; u32 flags = 0; u64 manifest_offset = 64; u64 manifest_length; u64 data_offset (multiple of 64, ≥ 64 + manifest_length, ≤ file length); u64 data_length = file length − data_offset; u64 file_length (exact); u64 reserved = 0. A reader rejects anything else, and checks every field against the real length before reading the manifest (limit 256 MiB, JSON depth 32).

Manifest: {"format":"onemodel-checkpoint","version":1,"model":[T...],"optimizer":{"kind","step","hyper":[{"name","bits"}...],"state":[T...]},"extra":[{"key","value"}...]} (optimizer only when saved). Each tensor entry T is {"name","dtype":"f32"|"i32"|"i64","shape":[...],"offset","length"}: offset from the start of the data region (a multiple of 64), length in bytes. Tensors appear in the state order (section 1), model first, then optimizer state in StateTensors order; hyper lists HyperParameters() in order, each float as the unsigned integer of its IEEE-754 bit pattern; extra holds integers sorted by key (ordinal). Tensor data is raw little-endian elements, row-major, floats by bit pattern (NaN payloads and signed zeros survive). The writer emits compact JSON in exactly this field order, so the same state writes the same bytes.

Loading checks, before changing anything: header; manifest shape; tensor names, types and shapes equal the target's, in order; every offset and length inside the data region and equal to the shape's size (so a hostile shape never allocates); the optimizer kind and its fixed settings (momentum, betas, ε, amsgrad). Then it copies the tensors, restores hyper-parameters (the learning rate included) and the step count. Errors are CheckpointException(Kind): BadHeader, BadManifest, BadTensor, Mismatch. A checkpoint with optimizer state loads into a model alone; asking for optimizer state from a model-only file is a Mismatch.

Hostile input. A checkpoint is untrusted: the reader reads the whole input, then checks in this order before it touches the model or optimizer: the header (magic, version 1, flags 0, reserved 0, file_length equal to the real length, manifest and data offsets inside the file and data_offset a multiple of 64) -> BadHeader; the manifest (valid UTF-8 JSON of depth at most 32, an object with the required fields, integers that fit, dtype one of f32/i32/i64, no negative or oversized dimension) -> BadManifest; each tensor entry against the target (names, types and shapes equal, in order, and the shape is compared before any size is computed from it, so a hostile shape never allocates) -> Mismatch; then each entry's offset and length (length equal to the shape's byte size, offset at least 0, both inside the data region, checked without overflow) -> BadTensor. Only then are values copied. No bad input may raise any other exception type, and a rejected file leaves the model and optimizer unchanged. The committed cases are in testdata/checkpoint/hostile/<kind>-<name>.ckpt (41 files; CheckpointHostileTests assembles each by hand from this section, regenerates them with ONEMODEL_WRITE_TESTDATA=1, and requires the typed error from memory and from the file). The fuzz target checkpoint (tools/run_fuzz.sh) tries every input against four model/optimizer pairs and again with its header length fields repaired, so that mutations of the manifest and data get past the length checks; seeds are valid checkpoints of each pair plus the hostile files. Its cov figure in libFuzzer's output is a constant of the SharpFuzz driver (3 for every target); the meaningful figure is ft.

Resume rule. Model, buffers (including dropout counters), optimizer state, hyper-parameters and the step count are the whole training state; schedules and dropout are functions of them. A run saved at step k and continued in a fresh process reproduces the uninterrupted run bit for bit in the same arithmetic mode and thread count (and, in exact mode, on any machine). The caller supplies the data order (store the sample counter in extra).

7. Determinism

Fast mode: forward and backward are the engine's (autograd.md, thread-count independent); the optimizer, clipping, schedules, losses' constants, dropout and the checkpoint code are scalar, fixed-order code. Exact mode: all of it is bit-identical across ISA paths, OS and thread counts; checked by pinned hashes of the full training state after 20 steps of the conv net (AdamW + cosine + clipping) and of the transformer classifier with dropout (TrainingDeterminismTests), under default, DOTNET_EnableAVX512=0, DOTNET_EnableAVX2=0 and DOTNET_EnableHWIntrinsic=0, at 1 and 8 threads, and by CI on every platform. Fast-mode state hashes are equal at 1, 3 and 8 threads.

8. What was measured against PyTorch 2.15.0 (CPU, black box)

tools/training_oracle.py runs seeded trainings in PyTorch and writes the initial weights, data, batch order, per-step losses and final parameters to testdata/training/; the tests load the same weights and batches into our modules. Tolerance accepted: per-step loss relative difference ≤ 1e-3 and final parameter scaled error (max|Δ| / max|torch|) ≤ 1e-3 — the curves overlay far inside it.

Parameters whose true gradient is zero (the bias of a convolution in front of batch norm, the key bias of attention, which softmax ignores) receive only rounding noise as gradient, and Adam turns noise into steps of the full learning rate, so these differ between any two implementations (2e-2 and 4e-3 scaled here); they are excluded from the parameter comparison, and the batch-norm running mean that absorbs the first one is not gated (the running variance is). Single-step checks (5 steps, given gradients; scaled error against PyTorch): SGD variants ≤ 8.1e-8, Adam/AdamW variants (L2, AMSGrad, other betas) ≤ 3.4e-8; clipping ≤ 1.7e-7 with equal norms; learning-rate schedules (100 steps) ≤ 1.2e-7 relative.

9. Reduced precision: measured, not built

Question: does bfloat16, float16 or loss-scaled training pay on a CPU, on the kernels we have? Method (Micikevicius et al. 2018 describe the schemes): bench/Micro command mixed [threads] [steps] runs the float32 AdamW step of three models, reports the median step time (steps 3 on), the peak pool bytes sampled during the step, the activation bytes held at the end of the forward pass, and the time of one float32 → bfloat16 → float32 pass over all activations and parameters (a lower bound on what any reduced-precision scheme adds per step). Indicative numbers, loaded hosts, not a benchmark.

Findings.

• Time does not pay. The .NET 10 intrinsics expose no bfloat16 dot product on x86 (Avx512Bf16 is absent; the listed classes end at Avx10v2, AvxVnni*) or on Arm (no Bf16), although the worker's CPU has AVX512_BF16 and the M2 has FEAT_BF16. Without it, low-precision operands are widened to float32 before the multiply, so a reduced-precision matrix product does the same arithmetic plus the conversion; the products are compute-bound at these sizes, so halving operand bytes buys nothing. The best case is therefore a step time ≥ 1.0× float32's, never the 1.2× gain required.
• Memory pays only for the transformer. Saved activations are 84% of its peak, so holding every activation as bfloat16 would cut peak memory by up to about 42% (above the 30% bar); in the conv net and CRNN they are 41–42% of the peak and gradients of the same size follow in backward, so the saving is about 20% (below the bar). Parameters, gradients and optimizer state are 0.8–26 MB: irrelevant.
• Not built. The saving needs bfloat16 tensors flowing between operators (each operator widens its inputs into scratch, narrows its output, and its record saves the narrow tensor and widens it again in backward). Holding a bfloat16 copy beside the float32 original saves nothing, because activations are released together at the end of the step and a caller may still hold the float32 tensor. That change touches every operator's entry point and saved-slot handling, and is a separate piece of work worth doing only if long-sequence transformer training on a CPU becomes a real use. Loss scaling is a float16 remedy and has nothing to scale while gradients stay float32, so it waits for the same work.

10. Fine-tuning a CRNN line recognizer (measured against PyTorch)

The reproduction check for the training stack: fine-tune a convolutional-recurrent line recognizer (six conv blocks with batch norm and max pooling, a two-layer bidirectional LSTM of 192, a linear layer to 622 classes, CTC) for 2000 steps from a released checkpoint, once with a PyTorch trainer and once with this library, and compare. tools/CrnnFinetune is the driver (train, eval, export); tools/crnn_ref/ runs the PyTorch side as a black box and compares the runs.

Recipe. AdamW (lr 6e-5, weight decay 0.01, betas 0.9/0.999), 200-step linear warm-up from lr/200 then a half cosine down to 2% of the peak over the remaining steps (WarmupCosineLr), gradient-norm clip 5, CTC (blank 0, mean over lines, zeroInfinity), batch statistics in training mode, 128 lines per batch of mixed sources. Lines are 32 px high and padded to the batch's widest line by repeating the last column; the model sees only each line's true width (lengths below).

Library additions. Ops.Lstm/Ops.Gru and Recurrent.Forward(input, lengths) take per-row sequence lengths with packed-sequence meaning: the forward direction needs no change; the reverse direction runs as a forward pass over each row reversed within its length (a gather by an involutive permutation), so it starts at step length-1 from the zero state; outputs at padded steps are exactly zero and carry zero gradient (the layer output is multiplied by a 0/1 mask). Tests: every padded row equals the row run alone at its length, for LSTM and GRU, uni- and bidirectional, two layers; gradients by finite differences. WarmupCosineLr is base·(s+1)/W for s < W, then base·(f + (1−f)·(1+cos(πp))/2), p = (s−W)/(T−W).

Data. The trainer's augmentation draws from a GPU generator, which cannot be reproduced here, so the PyTorch trainer dumps each step's batch (after augmentation, rounded to bytes, padded) and both sides train on that one stream. The stream depends only on the seed and step, except that the targeted-error source is weighted by a measurement of the starting model, so a reference run reuses the first run's measurement file to get the same stream. Augmentation is therefore not ported; that is the remaining gap for training without PyTorch.

Results (host: a 32-thread Zen 4 machine shared with other jobs; ours on CPU, 8 threads, fast mode; references on an RTX 4070; the generated-photo source of the original recipe is left out, which changes the mix, not the comparison).

• First step. Loss 0.166759 ours against 0.166759 for PyTorch with TF32 off (relative 4e-7); gradient norm within 1e-5. Steps 0–100 agree to 1e-7 relative in the loss.
• Later steps. Two fp32 trajectories of a chaotic training diverge: the per-step loss differs by 5e-3 (median relative) between ours and strict-fp32 PyTorch over steps 0–999, and by 5e-3 between PyTorch with TF32 on and PyTorch with it off, so ours is as close to PyTorch as PyTorch is to a re-rounded PyTorch. Means over 200-step windows agree within 3e-3 (at steps 0–199, 6e-6).
• Quality at 2000 steps. Character error rate on the twelve validation sets, in percent, against the shipped trainer (bf16 autocast): val 1.642 / 1.620, receipt-style 7.775 / 7.810, gaps 4.509 / 4.472, interface 0.380 / 0.383, tight 1.370 / 1.366, receipt photos 2.163 / 2.173, mask 13.403 / 13.542, quantity 12.389 / 12.463, figures 5.291 / 5.313, small text 15.277 / 15.239, reader 1.758 / 1.749, shot 0.329 / 0.337 (reference / ours). Every set is within 5% relative (the largest is 2.3%, shot), and all but mask (0.139 pp) and quantity (0.074 pp), the two sets with the highest error rates, are also within 0.05 percentage points.
• Export. CrnnFinetune export writes the weights as a model container from its specification (batch norm folded in double precision, binary16 weights, nearest even). The consumer's own parity test, run unchanged, passes on the exported file and fails when a byte of a tensor is flipped.
• Cost. See the table at the end of this section.

So a CPU step here takes about twice PyTorch's, mostly in the convolutions' forward and backward; peak memory is 1.15× PyTorch's. The 2000-step run took 7.7 h wall-clock with queue waits; the work itself is far below what a GPU does, and no speed tuning was attempted.