如何为大规模预训练设置批大小
人工智能
2026-01-12 v2
摘要
由 OpenAI 首创的临界批大小概念,长期以来一直作为大规模预训练的基础原则。然而,随着向 Warmup-Stable-Decay(WSD)学习率调度器的范式转变,我们观察到原有的理论框架及其潜在机制无法与新的预训练动态相吻合。为了弥合理论与实践之间的差距,本文推导了专为 WSD 调度器定制的修正 E(S) 关系,刻画了预训练期间训练数据消耗 E 与步数 S 之间的权衡。我们的理论分析揭示了基于 WSD 的预训练的两个基本属性:1)B_min,即达到目标损失所需的最小批大小阈值;2)B_opt,即通过最小化总 token 数来最大化数据效率的最优批大小。基于这些属性,我们提出了一种动态批大小调度器。大量实验表明,我们修正后的公式精确捕捉了大规模预训练的动态,由此产生的调度策略显著提升了训练效率与最终模型质量。
引用
@article{arxiv.2601.05034,
title = {How to Set the Batch Size for Large-Scale Pre-training?},
author = {Yunhua Zhou and Junhao Huang and Shuhao Xing and Yechen Zhang and Runyu Peng and Qiping Guo and Xipeng Qiu},
journal= {arXiv preprint arXiv:2601.05034},
year = {2026}
}