Pretrain a small language model

Pretrain a model in the 10M to 100M range. That is large enough to show instability and small enough to finish. You are not chasing a public leaderboard. You are building the replica of a pretrain loop: data, step, log, checkpoint, resume.

Coding: train_step

import torch

def train_step(model, batch, optimizer, max_norm: float = 1.0):
    model.train()
    optimizer.zero_grad(set_to_none=True)
    logits, loss = model(batch["input_ids"], batch["labels"])
    if not torch.isfinite(loss):
        raise RuntimeError("non-finite loss")
    loss.backward()
    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
    optimizer.step()
    return {
        "loss": float(loss.detach()),
        "grad_norm": float(grad_norm),
        "clipped": bool(grad_norm > max_norm),
    }

Call this in a loop with cosine warmup from week 2. Save a checkpoint every N steps with optimiser state. Resume once on purpose so you know the path works before a crash teaches you.

Instability you should expect

• Non-finite loss after a bad batch or a too-high learning rate. Abort, do not continue with NaNs.
• Grad norms that spike at the start of a new data shard. Clip and look at the texts.
• Loss that will not move. Check the causal mask, label shift, and whether you are training the embedding table.
• Throughput lies. Tokens per second must use the same token definition as your budget.

Parameter count is not a personality. A 10M model with a clean log teaches more than a 100M run you cannot resume.

Assignment

• Pick a width and depth that land between 10M and 100M parameters. Print the count.
• Train for a token budget you wrote in week 4. Log loss, grad norm, clip rate, and tokens per second.
• Resume from a checkpoint and show that step and loss continue. That is the production test.

Opinion checkpoint

Core project 3: a small LM checkpoint with a training log you trust.

Next: Distributed training simulation.