Mica 7.7: a training step on the card, in every float type

Mica 7.7 trains AI models on the GPU end to end, in every float type they train in. A model’s whole training step runs on an NVIDIA card in one Mica program: attention over its heads side by side on the tensor cores, the layer normalization, the softmax, the cross-entropy, the GELU and the dropout, the AdamW update and the gradient clip, each a verb of the language, with the gradients the tape records.

float16 and bfloat16 stand beside float32 and float64 everywhere a float goes, on the card and on the processor, and a model trains in bfloat16 while its weights stay float32. The numerics are the program’s to state, a GPU program carries its kernels precompiled for five card generations, and the one who starts a program chooses the card it runs on.

The release at a glance

A training step on the card

A transformer layer is a handful of statements:

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 queries’, keys’ and values’ projections side by side, as GPT-2’s c_attn holds them, so one product gives every row its query, key and value. CausalSelfAttention reads the heads side by side within each third, scales, masks, takes the softmax, draws the dropout and weighs the values, on NVIDIA’s tensor cores, as one kernel each way: the score matrix never exists in memory. CausalAttention, Attention and SelfAttention are its siblings, and every one of them gives on the processor exactly the bits its six hand-written lines give.

The layer normalization, the row softmax and the cross-entropy take each row with a lane group of the card, and where a dropout, the residual add of its answer and the layer normalization of the sum follow each other, the compiler runs them as one pass each way. Its remarks say where it did and, where it did not, what stood in the way.

The whole model in one call

nrm := ClipGradNorms(1.0, g1, g2, g3);          { one cap over the joint length }
AdamWSteps(step, α, β₁, β₂, ε, decay,
           w1, g1, m1, v1, w2, g2, m2, v2);     { every parameter tensor at once }

A model clips its gradients by their joint norm and updates every parameter in one call each. On a GPU each call is one kernel launch however many tensors it names, and each group of a weight, its gradient and its two moments is updated exactly as AdamWStep updates it alone. ScaleTensors multiplies every cell of several tensors by one factor, the unscaling a float16 step needs before its clip.

Values that stay on the card

w : tracked matrix[384, 384] of float32 on gpu;
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 }

A model’s weights may be tracked device variables, which the optimizer steps update between the windows that record them. A loss and a gradient norm declared on the card stay there until the program reads them with ToHost, so the processor waits for the card only where the program says. loss * scale scales a device loss on the card and on the tape, the loss scaling of a float16 step. The tape lends a gradient that reaches an operand unchanged, borrows the values its rules read instead of copying them, and gives each gradient’s memory back to the card once its walk has passed it.

float16 and bfloat16, whole

The two half types now answer every question a float type answers. Vectors, matrices and tensors of them compute on the processor’s lanes and on the card, record on both tapes, convert with as whole tensors at a time, fold at compile time, are guarded on the checked tier, and read as themselves in the debugger.

A model trains in bfloat16 the way the industry’s mixed precision does: its weights stay float32, and a product reads them as bfloat16 where it needs them.

tape
    zb := xb * (w as bfloat16);                 { the product reads the half }
    loss := CrossEntropy(zb as float32, targets);
end;

Backward(loss);
g := ToHost(Gradient(w));                       { float32, as the weight }

The conversion is recorded on the tape, so the weight’s gradient arrives in float32. Over half logits, CrossEntropy computes its loss in float32 without making a float32 copy of them.

The numerics a program states

Every choice that changes a program’s digits is the program’s to make, at build time or at its start:

• --optimize O2,fp_tensor_float32 lets float32 products read their inputs in TF32 on the tensor cores, nanoGPT’s setting, while an attention keeps float32’s precision.
• fp_half_intermediates lets a half operation hold its intermediates in the half, PyTorch’s reduced-precision reduction and FlashAttention’s numerics.
• LayerNorm, RmsNorm and BatchNorm take their epsilon as their last argument, 1e-5 where a program leaves it out, so a model normalizes as the same model written for PyTorch.
• +mica products=FILE -mica runs each float32 or half product by the routine measured fastest on the card for its shape and keeps the choice, so every later run repeats its digits.

Kernels ready at the start

A GPU program’s kernels are compiled at its link for the card generations sm_80, sm_86, sm_89, sm_90 and sm_120, and the program carries them. On a card of those generations it loads them when it starts, and NVIDIA’s runtime compiler never runs. +mica info -mica says whether it will:

"device":{"provider":"nvidia","card":0,"name":"NVIDIA GeForce RTX 4090","generation":"sm_89","precompiled":true}

An optimizing build queues its kernels and products on the card’s stream, as PyTorch and JAX queue theirs, and the debug tier waits for each, so a failure is named where it happened.

The card a program uses

./train +mica check -mica           # lists the machine's cards
./train +mica gpu=1 -mica           # trains on the second

+mica gpu=N -mica chooses the card when the program starts. check lists every card by the number the flag names it by, and info names the one in use. A number with no card behind it is refused at the start, before any of the program’s work. A release build admits the flag too, because which card a program uses is a fact of the machine it runs on, not of its build. The runtime flags page has the details.

On the processor

A float loop around a square root or a fused multiply-add keeps its values in registers, as a C compiler’s loop does, and a float constant reaches its register the way GCC and Clang bring it there: built in place where an instruction or two can, and loaded from a read-only pool where they cannot. A two-word integer shifts inline. A large function compiles in a third less time and half the memory, and a program whose static storage lies past what its code can reach is refused by the compiler in its own words, with the way out: allocate the largest arrays with new, or place them on the device.

The editor 5.2

The extension’s assembly and Spectra views know the instructions 7.7’s code emits, the double shifts, the conditional move and the half-precision conversions, and the Spectra view’s hover names the block barrier a device kernel waits at.

Upgrading

Three things a 7.6.5 program meets differently in 7.7:

• Every normalization adds an epsilon under its square root, 1e-5 unless the program states another as its last argument, so its answers are the same model’s in PyTorch.
• An optimizing build reports a failure of a queued GPU kernel at the next act that waits, with a note saying so; the debug tier names it where it happened.
• A float operation whose answer is a not-a-number runs instead of folding at compile time, so the processor chooses the not-a-number’s bits.

The packages and the commands that install them are on the install page.