大规模预训练该如何设置批量大小？

由OpenAI开创的临界批量大小（Critical Batch Size）概念长期以来一直充当大规模预训练的基础原则。然而，随着范式转向Warmup-Stable-Decay（WSD）学习率调度器，我们观察到原有理论框架及其底层机制已无法契合新的预训练动态。为弥合理论与实践之间的鸿沟，本文推导了为WSD调度器量身修正的E(S)关系，以刻画预训练期间训练数据消耗E与步数S之间的权衡。我们的理论分析揭示了基于WSD的预训练的两个基本性质：1）B_min，即达到目标损失所需的最小批量阈值；2）B_opt，即通过最小化总token来最大化数据效率的最优批量大小。基于这些性质，我们提出一种动态批量大小调度器。大量实验表明，我们修正后的公式精确刻画了大规模预训练的动态，且由此得到的调度策略显著提升了训练效率与最终模型质量。