ZClip: Adaptive Spike Mitigation for LLM Pre-Training
Authors: name: Abhay Kumar, name: Louis Owen, name: Nilabhra Roy Chowdhury, name: Fabian Güra

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.