Pretrain LLMs

What pretraining is

Take a very large corpus, tokenize it, and repeatedly ask the model to predict token t+1 from tokens up to t. Compare its distribution to the true next token with cross-entropy, backpropagate, step the optimiser. Repeat for billions of tokens.

The reason this produces something impressive is that predicting text well requires modelling everything text describes. To predict the last word of 'the capital of France is' you must have stored a fact. To close a bracket correctly you must track state. Nobody trains those abilities directly; they are what low loss requires.

SGD, Adam, AdamW

Why AdamW rather than Adam

Adam divides each parameter's step by the running root-mean-square of its own gradients, so parameters with small gradients still move. That adaptivity is why it trains transformers out of the box where SGD needs careful schedules.

The flaw is weight decay. In Adam, L2 regularisation is added into the gradient, so it then gets divided by that same running RMS — parameters with large gradients end up barely regularised. AdamW applies the decay directly to the weights, outside the adaptive step, which restores the intended uniform pull toward zero. On transformers the difference is consistently measurable.

One more standard practice: do not decay biases or layernorm gains. Shrinking them toward zero has no regularising benefit and actively hurts.

import torch, torch.nn as nn

@torch.no_grad()
def evaluate(model, loader, loss_fn, device):
    model.eval()                            # disables dropout
    total, n = 0.0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        logits = model(x)
        loss = loss_fn(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
        total += loss.item() * y.numel()
        n += y.numel()
    model.train()
    return total / n

def train(model, train_loader, test_loader, epochs=5, device="cuda"):
    model.to(device).train()
    opt = make_optimizer(model)
    loss_fn = nn.CrossEntropyLoss()

    for epoch in range(epochs):
        for x, y in train_loader:
            x, y = x.to(device), y.to(device)
            logits = model(x)
            loss = loss_fn(logits.reshape(-1, logits.size(-1)), y.reshape(-1))

            opt.zero_grad(set_to_none=True)     # cheaper than zeroing
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()

        print(epoch, "train", loss.item(), "test", evaluate(model, test_loader, loss_fn, device))

# The test set must be a contiguous held-out slice of the corpus, not random
# windows: overlapping windows leak the same text into both splits.

Note: Forgetting model.eval() leaves dropout active during evaluation, which makes test loss look worse and noisier than it is.

Weight initialisation

Initialise all weights to zero and every neuron in a layer computes the same thing and receives the same gradient forever — the symmetry never breaks. Initialise too large and activations explode through depth; too small and they vanish.

The principled schemes keep the variance of activations roughly constant across layers. Xavier scales by 1/√fan_in for symmetric activations; He scales by 2/√fan_in to compensate for ReLU discarding half the distribution. GPT-2 uses a normal with standard deviation 0.02, and additionally scales the residual-projection weights by 1/√(2·n_layers) so that the variance added to the residual stream does not grow with depth.

Starting from pretrained embeddings

A cheap and effective middle path: initialise your embedding table from GPT-2's rather than randomly, so the model begins with a sensible semantic geometry and spends its capacity on the harder parts. Loss drops faster and lower on small corpora.

The requirement is a matching tokenizer — the rows must correspond to the same strings. Use GPT-2's tokenizer and its embedding rows line up by construction. You can either freeze them, saving memory and preserving the pretrained structure, or let them adapt.

Dropout

Dropout zeroes each activation independently with probability p during training, and scales the survivors by 1/(1−p) so the expected magnitude is unchanged. At evaluation it does nothing.

It regularises by preventing any unit from being relied upon — the network must build redundant paths. Transformers apply it after attention weights, after the output projection, and inside the MLP. Typical values are 0.1 for large-corpus pretraining and up to 0.3 when fine-tuning on a small dataset where overfitting is the real risk.

The 1/(1−p) scaling is the part worth remembering: it is why you must call model.eval(), and why a hand-rolled dropout that forgets the rescale shifts every downstream activation.

Numerical scaling

float32 holds roughly seven significant digits and float16 about three, with a maximum around 65,504. Mixed-precision training runs the forward and backward passes in low precision for speed and keeps a float32 copy of the weights for the update, with a loss scaler multiplying the loss so small gradients do not flush to zero.

bfloat16 has the same exponent range as float32 with fewer mantissa bits, so it rarely overflows and needs no loss scaling. Where available it is the easier choice. When you see NaNs appear suddenly mid-training, the usual causes are an overflow in float16, a division by a near-zero denominator, or a log of zero.

Data, and tokens that never appear

Modern open pretraining uses large filtered web corpora — FineWeb is the well-documented example, derived from Common Crawl with deduplication, language filtering and quality heuristics. The filtering is most of the work; raw web text trains a markedly worse model than the same volume cleaned.

A curiosity with real consequences: tokens in the vocabulary that never appear in training keep their initialised embedding rows and receive no gradient. Their unembedding rows are similarly untrained, so their logits are effectively arbitrary. These are the 'glitch tokens' that make models behave bizarrely when one is forced into a prompt — a direct, visible artefact of the vocabulary being fixed before the data was seen.

The optimisation options you actually turn

• Learning rate — by far the most impactful. Warm up linearly, then decay with cosine.
• Batch size — larger gives less gradient noise; use gradient accumulation to simulate it when memory is short.
• Gradient clipping — clip the global norm to 1.0; cheap insurance against a single bad batch.
• Weight decay — around 0.1 on matrices only.
• Dropout — 0.0–0.1 for large corpora, higher when data is scarce.
• Precision — bf16 where supported, otherwise fp16 with a loss scaler.
• Context length — quadratic cost in attention, so this dominates memory.