Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Training

Thank you for your review!

Regarding weakness (1), our claim is that batch size warmup attains comparable (in some cases slightly better) loss compared to the small-batch run while using fewer gradient steps. On the other hand, the large-batch run degrades loss, as shown in Figure 4 and Table 2. The takeaway we want to highlight is that the warmup method is a way to achieve comparable loss with fewer gradient steps; while the large-batch method saves more gradient steps, we find it degrades loss.

Regarding Lines 71-73, this was a typo about the branched training method. Thanks for catching this! In fact, branched training is less expensive compared to training many full runs, and batch size warmup does save 43% gradient steps relative to the small-batch control. We will update lines 71-73 to correctly say “less expensive”.

Regarding data scale, we agree this is an interesting question, but outside the scope of this project, where we focus on the evolution of CBS over the course of training given a fixed data distribution. As mentioned in Section 2, studies [1] have suggested that CBS scales with data scale.

Regarding the comparison of gradient noise scale with the branched training approach, the contrast between Figures 2 and 3 does give a comparison of these approaches for estimating the CBS. We do agree that it could have been interesting to consider batch size warmup with a gradient noise schedule, but we did not have time to launch this experiment during the rebuttal period, and, while it would have been interesting to have, this experiment is not central to the main claims in the paper (see further discussion in our response to Reviewer Hdu9).

Regarding doubling the batch size, we chose this approach because of the practical constraint that we want the batch size to remain a multiple of the number of GPUs whenever we increase it. Because our number of GPUs during pretraining is a power of 2, doubling whenever we reach a new power of two was an easy way to maintain this invariant. In general, one could imagine more complex approaches, such as explicitly selecting the largest multiple of the number of GPUs less than the CBS. We will clarify our motivation and the potential alternatives around this hyperparameter choice.

[1] https://arxiv.org/abs/2410.21676

Thank you for your review and your appreciation of our work! We are very happy to hear that you appreciated the contributions and clarity of the paper.

Regarding your questions about the branched training method, we have some thoughts.

The authors write that the gradient noise estimate of the CBS "underestimates the CBS by several orders of magnitude". What is the ground truth CBS that the authors are using in such a statement? Figure 3 would suggest that the answer is 4096 (at least, towards the end of training) but I'm not sure why this is treated as the actual CBS. This number is closer to what their own methodology suggests is the correct CBS, but that doesn't make it the ground truth. For example, is it possible that the method in Section 3.1 overestimates the CBS?

This is a good question, and we will try to make our reasoning more explicit. Indeed, the comparison between Figures 2 and 3 suggests that the gradient noise scale underestimates the CBS relative to our method. A priori, it is possible that our method is an overestimate and that the gradient noise scale is correct. However, our pretraining results in Figure 4 provide evidence against this (i.e., they show that our method is not an overestimate), because we can train with the larger batch size predicted by our method and obtain no degradation in final loss. By definition of the CBS, this means the batch size produced by our method does not exceed the CBS (and thus the smaller gradient noise estimate is an underestimate).

Why not train with a gradient noise scale-based batch size schedule?

We think this could have been a nice, though not entirely central, comparison to have. The prediction based on the previous sections would be that we could still obtain decent pretraining loss (since we’re training below the CBS) but that the number of gradient steps would be significantly larger than training with the branched training batch size. The latter claim is true by construction and thus not worth testing, but it could have been interesting to measure the loss attained. On the other hand, it takes a while to launch the pretraining runs behind the results in Figure 3, and we were not able to launch this pretraining run during the short rebuttal period, especially given that the results would not be essential to our main findings.

How does the CBS estimate depend on $\Delta$?

Under the local recovery assumption (Assumption 3), increasing $\Delta$ up to the total pretraining token budget will only increase the estimate of the CBS, so CBS estimates with smaller $\Delta$ are still a lower bound. The choice of $\Delta$ can thus be understood as striking a tradeoff between finding the largest possible CBS lower bound and the compute budget one is willing to spend (since larger values of $\Delta$ require more training). We did not explore this tradeoff in further empirical detail since we chose to allocate our attention and compute towards the large scale batch size warmup pretraining runs instead. However, we’ll try to clarify the role of $Delta$ in the paper and add a note that understanding the role of $\Delta$ is an open question for future work.

I found Appendix D to be quite interesting … In the context of fitting such a power law, what would happen if you tried to fit your local CBS estimates to a power law?

We spent some time trying to fit a parametric form to our CBS plots, and we did not find any form that we were confident fit the shape of the function. This could be because there is a phase change between the first and second parts of the trend. Because we did not have any conclusive insights here, we did not make any claims about the functional form of the CBS trend in the paper, though we thought it would still be interesting to include the theoretical investigation related to this in the appendix since it might be relevant for future research.

Thanks for your review!

Regarding the downstream evaluations, we have found the loss-based evaluations as represented in https://arxiv.org/abs/2412.04403 to be some of the most reliable for evaluating pretraining interventions. Of course, this kind of evaluation is an unsolved problem with a lack of scientific consensus. If you have some specific suggestions for alternate types of evaluations we could use here, we would be curious for suggestions.

Indeed, the batch size warmup schedule is based on the results in Section 3. Because we only needed to determine two thresholds for increasing the batch size, we determined them manually based on Figure 2 before launching the batch size warmup pretraining runs. We did not have the resources to relaunch the pretraining runs using a more online approach, but we agree this would be a useful refinement of the approach and revise the paper to clarify this.

From Figure 2, it seems that the 1B model should reach the batch size of 4K at 100B tokens? However, the authors say that this is reached at 503B tokens in section 4. Not sure where I get it wrong?

We chose the thresholds manually and conservatively in Figure 4, with the goal of having the warmup batch size upper bounded by the curve in Figure 2. Based on a noisier initial version of Figure 2, we preferred a later threshold for increasing to batch size 4K, and once we produced the newer Figure 2 plot, we did not have the resources to relaunch the pretraining runs. We will clarify this manual choice in the paper and emphasize that there is an opportunity for future work to determine batch size thresholds more systematically.

Thanks for your review! We are glad to hear that you find our paper well-written and that its findings are convincing.

We’ll make sure to incorporate some comparison with [1] as you suggest.

Regarding the minor arbitrariness in setting the batch sizes thresholds for the warmup experiment, we set the threshold manually early on the research project based on a preliminary version of Figure 2 and then launched the batch size warmup pretraining runs. We agree that, all else equal, it would be nice to eliminate this arbitrariness. However, the pretraining runs for the warmup experiments were quite expensive, and we judged that the substantial cost for rerunning those experiments was not worth it to eliminate minor arbitrariness, especially since we believe the results are not especially sensitive to a precise choice of threshold. We will add some brief discussion of this limitation in the paper and highlight that making the choice of threshold more systematic would be appealing in future work.