How This LLM Works

Running python train.py calls train.main. It reads text, trains a tokenizer, trains the model and saves a checkpoint. The sections below follow that function in code order. Each one shows the real lines.

The code uses NumPy only for arrays. There is no automatic differentiation: every backward function is written by hand. The model has no batch dimension, so forward and loss_and_grads process one sequence at a time. Architecture shows the model itself.

Prepare the data

Seed the random numbers

One generator, seeded with 42, is used for initial weights, for shuffling and for sampling.

37    rng = np.random.default_rng(config.seed)

Read the corpus

load_corpus reads at most 2000 rows of one CSV column, drops missing values, strips each row and joins the rows with newlines. See Data Pipeline.

40    text = load_corpus(config.corpus_path, config.corpus_column, config.max_rows)

Train the tokenizer

build_tokenizer learns up to 500 BPE merges from the text. It works on Unicode characters, not bytes. See Tokenization.

41    tok = build_tokenizer(text, config.num_merges)

Encode the text as ids

encode_corpus turns the whole text into one 1-D array of token ids.

42    ids = encode_corpus(text, tok)

Split into training and validation

split_ids does not shuffle. The leading part of the stream is for training, and the last 0.1 of it is for validation. The split happens before windowing, so no window crosses it.

43    train_ids, val_ids = split_ids(ids, config.val_fraction)

Cut windows

make_windows slides a window of 16 tokens along each part, moving 1 token at a time. Targets are the inputs shifted by one token. The results are pandas DataFrames.

44    tr_x, tr_y = make_windows(train_ids, config.seq_len, config.stride)45    va_x, va_y = make_windows(val_ids, config.seq_len, config.stride)

Build the model and optimizer

ModelConfig takes its vocabulary size from the tokenizer, so V depends on the corpus. init_params creates the arrays listed in Parameters. Adam holds the optimizer state. It takes a weight decay of 0.01, applied separately from the gradient, so it behaves as AdamW.

48    cfg = ModelConfig(tok.vocab_size, config.context_len, config.emb_dim,49                      config.num_heads, config.ff_dim, config.num_layers)50    params = init_params(cfg, rng)51    opt = Adam(params, config.beta1, config.beta2, config.adam_eps, config.weight_decay)

ln_eps is not passed, so ModelConfig keeps its own default and config.ln_eps has no effect.

The number of steps follows from the number of training windows, so it depends on the corpus too. The learning-rate schedule needs this total.

54    steps_per_epoch = math.ceil(len(tr_x) / config.batch_size)55    total_steps = config.epochs * steps_per_epoch

The training loop

The loop runs 20 times over the training windows. Each pass shuffles them and yields batches of 16. A final short batch is kept.

59    for epoch in range(config.epochs):60        for xb, yb in iterate_batches(tr_x, tr_y, config.batch_size, rng):

Gradient

batch_loss_and_grads runs loss_and_grads once per window in a Python loop, then averages the losses and gradients. See How the Model Learns.

61            loss, grads = batch_loss_and_grads(params, xb, yb, cfg)

12def batch_loss_and_grads(params, xb, yb, cfg: ModelConfig):13    B = len(xb)14    grads = {k: np.zeros_like(v) for k, v in params.items()}15    total_loss = 0.01617    for x, y in zip(xb, yb):18        loss, g = loss_and_grads(params, x, y, cfg)19        total_loss += loss / B2021        for k in grads:22            grads[k] += g[k] / B2324    return total_loss, grads

Clip

clip_grad_norm scales all gradients together if their global norm is above 1.0. It returns the norm from before clipping. See AdamW / Learning Rate / Clipping.

62            norm = clip_grad_norm(grads, config.grad_clip)

Learning rate

step increases first, so the schedule sees steps starting at 1. The rate rises linearly for 50 steps to 1e-3, then falls along a cosine to 1e-4.

63            step += 164            lr = lr_schedule(step, config.warmup_steps, total_steps,65                             config.lr_max, config.lr_min)

AdamW

opt.step updates the parameters in place. Weight decay is applied inside this call and only to arrays with two or more dimensions. Decay therefore does not pass through clipping.

66            opt.step(params, grads, lr)

Then the loop prints the loss, the norm and the learning rate on the first step and at the interval in the code.

67            if step == 1 or step % 50 == 0:68                print(f"step {step:5d}  loss {loss:.3f}  |g| {norm:.2f}  lr {lr:.5f}")

After each epoch

Validation

validation_loss averages the cross-entropy over all validation windows. It runs forward only. perplexity is the exponential of this loss, and it is only used for this printed line.

70        v = validation_loss(params, cfg, va_x, va_y)71        print(f"epoch {epoch + 1}: val loss {v:.3f}  ppl {perplexity(v):.1f}")

26def validation_loss(params, cfg, inputs, targets):27    losses = []2829    for x, y in zip(inputs.to_numpy(), targets.to_numpy()):30        logits, _ = forward(params, x, cfg)31        loss, _ = cross_entropy_forward(logits, y)32        losses.append(loss)3334    return float(np.mean(losses))

Sample

The loop prints a sample: it takes the first tokens of the first validation window as a seed and generates from them, using the sampling settings in config.py. See Autoregressive Generation. This site shows no output from these lines, because no run of the pinned commit and its corpus is available.

72        seed_ids = [int(i) for i in va_x.iloc[0][:3]]73        out = generate(params, cfg, seed_ids, 30, config.seq_len,74                       config.temperature, config.top_k, config.top_p, rng)75        print("  sample:", repr(tok.decode(out)))

Save the checkpoint

Saving happens once, after the last epoch. save_checkpoint writes three files that share a prefix: .npz for weights, .json for the model shape and .tok.json for the tokenizer. See Checkpoints & Tests.

78    save_checkpoint(prefix, params, cfg, tok)

Using the trained model

talk.main is the other entry point. It is a prompt loop that completes text. It has no chat format, no message history and no stopping token.

load_checkpoint restores the weights, the model shape and the tokenizer.

8    params, cfg, tok = load_checkpoint(prefix)

Each prompt starts fresh and is encoded on its own.

22        ids = tok.encode(prompt)

generate appends 50 sampled tokens. It feeds forward only the last 16 ids. The window and sampling settings come from the current config.py, while the model shape comes from the checkpoint.

24        out = generate(params, cfg, ids, config.max_new_tokens, config.seq_len,25                       config.temperature, config.top_k, config.top_p, rng)

Only the new tokens are printed.

27        print(tok.decode(out[len(ids):])) # only the new tokens, the prompt is already on screen