提出直接测量语言模型训练中临界批量大小的新方法,提升大规模训练效率。
Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Training
- 通过实验直接测量临界批量大小,而非依赖理论假设。
- 发现临界批量随训练进程先快速上升后趋于平稳,适用于不同规模模型。
- 提出批量大小渐增策略,用更少梯度步数达到更好效果,适合大规模训练。
在大规模语言模型训练中,合适的批量大小至关重要:过大虽加速训练但降低词元效率。McCandlish 等(2018)提出基于梯度噪声的临界批量(CBS)估算方法,但需强假设支撑,可信度受限。本文提出一种简单、实证的直接测量方法,揭示了 CBS 随训练演化的规律:初始化时接近 0,初期迅速上升,随后趋于平稳。该趋势在 1B 与 7B 模型上均成立,表明小规模训练的 CBS 可指导大规模训练。据此提出批量大小渐增策略——从较小批量开始,随 CBS 增长逐步增大。在 OLMo 1B 上验证,仅用 43% 的梯度步数,即达比原训练略优的损失值,证明该框架可高效实现大批量训练,提升数据并行性而不牺牲性能。
原文摘要 · Abstract (English)
The right batch size is important when training language models at scale: a large batch size is necessary for fast training, but a batch size that is too large will harm token efficiency. To navigate this tradeoff, McCandlish et al. (2018) suggest that a critical batch size (CBS), below which training will not substantially degrade loss, can be estimated based on the gradient noise scale during training. While their method has been adopted in practice, e.g., when training GPT-3, strong assumptions are required to justify gradient noise as a proxy for the CBS, which makes it unclear whether their approach should be trusted in practice, limiting its applicability. In this paper, we introduce a simple, empirical approach to directly measure the CBS and show how the CBS evolves over training. Applying our approach to the OLMo models, we find that CBS is near 0 at initialization, increases rapidly at first, and then plateaus as training progresses. Furthermore, we find that this trend holds across different model sizes (1B and 7B), suggesting CBS from small training runs can inform larger-scale training runs. Our findings about how the CBS changes over training motivate batch size warmup as a natural way to reliably train language models at large batch size: start the batch size small and increase it as the CBS grows. To validate this claim, we use batch size warmup to train OLMo 1B to slightly better loss than the original training run with 43% fewer gradient steps. This shows how our framework can be applied to reliably train language models at larger batch sizes, increasing data parallelism without compromising performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。