提出新方法确定大模型预训练最优批量大小,提升训练效率与效果。
How to Set the Batch Size for Large-Scale Pre-training?
- 基于温控稳定衰减调度器,推导出新的训练数据消耗与步数关系
- 发现最小批量大小和最优批量大小可显著提升数据利用率
- 动态调整批量大小,适用于大规模模型高效训练
OpenAI提出的临界批量大小概念长期作为大规模预训练的基础原则。然而,随着温控-稳定-衰减(WSD)学习率调度器的兴起,原有理论框架与实际训练动态不再一致。本文针对WSD调度器推导出修正后的E(S)关系,刻画了预训练中训练数据消耗E与训练步数S之间的权衡。理论分析揭示了基于WSD预训练的两个基本特性:1)B_min——达到目标损失所需的最小批量大小;2)B_opt——通过最小化总令牌数实现数据效率最大化的最优批量大小。基于此,我们提出一种动态批量大小调度策略。大量实验表明,修正公式能精确捕捉大规模预训练动态,所提调度策略显著提升训练效率与最终模型质量。
原文摘要 · Abstract (English)
The concept of Critical Batch Size, as pioneered by OpenAI, has long served as a foundational principle for large-scale pre-training. However, with the paradigm shift towards the Warmup-Stable-Decay (WSD) learning rate scheduler, we observe that the original theoretical framework and its underlying mechanisms fail to align with new pre-training dynamics. To bridge this gap between theory and practice, this paper derives a revised E(S) relationship tailored for WSD scheduler, characterizing the trade-off between training data consumption E and steps S during pre-training. Our theoretical analysis reveals two fundamental properties of WSD-based pre-training: 1) B_min, the minimum batch size threshold required to achieve a target loss, and 2) B_opt, the optimal batch size that maximizes data efficiency by minimizing total tokens. Building upon these properties, we propose a dynamic Batch Size Scheduler. Extensive experiments demonstrate that our revised formula precisely captures the dynamics of large-scale pre-training, and the resulting scheduling strategy significantly enhances both training efficiency and final model quality.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。