Pretraining

How I trained a 126M model on 5 billion tokens in 19 hours, and then a 304M model on 10 billion tokens in about 90, on one RTX 5070 Ti.

By the end of stage 5 the trainer worked from start to finish on a small model: forward and backward passes in PTX, every fast kernel checked against a naive one, gradients checked with finite differences, AdamW, checkpoints. Pretraining is where all of that has to keep working for days in a row.

The presets

Every hyperparameter comes from a config file, so a different model is a different argument: bin\train main.

They all share the same 32,768-token vocabulary and the same Llama-style layout: RoPE, RMSNorm, SwiGLU, grouped-query attention, no biases, and the input embedding doubling as the output layer.

One step

An optimizer step is 524,288 tokens. The 304M only fits four sequences of 2,048 tokens on the card at once, which takes 14.9 of its 15.9 GiB, so one step is 64 micro-batches. Each one runs forward and backward, and the gradients pile up in fp32. Then AdamW updates the fp32 master weights once. The matmuls run on the tensor cores in bf16 with fp32 accumulation (mma.sync). The learning rate warms up for 1,000 steps to 5e-4, then follows a cosine down to a tenth of that.

Watching it

Every 10 steps the trainer redraws a progress bar in place in the console, and writes the same numbers to a log. Here are two lines from the 304M run’s log, one from halfway through and one from the very end, cut down to the main fields:

step=9530  loss=3.0273 val=2.9821 tok/s=29576 tflops=62.89
step=19070 loss=2.7597 val=2.8406 tok/s=31914 tflops=67.87

Checkpoints

A run this long gets interrupted, by me or by Windows. Ctrl+C lets the current step finish, saves and exits, and a second Ctrl+C quits on the spot. It also saves every 1,000 steps or 30 minutes and keeps the last three.

A checkpoint has everything an exact resume needs: the weights, both of Adam’s moments, the step, the learning rate schedule’s position, the random number generator and where the data loader was. It’s written to a temporary file first, and only once that’s flushed to disk does a single rename put it in place. A crash halfway through a save leaves the old checkpoint untouched:

mov rcx, rbx
call file_flush
mov rcx, rbx
call file_close
mov rcx, r13
mov rdx, r12
call file_replace
test eax, eax
jz .bad

There’s a test for it (test\resume): it stops a run partway, resumes it, and checks that the weights come out bit for bit the same as a run that never stopped.

The two runs

The 126M took 9,537 steps and 19.3 hours, at about 72,000 tokens a second (around 60 TFLOPS), and finished at a validation loss of 3.101.

The 304M took 19,074 steps and 90.4 hours, with one forced exit that cost about 94 steps, and one planned pause halfway to swap in a faster backward pass (more on that in chapter 9). It ran at 29,600 tokens a second on one CUDA stream and about 32,000 after the swap, which is 68 TFLOPS. By 5 billion tokens it had already passed where the 126M finished, and it ended at 2.841.

The last sample the 304M wrote on its own, from the prompt “The history of”:

The history of Canada, spanning more than 65 years, has been inextricably linked to the country’s development. A country beset by political instability, it has served as a hub for various international organizations, including the World Bank, the United…

It’s fluent and confidently wrong, which is about what you’d expect from a base model this size. Making it answer questions is the next chapter.

A stack that was 8 bytes off

The bug that cost me the most time wasn’t in a GPU kernel. One routine in the trainer pushed four registers and then subtracted 32 from rsp, so everything it called ran with the stack 8 bytes off the 16-byte alignment Windows expects. Nothing noticed for a long time. Then, 55 million tokens into a dev run, the data loader crossed into the next token file and called CreateFileA from there. The fault happened inside a system DLL, so my crash handler, which only reports crashes in my own code, stayed quiet, and the console window just vanished.

I found it by cutting the data into three tiny files, so the loader switched files every few steps and the crash came right away. Since then I check every function against one rule: pushes × 8 + the sub has to come to 8 modulo 16.