Between the tape and a full trainer sits a toolkit every training loop reaches for: activations, dropout, losses, update steps, a learning-rate schedule, and the seeded starts. Each is one call, each records where the tape needs it, and each prints the same digits every time it runs: the seeded starts, the masks, MseLoss and the update steps on every machine, and Gelu, GeluErf, Silu, CrossEntropy and CosineSchedule to the last digit of the C library whose tanh, erf, exp, log and cos they call.

The example is examples/LearningToolkit in the tutorial repository.

Getting the file

make -C mica-container/examples/LearningToolkit run
LearningToolkit — the pieces of a training loop

relu: s = 3.500000 g = 1.000000 0.000000 0.000000 1.000000
gelu at zero: g = 0.500000 0.500000
silu at zero: g = 0.500000 0.500000
mse: l = 2.500000 g = 1.000000 2.000000

mask (seed 42, rate one half): 2.000000 0.000000 0.000000 0.000000
dropout: s = 6.000000 g = 2.000000 0.000000 0.000000 0.000000

sgd: w = 0.950000 1.900000 v = 0.500000 1.000000
clip: norm was 5.000000, g = 1.500000 2.000000
cosine: 0.000000 0.500000 1.000000 0.000000

xavier: -0.269861 -1.183622 0.981659 0.203137
he:     -0.164746 -0.867995 -0.111014 -0.595560

The modern activations

s := Sum(Relu(x));      { negative cells drop; the derivative is the saved sign }
s := Sum(Gelu(z));      { the tanh spelling the GPT lineage trains with }
s := Sum(GeluErf(z));   { the exact erf spelling, PyTorch's nn.GELU() }
s := Sum(Silu(z2));     { the sigmoid-weighted input of the llama family }

Relu, Gelu, GeluErf, Silu and Sigmoid join Tanh in the recorded vocabulary, each an elemental routine that applies to a vector, a matrix and a higher-rank tensor exactly as Tanh does. Sigmoid is the logistic curve 1 / (1 + e^(−x)) of a classic network’s layers, so a forward pass reads as its formula, h := Sigmoid(w * a), and its rule is y (1 − y), read from its own saved answer. GeluErf is the gelu PyTorch computes by default, x times the Gaussian’s cumulative distribution 0.5 (1 + erf(x / √2)), in PyTorch’s own order, so a model ported from PyTorch trains the same activation. Each records with its derivative rule: relu’s is the saved sign (an input at exactly zero answers zero on both sides), gelu’s and silu’s are their exact slopes — both pass one half at zero, and the printed gradients say so. RmsNorm conditions rows the way modern language models do: LayerNorm without the mean subtraction and the bias, with a learned gain and its own recorded rule. Both add a small epsilon under their square root, 1e-5 unless the call names its model’s own as a last argument: RmsNorm(x, gain, 1e-6) is the first llama’s, LayerNorm(x, gain, bias, 1e-12) is BERT’s.

Dropout is a pair of visible values

r := Seeded(42);
mk := DropoutMask(0.5, address r);

tape
    s := Sum(Dropout(m, mk));
end;

The mask is a value: drawn once from the seeded generator, each kept cell carrying 1/(1-rate) (the inverted scaling that keeps the expectation put), each dropped cell zero. You can print it — the example does. The recorded product saves the mask beside the activation, and the derivative multiplies by exactly it: the printed gradient is the printed mask. Inference is the same program without those two lines — no evaluation mode, no hidden flag.

Same seed, same mask, same digits — on every machine at every tier, because the generator is the random unit’s bit-stable splitmix64.

Or one call that draws its own mask

r := Seeded(42);

tape
    s := Sum(Dropout(m, 0.5, address r));
end;

When a mask is only ever applied, Dropout(x, rate, address r) draws it where x lives, cell by cell from the same generator in the same row-major order. Those are the very numbers DropoutMask draws, so the answer and the gradient are the two-line form’s to the bit, and the generator ends where DropoutMask would leave it.

Splitmix64 computes its n-th number from its starting state and n alone, so every cell draws its own. On a GPU the kernel draws inside itself: the mask never exists in memory and nothing crosses from the processor. On the tape the recorded form keeps only the generator’s starting state and the rate, and draws the mask again for the derivative. A rate outside [0, 1) fails at the call with DropoutMask’s own contract.

Attention in one call

o := CausalAttention(q, k, v, 0.1, address r);   { training: scores, mask, softmax, dropout }
o := CausalAttention(q, k, v);                   { inference }
o := Attention(q, k, v);                         { every key read, as an encoder does }

CausalAttention is the six lines a transformer head writes by hand: the scores q * Transpose(k), their scale by one over the square root of the key width, the causal mask, RowSoftmax, the dropout and the weighted values. On the processor it gives exactly the bits those six lines give, and with a rate and a generator its dropout draws the very mask Dropout would. The operands are matrices or higher-rank tensors; every axis before the last two is a batch.

On an NVIDIA card with tensor cores the call runs as one kernel each way, and the score matrix never exists in memory. Its gradient recomputes the weights instead of storing them. The kernels keep float32’s accuracy whether or not the program is built with fp_tensor_float32: that flag lets the program’s matrix products read TF32, and the attention multiplies three times over each operand’s split parts, as PyTorch’s float32 attention does.

Many heads, one call

q3 := Reshape(a * Wq);                          { one product each for every head }
k3 := Reshape(a * Wk);
v3 := Reshape(a * Wv);
o := Reshape(CausalAttention(q3, k3, v3, Heads, 0.1, address r));
x := x + o * Wo;                                { one output product }

A head count after the values says the rows hold that many attentions side by side. With 384 columns and six heads, each head reads its own 64 columns of the queries, keys and values and writes its own 64 columns of the answer. That is exactly the layout one wide product gives, so a transformer layer needs three products for all its heads, one attention call and one output product. On a GPU a reshape of a value its statement computes takes that value’s memory over, so nothing is copied to split or join the heads; a reshape of a variable copies it, since the variable keeps its value. The call gives the bits of six calls one head at a time, dropout included. The count is a literal or a named constant, and it must divide the widths.

Queries, keys and values from one product

qkv := Reshape(a * Wqkv);                       { one product for all three }
o := Reshape(CausalSelfAttention(qkv, Heads, 0.1, address r));
x := x + o * Wo;

Wqkv holds the three projections side by side, 384 by 1152, as GPT-2’s c_attn holds them, so one product gives every row its query, key and value side by side. CausalSelfAttention reads each row’s first third as the query, the next as the key and the last as the value, the heads side by side within each third, and gives the bits CausalAttention gives over the three thirds taken apart. SelfAttention is its twin that reads every key. On a GPU the kernels read the thirds where they stand, so nothing is copied apart, and one product does the work of three. The recorded form writes the three gradients side by side into the one gradient.

The losses

CrossEntropy you know from the GPT. MseLoss stands beside it for regression: the mean over all cells of the squared difference, recorded, with the exact derivative 2·(a−b)/n flowing to both operands. Its views take a vector, a matrix, or a higher-rank tensor read flat.

The update family

SgdStep(address w, g, address v, 0.1, 0.5);         { v ← μ·v + g;  w ← w − α·v }
AdamWStep(address w, g, address m, address v,
          step, α, β₁, β₂, ε, decay);               { AdamStep + decoupled decay }
AdamWSteps(step, α, β₁, β₂, ε, decay,
           w1, g1, m1, v1, w2, g2, m2, v2);         { every parameter tensor in one call }
nrm := ClipGradNorm(address g, ceiling);            { cap the length, keep the direction }
nrm := ClipGradNorms(ceiling, g1, g2, g3);          { one cap over the joint length of all three }
ScaleTensors(inverse, g1, g2, g3);                  { every cell of all three times one factor }
α := CosineSchedule(step, warmup, total, peak);     { the warmup ramp + half cosine }

All seven are plain code over plain values — they run outside the window, on gradients you extracted with Gradient. AdamWStep at zero decay is bit-equal to AdamStep; ClipGradNorm answers the norm it measured before any clipping, the number the field logs; and the schedule is a function you call once per step — the learning rate is a variable in your loop, not an object’s hidden state.

ClipGradNorms is the clip a whole model takes before its update: it measures the joint norm of every gradient variable you name, scales them all in place by one factor when that norm exceeds the ceiling, and answers the norm. Each gradient is measured on the same reduction tree as ClipGradNorm, so a call naming one gradient answers exactly what ClipGradNorm answers for it. On a GPU the call waits for the card once, however many gradients it names, and not at all when its norm stays on the card, as below.

AdamWSteps is the update a whole model takes. Behind the step count and the five hyperparameters come groups of four variables, one group per parameter tensor: the weight, its gradient and its two moments. Each group is updated exactly as AdamWStep updates it, so the call gives the numbers of one call per tensor. A model whose gains take no decay makes two calls, one with the decay for its matrices and one without it for its gains. On a GPU the call is one kernel launch however many tensors it names, and the steps compute in float32 for float32 tensors.

ScaleTensors multiplies every cell of each tensor variable you name by one factor, in place, so each holds the numbers g1 := g1 * inverse would leave in it. A step that scaled its loss to keep float16 gradients from vanishing unscales them with it before the clip and the update read them, as PyTorch’s GradScaler does. On a GPU the call is one kernel launch however many tensors it names, where a statement per gradient launches once per gradient.

On a GPU a model’s weights may be tracked variables themselves: the tape records them directly, and the steps update them in place between the windows that record them.

w : tracked matrix[384, 384] of float32 on gpu;

AdamWStep(w, Gradient(w), m, v, step, α, β₁, β₂, ε, decay);

A device step names its weight, moments and velocity bare, the variables it updates in place, and ToHost(w) reads a tracked weight’s value for a checkpoint.

The loss and the norm on the card

loss : tracked float32 on gpu;
nrm : float32 on gpu;

loss := CrossEntropy(lg, targets);        { in the window, answered on the card }
nrm := ClipGradNorms(1.0, g1, g2, g3);    { measured, decided and scaled there }
lv := ToHost(loss);                       { the host waits here, and only here }

Reading a loss or a norm on the host makes the host wait for everything the card has queued. Declared on the card, as one cell of a float element, they stay there: CrossEntropy, ClipGradNorm and ClipGradNorms answer into a device scalar when the call stands whole on the right of the assignment, and a tracked device loss walks back as a host loss does. ToHost reads the value where the program chooses to wait, at a report step for instance. A device scalar assigned straight to a host variable is refused, naming ToHost.

Over bfloat16 or float16 logits, CrossEntropy answers into a float32 device scalar as well: the loss is computed in float32 over the half logits, as PyTorch’s autocast computes it, and no float32 copy of the logits is made. At GPT-2’s vocabulary that copy would be 1.6 GB a step.

scaled : tracked float32 on gpu;

scaled := loss * scale;                   { in the window, on the card }
Backward(scaled);                         { every gradient scaled by 'scale' }
g1 := Gradient(w1) * inverse;             { read back unscaled }

A device scalar takes one arithmetic on the card: its scaling by a host value of its own element type, with the value on either side of *. That is the loss scaling a float16 step needs, the multiply PyTorch’s GradScaler puts in front of backward(): the scaled loss walks back to gradients scaled by the same value, which the program reads back unscaled. Any other arithmetic reads the value with ToHost first.

The seeded starts

r := Seeded(7);
w1 := XavierInit(address r);    { uniform in ±sqrt(6/(R+C)) — Glorot's bound }
w2 := HeInit(address r);        { uniform in ±sqrt(6/R) — Kaiming's, fan-in = rows }

Both draw from the caller’s own generator in row-major cell order, so a training run’s starting point replays exactly — the reproducibility promise extended to the very first step. The fan-in is the row count because this library multiplies activations from the left (h * w): a weight’s rows face the incoming features.

The elementwise pair, by name

* between matrices is the linear-algebra product, so the elementwise product is written ⊙, g ⊙ u, with Hadamard as its name, and the quotient has a name alone: Divide — and, like every elemental verb, both admit a static higher-rank operand through the trailing-axis flattening. Recorded, the product’s derivative multiplies by the other side’s saved value, and the quotient carries the textbook rule. The llama family’s SwiGLU gate is exactly this pair of lines:

g := Silu(x * wg);              { the gate branch }
h := Hadamard(g, x * wu);       { gated elementwise - the SwiGLU sentence }

And the quotient is what keeps a batch normalization’s inference scale one visible sentence — s := Divide(gain, Sqrt(rv + eps)) — the line the Image networks chapter builds on.

Where this leads

The batched trainer uses every piece on this page in one program, and checks the whole run against an independent oracle.