ZClip: Adaptive Spike Mitigation for LLM Pre-Training

TL;DR

This paper tackles gradient instability and loss spikes during large-scale pre-training by introducing ZClip, an adaptive gradient clipping method guided by z-score anomaly detection. ZClip maintains lightweight EMA-based estimates of gradient-norm statistics and adjusts clipping strength via a reciprocal z-score function, enabling stable optimization across dynamic learning-rate regimes. Empirical results on a 1B-parameter LLaMA-like setup show ZClip reduces or eliminates spikes, expands the viable learning-rate space, and improves convergence speed and downstream task performance with minimal overhead. The approach offers practical benefits for large-scale pre-training and potentially broader applicability to other noisy training scenarios such as RL or multimodal learning.

Abstract

Training large language models (LLMs) presents numerous challenges, including gradient instability and loss spikes. These phenomena can lead to catastrophic divergence, requiring costly checkpoint restoration and data batch skipping. Traditional gradient clipping techniques, such as constant or norm-based methods, fail to address these issues effectively due to their reliance on fixed thresholds or heuristics, leading to inefficient learning and requiring frequent manual intervention. In this work, we propose ZClip, an adaptive gradient clipping algorithm that dynamically adjusts the clipping threshold based on statistical properties of gradient norms over time. Unlike prior reactive strategies, ZClip proactively adapts to training dynamics without making any prior assumptions on the scale and the temporal evolution of gradient norms. At its core, it leverages z-score-based anomaly detection to identify and mitigate large gradient spikes, preventing malignant loss spikes while not interfering with convergence otherwise. Our code is available at: https://github.com/bluorion-com/ZClip.

Figures (13)

• Figure 1: Training loss graph comparing 1) training without clipping, 2) clipping with fixed threshold 1.0, and 3) ZClip for a LLaMA 1B model. The learning rate for all three experiments is $3.0 \times 10^{-3}$. While both "no clipping" and "constant clipping" exhibit spiky behavior and diverge early, ZClip (with $z_{\text{thres}}$=2.5 and $\alpha=0.97$) remains stable and continues to optimize effectively throughout training. Details on the model configuration and other training hyperparameters are presented in Appendix Section \ref{['sec:llama_config']}.
• Figure 2: Training loss graph for a LLaMA 1B model trained with fixed-threshold clipping c = 1.0. Gradient norm spikes persist due to a mismatch between the static threshold and the running distribution. This reveals a key limitation of fixed-threshold clipping in dynamically changing training regimes.
• Figure 3: Three possible choices for the z-score adjustment function $\xi(z_t)$ illustrated for $z_{\text{thres}} = 2.5$. Note the discontinuity for $\xi(z_t) = 0$, and the reciprocal nature of $\xi(z_t) = z_{\text{thres}}^2/z_t$ leading to more aggressive clipping for more extreme outliers.
• Figure 4: Comparison of Test Loss between ZClip (lr=3e-3) and a baseline model (lr=5e-4) on a 50B token corpus. ZClip achieved the same final loss as the baseline while requiring 18.6B fewer tokens to get there. The plot uses log-scaled training loss for visibility, and smoothing has been applied to reduce noise. ZClip allows for faster convergence without compromising on final loss value.
• Figure 5: ZClip and fixed-threshold clipping at higher learning rates. Each row shows training loss (left), gradient norm before clipping (middle), and after clipping (right). For $3.0 \times 10^{-3}$ ZClip stabilized gradients and reduces post-clipping spikes, unlike fixed-threshold clipping which accumulates instability. For $5.0 \times 10^{-3}$ both clipping methods saturated.