Build a Language Model · A trained language model

Train a Language Model from Scratch

This page puts the series' parts together into one program: a character-level GPT with one transformer block, trained in the browser on a 16 KB excerpt of Shakespeare, from the data pipeline through hand-derived gradients, AdamW with warmup and cosine decay, sampling, and the arithmetic that separates this 15,000-parameter model from GPT-2's 124 million.

Time~60 minLevelIntermediatePrereqsTransformers for attention and the block; Training Deep Networks for AdamW, schedules and scaling; Automatic Differentiation for the backward passes derived here by hand.StackC++20 & Rust · browser demos

01Predict the next character

A assigns a probability to a sequence of tokens. Written with the chain rule, that probability is a product of one next-token distribution per position, so the model only ever has to answer one question: given everything so far, what comes next? Training minimizes the mean negative log-probability the model gives to the token that did come next, averaged over every position of every window in a batch.^[1]

L = −(1/T) Σ_t log p_θ(x_t+1 | x₁ … x_t) PPL = exp(L)

T is the window length, x_t the token at position t, p_θ the model's distribution. The loss is in nats; dividing by ln 2 gives bits per character.

Three yardsticks make the number meaningful. A model that predicts uniformly scores ln V nats, which is 4.060 for the 58 characters used here, and that is also the loss a correctly initialized model should show before any training, since its logits start near zero.^[1]^[2] A model that ignores context and predicts each character by its frequency can at best reach the loss of a smoothed unigram model, 3.215 on this page's validation split; one that sees a single previous character is bounded by a smoothed bigram model, 2.615. The page's readouts report exp of the loss: 4.060 nats is a perplexity of 58, 2.4 nats about 11.

02Characters, splits and batches

The tokenizer here is the simplest one that works: every distinct character is a token, numbered in sorted order, which is exactly how nanoGPT's character-level Shakespeare preparation does it (65 symbols over 1.1 million characters there, 58 over 15,957 here).^[3] The last 10% of the text, in document order, is held out for validation. Shuffling before splitting would put the first half of a sentence in training and the second in validation, and the model would be scored on text it has partly seen.

A training batch is B windows of T tokens starting at random offsets into the training split, and the target for each position is the token one place to its right. Random offsets rather than fixed partitions mean every window boundary is eventually seen; d2l describes the alternative of partitioning from one random offset per epoch.^[1]^[4] The number of target predictions per iteration, B × T, is what nanoGPT prints as tokens per iteration, and it is the unit scaling laws are written in.^[4]

Space dominates the counts, then e, t and o, with s, a, i, r, h, n and the newline next; the rarest characters appear a handful of times or once, too few for the model to learn confident predictions for them. In the batch panel each target row is the input row shifted one place left: position t of the input predicts position t+1. Showing the model the target as an input, an off-by-one in the slicing, is one of the silent bugs Karpathy's recipe lists.^[2] With the target visible, training loss falls far faster than it should because the model only has to copy, and the copying is useless at generation time, when the next character is not available. The excerpt is the opening of tinyshakespeare, the same public-domain text nanoGPT trains on.^[3]

Production models do not use characters. GPT-2 and its successors use byte-level BPE: start from the 256 byte values, merge the most frequent adjacent pair into a new token, repeat until the vocabulary is the target size, 50,257 for GPT-2.^[5] Text is pre-split by a regular expression so merges never cross a letter–digit or letter–punctuation boundary, and the result is reversible and compresses text to about four bytes per token on average.^[6]^[5] Nothing downstream changes: the model sees integer ids either way. What changes is that each prediction covers roughly four characters, so a BPE model's loss per token is roughly four times its loss per character, and a per-token loss and a character model's per-character loss are comparable only after converting both to bits per byte.

03The model, with every gradient written down

The architecture is nanoGPT's with one block and no bias vectors.^[7] A token id selects a row of the token embedding, the position selects a row of the position embedding, and their sum is the residual stream. The block applies LayerNorm and causal multi-head attention, adds the result to the stream, then applies LayerNorm and a two-layer MLP with GELU and adds that too. A final LayerNorm feeds the output head, which is the token embedding matrix reused: the logit for token v is the dot product of the final vector with v's embedding, .

x = E_tok[id] + E_pos[t]
x ← x + W_proj·Attn(LN₁(x)) x ← x + W₂·GELU(W₁·LN₂(x))
z = E_tok·LN_f(x) p = softmax(z)

Attn is causal multi-head attention from the Transformers page: queries, keys and values come from one C × 3C matrix, scores are scaled by 1/√(head width), positions after t are masked, and the heads' outputs are concatenated. L is the number of blocks, 1 here, 12 for GPT-2 124M.

Weights start as N(0, 0.02), the two projections that write into the residual stream as N(0, 0.02/√(2L)) so that the sum of 2L such contributions keeps the stream's variance in check, and the LayerNorm gains at 1.^[7] With 58 characters, context 32, width 32 and 4 heads of width 8, the model has 15,264 parameters, 1,856 of them in the token embedding and 1,024 in the position table.

The page's trainer has no autodiff. Each layer's backward pass is written out, which is what a tape does for you on the Automatic Differentiation page and what makes the listings below checkable against finite differences. Three derivations carry the whole model.

∂L/∂z_v = (p_v − 𝟙[v = y]) / (B·T)

Softmax followed by cross-entropy has the gradient "probabilities minus one-hot", divided by the number of positions because the loss is their mean. Its sign and magnitude are the first thing to check: a confident wrong answer pushes hard, a correct one barely at all.

∂L/∂s_j = w_j ( g_j − Σ_k w_k g_k )

The softmax Jacobian applied to a gradient vector: each weight's gradient minus the weight-averaged gradient, times the weight. Masked positions have w = 0 and drop out. The score gradient then splits into the query and key through the scaled dot product, and the value gradient is the weight times the output gradient.

∂L/∂x = (1/σ) ( ĝ − mean(ĝ) − x̂ · mean(ĝ ⊙ x̂) ) ∂L/∂γ = Σ_positions ∂L/∂y ⊙ x̂

LayerNorm's backward pass subtracts the mean of the incoming gradient and its projection onto x̂, because the forward pass removed the mean and fixed the norm: any gradient component that would change either is discarded. The variance is the biased estimate and ε sits inside the square root, as in PyTorch.^[8] GELU's derivative follows from its tanh form; the page uses the tanh approximation that PyTorch documents alongside the exact x·Φ(x).^[9]

04Train it

The loop is nanoGPT's at toy scale: draw a batch, forward, backward, clip the global gradient norm to 1, take an AdamW step with weight decay 0.1 on the matrices and none on the gains, at a learning rate that warms up linearly for 50 iterations and then follows a cosine down to a tenth of its peak.^[4]^[10] β₂ is 0.99 rather than nanoGPT's default of 0.95 because each iteration here sees only 256 tokens and the second-moment estimate needs the longer memory, which is the reasoning nanoGPT's character config gives for the same choice.^[11] Validation loss is the mean over two fixed batches of the held-out split, 512 characters, so it is noisy to a few hundredths; nanoGPT averages 200 batches for the same reason.^[4]

Each Step runs six iterations of batch 8 × 32 on this tab's CPU; Play runs all 900. The loss starts at ln 58 = 4.06, the uniform baseline, which confirms the initialization. On the 16 KB excerpt at peak rate 0.01 the validation loss passes the unigram line within 30 iterations, crosses the bigram model's 2.615 before iteration 600 and ends near 2.39, so the model has learned more than one character of context. Switch to the first 2 KB and the curves separate: training loss keeps falling, to about 0.6, while validation loss bottoms out near 2.6 after about 115 iterations, turns upward and ends above 5, worse than the uniform model, because a memorizing model is confidently wrong on text it has not seen. The samples are drawn at temperature 0.8 from a newline.

Two checks that cost nothing catch many common training bugs. The loss at initialization should be −ln(1/V) to within a few hundredths, which this model shows on its first render. A value far above it means the initial logits are too large. A value clearly below it means something is supplying the answer: feeding this tied-embedding model its targets as inputs drops the initial loss from 4.07 to 3.66 before any training. And a small model should be able to drive the loss on one fixed batch to nearly zero; the 2 KB run above is a slower version of that overfitting test, and with 15,000 parameters, a model that cannot memorize 1,800 characters has a bug.^[2]

05Sample from it

Generation runs the forward pass on the prompt, takes the distribution at the last position, draws one character, appends it and repeats, cropping the context to the window length once it exceeds T.^[7] Two knobs reshape the distribution before the draw. Temperature divides the logits: at 0.3 the most likely character takes most of the mass, at 1.5 the distribution flattens toward uniform. Top-k zeroes everything outside the k most likely characters and renormalizes, which removes the long tail of rare characters that a flat distribution would otherwise sample.

The widget reads whatever state the training widget has reached; it refreshes every ten steps of training and on any control change. After 900 iterations on the 16 KB text, the character after "First Citizen:⏎" is a capital letter with high probability, as it is in the text, and the entropy of the distribution falls from about 2.4 nats at temperature 0.8 to under 1 nat at 0.3. The greedy continuation in yellow, argmax at every step, falls into "the the the" or "an and the an" within a few words: once the most likely continuation of a state leads back to that state there is no randomness to leave it, which the Transformers generation documentation describes as greedy search repeating itself on longer outputs.^[12] The sampled line below it is less fluent and does not loop.

06What it learned

A loss number says how much the model learned; two further measurements say what. The first is the attention pattern of each head on a fixed prompt. The second is the loss as a function of how many characters of context the position has, which separates what the model gets from one character, from a few, and from the full window.

Before training every row is nearly flat, each position spreading its weight over all the positions it may see. After 900 iterations the heads differ: in the run above, the mean weight a position gives to the character immediately before it rises to between 0.17 and 0.29 across the four heads, against 0.10 for uniform attention over a 32-character prompt, and the last few rows of head 0 put a quarter to a half of their weight on one recent character. A one-block model with 8-wide heads cannot form the sharp previous-token and induction patterns reported in larger transformers; what it shows is the beginning of them. Patterns depend on the random seed and the learning rate, so the specific rows describe this run only.

Every validation window of 32 characters, starting every 8 characters, is scored, and the loss at position t is averaged across windows, so position 1 is the loss when the model has seen one character, position 32 when it has seen the whole window. Untrained, the curve sits on the uniform line at 4.06. After 900 iterations at rate 0.01, the loss with one character of context is about 2.76, close to the smoothed bigram model's 2.615; from two characters onward it fluctuates between about 2.3 and 2.5 with no further trend. The model has learned to use roughly two or three characters of history and nothing beyond, which is consistent with its samples: plausible fragments, no words that depend on what came a line earlier.

07What changes at scale

The same code, with the constants changed, is nanoGPT's Shakespeare run: 6 blocks of 6 heads at width 384, context 256, batch 64, dropout 0.2, 5,000 iterations at a peak rate of 1e-3, about 10.7 million parameters, 3 minutes on one A100, and a best validation loss of 1.4697 nats, 2.12 bits per character.^[11]^[13] The same code again is GPT-2 124M: 12 blocks of 12 heads at width 768, context 1,024, a 50,257-token BPE vocabulary, 124,439,808 parameters of which 38.6 million are the token embedding, trained by nanoGPT on OpenWebText for 600,000 iterations of 491,520 tokens on eight A100s in about four days to a loss near 2.85 nats per token.^[13]^[4]^[14] OpenAI's model card dates the GPT-2 family to February 2019; the 1.5B model was the fourth and largest size released, after 124M, 355M and 774M, and all four were trained on WebText, the text behind 45 million links posted to Reddit.^[14]

Three quantities govern the jump. Parameter count follows from the shape: every block holds 12C² weights in its attention and MLP matrices, so width dominates depth once the vocabulary term V·C stops mattering. Training compute is close to 6 floating-point operations per parameter per token, 2 for the forward pass and 4 for the backward, plus an attention term of 12·L·H·Q·T per token that nanoGPT includes in its estimate following the PaLM paper.^[7] And memory for training in mixed precision is about 16 bytes per parameter: 2 for the bfloat16 weights, 2 for their gradients, and 12 for the 32-bit master weights and Adam's two moments, which is why DeepSpeed's ZeRO partitions exactly those optimizer states across devices and reports 18 GB of Adam state for a 1.5-billion-parameter GPT-2.^[15] The compute-optimal data budget, from the Chinchilla fit on the Training Deep Networks page, is about 20 tokens per parameter; this page's model sees its 14,000 training characters about 16 times over.^[16]

Parameter counts are exact for the GPT-2 shapes, with biases: 124,439,808 for GPT-2 124M and 1,557,611,200 for XL, and 10,745,088 for nanoGPT's bias-free Shakespeare model, which nanoGPT reports as 10.65M after excluding the position table.^[7] For GPT-2 124M the token embedding is 31% of the total; for XL it is 5%. At 10¹¹ tokens, forty percent utilization of eight A100s puts the 124M run at about a day (23.8 hours) by the 6N + 12LHQT count, consistent with nanoGPT's four-day figure, which covers about 3 × 10¹¹ tokens at a lower realized utilization.^[13] The memory panel is arithmetic, not a measurement: activations, the KV cache and framework overhead come on top.

Data changes too. Character-level Shakespeare is one author in one register. GPT-2's WebText was filtered web text; later open corpora such as the Pile combine 22 sources with explicit weights, Common Crawl at 18%, PubMed Central at 14%, books at 12%, GitHub, arXiv and more, to about 1,254 GiB after repeating some sources for up to three epochs.^[17] At that scale deduplication, filtering and the held-out split's contamination with training text become engineering problems in their own right, which the Evaluation page covers.

08A complete trainer, checked

These programs are the browser model without the browser, in the spirit of llm.c's CPU reference, which trains GPT-2 in about a thousand lines of C with every kernel written out and fine-tunes the released 124M weights on tinyshakespeare from a validation loss of 5.25 to 4.11 in 40 steps.^[18] Here: the one-block GPT, its hand-written backward pass, AdamW with clipping, warmup and cosine decay, random-offset batches with shifted targets, and sampling with temperature. They generate a synthetic corpus of templated sentences so they need no data file. Before training they assert two things: that the initial loss is within 0.05 of ln V, and that every parameter group's analytic gradient agrees with a centered finite difference to a relative error below 10⁻⁴, which is the check that makes a hand-derived backward pass trustworthy. After 400 iterations they assert the validation loss has at least halved, print it with its perplexity, and sample 120 characters.

A one-block GPT with hand-written gradients, complete programs

// A one-block GPT trained on characters, with every gradient written by hand.
// Build: g++ -std=c++20 -O2 -Wall -Wextra -Wpedantic -Werror lm.cpp && ./a.out
#include <algorithm>
#include <cassert>
#include <cctype>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <numbers>
#include <string>
#include <vector>

struct Rng {  // xorshift32, fixed seed: identical numbers in C++ and Rust
    uint32_t state = 7;
    double next() { state ^= state << 13; state ^= state >> 17; state ^= state << 5; return state / 4294967296.0; }
    double normal() { double u = 0; while (u == 0) u = next(); return std::sqrt(-2 * std::log(u)) * std::cos(2 * std::numbers::pi * next()); }
};

// Hyperparameters of the model and the training run.
constexpr int kBlock = 32, kEmbed = 32, kHeads = 4, kHeadDim = kEmbed / kHeads, kHidden = 4 * kEmbed, kBatch = 8;

struct Model {
    int vocab;
    // Parameters. Every linear map is a plain [in][out] matrix with no bias, like nanoGPT with bias=False.
    std::vector<double> tokenEmbedding, positionEmbedding, gain1, attentionWeight, attentionProjection, gain2, mlpExpand, mlpProjection, gainFinal;
    std::vector<std::vector<double>*> parameters() { return {&tokenEmbedding, &positionEmbedding, &gain1, &attentionWeight, &attentionProjection, &gain2, &mlpExpand, &mlpProjection, &gainFinal}; }
    Model(int vocabSize, Rng& rng) : vocab(vocabSize) {
        auto gaussian = [&](std::size_t size, double std) { std::vector<double> values(size); for (double& value : values) value = rng.normal() * std; return values; };
        tokenEmbedding = gaussian(std::size_t(vocab) * kEmbed, 0.02);   // also the tied output head
        positionEmbedding = gaussian(std::size_t(kBlock) * kEmbed, 0.02);
        gain1.assign(kEmbed, 1.0); gain2.assign(kEmbed, 1.0); gainFinal.assign(kEmbed, 1.0);
        attentionWeight = gaussian(std::size_t(kEmbed) * 3 * kEmbed, 0.02);
        attentionProjection = gaussian(std::size_t(kEmbed) * kEmbed, 0.02 / std::sqrt(2.0));  // GPT-2 residual scaling, 1 layer
        mlpExpand = gaussian(std::size_t(kEmbed) * kHidden, 0.02);
        mlpProjection = gaussian(std::size_t(kHidden) * kEmbed, 0.02 / std::sqrt(2.0));
    }
};

// Gradients, Adam moments: same shapes as the parameters.
struct Shadow {
    std::vector<std::vector<double>> slots;
    explicit Shadow(Model& model) { for (auto* parameter : model.parameters()) slots.emplace_back(parameter->size(), 0.0); }
    void zero() { for (auto& slot : slots) for (double& value : slot) value = 0; }
};

struct NormCache { std::vector<double> output, normalized, inverseStd; };
// LayerNorm without a bias term over rows of width kEmbed.
NormCache layerNormForward(const std::vector<double>& input, const std::vector<double>& gain, int rows) {
    NormCache cache{std::vector<double>(input.size()), std::vector<double>(input.size()), std::vector<double>(rows)};
    for (int row = 0; row < rows; ++row) {
        double mean = 0, variance = 0;
        for (int c = 0; c < kEmbed; ++c) mean += input[row * kEmbed + c];
        mean /= kEmbed;
        for (int c = 0; c < kEmbed; ++c) variance += (input[row * kEmbed + c] - mean) * (input[row * kEmbed + c] - mean);
        cache.inverseStd[row] = 1.0 / std::sqrt(variance / kEmbed + 1e-5);
        for (int c = 0; c < kEmbed; ++c) {
            cache.normalized[row * kEmbed + c] = (input[row * kEmbed + c] - mean) * cache.inverseStd[row];
            cache.output[row * kEmbed + c] = cache.normalized[row * kEmbed + c] * gain[c];
        }
    }
    return cache;
}
// dL/dx = (1/σ) (g⊙dy − mean(g⊙dy) − x̂ · mean(g⊙dy⊙x̂)), accumulated into inputGrad.
void layerNormBackward(const std::vector<double>& outputGrad, const NormCache& cache, const std::vector<double>& gain, std::vector<double>& gainGrad, int rows, std::vector<double>& inputGrad) {
    for (int row = 0; row < rows; ++row) {
        double sumGrad = 0, sumGradNormalized = 0;
        for (int c = 0; c < kEmbed; ++c) {
            double scaled = outputGrad[row * kEmbed + c] * gain[c], normalized = cache.normalized[row * kEmbed + c];
            gainGrad[c] += outputGrad[row * kEmbed + c] * normalized;
            sumGrad += scaled; sumGradNormalized += scaled * normalized;
        }
        for (int c = 0; c < kEmbed; ++c) {
            double scaled = outputGrad[row * kEmbed + c] * gain[c], normalized = cache.normalized[row * kEmbed + c];
            inputGrad[row * kEmbed + c] += cache.inverseStd[row] * (scaled - sumGrad / kEmbed - normalized * sumGradNormalized / kEmbed);
        }
    }
}
// output[rows][outDim] = input[rows][inDim] · weight[inDim][outDim]
std::vector<double> linearForward(const std::vector<double>& input, const std::vector<double>& weight, int rows, int inDim, int outDim) {
    std::vector<double> output(std::size_t(rows) * outDim, 0.0);
    for (int row = 0; row < rows; ++row) for (int i = 0; i < inDim; ++i) {
        double x = input[row * inDim + i];
        for (int o = 0; o < outDim; ++o) output[row * outDim + o] += x * weight[i * outDim + o];
    }
    return output;
}
// weightGrad += inputᵀ · outputGrad; inputGrad += outputGrad · weightᵀ
void linearBackward(const std::vector<double>& outputGrad, const std::vector<double>& input, const std::vector<double>& weight, std::vector<double>& weightGrad, int rows, int inDim, int outDim, std::vector<double>& inputGrad) {
    for (int row = 0; row < rows; ++row) for (int i = 0; i < inDim; ++i) {
        double x = input[row * inDim + i], accumulated = 0;
        for (int o = 0; o < outDim; ++o) { double gy = outputGrad[row * outDim + o]; weightGrad[i * outDim + o] += x * gy; accumulated += gy * weight[i * outDim + o]; }
        inputGrad[row * inDim + i] += accumulated;
    }
}
double gelu(double x) { double u = std::sqrt(2 / std::numbers::pi) * (x + 0.044715 * x * x * x); return 0.5 * x * (1 + std::tanh(u)); }
double geluGrad(double x) { double u = std::sqrt(2 / std::numbers::pi) * (x + 0.044715 * x * x * x), t = std::tanh(u); return 0.5 * (1 + t) + 0.5 * x * (1 - t * t) * std::sqrt(2 / std::numbers::pi) * (1 + 3 * 0.044715 * x * x); }

// Everything the backward pass needs from the forward pass.
struct Activations {
    int batch, length, rows;
    std::vector<int> inputs, targets;
    std::vector<double> qkv, weights, attentionOut, hidden, activated, probabilities;  // the residual sums are not needed again
    NormCache norm1, norm2, normFinal;
    double loss;
};

Activations forward(const Model& model, const std::vector<int>& inputs, const std::vector<int>& targets, int batch, int length) {
    Activations act; act.batch = batch; act.length = length; act.rows = batch * length; act.inputs = inputs; act.targets = targets;
    int rows = act.rows, vocab = model.vocab;
    std::vector<double> residual0(std::size_t(rows) * kEmbed, 0.0);
    for (int b = 0; b < batch; ++b) for (int t = 0; t < length; ++t) for (int c = 0; c < kEmbed; ++c)
        residual0[(b * length + t) * kEmbed + c] = model.tokenEmbedding[inputs[b * length + t] * kEmbed + c] + model.positionEmbedding[t * kEmbed + c];
    act.norm1 = layerNormForward(residual0, model.gain1, rows);
    act.qkv = linearForward(act.norm1.output, model.attentionWeight, rows, kEmbed, 3 * kEmbed);  // queries | keys | values
    act.weights.assign(std::size_t(batch) * kHeads * length * length, 0.0);
    act.attentionOut.assign(std::size_t(rows) * kEmbed, 0.0);
    double scale = 1.0 / std::sqrt(double(kHeadDim));
    for (int b = 0; b < batch; ++b) for (int h = 0; h < kHeads; ++h) for (int t = 0; t < length; ++t) {
        double* row = &act.weights[((b * kHeads + h) * length + t) * length];
        double maximum = -1e300;
        for (int s = 0; s <= t; ++s) {  // causal: only positions s ≤ t are scored
            double score = 0;
            for (int d = 0; d < kHeadDim; ++d) score += act.qkv[(b * length + t) * 3 * kEmbed + h * kHeadDim + d] * act.qkv[(b * length + s) * 3 * kEmbed + kEmbed + h * kHeadDim + d];
            row[s] = score * scale; maximum = std::max(maximum, row[s]);
        }
        double total = 0;
        for (int s = 0; s <= t; ++s) { row[s] = std::exp(row[s] - maximum); total += row[s]; }
        for (int s = 0; s <= t; ++s) {
            row[s] /= total;
            for (int d = 0; d < kHeadDim; ++d) act.attentionOut[(b * length + t) * kEmbed + h * kHeadDim + d] += row[s] * act.qkv[(b * length + s) * 3 * kEmbed + 2 * kEmbed + h * kHeadDim + d];
        }
    }
    std::vector<double> projected = linearForward(act.attentionOut, model.attentionProjection, rows, kEmbed, kEmbed);
    std::vector<double> residual1(residual0.size());
    for (std::size_t i = 0; i < residual1.size(); ++i) residual1[i] = residual0[i] + projected[i];
    act.norm2 = layerNormForward(residual1, model.gain2, rows);
    act.hidden = linearForward(act.norm2.output, model.mlpExpand, rows, kEmbed, kHidden);
    act.activated.resize(act.hidden.size());
    for (std::size_t i = 0; i < act.hidden.size(); ++i) act.activated[i] = gelu(act.hidden[i]);
    std::vector<double> mlpOut = linearForward(act.activated, model.mlpProjection, rows, kHidden, kEmbed);
    std::vector<double> residual2(residual1.size());
    for (std::size_t i = 0; i < residual2.size(); ++i) residual2[i] = residual1[i] + mlpOut[i];
    act.normFinal = layerNormForward(residual2, model.gainFinal, rows);
    // Logits through the tied token embedding, then softmax and mean cross-entropy.
    act.probabilities.assign(std::size_t(rows) * vocab, 0.0);
    act.loss = 0;
    for (int row = 0; row < rows; ++row) {
        double maximum = -1e300;
        for (int v = 0; v < vocab; ++v) {
            double logit = 0;
            for (int c = 0; c < kEmbed; ++c) logit += act.normFinal.output[row * kEmbed + c] * model.tokenEmbedding[v * kEmbed + c];
            act.probabilities[row * vocab + v] = logit; maximum = std::max(maximum, logit);
        }
        double total = 0;
        for (int v = 0; v < vocab; ++v) { act.probabilities[row * vocab + v] = std::exp(act.probabilities[row * vocab + v] - maximum); total += act.probabilities[row * vocab + v]; }
        for (int v = 0; v < vocab; ++v) act.probabilities[row * vocab + v] /= total;
        if (!targets.empty()) act.loss -= std::log(act.probabilities[row * vocab + targets[row]]);
    }
    if (!targets.empty()) act.loss /= rows;
    return act;
}

void backward(const Model& model, const Activations& act, Shadow& grads) {
    int rows = act.rows, batch = act.batch, length = act.length, vocab = model.vocab;
    // Gradient slots in the order of Model::parameters().
    auto& tokenGrad = grads.slots[0]; auto& positionGrad = grads.slots[1]; auto& gain1Grad = grads.slots[2]; auto& attentionGrad = grads.slots[3]; auto& projectionGrad = grads.slots[4];
    auto& gain2Grad = grads.slots[5]; auto& expandGrad = grads.slots[6]; auto& mlpProjectionGrad = grads.slots[7]; auto& gainFinalGrad = grads.slots[8];
    // Softmax + cross-entropy: dL/dlogit = (p − onehot(target)) / rows.
    std::vector<double> logitGrad(act.probabilities);
    for (int row = 0; row < rows; ++row) { for (int v = 0; v < vocab; ++v) logitGrad[row * vocab + v] /= rows; logitGrad[row * vocab + act.targets[row]] -= 1.0 / rows; }
    std::vector<double> normFinalGrad(std::size_t(rows) * kEmbed, 0.0);
    for (int row = 0; row < rows; ++row) for (int v = 0; v < vocab; ++v) {
        double gl = logitGrad[row * vocab + v];
        for (int c = 0; c < kEmbed; ++c) { normFinalGrad[row * kEmbed + c] += gl * model.tokenEmbedding[v * kEmbed + c]; tokenGrad[v * kEmbed + c] += gl * act.normFinal.output[row * kEmbed + c]; }
    }
    std::vector<double> residual2Grad(std::size_t(rows) * kEmbed, 0.0);
    layerNormBackward(normFinalGrad, act.normFinal, model.gainFinal, gainFinalGrad, rows, residual2Grad);
    // residual2 = residual1 + mlpOut: the gradient flows unchanged along the skip and through the MLP.
    std::vector<double> activatedGrad(std::size_t(rows) * kHidden, 0.0);
    linearBackward(residual2Grad, act.activated, model.mlpProjection, mlpProjectionGrad, rows, kHidden, kEmbed, activatedGrad);
    std::vector<double> hiddenGrad(activatedGrad.size());
    for (std::size_t i = 0; i < hiddenGrad.size(); ++i) hiddenGrad[i] = activatedGrad[i] * geluGrad(act.hidden[i]);
    std::vector<double> norm2Grad(std::size_t(rows) * kEmbed, 0.0);
    linearBackward(hiddenGrad, act.norm2.output, model.mlpExpand, expandGrad, rows, kEmbed, kHidden, norm2Grad);
    std::vector<double> residual1Grad(residual2Grad);
    layerNormBackward(norm2Grad, act.norm2, model.gain2, gain2Grad, rows, residual1Grad);
    std::vector<double> attentionOutGrad(std::size_t(rows) * kEmbed, 0.0);
    linearBackward(residual1Grad, act.attentionOut, model.attentionProjection, projectionGrad, rows, kEmbed, kEmbed, attentionOutGrad);
    std::vector<double> qkvGrad(std::size_t(rows) * 3 * kEmbed, 0.0);
    double scale = 1.0 / std::sqrt(double(kHeadDim));
    std::vector<double> weightGrad(length);
    for (int b = 0; b < batch; ++b) for (int h = 0; h < kHeads; ++h) for (int t = 0; t < length; ++t) {
        const double* row = &act.weights[((b * kHeads + h) * length + t) * length];
        double dotted = 0;
        for (int s = 0; s <= t; ++s) {  // through the weighted sum of values
            double accumulated = 0;
            for (int d = 0; d < kHeadDim; ++d) {
                double go = attentionOutGrad[(b * length + t) * kEmbed + h * kHeadDim + d];
                accumulated += go * act.qkv[(b * length + s) * 3 * kEmbed + 2 * kEmbed + h * kHeadDim + d];
                qkvGrad[(b * length + s) * 3 * kEmbed + 2 * kEmbed + h * kHeadDim + d] += row[s] * go;
            }
            weightGrad[s] = accumulated; dotted += row[s] * accumulated;
        }
        for (int s = 0; s <= t; ++s) {  // softmax backward, then the scaled dot product
            double scoreGrad = row[s] * (weightGrad[s] - dotted) * scale;
            for (int d = 0; d < kHeadDim; ++d) {
                qkvGrad[(b * length + t) * 3 * kEmbed + h * kHeadDim + d] += scoreGrad * act.qkv[(b * length + s) * 3 * kEmbed + kEmbed + h * kHeadDim + d];
                qkvGrad[(b * length + s) * 3 * kEmbed + kEmbed + h * kHeadDim + d] += scoreGrad * act.qkv[(b * length + t) * 3 * kEmbed + h * kHeadDim + d];
            }
        }
    }
    std::vector<double> norm1Grad(std::size_t(rows) * kEmbed, 0.0);
    linearBackward(qkvGrad, act.norm1.output, model.attentionWeight, attentionGrad, rows, kEmbed, 3 * kEmbed, norm1Grad);
    std::vector<double> residual0Grad(residual1Grad);
    layerNormBackward(norm1Grad, act.norm1, model.gain1, gain1Grad, rows, residual0Grad);
    for (int b = 0; b < batch; ++b) for (int t = 0; t < length; ++t) for (int c = 0; c < kEmbed; ++c) {
        tokenGrad[act.inputs[b * length + t] * kEmbed + c] += residual0Grad[(b * length + t) * kEmbed + c];
        positionGrad[t * kEmbed + c] += residual0Grad[(b * length + t) * kEmbed + c];
    }
}

// AdamW: decoupled weight decay on matrices only, bias-corrected moments, after clipping the global gradient norm to 1.
void adamStep(Model& model, Shadow& grads, Shadow& moment, Shadow& second, int step, double rate) {
    double squared = 0;
    for (auto& slot : grads.slots) for (double value : slot) squared += value * value;
    double clip = std::min(1.0, 1.0 / (std::sqrt(squared) + 1e-6));
    const double beta1 = 0.9, beta2 = 0.99, epsilon = 1e-8, weightDecay = 0.1;
    auto parameters = model.parameters();
    for (std::size_t p = 0; p < parameters.size(); ++p) {
        bool isMatrix = parameters[p]->size() > std::size_t(kEmbed);  // gains are vectors: no decay
        for (std::size_t i = 0; i < parameters[p]->size(); ++i) {
            double gradient = grads.slots[p][i] * clip;
            moment.slots[p][i] = beta1 * moment.slots[p][i] + (1 - beta1) * gradient;
            second.slots[p][i] = beta2 * second.slots[p][i] + (1 - beta2) * gradient * gradient;
            double correctedMoment = moment.slots[p][i] / (1 - std::pow(beta1, step)), correctedSecond = second.slots[p][i] / (1 - std::pow(beta2, step));
            (*parameters[p])[i] -= rate * (correctedMoment / (std::sqrt(correctedSecond) + epsilon) + (isMatrix ? weightDecay : 0.0) * (*parameters[p])[i]);
        }
    }
}

int main() {
    // A synthetic corpus of templated sentences, so the program needs no data file.
    const std::vector<std::string> subjects{"the cat", "a dog", "the old owl", "my horse", "one fox"}, verbs{"sat on", "ran past", "looked at", "slept under"}, objects{"the mat", "a red barn", "the gate", "our wall"};
    Rng rng;
    std::string text;
    while (text.size() < 24000) {
        // Separate declarations fix the draw order: the operands of + are unsequenced in C++, so three rng.next() calls in one expression run in a compiler-chosen order.
        const std::size_t subject = std::size_t(rng.next() * subjects.size()), verb = std::size_t(rng.next() * verbs.size()), object = std::size_t(rng.next() * objects.size());
        std::string sentence = subjects[subject] + " " + verbs[verb] + " " + objects[object] + ".";
        sentence[0] = char(std::toupper(sentence[0]));
        text += sentence + (rng.next() < 0.3 ? "\n" : " ");
    }
    std::vector<int> charToId(256, -1); std::string vocabulary;
    for (unsigned char ch : text) if (charToId[ch] < 0) { charToId[ch] = int(vocabulary.size()); vocabulary.push_back(char(ch)); }
    int vocab = int(vocabulary.size());
    std::vector<int> data; for (unsigned char ch : text) data.push_back(charToId[ch]);
    std::size_t split = data.size() * 9 / 10;  // 90% train, 10% validation, in document order
    std::printf("%zu characters, vocabulary %d, ln V = %.3f\n", text.size(), vocab, std::log(double(vocab)));

    Model model(vocab, rng);
    Shadow grads(model), moment(model), second(model);
    auto sampleBatch = [&](std::size_t begin, std::size_t end, std::vector<int>& inputs, std::vector<int>& targets) {
        inputs.assign(std::size_t(kBatch) * kBlock, 0); targets.assign(std::size_t(kBatch) * kBlock, 0);
        for (int b = 0; b < kBatch; ++b) {
            std::size_t offset = begin + std::size_t(rng.next() * double(end - begin - kBlock - 1));  // random window
            for (int t = 0; t < kBlock; ++t) { inputs[b * kBlock + t] = data[offset + t]; targets[b * kBlock + t] = data[offset + t + 1]; }  // targets shifted by one
        }
    };
    std::vector<int> inputs, targets;

    // Check 1: at initialization the loss is close to ln V, because the logits are near zero.
    sampleBatch(0, split, inputs, targets);
    Activations act = forward(model, inputs, targets, kBatch, kBlock);
    assert(std::fabs(act.loss - std::log(double(vocab))) < 0.05);

    // Check 2: the hand-written backward pass agrees with centered finite differences on every parameter group.
    grads.zero(); backward(model, act, grads);
    auto parameters = model.parameters();
    double worstRelativeError = 0;
    for (std::size_t p = 0; p < parameters.size(); ++p) for (std::size_t i : {std::size_t(0), parameters[p]->size() / 2, parameters[p]->size() - 1}) {
        double original = (*parameters[p])[i], h = 1e-5;
        (*parameters[p])[i] = original + h; double lossPlus = forward(model, inputs, targets, kBatch, kBlock).loss;
        (*parameters[p])[i] = original - h; double lossMinus = forward(model, inputs, targets, kBatch, kBlock).loss;
        (*parameters[p])[i] = original;
        double numeric = (lossPlus - lossMinus) / (2 * h), analytic = grads.slots[p][i];
        worstRelativeError = std::max(worstRelativeError, std::fabs(numeric - analytic) / std::max(1e-6, std::fabs(numeric) + std::fabs(analytic)));
    }
    std::printf("gradient check: worst relative error %.2e\n", worstRelativeError);
    assert(worstRelativeError < 1e-4);  // double precision; the finite difference itself is accurate to roughly 1e-6

    // Train: warmup then cosine decay, as in nanoGPT's get_lr.
    const int iterations = 400, warmup = 40; const double maxRate = 3e-3, minRate = 3e-4;
    double initialLoss = act.loss, finalValidation = 0;
    for (int step = 1; step <= iterations; ++step) {
        double rate = step <= warmup ? maxRate * step / warmup : minRate + 0.5 * (1 + std::cos(std::numbers::pi * double(step - warmup) / (iterations - warmup))) * (maxRate - minRate);
        sampleBatch(0, split, inputs, targets);
        grads.zero();
        act = forward(model, inputs, targets, kBatch, kBlock);
        backward(model, act, grads);
        adamStep(model, grads, moment, second, step, rate);
        if (step % 100 == 0 || step == iterations) {
            double validation = 0;
            for (int k = 0; k < 5; ++k) { sampleBatch(split, data.size(), inputs, targets); validation += forward(model, inputs, targets, kBatch, kBlock).loss; }
            finalValidation = validation / 5;
            std::printf("step %4d  lr %.2e  train %.3f  val %.3f  val perplexity %.2f\n", step, rate, act.loss, finalValidation, std::exp(finalValidation));
        }
    }
    assert(finalValidation < initialLoss / 2);  // the model learned something: validation loss at least halved

    // Sample 120 characters with temperature 0.8 from a prompt, cropping the context to the block size.
    std::vector<int> context; for (char ch : std::string("The cat ")) context.push_back(charToId[(unsigned char)ch]);
    std::string sample;
    for (int n = 0; n < 120; ++n) {
        std::vector<int> window(context.end() - std::min<std::size_t>(context.size(), kBlock), context.end());
        Activations out = forward(model, window, {}, 1, int(window.size()));
        std::vector<double> logits(vocab);
        double maximum = -1e300;
        for (int v = 0; v < vocab; ++v) { logits[v] = std::log(out.probabilities[(window.size() - 1) * vocab + v]) / 0.8; maximum = std::max(maximum, logits[v]); }
        double total = 0; for (double& logit : logits) { logit = std::exp(logit - maximum); total += logit; }
        double u = rng.next() * total; int chosen = 0;
        while (chosen < vocab - 1 && (u -= logits[chosen]) > 0) ++chosen;
        context.push_back(chosen); sample.push_back(vocabulary[chosen]);
    }
    std::printf("sample: %s\n", sample.c_str());
    return 0;
}

// A one-block GPT trained on characters, with every gradient written by hand.
// Build: rustc --edition 2021 -O -D warnings lm.rs && ./lm
use std::f64::consts::PI;

struct Rng { state: u32 } // xorshift32, fixed seed: identical numbers in C++ and Rust
impl Rng {
    fn next(&mut self) -> f64 { self.state ^= self.state << 13; self.state ^= self.state >> 17; self.state ^= self.state << 5; self.state as f64 / 4294967296.0 }
    fn normal(&mut self) -> f64 { let mut u = 0.0; while u == 0.0 { u = self.next(); } (-2.0 * u.ln()).sqrt() * (2.0 * PI * self.next()).cos() }
}

// Hyperparameters of the model and the training run.
const BLOCK: usize = 32; const EMBED: usize = 32; const HEADS: usize = 4; const HEAD_DIM: usize = EMBED / HEADS; const HIDDEN: usize = 4 * EMBED; const BATCH: usize = 8;

// Parameters. Every linear map is a plain [in][out] matrix with no bias, like nanoGPT with bias=False.
// Order: token_embedding, position_embedding, gain1, attention_weight, attention_projection, gain2, mlp_expand, mlp_projection, gain_final.
struct Model { vocab: usize, slots: Vec<Vec<f64>> }
const TOKEN: usize = 0; const POSITION: usize = 1; const GAIN1: usize = 2; const ATTENTION: usize = 3; const PROJECTION: usize = 4; const GAIN2: usize = 5; const EXPAND: usize = 6; const MLP_PROJECTION: usize = 7; const GAIN_FINAL: usize = 8;
impl Model {
    fn new(vocab: usize, rng: &mut Rng) -> Model {
        let mut gaussian = |size: usize, std: f64| (0..size).map(|_| rng.normal() * std).collect::<Vec<f64>>();
        let slots = vec![
            gaussian(vocab * EMBED, 0.02),            // also the tied output head
            gaussian(BLOCK * EMBED, 0.02),
            vec![1.0; EMBED],
            gaussian(EMBED * 3 * EMBED, 0.02),
            gaussian(EMBED * EMBED, 0.02 / 2f64.sqrt()), // GPT-2 residual scaling, 1 layer
            vec![1.0; EMBED],
            gaussian(EMBED * HIDDEN, 0.02),
            gaussian(HIDDEN * EMBED, 0.02 / 2f64.sqrt()),
            vec![1.0; EMBED],
        ];
        Model { vocab, slots }
    }
    fn zeros_like(&self) -> Vec<Vec<f64>> { self.slots.iter().map(|slot| vec![0.0; slot.len()]).collect() }
}

struct NormCache { output: Vec<f64>, normalized: Vec<f64>, inverse_std: Vec<f64> }
// LayerNorm without a bias term over rows of width EMBED.
fn layer_norm_forward(input: &[f64], gain: &[f64], rows: usize) -> NormCache {
    let mut cache = NormCache { output: vec![0.0; input.len()], normalized: vec![0.0; input.len()], inverse_std: vec![0.0; rows] };
    for row in 0..rows {
        let slice = &input[row * EMBED..(row + 1) * EMBED];
        let mean = slice.iter().sum::<f64>() / EMBED as f64;
        let variance = slice.iter().map(|value| (value - mean) * (value - mean)).sum::<f64>() / EMBED as f64;
        cache.inverse_std[row] = 1.0 / (variance + 1e-5).sqrt();
        for c in 0..EMBED {
            cache.normalized[row * EMBED + c] = (slice[c] - mean) * cache.inverse_std[row];
            cache.output[row * EMBED + c] = cache.normalized[row * EMBED + c] * gain[c];
        }
    }
    cache
}
// dL/dx = (1/σ) (g⊙dy − mean(g⊙dy) − x̂ · mean(g⊙dy⊙x̂)), accumulated into input_grad.
fn layer_norm_backward(output_grad: &[f64], cache: &NormCache, gain: &[f64], gain_grad: &mut [f64], rows: usize, input_grad: &mut [f64]) {
    for row in 0..rows {
        let (mut sum_grad, mut sum_grad_normalized) = (0.0, 0.0);
        for c in 0..EMBED {
            let scaled = output_grad[row * EMBED + c] * gain[c];
            let normalized = cache.normalized[row * EMBED + c];
            gain_grad[c] += output_grad[row * EMBED + c] * normalized;
            sum_grad += scaled; sum_grad_normalized += scaled * normalized;
        }
        for c in 0..EMBED {
            let scaled = output_grad[row * EMBED + c] * gain[c];
            let normalized = cache.normalized[row * EMBED + c];
            input_grad[row * EMBED + c] += cache.inverse_std[row] * (scaled - sum_grad / EMBED as f64 - normalized * sum_grad_normalized / EMBED as f64);
        }
    }
}
// output[rows][out_dim] = input[rows][in_dim] · weight[in_dim][out_dim]
fn linear_forward(input: &[f64], weight: &[f64], rows: usize, in_dim: usize, out_dim: usize) -> Vec<f64> {
    let mut output = vec![0.0; rows * out_dim];
    for row in 0..rows { for i in 0..in_dim {
        let x = input[row * in_dim + i];
        for o in 0..out_dim { output[row * out_dim + o] += x * weight[i * out_dim + o]; }
    } }
    output
}
// weight_grad += inputᵀ · output_grad; input_grad += output_grad · weightᵀ
fn linear_backward(output_grad: &[f64], input: &[f64], weight: &[f64], weight_grad: &mut [f64], rows: usize, in_dim: usize, out_dim: usize, input_grad: &mut [f64]) {
    for row in 0..rows { for i in 0..in_dim {
        let x = input[row * in_dim + i];
        let mut accumulated = 0.0;
        for o in 0..out_dim { let gy = output_grad[row * out_dim + o]; weight_grad[i * out_dim + o] += x * gy; accumulated += gy * weight[i * out_dim + o]; }
        input_grad[row * in_dim + i] += accumulated;
    } }
}
fn gelu(x: f64) -> f64 { let u = (2.0 / PI).sqrt() * (x + 0.044715 * x * x * x); 0.5 * x * (1.0 + u.tanh()) }
fn gelu_grad(x: f64) -> f64 { let u = (2.0 / PI).sqrt() * (x + 0.044715 * x * x * x); let t = u.tanh(); 0.5 * (1.0 + t) + 0.5 * x * (1.0 - t * t) * (2.0 / PI).sqrt() * (1.0 + 3.0 * 0.044715 * x * x) }

// Everything the backward pass needs from the forward pass.
struct Activations {
    batch: usize, length: usize, rows: usize, inputs: Vec<usize>, targets: Vec<usize>,
    qkv: Vec<f64>, weights: Vec<f64>, attention_out: Vec<f64>, hidden: Vec<f64>, activated: Vec<f64>, probabilities: Vec<f64>,
    norm1: NormCache, norm2: NormCache, norm_final: NormCache, loss: f64,
}

fn forward(model: &Model, inputs: &[usize], targets: &[usize], batch: usize, length: usize) -> Activations {
    let (rows, vocab) = (batch * length, model.vocab);
    let mut residual0 = vec![0.0; rows * EMBED];
    for b in 0..batch { for t in 0..length { for c in 0..EMBED {
        residual0[(b * length + t) * EMBED + c] = model.slots[TOKEN][inputs[b * length + t] * EMBED + c] + model.slots[POSITION][t * EMBED + c];
    } } }
    let norm1 = layer_norm_forward(&residual0, &model.slots[GAIN1], rows);
    let qkv = linear_forward(&norm1.output, &model.slots[ATTENTION], rows, EMBED, 3 * EMBED); // queries | keys | values
    let mut weights = vec![0.0; batch * HEADS * length * length];
    let mut attention_out = vec![0.0; rows * EMBED];
    let scale = 1.0 / (HEAD_DIM as f64).sqrt();
    for b in 0..batch { for h in 0..HEADS { for t in 0..length {
        let base = ((b * HEADS + h) * length + t) * length;
        let mut maximum = f64::NEG_INFINITY;
        for s in 0..=t { // causal: only positions s ≤ t are scored
            let mut score = 0.0;
            for d in 0..HEAD_DIM { score += qkv[(b * length + t) * 3 * EMBED + h * HEAD_DIM + d] * qkv[(b * length + s) * 3 * EMBED + EMBED + h * HEAD_DIM + d]; }
            weights[base + s] = score * scale; maximum = maximum.max(weights[base + s]);
        }
        let mut total = 0.0;
        for s in 0..=t { weights[base + s] = (weights[base + s] - maximum).exp(); total += weights[base + s]; }
        for s in 0..=t {
            weights[base + s] /= total;
            for d in 0..HEAD_DIM { attention_out[(b * length + t) * EMBED + h * HEAD_DIM + d] += weights[base + s] * qkv[(b * length + s) * 3 * EMBED + 2 * EMBED + h * HEAD_DIM + d]; }
        }
    } } }
    let projected = linear_forward(&attention_out, &model.slots[PROJECTION], rows, EMBED, EMBED);
    let residual1: Vec<f64> = residual0.iter().zip(&projected).map(|(a, b)| a + b).collect();
    let norm2 = layer_norm_forward(&residual1, &model.slots[GAIN2], rows);
    let hidden = linear_forward(&norm2.output, &model.slots[EXPAND], rows, EMBED, HIDDEN);
    let activated: Vec<f64> = hidden.iter().map(|&value| gelu(value)).collect();
    let mlp_out = linear_forward(&activated, &model.slots[MLP_PROJECTION], rows, HIDDEN, EMBED);
    let residual2: Vec<f64> = residual1.iter().zip(&mlp_out).map(|(a, b)| a + b).collect();
    let norm_final = layer_norm_forward(&residual2, &model.slots[GAIN_FINAL], rows);
    // Logits through the tied token embedding, then softmax and mean cross-entropy.
    let mut probabilities = vec![0.0; rows * vocab];
    let mut loss = 0.0;
    for row in 0..rows {
        let mut maximum = f64::NEG_INFINITY;
        for v in 0..vocab {
            let mut logit = 0.0;
            for c in 0..EMBED { logit += norm_final.output[row * EMBED + c] * model.slots[TOKEN][v * EMBED + c]; }
            probabilities[row * vocab + v] = logit; maximum = maximum.max(logit);
        }
        let mut total = 0.0;
        for v in 0..vocab { probabilities[row * vocab + v] = (probabilities[row * vocab + v] - maximum).exp(); total += probabilities[row * vocab + v]; }
        for v in 0..vocab { probabilities[row * vocab + v] /= total; }
        if !targets.is_empty() { loss -= probabilities[row * vocab + targets[row]].ln(); }
    }
    if !targets.is_empty() { loss /= rows as f64; }
    Activations { batch, length, rows, inputs: inputs.to_vec(), targets: targets.to_vec(), qkv, weights, attention_out, hidden, activated, probabilities, norm1, norm2, norm_final, loss }
}

fn backward(model: &Model, act: &Activations, grads: &mut [Vec<f64>]) {
    let (rows, batch, length, vocab) = (act.rows, act.batch, act.length, model.vocab);
    // Softmax + cross-entropy: dL/dlogit = (p − onehot(target)) / rows.
    let mut logit_grad = act.probabilities.clone();
    for row in 0..rows { for v in 0..vocab { logit_grad[row * vocab + v] /= rows as f64; } logit_grad[row * vocab + act.targets[row]] -= 1.0 / rows as f64; }
    let mut norm_final_grad = vec![0.0; rows * EMBED];
    for row in 0..rows { for v in 0..vocab {
        let gl = logit_grad[row * vocab + v];
        for c in 0..EMBED { norm_final_grad[row * EMBED + c] += gl * model.slots[TOKEN][v * EMBED + c]; grads[TOKEN][v * EMBED + c] += gl * act.norm_final.output[row * EMBED + c]; }
    } }
    let mut residual2_grad = vec![0.0; rows * EMBED];
    layer_norm_backward(&norm_final_grad, &act.norm_final, &model.slots[GAIN_FINAL], &mut grads[GAIN_FINAL], rows, &mut residual2_grad);
    // residual2 = residual1 + mlp_out: the gradient flows unchanged along the skip and through the MLP.
    let mut activated_grad = vec![0.0; rows * HIDDEN];
    linear_backward(&residual2_grad, &act.activated, &model.slots[MLP_PROJECTION], &mut grads[MLP_PROJECTION], rows, HIDDEN, EMBED, &mut activated_grad);
    let hidden_grad: Vec<f64> = activated_grad.iter().zip(&act.hidden).map(|(g, &x)| g * gelu_grad(x)).collect();
    let mut norm2_grad = vec![0.0; rows * EMBED];
    linear_backward(&hidden_grad, &act.norm2.output, &model.slots[EXPAND], &mut grads[EXPAND], rows, EMBED, HIDDEN, &mut norm2_grad);
    let mut residual1_grad = residual2_grad.clone();
    layer_norm_backward(&norm2_grad, &act.norm2, &model.slots[GAIN2], &mut grads[GAIN2], rows, &mut residual1_grad);
    let mut attention_out_grad = vec![0.0; rows * EMBED];
    linear_backward(&residual1_grad, &act.attention_out, &model.slots[PROJECTION], &mut grads[PROJECTION], rows, EMBED, EMBED, &mut attention_out_grad);
    let mut qkv_grad = vec![0.0; rows * 3 * EMBED];
    let scale = 1.0 / (HEAD_DIM as f64).sqrt();
    let mut weight_grad = vec![0.0; length];
    for b in 0..batch { for h in 0..HEADS { for t in 0..length {
        let base = ((b * HEADS + h) * length + t) * length;
        let mut dotted = 0.0;
        for s in 0..=t { // through the weighted sum of values
            let mut accumulated = 0.0;
            for d in 0..HEAD_DIM {
                let go = attention_out_grad[(b * length + t) * EMBED + h * HEAD_DIM + d];
                accumulated += go * act.qkv[(b * length + s) * 3 * EMBED + 2 * EMBED + h * HEAD_DIM + d];
                qkv_grad[(b * length + s) * 3 * EMBED + 2 * EMBED + h * HEAD_DIM + d] += act.weights[base + s] * go;
            }
            weight_grad[s] = accumulated; dotted += act.weights[base + s] * accumulated;
        }
        for s in 0..=t { // softmax backward, then the scaled dot product
            let score_grad = act.weights[base + s] * (weight_grad[s] - dotted) * scale;
            for d in 0..HEAD_DIM {
                qkv_grad[(b * length + t) * 3 * EMBED + h * HEAD_DIM + d] += score_grad * act.qkv[(b * length + s) * 3 * EMBED + EMBED + h * HEAD_DIM + d];
                qkv_grad[(b * length + s) * 3 * EMBED + EMBED + h * HEAD_DIM + d] += score_grad * act.qkv[(b * length + t) * 3 * EMBED + h * HEAD_DIM + d];
            }
        }
    } } }
    let mut norm1_grad = vec![0.0; rows * EMBED];
    linear_backward(&qkv_grad, &act.norm1.output, &model.slots[ATTENTION], &mut grads[ATTENTION], rows, EMBED, 3 * EMBED, &mut norm1_grad);
    let mut residual0_grad = residual1_grad.clone();
    layer_norm_backward(&norm1_grad, &act.norm1, &model.slots[GAIN1], &mut grads[GAIN1], rows, &mut residual0_grad);
    for b in 0..batch { for t in 0..length { for c in 0..EMBED {
        grads[TOKEN][act.inputs[b * length + t] * EMBED + c] += residual0_grad[(b * length + t) * EMBED + c];
        grads[POSITION][t * EMBED + c] += residual0_grad[(b * length + t) * EMBED + c];
    } } }
}

// AdamW: decoupled weight decay on matrices only, bias-corrected moments, after clipping the global gradient norm to 1.
fn adam_step(model: &mut Model, grads: &[Vec<f64>], moment: &mut [Vec<f64>], second: &mut [Vec<f64>], step: i32, rate: f64) {
    let squared: f64 = grads.iter().flatten().map(|g| g * g).sum();
    let clip = (1.0 / (squared.sqrt() + 1e-6)).min(1.0);
    let (beta1, beta2, epsilon, weight_decay) = (0.9, 0.99, 1e-8, 0.1);
    for p in 0..model.slots.len() {
        let is_matrix = model.slots[p].len() > EMBED; // gains are vectors: no decay
        for i in 0..model.slots[p].len() {
            let gradient = grads[p][i] * clip;
            moment[p][i] = beta1 * moment[p][i] + (1.0 - beta1) * gradient;
            second[p][i] = beta2 * second[p][i] + (1.0 - beta2) * gradient * gradient;
            let corrected_moment = moment[p][i] / (1.0 - beta1.powi(step));
            let corrected_second = second[p][i] / (1.0 - beta2.powi(step));
            let decay = if is_matrix { weight_decay } else { 0.0 };
            model.slots[p][i] -= rate * (corrected_moment / (corrected_second.sqrt() + epsilon) + decay * model.slots[p][i]);
        }
    }
}

fn main() {
    // A synthetic corpus of templated sentences, so the program needs no data file.
    let subjects = ["the cat", "a dog", "the old owl", "my horse", "one fox"]; let verbs = ["sat on", "ran past", "looked at", "slept under"]; let objects = ["the mat", "a red barn", "the gate", "our wall"];
    let mut rng = Rng { state: 7 };
    let mut text = String::new();
    while text.len() < 24000 {
        let subject = (rng.next() * subjects.len() as f64) as usize; let verb = (rng.next() * verbs.len() as f64) as usize; let object = (rng.next() * objects.len() as f64) as usize;
        let mut sentence = format!("{} {} {}.", subjects[subject], verbs[verb], objects[object]);
        sentence.replace_range(0..1, &sentence[0..1].to_uppercase());
        text.push_str(&sentence); text.push(if rng.next() < 0.3 { '\n' } else { ' ' });
    }
    let mut char_to_id = vec![usize::MAX; 256]; let mut vocabulary: Vec<char> = Vec::new();
    for ch in text.bytes() { if char_to_id[ch as usize] == usize::MAX { char_to_id[ch as usize] = vocabulary.len(); vocabulary.push(ch as char); } }
    let vocab = vocabulary.len();
    let data: Vec<usize> = text.bytes().map(|ch| char_to_id[ch as usize]).collect();
    let split = data.len() * 9 / 10; // 90% train, 10% validation, in document order
    println!("{} characters, vocabulary {}, ln V = {:.3}", text.len(), vocab, (vocab as f64).ln());

    let mut model = Model::new(vocab, &mut rng);
    let (mut grads, mut moment, mut second) = (model.zeros_like(), model.zeros_like(), model.zeros_like());
    let sample_batch = |rng: &mut Rng, begin: usize, end: usize| -> (Vec<usize>, Vec<usize>) {
        let (mut inputs, mut targets) = (vec![0; BATCH * BLOCK], vec![0; BATCH * BLOCK]);
        for b in 0..BATCH {
            let offset = begin + (rng.next() * (end - begin - BLOCK - 1) as f64) as usize; // random window
            for t in 0..BLOCK { inputs[b * BLOCK + t] = data[offset + t]; targets[b * BLOCK + t] = data[offset + t + 1]; } // targets shifted by one
        }
        (inputs, targets)
    };

    // Check 1: at initialization the loss is close to ln V, because the logits are near zero.
    let (mut inputs, mut targets) = sample_batch(&mut rng, 0, split);
    let mut act = forward(&model, &inputs, &targets, BATCH, BLOCK);
    assert!((act.loss - (vocab as f64).ln()).abs() < 0.05);

    // Check 2: the hand-written backward pass agrees with centered finite differences on every parameter group.
    for slot in grads.iter_mut() { slot.iter_mut().for_each(|g| *g = 0.0); }
    backward(&model, &act, &mut grads);
    let mut worst_relative_error = 0.0f64;
    for p in 0..model.slots.len() { for i in [0, model.slots[p].len() / 2, model.slots[p].len() - 1] {
        let (original, h) = (model.slots[p][i], 1e-5);
        model.slots[p][i] = original + h; let loss_plus = forward(&model, &inputs, &targets, BATCH, BLOCK).loss;
        model.slots[p][i] = original - h; let loss_minus = forward(&model, &inputs, &targets, BATCH, BLOCK).loss;
        model.slots[p][i] = original;
        let (numeric, analytic) = ((loss_plus - loss_minus) / (2.0 * h), grads[p][i]);
        worst_relative_error = worst_relative_error.max((numeric - analytic).abs() / (numeric.abs() + analytic.abs()).max(1e-6));
    } }
    println!("gradient check: worst relative error {:.2e}", worst_relative_error);
    assert!(worst_relative_error < 1e-4); // double precision; the finite difference itself is accurate to roughly 1e-6

    // Train: warmup then cosine decay, as in nanoGPT's get_lr.
    let (iterations, warmup, max_rate, min_rate) = (400, 40, 3e-3, 3e-4);
    let initial_loss = act.loss; let mut final_validation = 0.0;
    for step in 1..=iterations {
        let rate = if step <= warmup { max_rate * step as f64 / warmup as f64 } else { min_rate + 0.5 * (1.0 + (PI * (step - warmup) as f64 / (iterations - warmup) as f64).cos()) * (max_rate - min_rate) };
        (inputs, targets) = sample_batch(&mut rng, 0, split);
        for slot in grads.iter_mut() { slot.iter_mut().for_each(|g| *g = 0.0); }
        act = forward(&model, &inputs, &targets, BATCH, BLOCK);
        backward(&model, &act, &mut grads);
        adam_step(&mut model, &grads, &mut moment, &mut second, step, rate);
        if step % 100 == 0 || step == iterations {
            let mut validation = 0.0;
            for _ in 0..5 { let (val_inputs, val_targets) = sample_batch(&mut rng, split, data.len()); validation += forward(&model, &val_inputs, &val_targets, BATCH, BLOCK).loss; }
            final_validation = validation / 5.0;
            println!("step {:4}  lr {:.2e}  train {:.3}  val {:.3}  val perplexity {:.2}", step, rate, act.loss, final_validation, final_validation.exp());
        }
    }
    assert!(final_validation < initial_loss / 2.0); // the model learned something: validation loss at least halved

    // Sample 120 characters with temperature 0.8 from a prompt, cropping the context to the block size.
    let mut context: Vec<usize> = "The cat ".bytes().map(|ch| char_to_id[ch as usize]).collect();
    let mut sample = String::new();
    for _ in 0..120 {
        let window = context[context.len().saturating_sub(BLOCK)..].to_vec();
        let out = forward(&model, &window, &[], 1, window.len());
        let mut logits: Vec<f64> = (0..vocab).map(|v| out.probabilities[(window.len() - 1) * vocab + v].ln() / 0.8).collect();
        let maximum = logits.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
        let mut total = 0.0; for logit in logits.iter_mut() { *logit = (*logit - maximum).exp(); total += *logit; }
        let mut u = rng.next() * total; let mut chosen = 0;
        while chosen < vocab - 1 { u -= logits[chosen]; if u <= 0.0 { break; } chosen += 1; }
        context.push(chosen); sample.push(vocabulary[chosen]);
    }
    println!("sample: {}", sample);
}

What's intentionally missing

Dropout, bias vectors, more than one block, the key-value cache that makes generation linear rather than quadratic, mixed precision, gradient accumulation, checkpoint saving, a real tokenizer, data loading from disk, and any kernel faster than three nested loops. The forward pass, loss, gradients, optimizer and schedule are the ones nanoGPT runs, except that GELU uses its tanh approximation where nanoGPT uses the exact form.

09Where small language models mislead

Perplexities are only comparable at the same tokenization on the same text

A character model's 2.4 nats and a BPE model's 2.85 nats describe different units: one is per character, the other per token of roughly four characters. Convert both to bits per byte before comparing, and compare only on the same evaluation text. The Evaluation page covers the further dependence on context length and stride.

A falling training loss on its own is not evidence of learning. The 2 KB run reaches a lower training loss than the 16 KB run while being a worse model of the language by every held-out measure. Report validation loss, state its noise, and look at samples.

Greedy samples look worse than the model is. A model whose argmax loops still assigns reasonable probabilities; the loop is a property of the decoding rule.^[12] Judge a language model by its loss and by samples at a stated temperature, not by its greedy output alone.

Scaling constants are shape-dependent. The 6 FLOPs per parameter per token count omits attention, which matters once T is large relative to C; the 16 bytes per parameter omits activations, which dominate at large batch and context. Treat the widget's numbers as lower bounds.

10What's next

This model predicts text. Making a model useful means changing what it optimizes for after pretraining, and the signal for that is a reward rather than a next token. Reinforcement Learning builds that machinery from a grid world up: value functions, policy gradients and the variance that makes them hard, which the fine-tuning page then applies to language models.

11Sources

The cited code defines the architecture, the training loop and the configurations quoted; the cited documents define the loss, the baselines and the decoding behaviour. Widget readouts report values computed in the page from the embedded excerpt and the model trained in the browser.

Aston Zhang, Zachary C. Lipton, Mu Li, Alexander J. Smola, 2023. Dive into Deep Learning, Section 9.3: Language Models, Cambridge University Press. Perplexity as the exponential of the mean cross-entropy, equal to 1 for a perfect model and to the vocabulary size for a uniform one, and partitioning a corpus into fixed-length windows from a random offset.
Andrej Karpathy, 2019. A Recipe for Training Neural Networks. Verifying the loss at initialization against −log(1/classes), overfitting a single batch, the off-by-one bug that lets an autoregressive model see its target, the input-independent baseline, and using gradients to check what each output depends on.
Andrej Karpathy. nanoGPT, data/shakespeare_char/prepare.py. Sorted character vocabulary (65 symbols over 1,115,394 characters), a 90/10 split in document order, and integer ids written as uint16.
Andrej Karpathy. nanoGPT, train.py. Random window offsets with targets shifted by one, tokens per iteration, loss estimated over eval_iters batches, gradient clipping at 1.0, weight decay 0.1 on tensors of rank 2 and above, and linear warmup followed by cosine decay to a tenth of the peak rate.
Andrej Karpathy. minbpe, README. Byte-level BPE over UTF-8 with the 256 byte values as the initial tokens, regex pre-splitting introduced with GPT-2, and special tokens that must be allowed explicitly at encode time.
OpenAI. tiktoken, README. Byte pair encoding as the tokenizer of the GPT models, reversible and compressing text to about four bytes per token on average.
Andrej Karpathy. nanoGPT, model.py. Token plus position embeddings, pre-LayerNorm blocks, GELU, tied input and output embeddings, N(0, 0.02) initialization with the residual projections scaled by 1/√(2L), generation with temperature and top-k, the parameter count excluding positions, and the FLOPs estimate 6N + 12·L·H·Q·T per token attributed to the PaLM paper's Appendix B against the A100's 312 TFLOP/s bfloat16 peak.
PyTorch contributors. torch.nn.LayerNorm, documentation. Normalization over the last dimensions with the biased variance estimate, ε inside the square root, and an elementwise affine gain and bias; cites Ba et al. (2016).
PyTorch contributors. torch.nn.GELU, documentation. GELU as x·Φ(x) and its tanh approximation.
PyTorch contributors. torch.optim.AdamW, documentation. Decoupled weight decay applied directly to the parameters rather than through the gradient, after Loshchilov and Hutter (2019).
Andrej Karpathy. nanoGPT, config/train_shakespeare_char.py. Batch 64, context 256, 6 layers of 6 heads at width 384, dropout 0.2, peak learning rate 1e-3 decaying to 1e-4 over 5,000 iterations with 100 warmup iterations, β₂ = 0.99.
Hugging Face. Generation strategies, Transformers documentation. Greedy search repeating itself on longer outputs, and sampling reducing repetition.
Andrej Karpathy. nanoGPT, README. The character-level Shakespeare run (6 layers, 6 heads, 384 channels, context 256; about 3 minutes on one A100; best validation loss 1.4697), the CPU configuration that reaches 1.88, and GPT-2 124M on OpenWebText reproduced on an 8×A100 node in about 4 days to a loss near 2.85.
OpenAI, 2019. GPT-2 model card. Model date February 2019, the 1.5B model as the fourth and largest version released after the 124M, 355M and 774M sizes, and WebText as text from 45 million outbound Reddit links.
DeepSpeed contributors. Zero Redundancy Optimizer, tutorial. Adam's optimizer state consisting of 32-bit weights and first and second moments, 16-bit gradients, and the 18 GB of Adam state for a 1.5-billion-parameter GPT-2.
Jordan Hoffmann et al., 2022. Training Compute-Optimal Large Language Models, arXiv:2203.15556. Compute-optimal training puts tokens and parameters in a roughly 20:1 ratio; the fit and its constants are discussed on the Training Deep Networks page.
Leo Gao et al., 2020. The Pile, repository README (paper: arXiv:2101.00027). Composition table of 22 sources with their weights, Pile-CC at 18.11%, PubMed Central at 14.40% and Books3 at 12.07%, totalling 1,254 GiB after up-weighting.
Andrej Karpathy. llm.c, README. GPT-2 training in plain C with a roughly 1,000-line CPU reference, and fine-tuning the 124M model on tinyshakespeare for 40 steps at batch 4 and context 64, validation loss 5.25 to 4.11.

Train a Language Model from Scratch

01Predict the next character

02Characters, splits and batches

Tokenize the excerpt and draw a batch

03The model, with every gradient written down

04Train it

Train the one-block GPT in the browser

05Sample from it

Reshape the next-character distribution

06What it learned

Attention weights of one head on a fixed prompt

Loss by how much context the position has