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

TL;DR

This work shows that the traditional Critical Batch Size and the OpenAI $E(S)$ framework do not capture the data-steps dynamics of modern large-scale pre-training under Warmup-Stable-Decay schedules. It introduces a novel $E(S)$ relationship tailored to the WSD regime, identifies two key metrics, $B_{min}$ and $B_{opt}$, and proves these increase as training loss decreases. A theoretical and empirical pipeline then yields a dynamic batch-size scheduler that grows the batch size over time, improving both training efficiency and final performance. The approach is validated across InternLM2-based fitting and Qwen3-based scheduling, with consistent gains on downstream benchmarks such as MMLU and CMMLU, underscoring its practical impact for large-scale pre-training with contemporary optimizers and LR schedules.

Abstract

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.

Figures (18)

• Figure 1: Loss curves for models trained with different batch sizes (Stable phase under WSD schedule). The red region denotes the regime where the $E(S)$ formula and Critical Batch Size theory remain valid. In the green region, the $E(S)$ relationship no longer holds, leading to a failure of the Critical Batch Size framework. Post-intersection, the partial ordering of data consumption among the various batch sizes is inverted.
• Figure 2: Fitting results of $E(S)$ for 1B model. We select the target loss interval as [2.93, 3.25] and perform fitting on the $E(S)$ curves for target losses within this interval.
• Figure 3: The variation of $B_{min}$ and $B_{opt}$ with respect to loss across different model sizes.
• Figure 4: Training loss curves for Qwen3 MoE using fixed and dynamic batch size strategies under a constant learning rate schedule.
• Figure 5: Comparison of downstream benchmark results for Qwen3 MoE under fixed vs. dynamic batch size scheduling at a constant learning rate.