# How to Set the Batch Size for Large-Scale Pre-training?

arXiv: https://arxiv.org/abs/2601.05034 · HTML: https://mlsanity.com/paper/2601.05034/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 30 references that are in the map.
- First posted on arXiv 8 Jan 2026.

## Builds on (most cited first, 23 of 23)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Qwen2 Technical Report](https://mlsanity.com/paper/2407.10671/) (2024)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [On the Variance of the Adaptive Learning Rate and Beyond](https://mlsanity.com/paper/1908.03265/) (2019)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [DeepSeek LLM: Scaling Open-Source Language Models with Longtermism](https://mlsanity.com/paper/2401.02954/) (2024)
- [CMMLU: Measuring massive multitask language understanding in Chinese](https://mlsanity.com/paper/2306.09212/) (2023)
- [InternLM2 Technical Report](https://mlsanity.com/paper/2403.17297/) (2024)
- [An Empirical Model of Large-Batch Training](https://mlsanity.com/paper/1812.06162/) (2018)
- [One weird trick for parallelizing convolutional neural networks](https://mlsanity.com/paper/1404.5997/) (2014)
- [On the SDEs and Scaling Rules for Adaptive Gradient Algorithms](https://mlsanity.com/paper/2205.10287/) (2022)
- [On the Generalization Benefit of Noise in Stochastic Gradient Descent](https://mlsanity.com/paper/2006.15081/) (2020)
- [Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models](https://mlsanity.com/paper/2501.11873/) (2025)
- [Surge Phenomenon in Optimal Learning Rate and Batch Size Scaling](https://mlsanity.com/paper/2405.14578/) (2024)
- [Scaling Law for Language Models Training Considering Batch Size](https://mlsanity.com/paper/2412.01505/) (2024)
- [A Multi-Power Law for Loss Curve Prediction Across Learning Rate Schedules](https://mlsanity.com/paper/2503.12811/) (2025)

## Abstract (arXiv)

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.