Spike No More: Stabilizing the Pre-training of Large Language Models
This addresses a critical stability issue for researchers and practitioners training large language models, though it appears incremental as it builds on existing pre-training frameworks.
The paper tackles the problem of loss spikes during pre-training of large language models, which degrade performance and waste computational resources, and finds that methods based on small sub-layers and large shortcuts effectively prevent these spikes.
Loss spikes often occur during pre-training of large language models. The spikes degrade the performance of large language models and sometimes ruin the pre-training. Since the pre-training needs a vast computational budget, we should avoid such spikes. Based on the assumption that the loss spike is caused by the sudden growth of the gradient norm, we explore factors to keep the gradient norm small through an analysis of the spectral norms of the Jacobian matrices for the sub-layers. Our findings suggest that stabilizing the pre-training process requires two conditions: small sub-layers and large shortcut. We conduct various experiments to empirically verify our theoretical analyses. Experimental results demonstrate that methods satisfying the conditions effectively prevent loss spikes during pre-training.