3.11 Training Stability, Loss Spikes & Debugging Large Runs¶

Pretraining a large language model is one of the most expensive experiments a human team can run. At the scale of hundreds of billions of parameters and trillions of tokens, a single unrecovered divergence can waste millions of dollars of compute and weeks of calendar time. Yet instability is not rare — it is an almost universal companion of large-scale training. Every serious pretraining team has stories of mysterious loss spikes, subtle data bugs discovered only after 10 days of training, and GPU clusters operating at 40 % throughput while engineers hunt a deadlock.

This chapter is the field manual for that experience. We cover the mechanisms behind instability, the mitigations built into modern architectures and training pipelines, the monitoring infrastructure that surfaces problems before they become catastrophes, and the decision-making playbook you need when things go wrong at 3 AM.

Related context that we assume you have read or will read alongside this chapter: Mixed Precision, bf16 & FP8 Training for floating-point root causes, Optimizers: SGD, Adam, Adafactor, Lion, Muon & Shampoo for optimizer dynamics, Learning Rate Schedules, Warmup, Batch Size & Hyperparameters for schedule interactions, and Checkpointing, Fault Tolerance & Long-Running Jobs for recovery mechanics. Everything here is exercised concretely in Part XIV: The Pretraining Run wires gradient clipping, a non-finite-loss skip, a fused CE + z-loss and a periodic QK-clip into the single-GPU loop that trains Stack-100M, and Optimizer & Schedule derives the clip threshold it uses.

What a Stable Run Looks Like — and When to Worry¶

Before diagnosing instability, you need a baseline for health. A well-configured pretraining run exhibits a smooth loss curve that decreases monotonically on a log scale, with mild stochastic noise but no sustained plateaus or upward excursions larger than roughly 0.05–0.10 nats over a few hundred steps.

Healthy run:                        Spike event:
 loss                                loss
  |                                   |
3.5|*                              3.5|*
   | **                                | **
3.0|   ***                         3.0|   ***
   |      *****                        |      *****
2.5|           *******             2.5|           ***  /\  ****
   |                 ****              |                \/
2.0|                     ****      2.0|
   +----------------------- step      +----------------------- step

Warning thresholds (guidelines, not hard rules):

These numbers come from practical experience across GPT-class, Llama-class, and similar training runs. The exact thresholds depend on model size, learning rate, and optimizer — but the relative ratios are robust.

Anatomy of a Loss Spike¶

A loss spike is a sudden increase in training loss, typically resolved by the optimizer over tens to hundreds of subsequent steps. Understanding what actually happens mechanically is important for choosing the right mitigation.

The causal chain¶

A spike almost always follows this causal chain:

The “bad state” can be a transient — the model recovers on its own — or it can be absorbing: the optimizer finds a new basin with higher loss, and the run effectively diverges. Whether a spike is transient or absorbing depends heavily on the magnitude relative to the Adam/AdamW second-moment estimate, which we discuss in the next section.

Adam’s role in amplifying spikes¶

Adam (and its variants) maintains a running estimate of the second moment \(\hat{v}_t\) for each parameter gradient. When a gradient is unusually large, the effective learning rate for that parameter is:

If the large gradient is new (not predicted by the historical \(\hat{v}_t\)), the denominator is small (accumulated from previous, smaller gradients), so the parameter update is disproportionately large. This is the primary amplification mechanism. After the spike, \(\hat{v}_t\) quickly absorbs the new large value, damping future updates — which is why spikes are usually transient. But if the first large update throws the model into a region of high curvature, recovery may take hundreds of steps or fail entirely.

Root Causes of Instability¶

Bad data batches¶

The single most common cause of spikes in practice is a batch containing anomalous content: repeated tokens, extremely long documents (padding or truncation bugs), encoding corruption, or near-duplicate toxic sequences that the model assigns very low probability to. The resulting loss is high, the gradient is large, and Adam amplifies it.

Canonical bad-data patterns:

• Repeated n-grams or copy-paste artifacts. A sequence like aaaa...aaaa (10 000 repetitions) has near-zero cross-entropy under a well-trained model but produces high loss on an in-progress model, and the gradient is sharply peaked at the repetition token.
• Mixed-language documents with incorrect tokenization. A Chinese document tokenized with a primarily-English BPE vocabulary produces absurdly long token sequences, often hitting the sequence-length limit mid-word. See Tokenization: BPE, WordPiece, Unigram & Byte-Level for how this manifests.
• Numeric/code sequences with large literal integers. Sequences like 0000000000000000000000001 produce degenerate token sequences that confuse positional encodings.
• HTML/boilerplate leaking through filtering. A batch with 30 % <div class="..."> repeating patterns is effectively adversarial input.

# Tool: detect anomalous batches before they enter the training loop.
# Run as a filter in the data loader or offline as a data audit.

import torch
from collections import Counter

def batch_anomaly_score(
    input_ids: torch.Tensor,  # (B, T)
    ngram_n: int = 4,
    repetition_threshold: float = 0.4,
) -> torch.Tensor:
    """
    Returns a per-example anomaly score in [0, 1].
    High score = potentially bad batch element.

    Two signals:
      1. Repetition: fraction of n-grams that are duplicates.
      2. Entropy: low token-level entropy suggests degenerate text.
    """
    B, T = input_ids.shape
    scores = torch.zeros(B)

    for b in range(B):
        tokens = input_ids[b].tolist()

        # Signal 1: n-gram repetition fraction
        ngrams = [tuple(tokens[i:i+ngram_n]) for i in range(T - ngram_n)]
        if ngrams:
            counts = Counter(ngrams)
            # fraction of positions that are a repeated n-gram
            repeated = sum(v - 1 for v in counts.values() if v > 1)
            rep_frac = repeated / len(ngrams)
        else:
            rep_frac = 0.0

        # Signal 2: unigram entropy (in bits)
        tok_counts = Counter(tokens)
        total = len(tokens)
        entropy = -sum(
            (c / total) * (torch.log2(torch.tensor(c / total)).item())
            for c in tok_counts.values()
        )
        # For a typical English document, entropy > 8 bits; < 3 is suspicious.
        entropy_score = max(0.0, 1.0 - entropy / 8.0)

        scores[b] = 0.5 * rep_frac + 0.5 * entropy_score

    return scores


def should_skip_batch(input_ids: torch.Tensor, threshold: float = 0.35) -> bool:
    """Return True if the batch contains too many anomalous examples."""
    scores = batch_anomaly_score(input_ids)
    # Skip if average score is high OR if any single example is very bad
    return bool(scores.mean() > threshold or scores.max() > 0.75)

Learning rate issues¶

An LR too high produces large gradients on every batch, not just anomalous ones. The tell-tale sign: the loss is healthy during warmup but explodes as soon as the LR reaches its peak value. A too-short warmup (reaching full LR before the Adam moments have stabilized) has the same symptom. See Learning Rate Schedules, Warmup, Batch Size & Hyperparameters for schedule mechanics.

Rule of thumb for initial LR selection: For AdamW, a reasonable peak LR scales with width as \(\alpha \approx C / \sqrt{d_\text{model}}\), with \(C\) landing empirically in the range \(6\times10^{-3}\) to \(2\times10^{-2}\) across published GPT-3-, Llama- and OLMo-class configurations. Sanity-check both ends: at \(d_\text{model} = 4096\), \(C = 2\times10^{-2}\) gives \(\alpha \approx 3\times10^{-4}\) — the value 7B-class runs typically publish; at \(d_\text{model} = 12\,288\), \(C = 6\times10^{-3}\) gives \(\alpha \approx 5\times10^{-5}\), the order of GPT-3 175B’s published peak. Two caveats. This is a width rule only — batch size is governed by the separate \(\alpha \propto \sqrt{B}\) heuristic below the critical batch size (see Learning Rate Schedules, Warmup, Batch Size & Hyperparameters). And it is the standard-parametrization compromise: the maximal update parametrization (μP) of Yang et al. prescribes the sharper \(\alpha \propto 1/\text{fan-in}\) for matrix-like parameters, which is what makes LR genuinely transferable across widths. Treat \(C/\sqrt{d_\text{model}}\) as the centre of a sweep, not a final answer.

Floating-point issues¶

At bf16, the representable range is roughly \(\pm 3.4 \times 10^{38}\) (same 8 exponent bits as fp32), but the significand is only 8 bits — about 2–3 significant decimal digits. Overflow is rare but not impossible; underflow to zero is more common and more insidious.

At fp16, overflow occurs above \(65\,504\), and activations can silently become inf or NaN during the forward pass if any intermediate value — typically in the attention softmax or MLP feedforward — exceeds this. See Mixed Precision, bf16 & FP8 Training for the full picture.

The attention logit overflow problem. Without QK normalization or attention bias clipping, the logits \(QK^\top / \sqrt{d_k}\) can grow arbitrarily large. For a model with \(d_k = 128\), the scale factor is \(1/\sqrt{128} \approx 0.088\). If query and key vectors both have L2 norm \(\approx 100\) (feasible in a large model at late training), the maximum logit magnitude is \(\approx 100 \times 100 \times 0.088 = 880\), which is within fp16 range but causes the softmax distribution to collapse to a single-token delta. The exponential sum underflows, producing NaN gradients.

Embedding table instabilities¶

The input embedding matrix is updated on every step via the language model head (tied weights) but each row is only updated when the corresponding token appears. Rare tokens accumulate large Adam second moments slowly, so their effective LR can be anomalously high when they do appear. The output embedding rows also receive gradients from every token in the vocabulary through the CE loss, but the magnitude varies wildly with token frequency.

Architectural Mitigations¶

Modern LLM architectures have baked in a portfolio of stability techniques. Understanding each one individually helps you reason about which to apply when you start from scratch or debug an existing run.

QK-Norm¶

Instead of relying on weight initialization to keep query and key norms bounded, QK-Norm explicitly normalizes the query and key projections before computing attention logits:

This ensures the logit magnitudes grow at most as \(O(\sqrt{d_k})\) regardless of the raw activations. Popularised at scale by Dehghani et al.’s ViT-22B and now used in production by Gemma 3, Qwen3 (which explicitly replaced Qwen2’s QKV bias with QK-norm) and OLMo 2, it is a default recommendation — and not only for fp16. bf16’s wider exponent range postpones overflow, but it does nothing to prevent softmax saturation into a one-hot with vanishing gradients, so QK-norm earns its keep in bf16 runs too. In HuggingFace transformers you can read the real implementations directly: the Qwen3 and Gemma-3 attention modules instantiate self.q_norm / self.k_norm RMSNorms over the head dimension and apply them to the reshaped Q and K, exactly as below.

One caveat the bound above hides: the learnable RMSNorm gains \(\gamma_q, \gamma_k\) are unconstrained, so the guarantee is \(O(\sqrt{d_k}) \cdot \lVert\gamma_q\rVert_\infty\lVert\gamma_k\rVert_\infty\) — structural, but not fixed. Log \(\max|\gamma|\) per layer alongside the max-logit monitor; slow upward drift is the precursor of the failure QK-clip (next subsection) exists to catch.

import torch
import torch.nn as nn
import torch.nn.functional as F

class QKNormAttention(nn.Module):
    """
    Multi-head attention with per-head QK normalization.
    Prevents attention logit overflow, a common source of training spikes.
    """
    def __init__(self, d_model: int, n_heads: int, eps: float = 1e-6):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

        # Learnable scale parameters, one per head dimension.
        # RMSNorm without learned scale would fix norm to 1.0;
        # a learnable scale restores representational flexibility.
        self.q_norm = nn.RMSNorm(self.d_k, eps=eps)
        self.k_norm = nn.RMSNorm(self.d_k, eps=eps)

    def forward(self, x: torch.Tensor, mask=None) -> torch.Tensor:
        B, T, D = x.shape

        # Project to Q, K, V and reshape to (B, n_heads, T, d_k)
        def split_heads(t):
            return t.view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        Q = split_heads(self.W_q(x))  # (B, H, T, d_k)
        K = split_heads(self.W_k(x))
        V = split_heads(self.W_v(x))

        # --- QK Norm: the key stability ingredient ---
        # Normalize along the head dimension; norms are now bounded.
        Q = self.q_norm(Q)
        K = self.k_norm(K)

        # Scaled dot-product attention (safe now that Q, K are normalized)
        scale = self.d_k ** -0.5
        attn = (Q @ K.transpose(-2, -1)) * scale  # (B, H, T, T)
        if mask is not None:
            attn = attn.masked_fill(mask == 0, float('-inf'))
        attn = F.softmax(attn, dim=-1)

        out = (attn @ V).transpose(1, 2).contiguous().view(B, T, D)
        return self.W_o(out)

QK-Clip (MuonClip)¶

QK-norm prevents logit runaway by construction. The alternative is to let the projections be, measure the damage, and correct it after each optimizer step. This is QK-clip, the active ingredient of MuonClip, introduced by the Kimi K2 technical report (Moonshot AI, 2025) to keep the Muon optimizer stable at trillion-parameter scale.

The motivation is optimizer-specific. Orthogonalizing optimizers such as Muon (see Optimizers: SGD, Adam, Adafactor, Lion, Muon & Shampoo) push with equal magnitude along every singular direction of a weight matrix, so \(\lVert W_Q \rVert\) and \(\lVert W_K \rVert\) can drift upward together in a way AdamW’s per-coordinate normalization tends to hide. Once \(\max_{ij} q_i^\top k_j / \sqrt{d_k}\) climbs into the hundreds, the softmax saturates, gradients through it vanish, and you get a spike or an outright NaN.

The fix exploits the fact that the logit is bilinear in \(W_Q\) and \(W_K\). If the observed per-head maximum logit is \(S_{\max} > \tau\), scaling both matrices by \(\gamma = \sqrt{\tau / S_{\max}}\) multiplies every logit in that head by exactly \(\gamma^2 = \tau / S_{\max}\), so the post-clip maximum is exactly \(\tau\):

@torch.no_grad()
def qk_clip_(attn: nn.Module, max_logit: float, tau: float = 100.0) -> bool:
    """QK-clip -- the active ingredient of MuonClip (Kimi K2, Moonshot 2025).

    Call AFTER optimizer.step(). `max_logit` is the largest pre-softmax
    attention logit observed for this module on a recent forward pass.
    Returns True if the clip fired (log this: it is the single most
    informative stability signal a Muon run produces).
    """
    if not (max_logit > tau):        # also covers NaN: NaN > tau is False
        return False
    gamma = (tau / max_logit) ** 0.5
    if getattr(attn, "q_norm", None) is not None:
        # QK-norm is ON: W_Q / W_K are pre-normalized, so rescaling them is a
        # NO-OP. The free knob is the learnable RMSNorm gain -- clip that.
        attn.q_norm.weight.mul_(gamma)
        attn.k_norm.weight.mul_(gamma)
    else:
        attn.W_q.weight.mul_(gamma)  # original Kimi K2 configuration
        attn.W_k.weight.mul_(gamma)
    return True

Two practical points that most summaries skip.

QK-norm and QK-clip are alternatives on the same knob, not layers that stack. Kimi K2 needed the clip on \(W_Q, W_K\) precisely because it does not QK-norm. If you do normalize, rescaling \(W_Q\) changes nothing — RMSNorm divides it right back out — so the clip must target the learnable gains instead, as the branch above does.

Measuring \(S_{\max}\) is not free. The fast attention path is F.scaled_dot_product_attention, which dispatches to a FlashAttention-style kernel whose entire point is that it never materializes the \((B, H, T, T)\) score matrix (see FlashAttention I: IO-Awareness & The Online Softmax). There is no .amax() to take. Either keep a cheap eager-attention probe path, or — since attention-logit drift moves at the speed of one decayed LR per step — take the reading every few hundred steps on a single small probe batch with post-step weights. Optimizer & Schedule implements exactly this for Stack-100M, including the GQA-aware per-head variant and the choice of \(\tau\) (Kimi’s \(\tau = 100\) as an inert backstop; \(\tau \approx 30\) when you want the trigger count as an early-warning sensor).

Z-loss¶

The z-loss is a small auxiliary penalty on the log-partition function of the softmax:

where \(z_{i,v}\) are the pre-softmax logits for position \(i\) and vocabulary token \(v\). If the logits grow large, the log-partition function grows, and the z-loss penalizes this. This is especially useful for Mixture-of-Experts models (see Mixture-of-Experts (MoE) Architectures) where the router softmax is a common source of collapse. A typical \(\beta_z = 10^{-4}\) adds negligible loss overhead but provides a gradient pressure that keeps logit norms bounded.

def z_loss(logits: torch.Tensor, beta: float = 1e-4) -> torch.Tensor:
    """
    Z-loss regularizer (Zoph et al., ST-MoE 2022).
    logits: (B, T, V) or (B, V) pre-softmax values.

    Penalizes large log-partition values to prevent logit explosion.
    Typical beta: 1e-4 for MoE router; 1e-5 for LM head.
    """
    # log(sum_v exp(z_v)) = log-partition function per position
    # torch.logsumexp is numerically stable
    log_z = torch.logsumexp(logits, dim=-1)  # (B, T) or (B,)
    return beta * (log_z ** 2).mean()

# Usage in training loop:
# lm_loss = cross_entropy_loss(logits, targets)
# aux_loss = z_loss(logits, beta=1e-5)
# total_loss = lm_loss + aux_loss

Careful initialization¶

The variance of activations through the network is set by initialization. The classical analysis by He et al. and the subsequent improvements (μP, Transformers specific scaling) give concrete recipes:

• Embedding matrix: Initialize with \(\mathcal{N}(0, \sigma^2)\) where \(\sigma = d_\text{model}^{-0.5}\) (this keeps embedding norms \(\approx 1\) immediately, avoiding early logit explosions).
• Residual projections (output of attention and MLP): Scale down by \(1/\sqrt{2L}\) where \(L\) is the number of layers. This is the “depth scaling” trick — each layer contributes \(1/\sqrt{2L}\) to the residual stream, keeping the cumulative norm growth bounded.
• QKV projections: Initialize so that the expected logit variance is \(\approx 1\). With head dimension \(d_k\), this means \(\sigma_{QK} = d_k^{-0.25}\) (so that \(QK^\top / \sqrt{d_k}\) has variance 1).

import torch.nn as nn

def init_transformer_weights(model: nn.Module, n_layers: int, d_model: int):
    """
    Stability-oriented weight initialization for a GPT-style transformer.
    Based on the GPT-2/NanoGPT pattern with depth scaling.
    """
    for name, param in model.named_parameters():
        if param.dim() < 2:
            # Biases, norms — leave as default (zeros / ones)
            continue

        if 'embedding' in name:
            # Embedding table: small init to keep logits bounded at step 0
            nn.init.normal_(param, mean=0.0, std=d_model ** -0.5)

        elif 'c_proj' in name or 'out_proj' in name:
            # Residual-path output projections (attn output + MLP output).
            # Scaled down by 1/sqrt(2 * n_layers) so that the residual stream
            # norm grows as O(1) rather than O(sqrt(L)) at initialization.
            std = (2 * n_layers) ** -0.5
            nn.init.normal_(param, mean=0.0, std=std)

        elif 'q_proj' in name or 'k_proj' in name:
            # Query/Key projections: initialize so logit std ≈ 1
            d_k = param.shape[0]  # assuming (d_k, d_model) layout
            nn.init.normal_(param, mean=0.0, std=d_k ** -0.25)

        else:
            # Default: Kaiming/He for everything else
            nn.init.normal_(param, mean=0.0, std=0.02)

Embedding norm clipping¶

Even with small init, embeddings can grow unbounded during training. A simple but effective technique is to project embeddings back to a unit ball (or ball of radius \(r\)) after each optimizer step:

@torch.no_grad()
def clip_embedding_norm(embedding: nn.Embedding, max_norm: float = 1.0):
    """
    After optimizer step: clip embedding row norms to max_norm.
    Prevents rare-token embedding rows from drifting far from the manifold.
    """
    norms = embedding.weight.norm(dim=-1, keepdim=True)  # (V, 1)
    # Only clip rows that exceed max_norm; leave smaller ones untouched.
    clipped = embedding.weight * (max_norm / norms.clamp(min=max_norm))
    embedding.weight.copy_(clipped)

Training-Time Mitigations¶

Gradient norm clipping¶

Gradient clipping is the first line of defense against spike propagation. The global gradient norm is:

and we rescale all gradients by \(\min(1, \tau / \|g\|_2)\) where \(\tau\) is the clip threshold. This prevents large gradients from causing large parameter updates but preserves their direction. The standard value is \(\tau = 1.0\); some works use \(\tau = 0.5\) for extra stability at the cost of slightly slower early learning.

Skip-batch logic¶

When a batch produces an anomalous gradient — detected either by norm threshold or by explicit batch quality scoring — you can skip the optimizer step for that batch. This is conservative: you still do the forward and backward pass (wasting compute), but you do not update the model state. The Adam moments are also not updated.

import math

def training_step(
    model,
    optimizer,
    batch: dict,
    scaler=None,      # torch.amp.GradScaler -- REQUIRED for fp16, must be None for bf16
    amp_dtype=torch.bfloat16,
    grad_clip: float = 1.0,
    grad_skip_threshold: float = 5.0,  # skip if norm > 5x clip threshold
    anomaly_threshold: float = 0.35,
) -> dict:
    """
    One training step with skip-batch and gradient norm monitoring.
    Returns a metrics dict for logging.
    """
    assert not (scaler is not None and amp_dtype is torch.bfloat16), (
        "bf16 has fp32's exponent range, so loss scaling is unnecessary and "
        "GradScaler's inf-detection logic is not what you want. Pass scaler=None."
    )
    # --- Optional: skip anomalous batch early (before forward pass) ---
    if should_skip_batch(batch['input_ids'], threshold=anomaly_threshold):
        return {'loss': float('nan'), 'skipped': True, 'reason': 'bad_data'}

    optimizer.zero_grad()

    # Forward pass under autocast
    with torch.autocast(device_type='cuda', dtype=amp_dtype):
        logits = model(batch['input_ids'])
        loss = F.cross_entropy(
            logits.view(-1, logits.size(-1)),
            batch['labels'].view(-1),
            ignore_index=-100,
        )

    # Backward pass. Loss scaling exists only to keep fp16 gradients out of
    # the subnormal range; under bf16 we go straight to .backward().
    if scaler is not None:
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)   # must unscale before clipping/inspection
    else:
        loss.backward()

    # clip_grad_norm_ returns the PRE-clip global norm -- exactly what we log.
    grad_norm = torch.nn.utils.clip_grad_norm_(
        model.parameters(),
        max_norm=grad_clip,
    ).item()

    # --- Skip-step logic: skip on a huge norm OR a non-finite one. ---
    # The isfinite() check is not optional: `float('nan') > x` is False in
    # Python, so a NaN gradient would sail past a bare threshold comparison
    # and be applied to the weights, corrupting every parameter at once.
    if not math.isfinite(grad_norm) or grad_norm > grad_skip_threshold * grad_clip:
        optimizer.zero_grad(set_to_none=True)
        if scaler is not None:
            scaler.update()          # let the scaler back off its scale factor
        return {
            'loss': loss.item(),
            'grad_norm': grad_norm,
            'skipped': True,
            'reason': 'grad_spike' if math.isfinite(grad_norm) else 'nonfinite',
        }

    # Normal step
    if scaler is not None:
        scaler.step(optimizer)
        scaler.update()
    else:
        optimizer.step()

    return {
        'loss': loss.item(),
        'grad_norm': grad_norm,
        'skipped': False,
    }

Loss spike recovery: rolling back vs. continuing¶

When a spike is detected (loss increases by more than, say, 0.2 nats over 20 steps), you have three options:

1. Continue and hope. Works for mild, transient spikes. The Adam moments adapt within 100–300 steps. Monitor closely.
2. Resume from the last good checkpoint. If the spike is severe or you can identify a batch to blame, roll back to the checkpoint before the bad batch and skip it. See Checkpointing, Fault Tolerance & Long-Running Jobs for checkpoint strategy.
3. Resume with lower LR. Load the last good checkpoint and continue with the LR reduced by 20–40 %. This treats the spike as a symptom of the LR being at its peak (common mid-run), and the reduced LR prevents recurrence while the moments stabilize.

Spike decision tree:

spike detected (loss > baseline + 0.15 nats for > 50 steps)
     |
     +-- grad_norm normal (< 2x clip)?
     |        YES: likely LR/optimizer issue → reduce LR by 20%, continue
     |        NO:  large grad norm → check batch data
     |
     +-- specific batch identifiable (from data logs)?
     |        YES: roll back to pre-spike checkpoint, skip that batch
     |        NO:  roll back to pre-spike checkpoint, reduce LR by 30%
     |
     +-- spike severity?
              MILD (< 0.3 nats, < 200 steps): continue + monitor
              SEVERE (> 0.5 nats, model not recovering): roll back

Monitoring Infrastructure¶

You cannot debug what you cannot see. A production pretraining run should log at minimum the following signals every N steps (typically every 5–50 steps, depending on cluster size and logging overhead).

Essential metrics¶

import math
import torch
import wandb  # Weights & Biases; swap for mlflow / torch.utils.tensorboard

class TrainingMonitor:
    """
    Collects and logs training health signals.
    Designed to add minimal overhead: most metrics are computed from
    tensors already in memory.
    """
    def __init__(self, log_every: int = 10, spike_window: int = 50):
        self.log_every = log_every
        self.spike_window = spike_window
        self.loss_history = []
        self.step = 0

    def log_step(
        self,
        loss: float,
        grad_norm: float,
        model: torch.nn.Module,
        lr: float,
    ):
        self.step += 1
        self.loss_history.append(loss)

        if self.step % self.log_every != 0:
            return

        metrics = {
            'train/loss': loss,
            'train/perplexity': math.exp(min(loss, 20)),  # clamp to avoid overflow
            'train/grad_norm': grad_norm,
            'train/lr': lr,
            'train/step': self.step,
        }

        # --- Activation statistics (sampled from a few layers) ---
        # Hook-based; only active when we log.
        act_stats = self._sample_activation_stats(model)
        metrics.update(act_stats)

        # --- Spike detector ---
        if len(self.loss_history) >= self.spike_window:
            window = self.loss_history[-self.spike_window:]
            baseline = sum(window[:self.spike_window // 2]) / (self.spike_window // 2)
            recent = sum(window[self.spike_window // 2:]) / (self.spike_window // 2)
            metrics['train/spike_delta'] = recent - baseline

        wandb.log(metrics)

    @torch.no_grad()
    def _sample_activation_stats(self, model: torch.nn.Module) -> dict:
        """
        Compute per-layer activation norms for the most recent forward pass.
        In practice, wire this to forward hooks registered on a few layers.
        Here we illustrate with a direct parameter-norm proxy.
        """
        stats = {}
        for name, param in model.named_parameters():
            if 'weight' in name and param.dim() >= 2:
                # Track weight matrix spectral norm proxy (Frobenius / sqrt(numel))
                rms = param.norm() / (param.numel() ** 0.5)
                short = name.replace('.weight', '').replace('model.', '')
                stats[f'weights/{short}_rms'] = rms.item()
        return stats

What to watch and why¶

Distributed monitoring considerations¶

In a multi-node run across hundreds of GPUs (see Distributed Training I: Data Parallelism, DDP, ZeRO & FSDP), you should:

• Log metrics only from rank 0 to avoid redundant writes.
• But compute gradient norms globally (all-reduce) before logging, since rank-0’s local norm may not represent the global norm.
• Alert on NaN/Inf using torch.isnan(loss).any() — in ZeRO-3, a NaN on any rank will propagate during the gradient all-reduce, so catching it early saves wasted compute.
• Reach for the stack’s own instrumentation before writing your own: torchtitan and Megatron-LM both log grad norm, loss scale and MFU out of the box; TORCH_DISTRIBUTED_DEBUG=DETAIL surfaces mismatched collectives and desynchronized ranks; NCCL_DEBUG=WARN catches the flaky-link class of hang; and torch.autograd.set_detect_anomaly(True) pinpoints the exact op that first produced a NaN (at a 2–5× slowdown, so enable it only on a replay, never on the real run).

def check_for_nan_distributed(loss: torch.Tensor, grad_norm: float) -> bool:
    """
    Check for NaN/Inf on all ranks and broadcast the result.
    Returns True if any rank has a problem.
    """
    import torch.distributed as dist

    # Flag: 1.0 if problematic, 0.0 if OK
    bad = torch.tensor(
        1.0 if (not torch.isfinite(loss) or not math.isfinite(grad_norm)) else 0.0,
        device=loss.device,
    )
    # Max-reduce: any rank with a problem sets the flag for all
    dist.all_reduce(bad, op=dist.ReduceOp.MAX)
    return bad.item() > 0.5

The Lived Experience: War Stories & Common Failure Modes¶

This section distills patterns from real large-scale training runs. No precise benchmark numbers are attributed, but the patterns are real and well-documented in public technical reports from teams including those behind Llama, PaLM, Gemini, and similar.

Story 1: The silent data corruption¶

A team ran a 70B model to 500B tokens with no obvious spikes, but evaluation metrics on a standard benchmark plateau early and never recover. Post-mortem: a deduplication bug in the data pipeline had removed 40 % of the math and code content, replacing it with duplicate web text. The loss curve was smooth because the model simply memorized the duplicates efficiently. Lesson: loss alone is not a sufficient health signal. Hold out a small fixed eval set covering each domain (code, math, multilingual, factual QA) and report it every 10B tokens.

Story 2: The creeping LR spike¶

Every modern LLM training run has seen this: things are fine for weeks, then at precisely the step where the warmup ends and the cosine peak is reached, loss jumps 0.4 nats and the grad norm hits 20×. The model recovers, but with a shifted loss baseline. Root cause: the Adam moments were well-calibrated for the warmup LR, not the peak LR. The effective LR at the peak is \(\alpha_\text{peak} / \sqrt{\hat{v}}\), and \(\hat{v}\) was too small. Fix: use a longer warmup (1–2 % of total steps rather than 0.1 %) or apply a sqrt-scaled warmup that grows \(\alpha\) slower than Adam’s \(\sqrt{t}\) moment term.

Story 3: The fp16 attention NaN cascade¶

In a 13B model trained in fp16 (not bf16), at step 42 000 the loss becomes NaN. Debugging with activation hooks reveals that the attention logits for the last layer’s head 7 are producing values near 60 000 before the softmax — approaching the fp16 max of 65 504. A batch with a 32 000-token nearly-identical sequence (a repeated copyright boilerplate) pushed query and key norms to an extreme. The softmax then produces Inf, the backward pass propagates NaN, and all parameters are corrupted. Fix: add QK-Norm (see above) and switch to bf16, which has a much larger dynamic range. The retrospective also added a maximum logit monitor to the activation hooks.

Story 4: The zombie GPU¶

During a 256-GPU training run, GPU #183 silently corrupts its computation starting at step 15 000 (the gradient it contributes is numerically wrong but finite). Because the corrupt gradient is averaged in during the DDP all-reduce, the model trains fine for another 20 000 steps — just slightly worse than it should. Discovered only when the team compared two runs that should have been identical under different partitioning and found a large discrepancy. Fix: implement periodic determinism checks (run a single fixed micro-batch through each GPU independently and compare; see also chapter Checkpointing, Fault Tolerance & Long-Running Jobs). This class of fault is exactly what NVIDIA’s DCGM is for — dcgmi diag -r 3 runs the long hardware/ECC/memory-bandwidth diagnostic, and a DCGM exporter scraping XID errors and uncorrectable ECC counts per node will usually name the sick GPU before your loss curve does.

Debugging Playbook: A Step-by-Step Checklist¶

When something goes wrong, follow this ordered checklist. Each step narrows the hypothesis space.

===== TRAINING STABILITY DEBUGGING CHECKLIST =====

STEP 1: ESTABLISH GROUND TRUTH
  [ ] Is the loss spike visible on ALL tracked metrics (val loss, ppl)?
  [ ] Is the spike visible from ALL ranks (or just one shard)?
  [ ] Is there a NaN or Inf anywhere in the loss or grad norm logs?
  [ ] What is the grad_norm at the spike step vs. baseline?

STEP 2: NARROW TO CATEGORY
  [ ] Grad norm normal, loss spiked → data issue (anomalous batch)
  [ ] Grad norm large (>5x), loss spiked → optimizer amplification
  [ ] Grad norm NaN → activation overflow (fp issue)
  [ ] Grad norm zero or near-zero → vanishing gradient or bad init
  [ ] Spike correlates with LR peak → warmup length issue

STEP 3: ISOLATE THE STEP
  [ ] Roll back to the checkpoint BEFORE the spike
  [ ] Replay forward+backward for the exact batch at the spike step
  [ ] Log per-layer gradient norms (which layer blows up first?)
  [ ] Log per-head attention logit max (which head has overflow?)

STEP 4: IDENTIFY ROOT CAUSE
  [ ] Inspect the batch: run batch_anomaly_score, look at raw tokens
  [ ] Check the data pipeline: was this batch from a specific shard/source?
  [ ] Check fp overflow: run the forward pass in fp32 and compare
  [ ] Check LR: what was the LR at the spike step?

STEP 5: MITIGATE AND RESUME
  [ ] Implement the fix (QK-norm, skip-batch, lower LR, fix data)
  [ ] Resume from pre-spike checkpoint
  [ ] Monitor for 500 steps before lowering alert thresholds
  [ ] Document the incident (what failed, why, what was fixed)

STEP 6: PREVENT RECURRENCE
  [ ] Add the root cause to the pre-training data quality audit
  [ ] Add the relevant monitoring signal permanently
  [ ] Consider adding skip-batch logic if not already present
  [ ] Consider adding QK-norm / z-loss if not already present

Here is a minimal but complete diagnostic script to run against a checkpoint and a batch:

"""
stability_probe.py — Diagnose a training spike post-hoc.

Usage:
  python stability_probe.py \
    --ckpt /path/to/checkpoint_before_spike.pt \
    --batch /path/to/spike_batch.pt \
    --model_config /path/to/config.json

Outputs: per-layer gradient norms, attention logit statistics,
         and a verdict on likely root cause.
"""
import torch
import json
import argparse
import math
from typing import Dict

def load_model_and_batch(ckpt_path: str, batch_path: str, config_path: str):
    """Load model from checkpoint and batch from saved tensor file."""
    # In practice: instantiate your model class, load state dict.
    # Here we use a placeholder to show the diagnostic logic.
    raise NotImplementedError("Wire to your model class")

@torch.no_grad()
def probe_activation_norms(
    model: torch.nn.Module,
    batch: Dict[str, torch.Tensor],
    dtype: torch.dtype = torch.float32,  # Always probe in fp32 for accuracy
) -> Dict[str, float]:
    """
    Register forward hooks to capture activation norms at each transformer block.
    Returns a dict: layer_name → max_activation_norm.
    """
    act_norms = {}
    hooks = []

    def make_hook(name):
        def hook(module, input, output):
            out = output[0] if isinstance(output, tuple) else output
            act_norms[name] = out.abs().max().item()
        return hook

    for name, module in model.named_modules():
        if 'attn' in name or 'mlp' in name:
            h = module.register_forward_hook(make_hook(name))
            hooks.append(h)

    model.to(dtype=dtype)
    with torch.autocast(device_type='cuda', enabled=False):
        _ = model(**batch)

    for h in hooks:
        h.remove()

    return act_norms


def compute_per_layer_grad_norms(
    model: torch.nn.Module,
    batch: Dict[str, torch.Tensor],
    labels: torch.Tensor,
) -> Dict[str, float]:
    """Run a backward pass and report per-parameter gradient norms."""
    model.train()
    logits = model(**batch)
    loss = torch.nn.functional.cross_entropy(
        logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-100
    )
    loss.backward()

    grad_norms = {}
    for name, param in model.named_parameters():
        if param.grad is not None:
            grad_norms[name] = param.grad.norm().item()
    return grad_norms


def diagnose(act_norms: Dict[str, float], grad_norms: Dict[str, float]) -> str:
    """Heuristic verdict based on activation and gradient diagnostics."""
    max_act = max(act_norms.values()) if act_norms else 0
    max_grad = max(grad_norms.values()) if grad_norms else 0

    if max_act > 60000:
        return "VERDICT: fp16 overflow likely (max_act={:.0f}). Add QK-Norm or switch to bf16.".format(max_act)
    elif max_act > 1000:
        return "VERDICT: activation explosion (max_act={:.0f}). Check init & norms.".format(max_act)
    elif max_grad > 50:
        layer = max(grad_norms, key=grad_norms.get)
        return "VERDICT: gradient explosion at '{}' (norm={:.1f}). Check data batch or reduce LR.".format(layer, max_grad)
    elif max_grad < 1e-6:
        layer = max(grad_norms, key=grad_norms.get)
        return "VERDICT: vanishing gradient (max_grad={:.2e}). Check depth or init.".format(max_grad)
    else:
        return "VERDICT: no clear activation/gradient anomaly. Check data content and LR schedule."

Pre-Run Stability Hardening Checklist¶

Before starting a large, expensive pretraining run, validate every item on this list. Discovering a bug at step 5 000 of a 500B-token run is far cheaper than discovering it at step 500 000.

===== PRE-RUN HARDENING CHECKLIST =====

ARCHITECTURE
  [ ] QK-Norm applied to all attention layers
      (if not: QK-clip / MuonClip enabled -- mandatory under Muon)
  [ ] Residual output projections initialized with 1/sqrt(2L) scaling
  [ ] Embedding init std ≤ 1/sqrt(d_model)
  [ ] Z-loss enabled (beta ≈ 1e-5 for LM head, 1e-4 for MoE router)
  [ ] Pre-norm (RMSNorm before attention/MLP) rather than post-norm

OPTIMIZER
  [ ] Gradient clipping at 1.0 (or 0.5 for extra stability)
  [ ] Adam epsilon: 1e-8 is the usual default, but INCREASE it (1e-6..1e-5)
      if you see spikes -- a larger eps floors the denominator, so a
      collapsed v_hat can no longer produce an unbounded effective LR
      (Molybog et al., 2023, recommend exactly this)
  [ ] Warmup ≥ 1% of total steps (for 1T-token run: ≥ 1B tokens)
  [ ] Weight decay ≠ 0 (0.1 is standard; reduces weight growth)
  [ ] No weight decay on embeddings / normalization parameters

PRECISION
  [ ] Using bf16 (not fp16) for forward/backward passes
  [ ] Loss scaling disabled when using bf16 (not needed)
  [ ] Master weights in fp32 (bf16 master weights can accumulate error)

DATA PIPELINE
  [ ] Batch anomaly detection enabled (or offline data audit complete)
  [ ] Sequence length distribution checked (no extreme outliers)
  [ ] Dataset mixture ratios validated against intent
  [ ] Fixed held-out eval batches per domain (code, math, web, multilingual)

MONITORING
  [ ] Grad norm logged every N steps (N ≤ 50)
  [ ] Loss spike detector with alerting (threshold: > 0.1 nat window delta)
  [ ] NaN/Inf detection with distributed reduce
  [ ] Checkpoint every K steps with verified restore test
  [ ] Compute and log MFU (Model FLOP Utilization) — sudden drops signal hardware issues

Key Takeaways¶

Further Reading¶

• Zoph et al., “ST-MoE: Designing Stable and Transferable Sparse Expert Models” (2022) — Introduces z-loss and provides a thorough analysis of instability in MoE training, with ablations on every stability technique discussed in this chapter.
• Wortsman et al., “Small-scale proxies for large-scale Transformer training instabilities” (2023) — Systematically studies which instabilities are predictable at small scale, providing a framework for the pre-run hardening approach.
• Yang et al., “Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer” (2022) — The μP framework that underpins principled LR and init scaling, explaining why \(\alpha \propto 1/\sqrt{d}\) keeps training stable across widths.
• Grattafiori et al., “The Llama 3 Herd of Models” (Meta AI, 2024) — The technical report contains candid discussion of training instabilities encountered during Llama 3 pretraining and the mitigations applied.
• Anil et al., “PaLM 2 Technical Report” (Google, 2023) — Documents the training stability experience for a series of large models including the use of bf16, careful init, and monitoring infrastructure.
• Molybog et al., “A Theory on Adam Instability in Large-Scale Machine Learning” (Meta AI, 2023) — Analytical backing for the Adam amplification model described in this chapter, including the argument that increasing Adam’s epsilon damps the instability.
• Kimi Team, “Kimi K2: Open Agentic Intelligence” (Moonshot AI, 2025) — The technical report that introduced MuonClip/QK-clip and documents a trillion-parameter run kept spike-free with it.
• nanoGPT (Andrej Karpathy, GitHub) — The canonical minimal GPT implementation. The train.py file is a useful starting point for understanding gradient clipping, skip-batch, and monitoring in a single-file, readable codebase.

Exercises¶

1. (Conceptual) The chapter states that loss spikes under AdamW are usually transient — the model recovers on its own after tens to hundreds of steps — yet a small fraction become absorbing and effectively diverge. Using the update rule \(\Delta\theta = -\alpha\,\hat{m}_t / (\sqrt{\hat{v}_t} + \epsilon)\) and the roles of \(\beta_1\) and \(\beta_2\), explain (a) the self-healing mechanism that makes most spikes transient, and (b) what distinguishes a spike that heals from one that absorbs.

2. (Quantitative) A parameter has historical gradient RMS \(g_\text{rms} = 0.02\), so \(\hat{v} \approx g_\text{rms}^2 = 4\times10^{-4}\) and \(\sqrt{\hat{v}} \approx 0.02\) (take \(\epsilon\) negligible). Training uses AdamW with \(\alpha = 4\times10^{-4}\), \(\beta_1 = 0.9\), and the fresh-moment model \(\hat{m} \approx (1-\beta_1)\,g\) with the denominator held at its stale value \(\sqrt{\hat{v}} = 0.02\) for the arriving step, so \(|\Delta\theta| \approx \alpha\,(1-\beta_1)\,g/\sqrt{\hat{v}}\).

(a) Compute \(|\Delta\theta|\) for a normal step with \(g = 0.02\). (b) A bad batch produces \(g = 0.5\) on this parameter. Compute \(|\Delta\theta|\) and the spike-to-normal ratio. © Global gradient clipping is enabled at \(\tau = 1.0\), and on the spike step the global gradient norm is \(\|g\| = 5.0\). Recompute \(|\Delta\theta|\) for the spike and the new ratio to a normal step. How much of the amplification did clipping remove?

3. (Quantitative) A model uses head dimension \(d_k = 64\), so the attention scale is \(1/\sqrt{d_k} = 1/8 = 0.125\). Late in training, a repeated-boilerplate batch drives one head’s query and key vectors to L2 norm \(\|q\| = \|k\| = 250\).

(a) What is the maximum possible attention logit \(q\cdot k / \sqrt{d_k}\) for this head? Is it within fp16 range (max \(65\,504\))? Why is it still dangerous even if it does not overflow? (b) The chapter’s QKNormAttention applies RMSNorm to \(Q\) and \(K\) before the dot product. After RMSNorm (RMS \(= 1\) per vector), what is the L2 norm of each normalized vector, and hence the new maximum logit? Confirm this matches the claimed \(O(\sqrt{d_k})\) bound.

4. (Conceptual) In “Story 2: The creeping LR spike,” a run is stable for weeks, then loss jumps 0.4 nats and the gradient norm hits 20x precisely at the step where warmup ends and the cosine peak LR is reached. (a) Explain mechanistically why the end of warmup is the dangerous moment, in terms of the effective learning rate \(\alpha_\text{peak}/\sqrt{\hat{v}}\). (b) Why does lengthening the warmup fix it, and why is a too-short warmup (e.g. 0.1% of steps) a classic cause?

5. (Implementation) The chapter’s QKNormAttention bounds logits to \(O(\sqrt{d_k})\). A complementary technique referenced in the SOTA section (Rybakov et al.) is softmax logit capping: pass the pre-softmax logits through \(c\cdot\tanh(\text{logit}/c)\), which smoothly saturates any logit to the range \((-c, c)\) while leaving small logits nearly unchanged. Implement a function attn_logit_softcap(attn, cap) and modify QKNormAttention.forward to apply the cap before the causal mask and softmax. Explain why the cap must be applied before, not after, masking.

import torch
import torch.nn.functional as F

def attn_logit_softcap(attn: torch.Tensor, cap: float = 50.0) -> torch.Tensor:
    """
    Soft-cap pre-softmax attention logits into (-cap, cap).
    attn: (B, H, T, T) raw logits (Q @ K^T * scale).
    Leaves |logit| << cap almost unchanged; saturates large logits.
    """
    return cap * torch.tanh(attn / cap)

def forward(self, x: torch.Tensor, mask=None, logit_cap: float = 50.0):
    B, T, D = x.shape

    def split_heads(t):
        return t.view(B, T, self.n_heads, self.d_k).transpose(1, 2)

    Q = split_heads(self.W_q(x))
    K = split_heads(self.W_k(x))
    V = split_heads(self.W_v(x))

    # QK-Norm: bounds logits to O(sqrt(d_k))
    Q = self.q_norm(Q)
    K = self.k_norm(K)

    scale = self.d_k ** -0.5
    attn = (Q @ K.transpose(-2, -1)) * scale        # (B, H, T, T)

    # --- Soft-cap BEFORE masking and softmax ---
    attn = attn_logit_softcap(attn, cap=logit_cap)

    if mask is not None:
        attn = attn.masked_fill(mask == 0, float('-inf'))
    attn = F.softmax(attn, dim=-1)

    out = (attn @ V).transpose(1, 2).contiguous().view(B, T, D)
    return self.W_o(out)