6.3 Large-scale Training Stability

A loss spike is an incident signal, not a root cause. The same curve can be produced by a malformed batch, a mask or label-shift bug, activation growth, low-precision overflow, an optimizer transition, a desynchronized rank, a bad checkpoint, or hardware/storage corruption. A useful stability system preserves enough evidence to distinguish these hypotheses before changing the run.

Provider reports can show that a particular training stack completed a long run without an unrecoverable spike—for example, the DeepSeek-V3 report describes its own FP8 training outcome [1]. That is evidence about the reported system, not a guarantee that QK normalization, FP8, or one recipe makes another model rollback-free.

1. What the Loss Can and Cannot Explain

For one target class $y$, softmax cross-entropy has logit derivative

Each component lies in $[-1,1]$. Therefore, one out-of-distribution target token does not directly create an “astronomical” gradient with respect to its final logit. Parameter gradients can still become large because the bounded logit gradient is multiplied by large activations and network Jacobians; accumulation, optimizer state, numerical overflow, or distributed corruption can amplify the effect. This distinction determines what to inspect.

Source: AI-generated illustration. Treat the shown trajectory as intuition, not a universal causal sequence.

Separate four incident classes:

• Data/objective: corrupted text, extreme repetition, wrong labels, lost EOS boundaries, attention or padding-mask mistakes, duplicated packs, unexpected mixture drift.
• Optimization: learning-rate discontinuity, insufficient warmup, clipping regime change, optimizer-state mismatch, excessive effective batch or stale gradients.
• Numerics/model: nonfinite activations, attention or LM-logit growth, unstable reductions, low-precision overflow/underflow, initialization or normalization defects.
• Systems: desynchronized ranks, collective failure, silent data corruption, incorrect checkpoint shards, kernel/compiler regression, storage retry returning wrong content.

Do not infer the class from the loss plot alone. Preserve batch IDs, artifact hashes, rank-level metrics, and a known-good checkpoint so the event can be replayed.

2. Stabilizers and Their Scope

QK normalization

Normalizing query and key vectors before their dot product can control attention-logit scale:

where $M$ contains causal and padding constraints. This may reduce attention saturation in some architectures, but it does not bound feed-forward activations, the final language-model logits, optimizer state, or bad data. Learning-rate claims must be re-established for the actual model and initialization rather than copied as a universal multiplier.

Attention-logit and LM-logit capping

A bounded transform such as

can be applied to attention scores or to final LM logits. These are different interventions with different gradients and quality trade-offs. Name the tensor, cap value, precision, masking order, and ablation result. Capping can hide a growth symptom; it does not replace finding why the uncapped tensor drifts.

The $z$-loss

The auxiliary $z$-loss penalizes the squared log partition,

and has been used to discourage logit-scale drift in large language models [2]. Compute it in adequate precision and only on valid causal target positions. Applying it to padding or ignored prompt positions changes the effective objective and makes its magnitude batch-shape dependent.

3. Mask-Correct Causal Loss

The following is a smoke-testable loss helper, not a complete training loop. It assumes logits has shape (batch, sequence, vocabulary), labels contain the input token IDs with padding or excluded positions set to ignore_index, and the model has already applied causal attention and the correct padding mask.

import torch
import torch.nn.functional as F

def causal_cross_entropy_with_z_loss(
    logits: torch.Tensor,
    labels: torch.Tensor,
    *,
    ignore_index: int = -100,
    z_loss_weight: float = 1e-4,
):
    if logits.ndim != 3 or labels.shape != logits.shape[:2]:
        raise ValueError("expected logits (B, L, V) and labels (B, L)")

    # Token t predicts token t+1.
    shift_logits = logits[:, :-1, :].float().contiguous()
    shift_labels = labels[:, 1:].contiguous()
    valid = shift_labels.ne(ignore_index)
    if not torch.any(valid):
        raise ValueError("batch has no supervised causal targets")

    ce_sum = F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1),
        ignore_index=ignore_index,
        reduction="sum",
    )
    ce = ce_sum / valid.sum()

    log_z = torch.logsumexp(shift_logits, dim=-1)
    z_loss = log_z.square()[valid].mean()
    total = ce + z_loss_weight * z_loss
    return total, {"cross_entropy": ce.detach(), "z_loss": z_loss.detach()}

# Shape, masking, dtype, and finite-value smoke test.
torch.manual_seed(7)
toy_logits = torch.randn(2, 5, 11, dtype=torch.bfloat16, requires_grad=True)
toy_labels = torch.tensor([[1, 2, 3, 4, 5], [6, 7, -100, -100, -100]])
toy_loss, parts = causal_cross_entropy_with_z_loss(toy_logits, toy_labels)
toy_loss.backward()
assert torch.isfinite(toy_loss)
assert toy_logits.grad is not None and torch.isfinite(toy_logits.grad).all()

In distributed training, normalize by the global valid-token count, not the average of unequal per-rank means. Otherwise padding and variable-length batches change rank weights. If using sequence/context parallelism, verify which process owns each target and which collective produces the numerator and denominator.

4. Observability Before Intervention

Log signals at a cadence that can resolve a spike without overwhelming the telemetry path:

Histograms and per-layer tails often reveal a localized failure hidden by global means. Keep telemetry overhead bounded and test that alerts still arrive during a collective or storage incident.

5. Loss-Spike Triage Runbook

1. Freeze evidence. Record the last known-good checkpoint, current artifact hashes, global step/token count, learning rate, sample/pack IDs, rank logs, and hardware events. Do not overwrite the suspect checkpoint or data shard.
2. Check synchronization and finiteness. Confirm every rank agrees on step, token count, checkpoint generation, and nonfinite status. Identify the first layer/rank where values diverge.
3. Replay from known good. Restore the complete checkpoint and run the suspect batch with identical code/config/world size. A non-reproducible event shifts suspicion toward nondeterminism, system state, or hardware.
4. Compare controlled branches. Run replay, skip, and clean-control batches. Change one factor at a time; do not simultaneously skip data, reduce the learning rate, and discard optimizer state.
5. Bisect the class. Validate shard checksums and masks; compare higher-precision or safer kernels; inspect optimizer and scheduler identity; move the workload away from suspect hardware.
6. Choose a repair and gate it. Resume only after a short continuation matches expected loss, update norms, throughput, and held-out metrics. Record why the repair is causal rather than merely correlated.

Skipping a batch can hide a deterministic model or mask bug. Discarding optimizer state changes the training trajectory and may destabilize the restart. Both are experiments, not default remedies.

6. Checkpoint and Recovery Signals

A recovery checkpoint contains sharded model and optimizer state, scheduler, gradient scaler or FP8 state, global step and consumed tokens, RNG states, sampler and exact data cursor, plus dataset/tokenizer/model/code/config hashes. Publish it with checksums and a completion marker; load only completed generations.

Choose checkpoint cadence from failure economics. If checkpoints take $C$ minutes and the observed mean time between recoverable failures is $F$, the best interval depends on write interference, restart time, and expected lost compute—not a universal number of steps. Measure background checkpoint impact and restore bandwidth.

Run periodic drills:

• restore at the original world size and compare next batch IDs and updates;
• restore at a changed world size only if the format and data sampler support resharding;
• corrupt or omit a checkpoint shard and confirm the loader fails closed;
• restore the last-known-good artifact and execute the production rollback path;
• distinguish bitwise replay from statistically equivalent continuation under the actual kernels and framework version.

7. MoE and Low-Precision Stability

MoE training adds routing imbalance, token drops, expert-parallel collectives, and router-objective interactions. Bias-based load-balancing methods reported by DeepSeek-V3 are one design point [1]; “auxiliary-loss-free” does not mean balance is perfect or that the router has no operational state. Measure quality and load together and keep the implementation tied to the cited algorithm.

FP8 adds tensor-specific range and scaling state. The E4M3 maximum and chosen block/tensor scaling scheme alone do not determine stability. Log amax histories, saturation/underflow, scaling delays, accumulation/reduction dtype, fallback tensors, and kernel version. Save scaling state in checkpoints and compare BF16 controls during pilots. A numerically stable run must also meet quality and throughput gates.

8. Interactive: The Loss Spike Simulator

The simulator illustrates how internal norm growth and stabilizers can affect a hypothetical run. It is not an incident classifier; use the evidence and replay workflow above on real training jobs.

Stability engineering is successful when the team can detect, localize, reproduce, repair, and safely resume an incident—not merely when a dashboard looks smooth.

Quizzes

References

1. DeepSeek-AI. (2024). DeepSeek-V3 Technical Report. arXiv:2412.19437.
2. Chowdhery, A., et al. (2022). PaLM: Scaling Language Modeling with Pathways. arXiv:2204.02311.
3. PyTorch. Numerical Accuracy. PyTorch developer notes.
4. PyTorch. Reproducibility. PyTorch developer notes.