Chapter 30 Designing a Pre-Training Run from Scratch

30.1 From Compute Budget to Architecture

30.1.1 Choosing Model Size, Token Count, and Schedule

You have a fixed compute budget - say $10^{23}$ FLOPs - and you need to decide, before writing a single line of training code, how many parameters your model should have and how many tokens to train on. Get this wrong in either direction and you either leave capability on the table or burn money training a model that costs too much to serve. This decision comes up in system design interviews at every frontier lab, typically framed as “given $X$ H100-days, design the training run.” The answer is not a guess; it follows directly from isoFLOP analysis and the Chinchilla scaling law, two tools you can derive from small proxy runs well before committing to the full job.

Every transformer training run consumes roughly

floating-point operations, where $N$ is the number of non-embedding parameters and $D$ is the number of training tokens. The factor of 6 accounts for the forward pass ($\approx 2 N D$) plus the backward pass ($\approx 4 N D$ including the parameter gradient and the activation gradient). This approximation holds within 10% for most dense decoder architectures in the regime $D \gg \sqrt{N}$.

The Chinchilla paper (Hoffmann et al., 2022) posits that cross-entropy loss follows a joint scaling law of the form

where $L_{\infty}$ is the irreducible error of the data distribution and $\alpha , \beta \approx 0 . 34$ empirically. Both the model-size term and the data term decay as independent power laws; the optimal allocation sets their marginal contributions equal.

The practical procedure is the isoFLOP sweep: fix a compute budget $C$, then train a grid of models with different $N$ while adjusting $D = C / ( 6 N )$ to keep total FLOPs constant. Each model reaches a different terminal loss. Fitting a quadratic to those losses and reading off the minimum gives you $N^{\ast}$ and $D^{\ast} = C / ( 6 N^{\ast} )$ simultaneously. Repeating this at several compute scales reveals that the optimal ratio $D^{\ast} / N^{\ast}$ follows its own power law, and at the scales studied by Chinchilla the ratio stabilizes near 20 tokens per parameter.

The punchline that every engineer should be able to state under pressure: At a fixed compute budget, training-optimal model size and token count scale as $C^{0 . 5}$ each; the ratio of tokens to parameters determines how the budget is split, and that ratio has climbed steadily from 20:1 (Chinchilla) toward 40-700:1 in modern inference-optimized runs.

The reason the power law form is so robust is that language modeling at scale is essentially a high-dimensional nonparametric regression problem. Classical theory predicts error should decay as $n^{- 1 / d}$ where $d$ is the intrinsic dimensionality of the task. For natural language, $d$ is large, which yields exponents much smaller than 0.5 and explains why empirical slopes around 0.1-0.3 are consistently observed rather than the $n^{- 1}$ or $n^{- 0 . 5}$

rates you would expect for simple function classes.

The training-optimal versus inference-optimal distinction is where modern practice diverges most sharply from the original Chinchilla result. When the Chinchilla paper was written, LLMs were research artifacts and inference cost was irrelevant. The training-optimal answer minimizes loss per FLOP spent training. The inferenceoptimal answer minimizes loss per FLOP spent serving the model, which means accepting higher training cost in exchange for a permanently smaller model. Once models become products, inference cost dominates because you pay training cost once and serving cost for every query indefinitely. This is why Llama 3 trained at roughly 39 tokens per parameter, why recent Qwen models trained on 30 trillion tokens, and why the modern rule of thumb has moved well above 20:1.

The Kaplan (2020) paper estimated a lower ratio - closer to 5-10 tokens per parameter - because its analysis used models trained to a fixed step count rather than to full convergence with a cosine cooldown. A cosine learning rate schedule cannot be truncated early; the cooldown phase is responsible for a substantial fraction of the final loss reduction. Models stopped mid-cosine are not valid samples on the training-optimal frontier, which caused Kaplan to underestimate the token contribution systematically. Chinchilla corrected this by training every model to its natural termination point.

A Worked Example

Assume a compute budget of $C = 2 \times 10^{23}$ FLOPs, roughly what you get from 1,000 H100s running for 30 days at 40% model FLOP utilization.

Step 1 - Training-optimal allocation. From $C = 120 N^{2}$ (substituting $D = 20 N$ into $C = 6 N D$):

Step 2 - Inference-optimal allocation (7B target). If serving cost constraints require a 7B model, spend all of $C$ on tokens:

The token-to-parameter ratio becomes $4 . 8$T$/ 7$B $\approx 686$:1. The model trains for far longer than Chinchilla would prescribe; each additional token past the Chinchilla point delivers diminishing returns on training loss but still improves the model because the irreducible floor has not been reached.

Step 3 - Validate with a proxy isoFLOP sweep. Before committing to either configuration, run five to seven models ranging from 1B to 10B parameters, each at fixed $C_{\text{proxy}} = C / 100 = 2 \times 10^{21}$ FLOPs. Fit a quadratic to the terminal losses and confirm the minimum. If the minimum shifts consistently toward smaller $N$, the data suggests an inference-optimal regime; if it aligns near the Chinchilla ratio, training-optimal dominates. This proxy sweep costs about 1% of the full budget and de-risks the entire run.

What You Actually Decide in Practice

• Pick the token-to-parameter ratio based on deployment, not theory. If you will serve the model at scale, lean toward 40-200:1 or higher. If this is a research artifact unlikely to see heavy traffic, 20:1 is fine. Never default to Chinchilla’s 20:1 for a production model without thinking through the serving economics.

• Use the WSD (warmup-stable-decay) learning rate schedule. A cosine schedule requires committing to a fixed endpoint before training starts; early checkpoints are incomparable because they have different cooldown trajectories. WSD holds the learning rate flat during the stable phase, letting you branch off cooldown ramps at arbitrary token counts. This turns the Chinchilla isoFLOP analysis from an $O ( n^{2} )$ training experiment into roughly an $O ( n )$ one.

• Exclude embedding parameters from $N$ when fitting scaling laws. Embedding layers have a different scaling behavior than the attention and MLP blocks. Including them distorts the log-linear fit and produces a visible kink in the scaling curve. Always report and fit non-embedding parameter counts.

• The critical batch size grows as your target loss shrinks. As training progresses and loss decreases, gradient signal becomes more precise and you can productively use larger batches. Llama 3 increased batch size mid-run for exactly this reason. Plan for at least one batch-size ramp in your training schedule.

• If you are tuning learning rate and batch size without mu-P, fit a scaling law on them. Without maximal update parameterization, the optimal learning rate shifts roughly as $1 /$width as the model grows. Train five-plus proxy runs at different scales, find the optimal learning rate at each, and fit the resulting scaling law to extrapolate the optimal rate for your target model. Do not transplant the learning rate from a small run without this correction; it will be too high and the large model will diverge or converge to a suboptimal loss.

How This Shows Up in Interviews

• “You have a compute budget of $10^{24}$ FLOPs. Walk me through how you choose model size and token count.” Strong answer: state $C = 6 N D$, describe the isoFLOP sweep, report the training-optimal answer ($N \approx \sqrt{C / 120}$, $D = 20 N$), then immediately flag the inference-optimal alternative and ask the interviewer which objective matters.

• “Why did Kaplan’s scaling law predict a different optimal ratio than Chinchilla?” Strong answer: cosine LR cannot be truncated; Kaplan’s models were not all trained to their natural termination points, so the sample of runs on the optimal frontier was biased toward shorter, data-light runs. Chinchilla fixed the experimental design.

• “Why do modern models like Llama 3 train at 40:1 or higher tokens per parameter instead of Chinchilla’s 20:1?” Strong answer: Chinchilla optimizes for lowest loss per training FLOP; once models are products, you want the smallest model with acceptable intelligence because inference cost scales with every query. Over-training a small model trades one-time training cost for permanent serving savings.

• Senior follow-up: “How would you modify the Chinchilla analysis if you’re training a mixture-of-experts model?” Sketch of the senior answer: active parameters $N_{\text{active}} \ll N_{\text{total}}$ change both the compute formula ($C \approx 6 N_{\text{active}} D$) and the scaling law exponents. You must redo the isoFLOP sweep using active parameters, not total parameters, to find the right allocation. Hunyuan-Large reported a ratio of $\sim$96 active-parameter tokens per active parameter after doing exactly this replication.

Key Takeaways

• Training compute is $C \approx 6 N D$; fixing $C$ and sweeping $N$ (isoFLOP analysis) is the canonical method for finding the compute-optimal model size and token count.

• Chinchilla’s 20 tokens per parameter is training-optimal, not inference-optimal. Modern production runs use 40-700:1 to reduce serving cost at the expense of training cost.

• Never truncate a cosine LR run early and compare checkpoints; use a WSD schedule to reuse stable-phase checkpoints across different token-count targets.

• Always fit scaling laws on non-embedding parameters; embedding layers follow different scaling dynamics and distort the log-linear fit.

• Validate every allocation decision with a cheap proxy isoFLOP sweep ($\sim$1% of the full compute budget) before launching the full run.

30.1.2 Picking a Parallelism Strategy for Your Cluster

Once you have fixed your model size and token count, the next design decision is how to distribute the work across your GPU cluster. This question shows up in system design interviews as “you have a 70B model and 512 H100s, how do you parallelize training?” The answer requires knowing which axes of parallelism exist, what hardware constraints bound each one, and how to compose them into a configuration that fits in memory and keeps GPU utilization high. Getting this wrong means either hitting an OOM before the first forward pass or spending half your compute budget waiting on communication rather than doing math.

Modern LLM training uses a combination of three parallelism axes, typically composed multiplicatively:

• Tensor parallelism (TP): splits individual weight matrices across GPUs within a layer. An all-reduce is required after every matrix multiply, so TP ranks must be connected by NVLink or an equivalently high-bandwidth interconnect. TP degree is almost always set to the number of GPUs per node (typically 8).

• Pipeline parallelism (PP): partitions the model’s layers into sequential stages, one stage per PP rank. GPUs in different stages communicate only at stage boundaries via point-to-point sends. This communication uses slower inter-node links (InfiniBand) but is much less frequent than TP all-reduces.

• Data parallelism (DP): replicates the full model (or a TP$\times$PP shard of it) across independent groups that each process a different micro-batch. Gradients are aggregated with an all-reduce across DP ranks after each backward pass.

The total GPU count decomposes cleanly:

Within the DP dimension, ZeRO (Zero Redundancy Optimizer) shards optimizer states, gradients, and optionally parameters across the DP group rather than replicating them. ZeRO stage 1 shards only optimizer states, stage 2 adds gradient sharding, and stage 3 shards parameters as well. Each stage reduces memory proportionally to the DP degree at the cost of additional all-gather communication.

The standard recipe is: fill TP within a node (TP $=$ GPUs per node), fill PP across nodes to get the model into memory, then fill the rest of your cluster with DP to increase throughput.

The hierarchy follows the hardware topology directly. NVLink bandwidth inside a node is typically 600- 900 GB/s all-reduce aggregate, whereas InfiniBand between nodes is one to two orders of magnitude slower per link. Tensor parallelism inserts an all-reduce after every matrix multiply in every transformer layer - for a 70B model that is hundreds of all-reduces per forward-backward pass. Those collectives must complete in microseconds to not bottleneck compute, which means they can only run efficiently over NVLink. Placing TP ranks on separate nodes would stall the entire pipeline waiting for slow cross-node traffic.

Pipeline parallelism, by contrast, only sends a single activation tensor at each stage boundary once per micro-batch. The communication volume is small (hidden dimension $\times$ sequence length $\times$ micro-batch size in BF16) and the timing is predictable, so it tolerates the higher latency of inter-node links.

The main cost of pipeline parallelism is the pipeline bubble: at the start and end of each pipeline flush, some stages are idle. With $p$ stages and $m$ micro-batches per flush, the bubble fraction is approximately $( p - 1 ) / m$. The interleaved pipeline schedule (used by Megatron-LM) reduces this to $( p / v - 1 ) / m$ by assigning $v > 1$ non-contiguous layer chunks to each rank, at the cost of $v \times$ more point-to-point messages per micro-batch. As a rule of thumb, keep $m \geq 4 p$ to hold the bubble fraction below 25%; otherwise pipeline parallelism wastes too much GPU time on idle fills and drains.

A Worked Example

Setup: a 7B parameter dense decoder (BF16 weights), trained on 512 H100 SXM5 GPUs organized as 64 nodes $\times$ 8 GPUs per node, 80 GB HBM per GPU.

Memory accounting for a single GPU without parallelism:

• Parameters (BF16): $7 \times 10^{9} \times 2$ bytes $= 14$ GB.

• Adam optimizer states (FP32 master weights, first and second moments): $7 \times 10^{9} \times 12$ bytes $= 84$ GB.

• Total static memory: $\approx 98$ GB - exceeds the 80 GB HBM capacity before accounting for activations.

Target configuration: TP $= 8$, PP $= 8$, DP $= 8$. Check: $8 \times 8 \times 8 = 512$ GPUs. $\checkmark$

Memory after parallelism (per GPU):

• Parameters sharded by TP$\times$PP: $14$ GB$/ ( 8 \times 8 ) = 219$ MB.

• Optimizer states with ZeRO stage 1 across DP$= 8$: $84$ GB$/ ( 8 \times 8 \times 8 ) = 164$ MB.

• Activations (32 layers, sequence length 4096, micro-batch 1): roughly 8-12 GB with recomputation, depending on remat strategy.

Total per-GPU footprint: $\approx 10$-15 GB, well within 80 GB. The headroom accommodates gradient buffers and CUDA workspace.

Pipeline bubble check: with $p = 8$ pipeline stages, targeting $m = 32$ micro-batches per flush:

Enabling the interleaved schedule with $v = 2$ interleaved chunks reduces this to:

The global batch size at these settings is DP $\times m \times$ micro-batch $= 8 \times 32 \times 1 = 256$ sequences. If the target global batch exceeds 256, increase gradient accumulation steps or the micro-batch size until the batch matches the critical batch size derived from the Chinchilla analysis.

What You Actually Decide in Practice

• Set TP equal to GPUs per node, almost always. For H100 nodes with 8 GPUs, TP $= 8$. Going cross-node with TP (TP $= 16$, $32$) is technically possible but the NVLink advantage disappears and latency on TP all-reduces typically kills MFU. Only attempt cross-node TP if your model’s hidden dimension is so large (e.g., $d_{\text{model}} \geq 16384$) that single-node TP cannot absorb the compute.

• Size PP to fit the model, then stop. Choose the smallest PP degree such that per-GPU memory (parameters, optimizer states, activations) fits in HBM. Every doubling of PP adds pipeline bubble overhead and doubles the number of inter-node point-to-point messages. If ZeRO-3 across DP gets the model into memory without PP, prefer that; the all-gather communication in ZeRO-3 is typically lower latency than managing pipeline schedules.

• Keep pipeline bubbles below 15% with the interleaved schedule. Budget $m \geq 2 v ( p - 1 )$ microbatches per flush when using $v$ interleaved chunks. If the required $m$ pushes your global batch size above the critical batch size, reduce PP or reduce $v$.

• For MoE models, add expert parallelism (EP) as a fourth axis. EP routes each token’s experts to a separate GPU group. The all-to-all dispatch and combine collectives replace the DP all-reduce for the MoE layers. Set EP degree equal to the number of experts (or a divisor of it) such that each EP rank holds exactly one expert. The total GPU count becomes TP $\times$ PP $\times$ DP $\times$ EP, but DP and EP share the same physical dimension in most implementations (EP is implemented as DP over the expert dimension).

• Match your global batch size to the critical batch size. The critical batch size scales predictably with the inverse of the target loss. For a 7B model targeting a terminal loss around 2.0 nats, empirical estimates put the critical batch size near 1-4M tokens. Below that threshold, adding more DP replicas (more GPUs) is nearly as efficient as taking more steps. Above it, extra DP replicas produce diminishing gradient signal. Design your TP $\times$ PP $\times$ DP product so that the resulting global batch lands at or just below the critical batch size; then grow the batch size in stages as loss improves mid-training.

How This Shows Up in Interviews

• “You have a 70B model and 512 H100s. Walk me through your parallelism strategy.” Strong answer: compute static memory (70B $\times$ 2 bytes $+$ optimizer states), determine TP $= 8$ (per node), choose PP such that per-GPU memory fits, fill remaining GPUs with DP. State the resulting global batch size and check it against the critical batch size. Quote bubble fraction.

• “What is the pipeline bubble and how do you reduce it?” Strong answer: bubble fraction $\approx ( p - 1 ) / m$ for the simple 1F1B schedule; the interleaved schedule with $v$ chunks reduces it to $( p / v - 1 ) / m$ at the cost of $v \times$ more P2P sends per micro-batch. Increase $m$ (more micro-batches per flush) to reduce the bubble without changing $v$.

• “When would you use ZeRO-3 instead of pipeline parallelism for a 70B model?” Strong answer: ZeRO-3 shards parameters across DP ranks, so it works only if the DP group is large enough and the all-gather latency per step is acceptable. For a 70B model on 8 nodes, ZeRO-3 with DP $= 64$ shards parameters to $\sim$2 GB per GPU - affordable, and you avoid pipeline bubble overhead entirely. The tradeoff is more frequent all-gather communication on the critical path. ZeRO-3 wins when you have a large DP group within a fat-tree fabric; PP wins when pipeline stages fit on separate chassis with limited cross-chassis bandwidth.

• Senior follow-up: “How does sequence parallelism change this analysis for a 1M-token context model?” Sketch of the senior answer: at long sequence lengths, activations (which scale as sequence length $\times$ batch $\times d_{\text{model}}$) dominate over parameter memory. Sequence parallelism (SP) shards the sequence dimension across TP ranks, reducing activation memory by the TP degree. You need all-gather before the QKV projection and reduce-scatter after the output projection, exactly replacing the TP allreduces. For a 1M-token context, SP is mandatory: activations without SP can exceed 10 GB per layer per GPU even at micro-batch $= 1$, making the model impossible to fit at TP $= 8$ without aggressive activation recomputation.

Key Takeaways

• Total GPUs $=$ TP $\times$ PP $\times$ DP; set TP to the NVLink-connected group size (typically 8), use PP to get the model into memory, fill remaining capacity with DP.

• Pipeline bubble fraction is $( p - 1 ) / m$ for the basic schedule and $( p / v - 1 ) / m$ for the interleaved schedule; keep it below 15% by ensuring $m \gg p$.

• ZeRO-3 is often preferable to PP for single-node-scale models (up to $\sim$30B) when the DP group is large; PP is necessary when parameters plus optimizer states exceed per-node memory even after TP sharding.

• The global batch size produced by your parallelism configuration must stay at or below the critical batch size; ramp the batch size mid-training as the terminal loss decreases.

• For MoE models, add expert parallelism as a fourth axis; for ultra-long-context models, add sequence parallelism within the TP group to manage activation memory.

30.2 Operations and Failure Modes

30.2.1 Checkpointing, Fault Tolerance, and Silent Data Corruption

When Meta trained Llama 3, they logged over 50 job interruptions from hardware failures across the full run - roughly one every few days at production scale. Training a frontier model across thousands of GPUs is a distributed system in the reliability-engineering sense: components will fail, and your design must make the run survivable. Checkpointing strategy appears in senior-level design interviews precisely because the right answer is not simply “save every N steps.” It requires async I/O pipelines, rolling retention policies, cross-replica validation, and a concrete plan for the failure mode that inflicts the most damage: a GPU that silently produces wrong results without crashing or throwing an error.

A checkpoint must capture everything needed to resume training from an exact point: model parameters, optimizer state (momentum and variance tensors for Adam), the data loader cursor, the random number generator state for both Python and CUDA, and the current global step counter. Miss any of these and a resumed run diverges from the original trajectory - sometimes immediately, sometimes subtly over thousands of steps.

At scale, a synchronous checkpoint blocks the entire training loop. For a 70B-parameter model in ZeRO-3 across 512 H100s, the full checkpoint size is approximately:

Sharded across 512 workers that is about 1.9 GB per rank. An NVMe RAID writing at 10 GB/s can flush each shard in under 200 ms - but only if you do not wait for every rank serially. The standard solution is asynchronous checkpointing: each GPU serializes its shard into host DRAM (fast, under 100 ms), the training loop resumes immediately, and a background CPU thread drains that buffer to persistent storage. The GPU never stalls waiting for disk.

Rolling checkpoints address a second failure mode: a checkpoint file that is itself corrupted mid-write. Keep the three most recent sharded saves (at steps $N$, $N - \delta$, $N - 2 \delta$) and delete the oldest only after the newest has been checksum-validated. This gives you two rollback opportunities if the most recent save is bad.

A checkpoint is not a save file - it is the full resumable state of a distributed computation, and its reliability is a product of every component in that chain.

Asynchronous saves overlap I/O with compute so that checkpointing overhead approaches zero at reasonable intervals (every 200-1,000 steps). The host-RAM staging buffer absorbs the burst, and the background thread handles the slower NVMe write while training continues. Blocking the GPU until the file is flushed wastes expensive accelerator time for no algorithmic benefit.

Rolling policies with post-write validation address the corruption-during-write failure mode. A naively overwritten single checkpoint is worse than no checkpoint at all: you lose your fallback. A three-generation rolling policy with per-shard checksum provides both recency and a verified fallback, at the cost of roughly $3 \times$ the checkpoint storage.

Silent data corruption (SDC) is the harder case and the one that teams are often underprepared for. ECC memory catches single-bit DRAM errors, but multi-bit errors, PCIe link faults, and rare CUDA kernel bugs can produce numerically plausible but incorrect gradient values. The training loss does not spike - the model simply learns the wrong thing, gradually. The detection strategy that works in practice is maintaining a frozen canary batch: a small fixed set of held-out examples whose loss is evaluated separately every 100 steps. A canary loss that drifts upward while training loss looks healthy is an early SDC signal. For ZeRO-1 or ZeRO-2, where all replicas hold identical parameters after each reduce, you can also compute the global gradient norm independently on two separate ranks and compare; a sustained disagreement beyond numerical tolerance flags a corrupted reduce.

A Worked Example

Consider a 7B model training on 64 H100s under FSDP (ZeRO-3). Each GPU holds roughly:

Checkpointing every 500 steps with async saves adds negligible overhead. At 800K tokens per step, 500 steps equals 400M tokens of training - the maximum exposure window for any hardware failure.

At step 8,200 a single GPU develops a flaky NVLink lane that silently biases its contribute to the reducescatter by $\sim$0.1% per step. No NaN appears; the training loss trends normally. At step 10,300, the canary batch loss is 0.09 nats above its expected trajectory - 3.2 standard deviations from the rolling baseline. The run is paused.

Rolling back to the step-8,000 checkpoint (the last clean save before the anomaly), swapping the affected GPU, and resuming costs 300 steps of recompute across 64 GPUs. Without the canary, the team would have continued for weeks and discovered the problem only during post-training evaluation, losing the entire run.

What You Actually Decide in Practice

• Checkpoint interval. Every 200-500 steps is the practical range for large runs. Too frequent and async I/O cannot keep up with training throughput; too infrequent and a single GPU failure costs hours of recompute. At 500 steps, 1M tokens/step, and 512 H100s at $3/hour, one failure costs roughly $3,000 of wasted compute.

• Rolling vs. permanent snapshots. Keep three rolling checkpoints for fast rollback, plus permanent snapshots at fixed milestones (e.g., every 10K steps). Permanent snapshots let you restart mid-training experiments - if you changed the data mix at step 50K, you can re-evaluate from that exact checkpoint without rerunning from scratch.

• Validate every save before deleting the previous generation. After the background thread flushes to disk, read back the shard header and compare a checksum. A corrupted save that silently passes is worse than a failed save that triggers a retry. Track per-rank validation status in a distributed metadata store.

• Canary batch sizing. 512-2,048 held-out examples is sufficient. Evaluate every 100 training steps. Alert if the canary loss deviates by more than two standard deviations from a 500-step rolling mean. This adds roughly 0.1% overhead at a 128K-token canary batch.

• Save the full data cursor. Record the shard index, file offset, and shuffle seed alongside model state. Without these, a resumed run either replays already-seen data or skips data, introducing distribution shifts that are nearly impossible to diagnose later. This is the most commonly omitted element in naive checkpoint implementations.

How This Shows Up in Interviews

• “Walk me through how you would checkpoint a 70B model training run across 512 GPUs with FSDP.” Strong answer: describe ZeRO-3 sharding, async host-RAM staging, background flush to NVMe, per-rank checksum validation, and a three-generation rolling retention policy. Mention that optimizer state dominates checkpoint size and that the data loader cursor and RNG state must be saved alongside model weights.

• “What is silent data corruption and why is it particularly dangerous during training?” Strong answer: SDC produces numerically plausible but incorrect gradient values - no crash, no NaN - so the training loss can look healthy while the model learns from a corrupted signal for thousands of steps. Detection requires out-of-band validation (canary batch or cross-replica gradient comparison), not in-band loss monitoring alone.

• Senior follow-up: “You are 20 days into a 30-day run and discover a GPU has been silently corrupting gradients since day 5. What do you do?” Senior answer: quantify the contamination window by comparing canary loss to its expected trajectory - if divergence is small and within two standard deviations of noise, the model may be salvageable from a day-5 checkpoint with a targeted rerun. If the drift is large, a 15-day rollback may be unavoidable. In either case: isolate and quarantine the affected hardware, verify the day-5 checkpoint with an independent forward pass on a clean reference node, add cross-replica gradient audits before resuming, and file a hardware incident report so the vendor can diagnose the root cause.

Key Takeaways

• Asynchronous checkpointing stages weights to host RAM first so training resumes before the file hits disk - overhead approaches zero at 200-500 step intervals.

• Keep three rolling checkpoints plus permanent milestone snapshots; validate each save with a checksum before deleting the previous generation.

• Save optimizer state, data loader cursor, and RNG seeds alongside model parameters - a resume without these is not a true resume.

• Silent data corruption produces no NaN and no crash; detect it with a frozen canary batch evaluated every ~100 steps, or via cross-replica gradient comparison on ZeRO-1/2 runs.

• At Llama-3 scale, hardware failures are a statistical certainty; a fault-tolerant checkpoint strategy limits blast radius to hundreds of steps rather than days of recompute.

30.2.2 Monitoring Loss Spikes and Divergence

Three hours into a 21-day training run, the per-step loss plot jumps from 2.8 to 11.3. The gradient norm has been climbing for 50 steps. Do you roll back and resume, wait for self-recovery, or stop the run entirely? Getting this decision wrong in either direction wastes GPU-months. The judgment call itself - and the instrumentation that makes it reliable - comes up in system design interviews at every lab that trains models at scale, because it tests whether you understand training dynamics at the level of individual steps rather than epoch averages.

A training run produces two qualitatively different instability modes. A spike is a transient loss increase that recovers within a few hundred steps; the optimizer has encountered an anomalous gradient and the momentum terms absorb the shock. Divergence is a non-recovering increase - the optimizer has entered a region from which gradient descent cannot escape, usually because the gradient scale has overwhelmed the learning rate or the Adam second-moment estimate has underflowed.

Distinguishing these two in the first minutes of a suspected event requires per-step instrumentation, not epoch-level averages. The canonical metrics to log at every training step are:

• Training loss (cross-entropy in nats)

• Global $\ell_{2}$ gradient norm before clipping

• Effective learning rate (scheduler output at this step)

• Hardware throughput (tokens per second per GPU, as a hardware health proxy)

And, evaluated on a fixed held-out set every $K = 100$ steps:

• Canary batch loss (orthogonal signal independent of in-distribution training data)

Gradient norm is the leading indicator. During stable pre-training on a well-filtered dataset, the unclipped global norm rarely exceeds $2 \times$ the gradient clipping threshold. When the unclipped norm crosses $5$-$10 \times$ the clipping value, a bad batch or a numerical instability has already started to propagate - the loss spike will follow 10-50 steps later. This lead time is your intervention window.

Per-step gradient norm is the earliest warning signal; training loss is the lagging indicator.

Logging at step granularity rather than epoch granularity catches transient events before they compound. A single corrupted document that appears once may spike the loss for five steps and then recover; logged only at epoch boundaries, you never see it, and you cannot correlate the event with a data cursor position to filter the offending sample from future training.

The canary batch provides an orthogonal signal that distinguishes two very different root causes. If training loss spikes but canary loss stays flat, the model is healthy and the event is almost certainly a data artifact - something in the current batch that is anomalous relative to the rest of the distribution. If both training loss and canary loss spike, or if canary loss drifts upward over hundreds of steps while training loss looks normal, the model’s learned representations are degrading, which points toward a more serious problem: silent data corruption, a numerical instability in a specific layer, or a bug in the optimizer implementation.

Epoch-level averaging destroys this signal entirely. A brief loss spike that resolves within 200 steps is invisible in a per-epoch summary, yet the data cursor position that caused it can still be identified and the bad document filtered, preventing recurrence.

A Worked Example

A 13B model is training on 128 H100s. The stable phase runs from step 12,000 to 14,900: training loss $\approx 2 . 2$ nats, gradient norm $\approx 1 . 3$, clipping threshold 1.0.

At step 14,950 the unclipped gradient norm climbs to 4.1 - roughly $3 \times$ above the rolling baseline. At step 14,960 it reaches 9.7. Gradient clipping fires, but the Adam second-moment estimate has already been updated with the squared large gradient, depressing the effective learning rate for affected parameters for the next several steps. At step 14,963 the training loss reads 5.4; at step 14,967 it reads 9.1.

Decision protocol:

1. Check the canary batch loss at step 14,950 through 14,967: it reads $2 . 21$ - completely flat. This is a data artifact, not model divergence.

2. Pull the data loader cursor from the step-14,900 checkpoint: shard index 312, document offset 8,192.

3. Inspect documents at that offset. Find a 380,000-token sequence that was incorrectly chunked: a single training sample with no EOS token, heavy Unicode encoding errors, and near-zero perplexity on most tokens (a repeated byte pattern).

4. Roll back to the step-14,900 checkpoint. Add the document hash to the dataset blocklist. Reprocess the shard. Resume.

5. Recovery: within 60 steps, loss returns to $2 . 2$. The canary loss never deviated.

Total cost: 100 steps of recompute plus 30 minutes of investigation. The alternative - running forward 10,000 steps hoping for self-recovery - would have consumed 80 GPU-hours and produced a model degraded on the specific domain of that corrupted document.

What You Actually Decide in Practice

• Spike vs. divergence threshold. If training loss has not returned to within 5% of its pre-spike value within 200 steps, roll back. Self-recovery beyond 200 steps is rare for data-artifact spikes and almost never happens for true divergence. Waiting longer only compounds the optimizer state corruption from the large Adam second-moment update.

• Gradient norm alert threshold. Set an automated alert at unclipped grad norm exceeding $5 \times$ the 1,000- step rolling mean. At this level you have an intervention window before the loss visibly spikes. At $100 \times$ the rolling mean, you have a bad batch; pause immediately and inspect the data cursor before resuming.

• Do not adjust the learning rate as a recovery mechanism. Reducing LR mid-spike to “stabilize” the run introduces a new instability: the Adam momentum vectors accumulated under the original LR are now mismatched. The preferred recovery is a checkpoint rollback to before the anomaly. LR adjustments belong in the schedule, not in ad hoc responses to individual spikes.

• Build a bad-document filter as a continuous background job. Every time a spike investigation identifies a pathological document, add it to a blocklist and re-process the affected data shard offline. After three or four spike investigations, patterns emerge (mislabeled code samples, HTML artifacts, very long sequences with zero semantic content) that can be caught by a rule-based pre-filter, preventing the spike entirely rather than just recovering from it.

• Alert on canary divergence, not just training loss. If canary loss deviates by more than two standard deviations from a 500-step rolling mean, page the on-call engineer. This catches slow model degradation - from silent data corruption, optimizer bugs, or data distribution drift - that step-level training loss monitoring misses because the signal is spread across thousands of steps.

How This Shows Up in Interviews

• “You are training a 30B model and the loss spikes mid-run. How do you decide whether to roll back or continue?” Strong answer: check gradient norm trajectory (leading indicator - was it climbing before the loss jumped?), check canary batch loss (data artifact vs. model degradation), correlate spike timestamp with data loader cursor position to identify a candidate document, wait at most 200 steps for self-recovery, then roll back to the most recent clean checkpoint if loss has not recovered.

• “What metrics do you log during pre-training and at what granularity?” Strong answer: per-step training loss, unclipped global gradient norm, effective learning rate, and hardware throughput; canary batch loss every 100 steps. Epoch-level averages are insufficient because they mask transient spikes that are often correlated with specific, identifiable data artifacts.

• Senior follow-up: “Your gradient norm alert fires, but when you inspect the documents at the data cursor position you find nothing obviously wrong. What else could be causing the spike?” Senior answer: the list of candidates in rough order of frequency: (1) silent data corruption where the gradient is numerically plausible but wrong - cross-replica gradient norm comparison would reveal a discrepancy; (2) a learning rate scheduler bug that momentarily emitted an incorrect LR value - check the per-step LR log; (3) an out-of-memory event that silently fell back to lower precision in one kernel - check GPU memory logs around the affected steps; (4) a numerical instability in a specific attention head triggered by an unusually long sequence - profile the attention patterns on the documents near the spike. Work from most common to most exotic and do not assume the data is the problem until you have ruled out the infrastructure.

Key Takeaways

• Log training loss, unclipped gradient norm, effective learning rate, and GPU throughput at every step - epoch-level averages hide transient spikes that are correlated with specific, filterable bad documents.

• Gradient norm is the leading indicator: it typically begins rising 10-50 steps before the training loss visibly spikes, giving you an intervention window.

• A fixed canary batch evaluated every 100 steps distinguishes data artifacts (canary unaffected) from true model degradation (canary also drifts), which determines the correct response.

• Correlate every spike with the data loader cursor position to identify and filter the offending document - this is almost always faster and cheaper than rolling back without root-cause investigation.

• If training loss has not recovered within 200 steps, roll back to the most recent checkpoint; do not reduce the learning rate mid-spike, as this corrupts the optimizer momentum state accumulated under the original schedule.