Training a GPT-2-style model in Lean

Robert Joseph George · TorchLean Verified Examples

Okay, let's train a GPT in Lean :) I wanted to take TorchLean beyond small examples and put the model, data loader, autograd, optimizer, and checkpoint code to work together. For this experiment, I trained all 124,412,160 stored parameters of a GPT-2-style model on one A100.

minGPT and nanoGPT provided the recipe: a causal transformer, next-token prediction, AdamW, and a warmup-and-cosine learning-rate schedule. The experiment uses the main dimensions and tied token embeddings from Radford and colleagues' GPT-2 report. It was trained on FineWeb-Edu, with its own initialization and weights; our current model also uses bias-free query, key, and value projections. So this is a GPT-2-style architecture trained from scratch, not OpenAI's released GPT-2 checkpoint.

The model learned: validation loss fell during pre-training and improved with more training. I also tried turning it into an assistant, with much less success. The instruction-tuning losses went down, but it still got elementary questions wrong. We'll look at both, then compare 50 matched updates against PyTorch and go through the proofs for data alignment, resuming a run, and cached generation.

Karpathy describes minGPT as trying to be:

small, clean, interpretable and educational

I want to keep that quality in the Lean version too. We should be able to follow a batch from its token IDs to the loss and update, then open the proof of a specific property of that path. The large run tests whether those pieces work together at scale; the smaller calculations below help us understand what they are doing.

The model and one training step

The large preset has 12 transformer blocks, hidden width 768, 12 attention heads, and a 50,257-token vocabulary. Each head works with 64 coordinates, since 768 divides evenly among the 12 heads. The feed-forward part expands to four times the hidden width before projecting back. These dimensions describe the model independently of how many sequences we put in one batch.

Start with one sequence of token IDs. Embedding lookup turns each ID into a width-768 vector, and a learned position vector is added. Each block normalizes those vectors, computes causal attention, and adds its output to the residual stream. A second normalization and a feed-forward network produce another residual update. After the final block and normalization, the language-model head produces one score for every vocabulary item at each position.

The residual stream is the hidden sequence carried between blocks. A block adds an update to it rather than replacing it outright. Attention mixes information between allowed token positions; the feed-forward network transforms each position's vector. Normalization rescales a vector before these operations using its mean, variance, and learned scale and shift. Those operations repeat throughout the stack, with different learned weights at each block.

The input embedding and output projection share a parameter table. The table's rows provide input vectors; its transpose maps the final hidden vector to vocabulary scores. Training can therefore change the same parameters through either use. Keeping one shared table in the model definition avoids treating the two copies as independent weights in the optimizer or checkpoint.

For a next-token target yₜ, the cross-entropy loss at position t is the negative log probability assigned to that token:

Here zₜ,ᵥ is the raw score for vocabulary item v. Increasing the target's score relative to its competitors lowers the loss. A correct top-ranked token can still incur loss if its probability is small; conversely, the average loss can improve while particular questions remain wrong. This becomes relevant when we inspect the instruction-tuned model.

The gradient of this real-valued loss with respect to one vocabulary score is \(p_{t,v}-\mathbf{1}_{v=y_t}\), where \(p\) is the softmax probability and the indicator is one only for the target. If the target has probability 0.6, its score derivative is −0.4; a competing token with probability 0.1 has derivative 0.1. Gradient descent would raise the target score and lower that competing score if these scores were independently adjustable parameters.

In the model, they aren't independent parameters: they come from the shared embeddings and transformer blocks. Backward follows those dependencies to find parameter derivatives. Changing a weight can affect many positions and many scores at once. This is why a sensible local loss derivative doesn't guarantee that every question improves after an optimizer step. The formula describes exact real arithmetic; the training run computes rounded versions.

One update begins by choosing token windows and their targets. The model computes logits, the objective reduces them to a scalar loss, and our automatic differentiation (autograd) tape propagates derivatives back to the parameters. Lean controls the training loop and backward traversal; on GPU, the operation adapters call LibTorch's numerical primitives and local gradient operations with its autograd engine disabled.

AdamW, from Loshchilov and Hutter, uses the resulting gradient together with two running statistics for each parameter: a first moment, which averages recent gradients, and a second moment, which averages their squares. Bias correction accounts for their initialization at zero. The learning rate sets the update size, and decoupled weight decay contributes a separate shrinkage of the parameters. Those running statistics are why a training checkpoint needs more than the current weights.

The learning rate first rises during warmup, then decreases on a cosine schedule. Its input is the optimizer-step index. Restarting that index after loading a checkpoint would repeat the early schedule at a different parameter state. We'll return to the proof about global indices after looking at the completed runs.

Pre-training on FineWeb-Edu

The first run used a 1,024-token context and batch size 8. It processed 1,000,005,632 scheduled tokens over 122,071 updates, taking about 46 hours and 35 minutes on one A100. Validation loss fell from 10.963 at initialization to a best recorded value of 3.308 at step 107,000.

I continued from the saved checkpoint at step 100,000. That was the best checkpoint available on disk, with validation loss 3.333; the lowest evaluation at step 107,000 fell between saves. The distinction matters when counting the training history. The checkpoint at step 100,000 had seen 819,200,000 tokens. Adding 1,500,008,448 continuation tokens gives a total of 2,319,208,448 for the parameters along this history.

The continuation used a new FineWeb-Edu shard, batch size 12, and a lower learning rate. Its best recorded validation loss was 3.135083 at continuation step 122,000. The final evaluation, 71 updates later, was 3.273021. The best measurement again fell between checkpoint boundaries; the final parameter file contains the later state.

All of those measurements belong to the archived strict-FP32 runs. The current runtime uses LibTorch and retains Lean Float's binary64 precision, including optimizer and cache state. The old loss curves and timings describe the completed experiments; a new run with the current adapter has different precision and may have different rounding and speed. The configurations and records are in the development notes.

The token budget counts scheduled training positions. With batch size 8 and context length 1,024, one update schedules 8,192 positions. Multiplying by 122,071 gives the first run's 1,000,005,632 tokens. The runner chooses a whole number of updates, so a requested budget can be exceeded by less than one batch. This total measures positions processed, not a claim that every token came from a different document or appeared exactly once.

The checkpoint history is equally concrete. Updates after step 100,000 in the first run did not contribute to the continuation initialized from that earlier checkpoint. Adding the complete first-run total to the continuation would count updates that the continued parameters never received. I use the saved step's token count instead, which is why the total along this parameter history is about 2.32 billion rather than 2.50 billion.

Validation evaluates held-out windows without updating the model. A recorded loss is still a measurement on those chosen windows. In the instruction-tuning runs below, validation uses a fixed deterministic sample of eight batches, not every record in the shard. Checkpoint saving and evaluation also happen at different intervals. A best observed loss can identify a useful point in the curve without leaving a parameter file for that exact step.

Lower loss, unreliable answers

The next stage trained the model on a User:/Assistant: transcript format. The first round scheduled about ten million token rows from SmolTalk. The second started from that round's best saved checkpoint and scheduled about fifty million rows from a SmolTalk2 mixture. Loss was charged only to assistant targets, while the prompt remained available as context.

Within each round, the sampled validation loss fell: from 2.673017 to 2.135422 in the first, and from 2.818384 to 2.216872 in the second. The validation shards differ, so comparing the two final values would mix a model change with a dataset change.

The lower validation loss did not give us reliable answers to simple prompts. The prompt suite contains the recorded outputs. The checkpoint named Paris as the capital of France and summarized a sentence about freezing water correctly. It got 7 + 5 wrong, gave an incorrect explanation of a blue sky, refused a harmless rewrite, and produced malformed Python. Recognizing the positive sentiment of a film review did not make it obey the requested one-word answer format.

I also checked an actual validation record containing the question about Bartolomeo Cristofori. Its prefix matched the generator's transcript and its target was Piano; the model still answered incorrectly. That check ruled out a transcript mismatch for this case. The objective had improved, but this pretraining and instruction-tuning recipe did not produce a dependable assistant at this model size.

For training, a supplied answer gives us all of its preceding tokens as context while predicting the next one. During free generation, the model has to continue from its own predictions. One poor choice can lead to a prefix unlike the validation examples. The loss and prompt suite therefore inspect different behavior, even before we consider that a numeric answer or a valid Python function may require several correct tokens in a row.

I don't yet know which changes would fix each failed answer; this experiment doesn't isolate their causes. For now, I want to keep those outputs beside the loss curves. The curves improved, and the model still got 7 + 5 wrong. So... I'd keep the actual answers beside the loss plot, rather than treat either one as the whole evaluation.

Getting the targets and loss mask right

A training row predicts one token ahead. A window contains input tokens, and its targets are the same sequence shifted by one:

That shift becomes easy to get wrong when the loss is restricted to assistant responses. If row t predicts corpus token t + 1, the mask for that row must say whether token t + 1 belongs to the answer. Using the mask attached to the input token would train a different set of predictions. CausalTraining.lean proves the next-token and mask indexing equalities for the Lean definitions.

Written in exact arithmetic, the weighted objective has the form

For assistant-only training, prompt and padding targets receive zero weight, and active assistant targets share the normalized weight. A zero direct loss weight does not mean the prompt has no effect on learning. Attention reads its representations when predicting the answer, and gradients can flow through those uses. The mask controls which predictions we ask the model to get right, not which preceding tokens it can see.

The loader gives every active target in the batch the same weight. If one response has two active targets and another has six, each target receives 1/8, so the responses contribute one quarter and three quarters of the loss respectively. This is a token average, rather than an equal average over responses. The implementation counts the enabled rows, forms a floating-point reciprocal, and converts those weights to the model's scalar type. That numerical calculation includes rounding. Padding receives zero weight regardless of how many padding rows were needed to form the tensor.

Suppose an input row ends at the final prompt token and its target is the first answer token. That prediction should count toward the answer loss. If we used the input token's prompt mask, we would discard it. At the other end, an answer input whose target lies beyond the answer interval should not acquire weight merely because the input belonged to the answer. Shifting the mask with the target handles both boundaries.

There was a concrete packing bug in an earlier instruction-tuning run. Sampling arbitrary windows from a concatenated dialogue tape allowed a window to begin inside an answer or cross into another conversation. I replaced that sampler with bounded dialogue records and reran both stages. The losses above come from the repaired runs. Each record describes a token window and an assistant-target interval inside it; long answers can use several records with overlapping context. The loader checks that target intervals cover active targets once and stay within their own windows.

The Lean results in DialogueRecords.lean connect the executable target selector to that interval definition and exclude padded rows after a record ends. In Batching.lean, we use Lean's vcgen tactic to prove that the packing loop appends each input, target, and mask bit in row order. These proofs cover the pure index and packing functions; byte decoding and device execution still need their own checks.

The record has two roles. Its window provides the context available to the model; its target interval says which predictions belong to this training example. A long answer can need several windows because one context cannot hold it all. Overlapping context lets a later window see earlier answer tokens, while disjoint target intervals avoid charging loss twice for the same answer token. The loader validates those intervals before it assembles the batch.

The packing proof follows the assembled contents, not only their lengths. Starting from arbitrary initial arrays, appending rows preserves the input, successor target, and mask bit at each appended position. A shape check alone could accept a batch with two target rows exchanged. The content invariant would require those exchanged values to equal the intended row functions, which is the stronger property we need here.

Causal attention handles the corresponding constraint inside the model. An input position must not read the future target during training. For the hard-mask specification, the Lean theorem gives every strict-future attention coordinate zero forward weight and zero score derivative, for arbitrary scores and incoming weight gradients. That is a local statement about the attention equations, with the numerical backend responsible for implementing them.

In the attention figure, select “sat.” Its position may read “The,” “cat,” and itself, but not “on.” All four positions can be computed in a training batch because this mask enforces their different allowed contexts. Parallel computation does not mean the prediction at an earlier position receives the later target. The real-valued hard-mask theorem also rules out a backward score contribution at those forbidden coordinates.

Fifty matched updates with PyTorch

I wanted to check the training calculation against PyTorch too. I reconstructed the Week 3 model in PyTorch from the same TorchLean checkpoint and ran both trainers on the same 50 batches. Dialogue records, assistant-row weights, learning-rate schedule, and AdamW settings matched. Both runs used strict FP32 on the same A100, with TF32 disabled; PyTorch used scaled-dot-product attention (SDPA) and fused AdamW. TensorFloat-32 (TF32) is an NVIDIA matrix-computation mode with less input precision than full FP32; we disabled it to avoid mixing those arithmetic choices into this comparison.

A matched comparison needs agreement before the first update. That includes the parameter ordering, shared embedding table, normalization constants, and sampled data. It also needs agreement on the state carried between updates, especially AdamW's moments and step count. Matching only the model dimensions would leave many possible reasons for the curves to differ.

Precision has to be specified too. The archived comparison used FP32 on both paths and disabled TF32. The current scripts match the newer binary64 runtime. Week 3 also retains the original checkpoints' different LayerNorm epsilon settings on CPU and CUDA, so a cross-device comparison must account for that choice. The archived graph measures the configuration recorded with it; it should not be relabeled as a result for the new adapter.

The largest absolute training-loss difference was about 0.00001193. Final validation losses were 2.431321 in TorchLean and 2.431322932 in PyTorch. This differential test exercises the checkpoint layout, forward pass, target selection, loss, backward pass, and optimizer update together. The close curves give me evidence that those pieces agreed on the measured run. They do not establish agreement for every batch or explain the checkpoint's remaining errors.

PyTorch also ran faster. Its median update throughput was about 19,207 scheduled tokens per second, compared with 5,596 for TorchLean, a ratio of 3.43 on this setup. Both timings include the full training update. PyTorch used SDPA and fused AdamW, so several implementation choices changed together. We would need a backend trace to identify its selected attention kernel, and timings of the individual operations to work out where the speed difference comes from. The comparison record contains these values.

Resuming with the optimizer state intact

A long run needs more than a parameter file. AdamW carries first and second moments; the learning-rate schedule and sampled data windows depend on the step number. A stochastic continuation also needs its random-operation state. Reloading the weights while restarting those other fields changes the sequence of updates.

TrainingCorrectness.lean gives the indexed loop a small pure semantics. An update receives both the state and global step number. Split the run after m steps, restore that exact state, and continue for n steps with the same global indices:

resume_eq_uninterrupted proves this for any step-indexed update function. It tells me what the checkpoint path must preserve. The serializer and restoration of device buffers are outside that pure theorem. A separate small CUDA resume check, made before the LibTorch migration, compared the final parameter and optimizer files byte for byte after a split run with dropout; both matched the uninterrupted run.

The underlying definition is short. This excerpt from TrainingCorrectness.lean leaves the state type and update function general:

/-- Execute `count` deterministic training updates beginning at global step `start`. -/
def runSteps (update : Nat → State → State) : Nat → Nat → State → State
  | _start, 0, state => state
  | start, count + 1, state =>
      runSteps update (start + 1) count (update start state)

The zero-step case returns the state unchanged. The other case applies update at start, then runs the remaining steps beginning at start + 1. Induction over the number of steps in the first part proves the split-run equation: peel off one update, then apply the induction hypothesis to the remaining first part and the second part.

State is deliberately general because the equation concerns how we compose updates. For a training run, it must represent everything that influences those updates. Merely proving this equation with a state containing only weights would not justify forgetting the optimizer moments. Likewise, any random stream that changes the updates has to be included in the restored state or otherwise fixed by the update function's inputs.

Let's make the step index visible with a small update: add step + 1 to a natural-number state. Four updates give 1 + 2 + 3 + 4. Resuming after the first two gives the same answer, provided the suffix starts at step two.

import TorchLeanGPT.TrainingCorrectness

open TorchLeanGPT.Training

-- A small indexed update makes a reset of the global step visible.
def addStep (step value : Nat) : Nat := value + step + 1

#eval runSteps addStep 0 4 0
#eval runSteps addStep 2 2 (runSteps addStep 0 2 0)
#check runSteps_add

10

10

TorchLeanGPT.Training.runSteps_add.{u_1} {State : Sort u_1} (update : ℕ → State → State) (start first second : ℕ)
  (state : State) :
  runSteps update start (first + second) state =
    runSteps update (start + first) second (runSteps update start first state)

The optimizer memory can change the very next update even when the restored weights are identical. A one-dimensional momentum example makes that visible. Take the rule

At \(w_t=1\), \(m_t=2\), and next gradient three, the new momentum is 4.8 and the weight is 0.52. Reload the same weight but reset momentum to zero, and the next weight is 0.7. Both checkpoints begin with the same model output, yet their training continuations differ. This toy rule illustrates state restoration; the GPT run uses AdamW, whose two moment buffers and bias-correction index create their own restoration requirements.

The resume theorem works with the complete state without expanding every AdamW coordinate to prove how adjacent update sequences compose. It gives the serializer a precise obligation: the loaded state must equal the state reached at the split point, and the continuation must use the same global indices. The earlier byte-for-byte experiment checks one serializer and runtime configuration; the general equation explains why restoring those values preserves the continuation for any indexed update function.

Reusing the KV cache

The key/value (KV) cache stores earlier attention projections. The full-prefix generator reevaluates the complete padded prefix after each token. Most of that work repeats: earlier tokens have already produced their attention keys and values. The cached generator retains those rows and adds the projections for the newly consumed token.

The cache proof states exactly what must be stored. A represented cache contains the key and value projections of the consumed history, in order. One incremental step appends the new rows, returns the reference output, and preserves that invariant. Induction over a suffix gives CachedDecode.run_correct; the layered result carries it through the complete causal stack. These results use the same abstract attention function on both paths.

At one layer, the new token has a query and contributes a new key and value. Its query reads all the retained keys; the resulting weights combine the retained values. Appending the new key/value pair in the correct order gives the next step the same history that the full-prefix calculation would reconstruct. The cache stores the projections rather than the final vocabulary logits because those are the values future attention needs.

A multi-layer cache contains separate projections for each layer. A key from the first block cannot stand in for a key from the twelfth: those blocks have different parameters and hidden inputs. The proof carries the representation relation through the stack, so each incremental block sees the cached values for its own reference calculation.

The causal condition explains why retaining earlier entries is valid. An earlier hidden position cannot depend on a token that has not arrived yet. In deterministic inference, with fixed parameters and the same positional convention, processing a new suffix does not require recomputing those earlier projections. This argument applies to the represented history; dropping old tokens or changing their position indices needs its own cache relation. The model's 1,024-token context does not become unlimited because a cache exists.

Different numerical kernels can still produce different logits from that logical computation. The generation theorem handles this with a margin condition. If the cached logits are within ε of the reference logits and the reference winner has margin greater than 2ε, both choose the same next token. Applying that condition at each prefix reached by the reference continuation proves equality of the whole greedy suffix. The current runtime checker reports observed errors and margins; it does not construct the Lean certificates required by that theorem.

We only need certificates along the reference continuation we reach. After certifying one choice, both generators append the same token and receive the same next history. We can then use the certificate for that history. There is no need to enumerate all vocabulary branches of an enormous generation tree, although obtaining the numerical bound at each reached prefix remains a separate task.

This connects the cache result to the first post's margin argument. The cache proof checks what the saved keys and values mean. A numerical error bound and winning margin then connect possibly different implementations of that attention calculation to the same greedy tokens. Neither a well-shaped cache nor one agreeing output alone establishes the complete suffix.

In three archived warm runs, both generation paths emitted the same eleven token IDs before ending. Cached prefill and generation took about 168–172 milliseconds, versus 9.43–9.74 seconds for the full-prefix path. That roughly 56-fold speedup belongs to this checkpoint, prompt, short continuation, and A100 setup. A benchmark of the current LibTorch adapter needs a new run.

Conclusion

The assistant still needs work. Lower validation loss did not make its answers dependable, and I want to keep measuring that behavior as we improve the training recipe. It's useful to have the training curves and the actual answers together, especially when they tell different stories.

I'd also like to rerun the matched comparison with the current LibTorch adapter. The archived FP32 run was 3.43 times slower than PyTorch; that gives us a specific performance problem to work on. We can retain the target-alignment, resume, and cache proofs while changing the implementation, and check that the new run still meets their assumptions.

References

• Alec Radford and colleagues. Language Models are Unsupervised Multitask Learners, 2019. GPT-2's architecture and next-token objective.
• Andrej Karpathy. minGPT and nanoGPT. The training implementations I used as a starting point.
• Ilya Loshchilov and Frank Hutter. Decoupled Weight Decay Regularization, 2019. AdamW.
• Hugging Face. FineWeb-Edu, SmolTalk, and SmolTalk2. Dataset cards for the pre-training and instruction-tuning data.
• Lean. Weakest-precondition semantics and vcgen. The program-verification tools used for the batch-packing proof.
• PyTorch. Numerical accuracy. The TF32 arithmetic choice disabled in the archived matched comparison.

@misc{george2026gpt,
  author = {Robert Joseph George},
  title = {{Training a GPT-2-style model in Lean}},
  year = {2026},
  howpublished = {TorchLean Verified Examples},
  url = {https://robertboy18.github.io/TorchLean-Verified-Examples/week-03-gpt-training/}
}