Causal Pretraining and Training Infrastructure

Train a small causal language model end to end, inspect loss curves, and verify checkpoint resumption.

A model becomes a pretraining artifact only when the data path, objective, update count, validation boundary, and checkpoint state are all reproducible. This chapter keeps a NumPy loop as a transparent reference, then runs the same contracts through the PyTorch ProofLM decoder and the reusable trainer. The smoke path records training and validation loss, stores a checkpoint, and reloads the state needed for a reproducible continuation.

The CPU smoke profile is deliberately small enough to run locally. Its purpose is to validate the artifact lineage and resume behavior before a larger CUDA run is authorized.

Causal language modeling

For input tokens x_{t-C:t-1} and target x_t, a fixed-context causal model produces logits z_t. Across a batch of windows, the training objective is

\mathcal{L}_{\mathrm{LM}} =-\frac{1}{N}\sum_{i=1}^{N}\log p_\theta(y_i\mid x_i).

The target is one position to the right of the final input position. Validation must construct its windows from a separate text boundary and must not update parameters. A loss curve is meaningful only when these two roles remain separate.

The NumPy loop below makes the update and checkpoint ideas inspectable. The end-to-end smoke cell later uses the shared byte-level tokenizer, packed boundary masks, PyTorch decoder, trainer, and checkpoint artifact.

import io
import pickle

import numpy as np


SEED = 51
CONTEXT_LENGTH = 6


def log_softmax(logits):
    shifted = logits - np.max(logits, axis=-1, keepdims=True)
    return shifted - np.log(np.exp(shifted).sum(axis=-1, keepdims=True))


def softmax(logits):
    return np.exp(log_softmax(logits))


def make_examples(text, vocabulary, context_length):
    token_ids = np.asarray([vocabulary[character] for character in text], dtype=np.int64)
    if len(token_ids) <= context_length:
        return np.empty((0, context_length), dtype=np.int64), np.empty(0, dtype=np.int64)
    inputs = np.stack([
        token_ids[start:start + context_length]
        for start in range(len(token_ids) - context_length)
    ])
    targets = token_ids[context_length:]
    return inputs, targets


def initialize_model(vocab_size, context_length, embedding_dim=6, hidden_dim=32, seed=0):
    random_generator = np.random.default_rng(seed)
    return {
        "embedding": random_generator.normal(scale=0.25, size=(vocab_size, embedding_dim)),
        "w1": random_generator.normal(
            scale=1.0 / np.sqrt(context_length * embedding_dim),
            size=(context_length * embedding_dim, hidden_dim),
        ),
        "b1": np.zeros(hidden_dim),
        "w2": random_generator.normal(scale=1.0 / np.sqrt(hidden_dim), size=(hidden_dim, vocab_size)),
        "b2": np.zeros(vocab_size),
    }


def model_forward(parameters, inputs):
    embeddings = parameters["embedding"][inputs]
    flattened = embeddings.reshape(len(inputs), -1)
    hidden = np.tanh(flattened @ parameters["w1"] + parameters["b1"])
    logits = hidden @ parameters["w2"] + parameters["b2"]
    return logits, (inputs, embeddings, flattened, hidden)


train_text = "the cat sat. the dog sat. the cat ran. " * 5
validation_text = "dog ran. cat dog ran. " * 2
vocabulary_symbols = sorted(set(train_text + validation_text))
vocabulary = {symbol: index for index, symbol in enumerate(vocabulary_symbols)}
train_inputs, train_targets = make_examples(train_text, vocabulary, CONTEXT_LENGTH)
validation_inputs, validation_targets = make_examples(
    validation_text, vocabulary, CONTEXT_LENGTH
)
train_window_keys = {
    tuple(row) + (int(target),)
    for row, target in zip(train_inputs, train_targets)
}
validation_window_keys = {
    tuple(row) + (int(target),)
    for row, target in zip(validation_inputs, validation_targets)
}
parameters = initialize_model(
    len(vocabulary), CONTEXT_LENGTH, seed=SEED
)
print("vocabulary:", "".join(vocabulary_symbols))
print("train / validation windows:", train_inputs.shape, validation_inputs.shape)
print("shared train/validation windows:", len(train_window_keys & validation_window_keys))
assert train_inputs.shape[1] == CONTEXT_LENGTH
assert len(validation_inputs) == len(validation_targets)
assert not train_window_keys & validation_window_keys

vocabulary:  .acdeghnorst
train / validation windows: (189, 6) (38, 6)
shared train/validation windows: 0

The training and validation strings share a vocabulary but are separate sequences. That distinction is enough for this small demonstration: a validation window is never sampled from the training array. The MLP flattens six embeddings, applies one nonlinear hidden layer, and returns one logit vector for the next character. Its limited receptive field makes the token shift easy to inspect.

One training update

The backward pass follows the cached path in reverse. For h=\tanh(a), the local derivative is 1-h^2. After the loss gradient reaches all parameter arrays, global-norm clipping applies one common scale

\tilde g = g\min\left(1,\frac{\tau}{\lVert g\rVert_2+10^{-12}}\right),

so it changes the step size without changing the direction of the concatenated gradient. The training loop below samples a batch with a supplied generator, evaluates a loss, clips, and updates every parameter exactly once.

def loss_and_gradients(parameters, inputs, targets):
    logits, cache = model_forward(parameters, inputs)
    loss = -np.mean(log_softmax(logits)[np.arange(len(targets)), targets])
    probabilities = softmax(logits)
    dlogits = probabilities.copy()
    dlogits[np.arange(len(targets)), targets] -= 1.0
    dlogits /= len(targets)
    _, embeddings, flattened, hidden = cache
    gradients = {
        "w2": hidden.T @ dlogits,
        "b2": dlogits.sum(axis=0),
    }
    dhidden = dlogits @ parameters["w2"].T
    dhidden_pre = dhidden * (1.0 - hidden ** 2)
    gradients["w1"] = flattened.T @ dhidden_pre
    gradients["b1"] = dhidden_pre.sum(axis=0)
    dembeddings = (dhidden_pre @ parameters["w1"].T).reshape(embeddings.shape)
    gradients["embedding"] = np.zeros_like(parameters["embedding"])
    for position in range(inputs.shape[1]):
        np.add.at(gradients["embedding"], inputs[:, position], dembeddings[:, position])
    return float(loss), gradients


def gradient_norm(gradients):
    return float(np.sqrt(sum(np.sum(value ** 2) for value in gradients.values())))


def clip_gradients(gradients, maximum_norm):
    norm = gradient_norm(gradients)
    scale = min(1.0, maximum_norm / (norm + 1e-12))
    return {name: value * scale for name, value in gradients.items()}, norm


def update(parameters, gradients, learning_rate):
    for name in parameters:
        parameters[name] -= learning_rate * gradients[name]


initial_loss, initial_gradients = loss_and_gradients(parameters, train_inputs[:12], train_targets[:12])
clipped, unclipped_norm = clip_gradients(initial_gradients, maximum_norm=0.5)
print("initial batch loss:", round(initial_loss, 5))
print("gradient norm before clipping:", round(unclipped_norm, 5))
print("gradient norm after clipping:", round(gradient_norm(clipped), 5))
assert gradient_norm(clipped) <= 0.5 + 1e-10
assert all(value.shape == parameters[name].shape for name, value in clipped.items())

initial batch loss: 2.63168
gradient norm before clipping: 0.77634
gradient norm after clipping: 0.5

Clipping is inactive when the norm is already below the threshold and scales every array together when it is active. Per-parameter clipping would change the direction of the combined update and would represent a different optimizer. The explicit shape check protects the embedding scatter-add and matrix derivatives before the loop runs repeatedly.

Checkpoint state

A resumable run needs more than parameter arrays. The update count identifies the schedule position, and the random-generator state identifies the next batch draw. Serialize both with the parameters. An in-memory byte stream is used here so the notebook obeys its no-file fixture constraint; a production trainer would write the same payload to a versioned checkpoint path.

Generation receives a separate generator. Evaluating a sample must not advance the generator used to choose future training batches.

def copy_parameters(parameters):
    return {name: value.copy() for name, value in parameters.items()}


def save_checkpoint(parameters, step, random_generator):
    payload = {
        "parameters": copy_parameters(parameters),
        "step": int(step),
        "rng_state": random_generator.bit_generator.state,
    }
    buffer = io.BytesIO()
    pickle.dump(payload, buffer)
    return buffer.getvalue()


def load_checkpoint(serialized):
    payload = pickle.loads(serialized)
    restored_parameters = copy_parameters(payload["parameters"])
    restored_rng = np.random.default_rng()
    restored_rng.bit_generator.state = payload["rng_state"]
    return restored_parameters, payload["step"], restored_rng


def sample_next(logits, random_generator, temperature=0.8):
    if temperature <= 0.0:
        raise ValueError("temperature must be positive")
    probabilities = softmax(np.asarray(logits) / temperature)
    return int(random_generator.choice(len(probabilities), p=probabilities))


def generate(parameters, prompt, steps, random_generator, temperature=0.8):
    generated = list(prompt)
    for _ in range(steps):
        if len(generated) < CONTEXT_LENGTH:
            raise ValueError("prompt must contain at least CONTEXT_LENGTH symbols")
        context = np.asarray(generated[-CONTEXT_LENGTH:], dtype=np.int64)[None, :]
        logits = model_forward(parameters, context)[0][0]
        generated.append(sample_next(logits, random_generator, temperature))
    return "".join(vocabulary_symbols[index] for index in generated)


checkpoint_rng = np.random.default_rng(SEED + 1)
checkpoint_bytes = save_checkpoint(parameters, step=7, random_generator=checkpoint_rng)
restored_parameters, restored_step, restored_rng = load_checkpoint(checkpoint_bytes)
print("serialized checkpoint bytes:", len(checkpoint_bytes), "step:", restored_step)
for name in parameters:
    np.testing.assert_array_equal(parameters[name], restored_parameters[name])
assert restored_step == 7
assert restored_rng.integers(0, 1000) == checkpoint_rng.integers(0, 1000)

serialized checkpoint bytes: 13948 step: 7

The equality check consumes one draw from each restored generator only after loading, so both generators produce the same next integer. If the checkpoint stored only weights, a resumed run with random mini-batches would use a different data order even when every hyperparameter matched. The sample generator is intentionally independent, which keeps evaluation side-effect free with respect to training randomness.

Tiny pretraining run

Record metrics by update count and save a checkpoint at fixed intervals. Each recorded validation loss is computed without changing parameters. Example generations use the checkpoint’s parameter state, so the output is an observable artifact of the training trajectory rather than a sample from only the final model.

def train_updates(parameters, inputs, targets, steps, random_generator, learning_rate=0.15, batch_size=24, maximum_norm=3.0):
    history = []
    for _ in range(steps):
        indices = random_generator.integers(0, len(inputs), size=batch_size)
        loss, gradients = loss_and_gradients(parameters, inputs[indices], targets[indices])
        gradients, norm = clip_gradients(gradients, maximum_norm)
        update(parameters, gradients, learning_rate)
        history.append({"loss": loss, "gradient_norm": norm})
    return history


def train_with_checkpoints(parameters, random_generator, steps, checkpoint_steps):
    history = []
    snapshots = {}
    prompt = [vocabulary[character] for character in train_text[:CONTEXT_LENGTH]]
    for step in range(steps + 1):
        if step in checkpoint_steps:
            train_loss = loss_and_gradients(parameters, train_inputs, train_targets)[0]
            validation_loss = loss_and_gradients(
                parameters, validation_inputs, validation_targets
            )[0]
            snapshots[step] = save_checkpoint(parameters, step, random_generator)
            sample = generate(
                parameters,
                prompt,
                steps=18,
                random_generator=np.random.default_rng(SEED + 100 + step),
            )
            history.append((step, train_loss, validation_loss, sample))
        if step == steps:
            break
        train_updates(
            parameters,
            train_inputs,
            train_targets,
            steps=1,
            random_generator=random_generator,
        )
    return history, snapshots


run_parameters = initialize_model(len(vocabulary), CONTEXT_LENGTH, seed=SEED + 2)
run_history, run_snapshots = train_with_checkpoints(
    run_parameters,
    np.random.default_rng(SEED + 3),
    steps=60,
    checkpoint_steps=(0, 20, 40, 60),
)
for step, train_loss, validation_loss, sample in run_history:
    print(f"step={step:>2} train={train_loss:.4f} valid={validation_loss:.4f} sample={sample}")
assert len(run_snapshots) == 4
assert run_history[-1][1] < run_history[0][1]
assert all(np.isfinite(row[2]) for row in run_history)
assert len(set(row[3] for row in run_history)) > 1

step= 0 train=2.6281 valid=2.5284 sample=the caheoneg rgg tgdatns
step=20 train=1.8201 valid=2.3127 sample=the ca.ne.ogeasgashd.tht
step=40 train=1.0110 valid=1.9845 sample=the cat  ate tat sctd aa
step=60 train=0.5701 valid=1.8225 sample=the cat eaetsnaro the eo

The checkpoint table ties each sample to the same update axis as the losses. A decreasing training curve confirms that the objective is being optimized; the validation curve shows whether those updates transfer to a separate string. A changing sample is weaker evidence than either loss, but it verifies that the saved parameter state affects generation and that the sampler is not returning a constant fixture.

Tiny-batch overfitting

Overfitting a few consistent windows is a unit test for the entire forward and backward path. It should be run before interpreting validation behavior. Use four early windows whose contexts are distinct in this corpus, train them with full-batch updates, and require a large loss reduction.

tiny_parameters = initialize_model(
    len(vocabulary), CONTEXT_LENGTH, embedding_dim=8, hidden_dim=40, seed=SEED + 4
)
tiny_inputs = train_inputs[:4]
tiny_targets = train_targets[:4]
tiny_initial = loss_and_gradients(tiny_parameters, tiny_inputs, tiny_targets)[0]
for _ in range(450):
    tiny_loss, tiny_gradients = loss_and_gradients(tiny_parameters, tiny_inputs, tiny_targets)
    tiny_gradients, _ = clip_gradients(tiny_gradients, maximum_norm=10.0)
    update(tiny_parameters, tiny_gradients, learning_rate=0.2)
tiny_final = loss_and_gradients(tiny_parameters, tiny_inputs, tiny_targets)[0]
print("tiny-batch loss:", round(tiny_initial, 4), "->", round(tiny_final, 6))
assert tiny_final < 0.08

tiny-batch loss: 2.675 -> 0.001468

The tiny-batch loss is expected to approach zero because the selected windows have consistent targets and the model has more parameters than examples. This test does not estimate generalization. It establishes that a failure in the full run should be investigated as a data, gradient, or update problem before changing the validation split.

Checkpoint resumption

Run the same initial state for a fixed number of updates in two ways: uninterrupted, and split at a checkpoint. Restore both the parameter arrays and the generator state before continuing. Exact equality is stronger than matching rounded losses and catches a missing piece of serialized state.

resume_initial = initialize_model(len(vocabulary), CONTEXT_LENGTH, seed=SEED + 5)
continuous = copy_parameters(resume_initial)
continuous_rng = np.random.default_rng(SEED + 20)
train_updates(continuous, train_inputs, train_targets, steps=30, random_generator=continuous_rng)

split_parameters = copy_parameters(resume_initial)
split_rng = np.random.default_rng(SEED + 20)
train_updates(split_parameters, train_inputs, train_targets, steps=12, random_generator=split_rng)
serialized = save_checkpoint(split_parameters, step=12, random_generator=split_rng)
resumed, resumed_step, resumed_rng = load_checkpoint(serialized)
train_updates(resumed, train_inputs, train_targets, steps=30 - resumed_step, random_generator=resumed_rng)

for name in continuous:
    np.testing.assert_array_equal(continuous[name], resumed[name])
print("resume equivalence: exact parameter match after 30 updates")
assert resumed_step == 12

resume equivalence: exact parameter match after 30 updates

CPU smoke lineage

The NumPy loop proves the mechanics locally, but the reusable lineage must include the same dataset boundary, tokenizer identity, model configuration, optimizer state, and masked objective that later chapters consume. The following smoke run trains ProofLM for a bounded number of updates, evaluates a separate packed validation stream, saves a checkpoint under the project artifact root, and reloads it before reporting the result.

import json
from pathlib import Path

import torch

from proof_lm.data import batch_packed_examples, pack_causal_examples
from proof_lm.evaluation import evaluate_language_model
from proof_lm.model import DecoderConfig, ProofLM
from proof_lm.trainer import (
    load_training_checkpoint,
    save_training_checkpoint,
    train_batches,
)

smoke_pad_id = len(vocabulary) + 1
smoke_eos_id = len(vocabulary)
train_ids = [vocabulary[character] for character in train_text]
validation_ids = [vocabulary[character] for character in validation_text]
packed_train = pack_causal_examples(
    [train_ids],
    eos_id=smoke_eos_id,
    pad_id=smoke_pad_id,
    context_length=CONTEXT_LENGTH,
)
packed_validation = pack_causal_examples(
    [validation_ids],
    eos_id=smoke_eos_id,
    pad_id=smoke_pad_id,
    context_length=CONTEXT_LENGTH,
)
smoke_batches = batch_packed_examples(packed_train, batch_size=4)
validation_batches = batch_packed_examples(packed_validation, batch_size=4)
smoke_config = DecoderConfig(
    vocab_size=len(vocabulary) + 2,
    context_length=CONTEXT_LENGTH,
    n_layers=1,
    d_model=32,
    n_heads=4,
    d_ff=64,
)
torch.manual_seed(SEED)
smoke_model = ProofLM(smoke_config)
smoke_optimizer = torch.optim.AdamW(smoke_model.parameters(), lr=0.02)
smoke_state, smoke_losses = train_batches(
    smoke_model,
    smoke_batches,
    smoke_optimizer,
    max_steps=8,
)
validation_report = evaluate_language_model(smoke_model, validation_batches)
checkpoint_path = Path(
    "projects/proof-lm/artifacts/checkpoints/smoke/base-smoke.pt"
)
save_training_checkpoint(
    checkpoint_path, smoke_model, smoke_optimizer, smoke_state
)
reloaded_model = ProofLM(smoke_config)
reloaded_optimizer = torch.optim.AdamW(reloaded_model.parameters(), lr=0.02)
loaded_state = load_training_checkpoint(
    checkpoint_path, reloaded_model, reloaded_optimizer
)
probe = torch.tensor([packed_validation[0].input_ids])
torch.testing.assert_close(smoke_model(probe)[0], reloaded_model(probe)[0])
smoke_report = {
    "profile": "smoke",
    "checkpoint": str(checkpoint_path),
    "step": smoke_state.step,
    "processed_tokens": smoke_state.processed_tokens,
    "train_loss_start": smoke_losses[0],
    "train_loss_end": smoke_losses[-1],
    "validation": validation_report,
}
report_path = checkpoint_path.with_name("smoke-report.json")
report_path.write_text(json.dumps(smoke_report, indent=2, sort_keys=True) + "\n")
print("smoke checkpoint:", checkpoint_path)
print("updates / tokens:", smoke_state.step, smoke_state.processed_tokens)
print(
    "train loss:",
    round(smoke_losses[0], 4),
    "->",
    round(smoke_losses[-1], 4),
)
print("validation perplexity:", round(float(validation_report["perplexity"]), 4))
assert smoke_state.step == loaded_state.step == 8
assert smoke_losses[-1] < smoke_losses[0]
assert checkpoint_path.exists()
assert report_path.exists()

smoke checkpoint: projects/proof-lm/artifacts/checkpoints/smoke/base-smoke.pt
updates / tokens: 8 192
train loss: 2.7896 -> 1.7561
validation perplexity: 7.8602

The uninterrupted and resumed arrays match exactly because the checkpoint captured the state before update 13 and restored the next batch draw. Saving a rounded loss, a seed value, or weights without the generator state would not guarantee this result. The same principle applies to shuffled data order, dropout masks, and any learning-rate schedule state in a larger trainer.

Summary

• Causal pretraining minimizes next-token cross-entropy on shifted windows from a declared training split.
• The NumPy loop computes gradients, applies global clipping, logs update-indexed metrics, and evaluates validation data without updates.
• Checkpoints contain parameters, update position, and random-generator state; in-memory serialization makes the state testable without creating repository files.
• A tiny-batch overfit test validates the training path before validation curves are interpreted.
• Exact resume equivalence requires restoring every source of trajectory state, not only the model weights.

Chapter 06 isolates optimizer updates, schedules, clipping thresholds, and token accounting so the training budget can be compared across runs.

Exercises

Use the exercises to test the chapter’s invariants and connect the derivations to the reusable implementation. Solutions are hidden in the notebook source and are available through the course tooling when needed.

[P5.1] Exact checkpoint resume

Checkpoint state. List the state required for an exact resume when batches are sampled randomly. Explain why restoring only parameter arrays is insufficient, then identify the corresponding fields in this chapter.

[P5.2] Pretraining sanity checks

Pretraining sanity test. Design two checks that distinguish a decreasing training loss from a valid validation measurement. Include one check for target alignment and one check showing that validation examples are not drawn from the training array.

[P5.3]

Masked pretraining batches. A packed batch contains inputs, shifted targets, a validity mask, and a loss mask. Explain why the target immediately after an EOS token is present in the packed stream but should have a false loss-mask entry. Give one invariant relating the number of loss tokens to valid tokens.

def boundary_loss_invariant(valid_mask, loss_mask):
    # Return whether the loss mask is a subset of valid target positions.
    pass