Pre-training systems and optimization

Connect the causal objective, cross-entropy, backpropagation, optimizer, mixed precision, parallelism, and checkpoints.

Applied AI · advanced · Session 11

Session contract

  • Compute cross-entropy loss.
  • Trace gradient, optimizer, and update.
  • Explain reliable checkpoint recovery.

Mechanism map

 batch (B,T) ─▶ ┌──────────────┐ logits z (B,T,V)
                │ forward pass │────────────────┐
                └──────┬───────┘                ▼
                       │ saved           ┌────────────────────┐
                       │ activations     │ CE = −z_y + logΣe^z│  L = 0.408
                       │ (memory)        └─────────┬──────────┘
                       ▼                           │
                ┌──────────────┐   chain rule      │
                │ backward     │◀──────────────────┘
                └──────┬───────┘  g = ∂L/∂θ
                       ▼
                ┌──────────────┐  m, v per parameter
                │    AdamW     │  θ ← θ − lr·m̂/(√v̂+ε) − lr·λ·θ
                └──────┬───────┘  BF16 compute │ FP32 master
                       ▼
 CHECKPOINT = θ + (m,v) + scheduler + scaler + data position + RNG states
              └─ weights only = APPROXIMATE resume, not exact ─┘

1. Causal objective

The model maximizes next-token likelihood at every position permitted by the causal mask. Mean loss aggregates valid positions and examples.

L = −Σ log p(x_t | x_<t)

Check — Causal objective

L = −Σ log p(x_t | x_<t) on a batch B = 2, T = 4 with 3 padded positions. How many terms do you divide by, and what happens to the curve if you divide by 8?

2. Cross-entropy from logits

Stable log-softmax subtracts log-sum-exp. Loss then selects the target log-probability. Larger logits matter only relative to others.

CE(z,y)=−z_y+log Σ exp(z_j)

Check — Cross-entropy from logits

CE(z,y) = −z_y + log Σ exp(z_j). Add +10 to each of the logits [2,1,0] and recompute: what is the loss, and which log-softmax property have you just demonstrated?

3. Backpropagation

The chain rule computes how each parameter contributed to loss. Saved activations consume memory; activation checkpointing trades recomputation for memory.

Check — Backpropagation

Activation checkpointing frees the memory of activations saved between forward and backward. What do you pay in exchange, and on which axis do you measure the net gain?

4. Optimizer

AdamW combines gradient moments, learning rate, and weight decay. Clipping can bound extreme gradients but does not repair faulty data or architecture.

Check — Optimizer

AdamW keeps m and v per parameter. For 7 billion parameters in FP32, quantify optimizer-state memory alone, and say why counting weights only underestimates the requirement by a large factor.

5. Precision and parallelism

BF16 reduces tensor memory without storing every state at full precision. Data parallelism replicates weights; tensor and pipeline parallelism split other dimensions with communication.

Check — Precision and parallelism

In BF16 compute is 16-bit but a master copy stays FP32. Why not move everything to BF16? Name what degrades first: the forward pass or the accumulation of updates.

6. Complete checkpoint

Exact recovery requires weights, optimizer state, scheduler, optional scaler, data position, and random states. A weights-only file is not a complete training checkpoint.

Check — Complete checkpoint

You resume training at step 12,000 from a weights-only file. Name three missing checkpoint fields and the observable effect of each on the loss curve over the next 200 steps.

Worked case — data

For logits [2,1,0] and target 0, softmax ≈ [0.665,0.245,0.090], so CE ≈ 0.408. The lab changes learning rate, gradient, and weight, then shows fields required for recovery.

Worked case — full trace

logits z = [2, 1, 0], target y = 0

exp: e² = 7.389   e¹ = 2.718   e⁰ = 1.000   Σ = 11.107
softmax       ≈ [0.665, 0.245, 0.090]       (sum = 1.000 ✅)
CE = −z_y + log Σ e^z = −2 + log(11.107) = −2 + 2.408 = 0.408   ✅

SHIFT INVARIANCE: z + 10 = [12, 11, 10]
  −12 + log(e¹² + e¹¹ + e¹⁰) = −12 + 12.408 = 0.408   ✅ identical loss

NAIVE SOFTMAX, z = [1000, 999, 998]
  exp(1000) → inf;  inf/inf → NaN                      ❌ loss destroyed
  stable version: subtract max(z) → [0,−1,−2] → 0.408  ✅

RESUME FROM WEIGHTS ONLY (step 12,000)
  m = 0, v = 0 → the first AdamW step is badly calibrated
  lr restarted at the top of the scheduler → transient loss spike
  (e.g. 0.4 → ~1.9 in the lab), then recovery  ❌

SHAPE CHECK: z is (B,T,V) and y is (B,T). The mean divides by the number of
UNMASKED positions, not by B×T.

Memory/throughput levers in pre-training: what each one costs

Lever What it buys What it actually costs
BF16 (compute) Half-size tensors, faster matmuls FP32 master copy kept; accumulations stay sensitive
Activation checkpointing Activation memory sharply reduced Forward recomputation: roughly 30% more time
Data parallelism Near-linear throughput in GPU count Weights and optimizer state replicated on every rank
Gradient clipping Bounds norm spikes, avoids NaNs Fixes neither dirty data nor an unstable architecture

Causal lab

Predict → change one variable → run → explain the delta

/interactives/curriculum/optimizer-checkpoint.html?lang=en

Common error 1

“Larger logits mean lower loss.”

The trace refutes it in one line: [2,1,0] and [12,11,10] both give exactly 0.408. Only logit differences matter; absolute scale cancels in the log-sum-exp.

Common error 2

“The .safetensors file is my training checkpoint.”

That is an inference checkpoint. Without m, v, scheduler, scaler, data position, and RNG states, resuming at step 12,000 restarts with m = v = 0, producing a loss spike plus repeated or skipped samples.

Validity boundary

Distributed performance depends on hardware, network, model size, and implementation; no lab estimate is a benchmark.

Evidence status

Established mechanisms; numerical simplifications are pedagogical.

Sources to inspect

  • Kingma & Ba, “Adam: A Method for Stochastic Optimization”, ICLR (2015).
  • Loshchilov & Hutter, “Decoupled Weight Decay Regularization” (AdamW), ICLR (2019).
  • Micikevicius et al., “Mixed Precision Training”, ICLR (2018).
  • Course source packet supplied by the owner; named-product details remain source-reported until primary verification.

Transfer challenge

Your 7B run dies at step 12,000. Two possible resumes: A) a weights-only file at the right step; B) a complete checkpoint that is 2,000 steps old.

  1. Predict the loss curve over the first 200 steps for each option.
  2. Cost each option (full checkpoint size; steps to replay).
  3. Choose, and set the threshold that would change your mind.

Synthesis

  • Causal objective
  • Cross-entropy from logits
  • Backpropagation
  • Optimizer
  • Precision and parallelism
  • Complete checkpoint

Exit ticket

Mechanism · trace · observation · boundary · evidence · next experiment

Instructor notes: Frame the problem before naming the mechanism. Collect an initial prediction and retain it for the exit ticket.

Instructor notes: These outcomes are observable: trace, calculation, comparison. A recited definition closes none of them.

Instructor notes: Connect every step to the next with a causal verb. Flag any merely decorative arrow.

Instructor notes: Start by writing shapes (B,T,V) and (B,T) on the board, then ask for the mean’s denominator. Those answering B×T are holding the most common production counting bug — let it live for a minute.

Instructor notes: Answer: divide by 5 — the unmasked positions — never by 8. Dividing by 8 flattens the curve optimistically while no prediction improves. Expected wrong answer: B×T = 8. Useful callback: the same denominator bug as intermediate session 11 (18 vs 13).

Instructor notes: Have them compute 0.408 by hand, then rerun with z + 10 without announcing the expected result. The surprise of equality beats an algebraic proof handed out in advance.

Instructor notes: Answer: the loss stays exactly 0.408 — shift invariance of the log-softmax; only logit differences matter. Expected wrong answer: “bigger logits, so lower loss”. Have them redo the trace line −12 + 12.408 before correcting.

Instructor notes: Have them draw activation memory as a stack growing through the forward pass. Ask where to cut the stack: checkpointing becomes a choice they make, not a library flag.

Instructor notes: Answer: you pay a forward recomputation, roughly 30% more time; the net gain is measured in freed activation memory, hence in feasible batch or model size. Expected wrong answer: “it is free since memory goes down”.

Instructor notes: Have the room quantify the full AdamW memory budget for 7B before showing the table. The gap between their estimate and the real total is the actual content of this slide.

Instructor notes: Answer: m + v in FP32 = 2 × 7 × 10⁹ × 4 bytes ≈ 56 GB — twice the FP32 weights (28 GB); weights + states ≈ 84 GB, i.e. 3× weights alone. Expected misses: counting one moment only, or forgetting FP32’s 4 bytes. The gap between their estimate and 56 GB is this slide’s content.

Instructor notes: Frame parallelism as a communication trade-off: ask what crosses the network each step in each of the three schemes. Reject any answer of the form “it is faster”.

Instructor notes: Answer: the accumulation of updates degrades first — tiny increments vanish in BF16’s short mantissa; the FP32 master copy preserves them. Expected wrong answer: “the forward pass goes wrong” — it tolerates reduced precision far better.

Instructor notes: Hand out a checkpoint field list with three fields missing and have them diagnose the expected symptom of each. This is the exercise that transfers best to a real incident.

Instructor notes: Answer — three fields and their symptoms: missing m,v → badly calibrated first AdamW step; missing scheduler → lr restarts, transient loss spike; missing data position → repeated or skipped samples; missing RNG states → non-reproducible run. Require the observable 200-step effect for every field named.

Instructor notes: Hide the final result. Elicit sign, shape, and order of magnitude before every operation.

Instructor notes: Walk line by line. Locate an inconsistency at the first faulty step, not only on the final line.

Instructor notes: Have learners fill the final row before revealing it: that trade-off is what decides in production.

Instructor notes: Retain initial and final values. Do not allow simultaneous changes that make the delta impossible to attribute.

Instructor notes: Have the group produce the smallest counterexample before giving the correction.

Instructor notes: Have the group produce the smallest counterexample before giving the correction.

Instructor notes: The boundary is not a footnote: it defines the cases where the mechanism no longer suffices.

Instructor notes: Separate verifiable mechanism, reported implementation choice, and experimental result. Evidence precision must match claim precision.

Instructor notes: Attribute each product choice to the supplied packet and retain reported status until an independent primary source confirms it.

Instructor notes: Expected: A resumes at the right step but with m = v = 0 and a restarted scheduler → transient spike then recovery; B replays 2,000 steps (pure GPU cost) with a healthy curve. Costing: full checkpoint ≈ 28 GB FP32 weights + 56 GB moments; B costs 2,000 × cost/step. Misconception to harvest: “weights are enough, the optimizer recalibrates quickly” — true only at small lr. Switch to B once the spike exceeds ~2× current loss. Eight minutes, pairs.

Instructor notes: Rebuild the chain without looking at the slides. Reopen only the first break.

Instructor notes: Six lines maximum. Compare with the initial prediction and name what actually changed. Cite one retained trace that lets a peer verify the conclusion.