Training Stability › Stability in Practice: A Reference Checklist

Bringing together the techniques from this chapter, here is a practical reference for building a stable training configuration:

Architecture choices:

• Use pre-norm (LayerNorm before sub-layers) rather than post-norm for transformers with more than 12 layers
• Scale down output projections by $1/\sqrt{2L}$ at initialization to prevent residual stream variance growth
• Consider QK normalization for models with very large heads or trained on very long sequences

Optimizer choices:

• Use AdamW with $\beta_1 = 0.9$, $\beta_2 = 0.95$ or $0.999$
• Set $\epsilon = 10^{-8}$ for float32 training; increase to $10^{-6}$ for mixed-precision
• Apply gradient clipping at $c = 1.0$ as a starting point; tune based on observed clip frequency
• Use non-zero weight decay (typically 0.01-0.1) to prevent unbounded parameter growth

Scheduling choices:

• Warm up the learning rate linearly for 1000-4000 steps depending on model size
• Use cosine decay after warmup for smooth long-term stability
• Avoid sudden learning rate changes after warmup (no step functions)

Monitoring choices:

• Log gradient norm at every step
• Log per-layer gradient norms every 100 steps for architectural debugging
• Log parameter norms every 200 steps to detect weight decay failure
• Set up automated spike detection with rolling window statistics
• Checkpoint every 500-1000 steps with full optimizer state

Data choices:

• Filter documents with extreme token repetition, very high loss, or encoding artifacts
• Consider curriculum learning (easier examples first) to avoid outlier batches in early training
• Log the data indices of batches that precede loss spikes for post-hoc analysis