← Back to Table of Contents

Chapter 10 — Pre-Training at Scale

“The remarkable thing about LLMs isn’t the architecture — it’s that next-token prediction on internet-scale data produces intelligence.”

The Pre-Training Pipeline

Pre-training is where the bulk of compute and data investment goes. A modern LLM pre-training run involves months of GPU time, trillions of tokens, and careful orchestration.

Training Objective

The training objective is deceptively simple — minimize cross-entropy loss at every position:

where $P_\theta$ is the model’s predicted probability for the correct next token.

This is computed efficiently: a single forward pass processes all T tokens in parallel (thanks to causal masking), and the loss is computed at every position simultaneously.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import torch
import torch.nn.functional as F

def compute_lm_loss(model, input_ids):
    """Standard causal LM loss computation."""
    # input_ids: [B, T]
    outputs = model(input_ids[:, :-1])          # predict from tokens 0..T-2
    logits = outputs.logits                      # [B, T-1, vocab_size]

    targets = input_ids[:, 1:]                   # targets are tokens 1..T-1
    loss = F.cross_entropy(
        logits.reshape(-1, logits.size(-1)),     # [B*(T-1), vocab_size]
        targets.reshape(-1),                      # [B*(T-1)]
    )
    return loss

Learning Rate Schedule

Almost all modern LLMs use a warmup + cosine decay schedule:

1
2
3
4
5
6
7
8
import math

def cosine_lr_schedule(step, warmup_steps, total_steps, max_lr, min_lr):
    """Warmup + cosine decay learning rate schedule."""
    if step < warmup_steps:
        return max_lr * step / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress))

Typical Hyperparameters

Pre-Training Datasets

Data Mix

Models aren’t trained on just web text. The training mix typically includes:

Training Loop with Gradient Accumulation

Real LLM training uses large effective batch sizes (millions of tokens) achieved through gradient accumulation across micro-batches:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def train_step(model, optimizer, dataloader, accumulation_steps, max_grad_norm=1.0):
    """Training loop with gradient accumulation and mixed precision."""
    model.train()
    optimizer.zero_grad()
    total_loss = 0.0

    for micro_step in range(accumulation_steps):
        batch = next(dataloader)  # [micro_batch_size, seq_len]
        input_ids = batch["input_ids"].to(device)

        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
            outputs = model(input_ids[:, :-1], labels=input_ids[:, 1:])
            loss = outputs.loss / accumulation_steps

        loss.backward()
        total_loss += loss.item()

    # Gradient clipping
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)

    optimizer.step()
    return total_loss

Example: micro_batch_size=4, seq_len=8192, accumulation_steps=32, 8 GPUs with DDP → effective batch = 4 × 8192 × 32 × 8 = ~8M tokens per step.

Training Infrastructure

Checkpointing and Recovery

Training runs crash. Hardware fails. Pre-training runs need robust checkpointing:

• Periodic checkpoints: Save every 500–1000 steps (model + optimizer state + RNG state)
• Async checkpointing: Save to storage without blocking training
• Loss spike detection: Monitor for loss spikes indicative of instability
• Automatic restart: Resume from last checkpoint on failure
• Deterministic recovery: Save RNG states for exact reproducibility

Loss Curves

A healthy pre-training loss curve shows rapid initial improvement followed by slow, steady decline:

Typical final pre-training loss: ~1.5–2.0 cross-entropy on general web text for a well-trained 7B model.

Compute Estimation

The compute needed for one forward+backward pass through the full dataset:

where $N$ = parameters, $D$ = tokens, and the factor 6 accounts for forward (2N FLOPs per token) + backward (4N FLOPs per token).

What’s Next

Pre-training produces a raw “base model” — capable but not aligned to human instructions or specialised domains. Before exploring the training stages further, the next chapter covers the optimisers and loss functions that drive all LLM training runs.

← Previous: Chapter 9 — Data for LLMs · Next: Chapter 11 — Optimizers & Loss Functions →

Last updated: April 2026