別丟棄 Dropout：優化層級稀疏性以實現高效的 LLM 訓練與推論

摘要

層丟棄（又稱隨機深度）已被證實能在語言與視覺Transformer中實現更快的訓練、更高的準確度，以及對零樣本層剪枝的穩健性。然而，隨著模型與資料集規模擴大，丟棄——尤其是層丟棄——已大幅從大型語言模型（LLM）的預訓練方案中消失。雖然先前部分研究曾報告丟棄可能降低準確度，但尚無全面性研究量化此影響，更遑論減輕此影響。在本研究中，我們證明層丟棄應運用於最先進的LLM訓練中，並為訓練與訓練後效益確立最佳實踐與規模分析。具體而言，在最佳的層分佈、時間排程與最佳化器超參數下，我們觀察到在相同訓練FLOPs下，層丟棄能帶來更低的損失。在給定訓練步數下，LLM可達成更低或相近的驗證損失，同時節省高達25%的訓練FLOPs。此外，層丟棄能實現顯著的訓練後優化，例如早期退出、中間層跳過與自推測解碼，在可忽略的準確度損失下提供最高1.5倍的推論加速。在超過2400次訓練實驗中，涵蓋參數量從2.71億到82億的模型，以及多達1600億詞元的資料集，我們證明這些發現能可靠地延伸至大規模訓練情境。所有預訓練實驗均在Cerebras CS-3系統上執行。

English

Layer dropout (a.k.a. stochastic depth) has been shown to enable faster training, higher accuracy, and robustness to zero-shot layer pruning in both language and vision transformers. However, as models and datasets have scaled, dropout - particularly layer dropout - has largely disappeared from large language models (LLMs) pre-training recipes. While some prior work has reported that dropout can degrade accuracy, no comprehensive study has quantified, let alone mitigated, this effect. In this study, we show that layer dropout should be used in state-of-the-art LLM training, establishing best practices and scaling analysis for both training and post-training benefits. Concretely, with optimal layer distribution, time schedule, and optimizer hyperparameters, we observe that at the same training FLOPs layer dropout leads to lower loss. For a given number of training steps, LLMs can achieve lower or similar validation loss while saving upto 25% of training FLOPs. Moreover, layer dropout enables significant post-training optimizations, such as early exit, intermediate-layer skipping, and self-speculative decoding, yielding up to 1.5x inference speedup with negligible accuracy loss. Across more than 2400 training experiments, spanning models from 271M to 8.2B parameters and datasets up to 160B tokens, we demonstrate that these findings extend reliably to large-scale training regimes. All pre-training experiments were run on Cerebras CS-3 systems.