Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation is Wasteful
Martin Marek, Sanae Lotfi, Aditya Somasundaram, Andrew Wilson, Micah Goldblum
Conventional wisdom dictates that small batch sizes make language model pretraining unstable, motivating gradient accumulation which trades off the number of optimizer steps for a proportional increase in batch size. While it is common to decrease learning rate for smaller batch sizes, other hyperparameters are often held fixed. In this work, we revisit small batch sizes, all the way down to batch size one, with appropriate hyperparameters. We find that small batch sizes (1) train stably, (2) are consistently more robust to hyperparameter choices, (3) achieve equal or better per-FLOP performance as larger batch sizes, and (4) notably enable stable language model training with vanilla SGD, even without momentum, despite storing no optimizer state.  Building on these results, we provide practical recommendations for selecting a batch size and setting optimizer hyperparameters. We further recommend against gradient accumulation in settings where throughput is not bound by inter-node bandwidth.