Chapter 95 — LLM Pretraining, Scaling & Alignment

📖 Definition

Building a chat LLM happens in stages:

🔍 Pretraining

• Objective: minimise cross-entropy of the next token = maximise the probability of the real text.
• Perplexity = exp(average negative log-likelihood) — "how many choices the model is confused between"; lower is better, comparable only with the same tokenizer.
• Data pipeline matters most: deduplication, quality filtering (classifiers, heuristics), toxicity/PII filtering, language and domain mixing, removing benchmark test sets (decontamination).
• Compute estimate: training FLOPs ≈ 6 × N × D (N = parameters, D = training tokens).

Example: 7B params × 2T tokens × 6 ≈ 8.4 × 10²² FLOPs
At ~40% utilisation of an H100 (~400 TFLOP/s effective bf16 dense) ≈ 2.1 × 10⁸ GPU-seconds
≈ 58,000 GPU-hours ≈ 1,000 GPUs for ~2.5 days

🔍 Scaling Laws

• Loss falls smoothly (as a power law) as you increase parameters, data and compute together.
• Chinchilla (2022): for a fixed compute budget, the optimal model uses about 20 training tokens per parameter. Many earlier models (e.g. GPT-3: 175B params, 300B tokens) were under-trained.
• In practice teams now train smaller models on far more tokens than Chinchilla-optimal (e.g. 8B model on 15T tokens) because inference cost depends on model size — a smaller, over-trained model is cheaper to serve forever.
• Emergent abilities debate: some skills appear to jump with scale, but part of that is the metric (exact-match) rather than the underlying ability.

🔍 Distributed Training (Senior Must-Know)

A 70B model's weights alone are 140 GB in bf16 — no single GPU can train it. Strategies:

Large runs combine them ("3D parallelism"): tensor parallel inside a node, pipeline across nodes, data/FSDP across groups.

Memory per parameter with AdamW + mixed precision ≈ 16 bytes: bf16 weights (2) + bf16 grads (2) + fp32 master weights (4) + Adam m (4) + Adam v (4). So a 7B model needs ~112 GB before activations — why ZeRO/FSDP exist.

Other realities: checkpoint often (hardware fails), watch for loss spikes (skip batches, rewind), monitor MFU (model FLOPs utilisation, 35–55% is good).

🔍 Supervised Fine-Tuning (SFT)

• Format examples with the chat template (system/user/assistant turns, special tokens).
• Mask the loss on prompt tokens — train only on the assistant's response.
• Quality beats quantity: a few thousand excellent, diverse examples can beat a million noisy ones (LIMA finding).
• Risks: overfitting to style, catastrophic forgetting of general skills, teaching the model to hallucinate if responses contain facts it doesn't know.

🔍 RLHF (Reinforcement Learning from Human Feedback)

1. Collect comparisons: for a prompt, show humans 2+ model responses; they pick the better one.
2. Train a reward model (RM): an LLM with a scalar head, trained so score(chosen) > score(rejected) using the Bradley–Terry loss −log σ(r_chosen − r_rejected).
3. Optimise the policy with PPO: generate responses, score them with the RM, update the model to increase reward minus a KL penalty that keeps it close to the SFT model.

objective = E[ reward(x, y) ] − β · KL( π_θ(y|x) || π_ref(y|x) )

Why the KL penalty? Without it the model finds reward hacking tricks — outputs the reward model loves but humans don't (overly long, flattering, repetitive).

Downsides of PPO-RLHF: four models in memory (policy, reference, reward, value), unstable, many hyperparameters.

🔍 DPO (Direct Preference Optimization)

DPO skips the reward model and RL loop. It uses the same preference pairs and a simple classification-style loss directly on the policy:

loss = −log σ( β · [ (log π_θ(y_w|x) − log π_ref(y_w|x)) − (log π_θ(y_l|x) − log π_ref(y_l|x)) ] )
          y_w = preferred answer, y_l = rejected answer, π_ref = frozen SFT model

💻 Code Example — DPO Loss in PyTorch

import torch
import torch.nn.functional as F

def sequence_logprob(model, input_ids, labels, mask):
    """Sum of log-probs of the response tokens only (mask = 1 on response tokens)."""
    logits = model(input_ids).logits[:, :-1]
    logp = logits.log_softmax(-1).gather(-1, labels[:, 1:].unsqueeze(-1)).squeeze(-1)
    return (logp * mask[:, 1:]).sum(-1)

def dpo_loss(policy, ref, batch, beta=0.1):
    pw = sequence_logprob(policy, batch["chosen_ids"], batch["chosen_ids"], batch["chosen_mask"])
    pl = sequence_logprob(policy, batch["rejected_ids"], batch["rejected_ids"], batch["rejected_mask"])
    with torch.no_grad():
        rw = sequence_logprob(ref, batch["chosen_ids"], batch["chosen_ids"], batch["chosen_mask"])
        rl = sequence_logprob(ref, batch["rejected_ids"], batch["rejected_ids"], batch["rejected_mask"])
    margin = beta * ((pw - rw) - (pl - rl))
    return -F.logsigmoid(margin).mean()

📊 Alignment Methods Compared

Reasoning models are trained with RL to produce long chains of thought before answering; spending more tokens at inference ("test-time compute") improves accuracy on hard problems.

🔍 Synthetic Data & Distillation

• Distillation: a large "teacher" model generates answers (or probability distributions) and a small "student" learns from them — cheaper to serve with much of the quality.
• Synthetic data: generate instructions, rephrase documents, create hard negatives, self-instruct. Filter heavily (verifiers, judges, dedup) or quality collapses.
• Model collapse risk: training repeatedly on model-generated text without fresh human data reduces diversity.

🧠 What a Senior Answer Adds

• Can estimate compute (6ND) and memory (16 bytes/param) on a whiteboard.
• Explains why DPO replaced PPO in many teams (simplicity, stability) and what each loses.
• Talks about data quality, decontamination and evaluation as the real levers, not just architecture.

🎯 Likely Interview Questions

1. Walk through how a chat model is built. — Pretraining → SFT → preference tuning (RLHF/DPO) → optional RL for reasoning; plus safety training and evals.
2. Base model vs instruct model? — Base continues text; instruct follows instructions in a chat format.
3. What do scaling laws say? What did Chinchilla change? — Loss falls predictably with scale; for fixed compute use ~20 tokens/param; earlier models were under-trained.
4. Why train small models on more tokens than "optimal"? — Inference cost dominates over the model's life.
5. Estimate training compute for a model. — FLOPs ≈ 6 × params × tokens.
6. Explain data, tensor and pipeline parallelism, and ZeRO. — See table; ZeRO shards optimizer states / grads / params across data-parallel ranks.
7. How much memory does full fine-tuning of a 7B model need? — ~16 bytes/param ≈ 112 GB + activations → multiple GPUs or PEFT.
8. Explain RLHF. — Preference data → reward model → PPO with KL penalty to the SFT reference.
9. What is reward hacking? — Policy exploits reward-model flaws (length, flattery); mitigated with KL penalty, better RMs, length control.
10. DPO vs PPO? — DPO optimises preferences directly with a classification loss, no reward model or sampling loop; simpler and more stable.
11. Why mask prompt tokens in SFT? — We want to learn to produce answers, not to predict the user's text.
12. What is perplexity? — exp(mean NLL); measures how well the model predicts text; only comparable across the same tokenizer.
13. What is distillation? — Train a small student on a big teacher's outputs to get most of its quality at lower cost.

← Transformer Architecture | Index | Next: Fine-Tuning & LoRA →