Large Batch Training › Worked Example: Training with Different Batch Sizes

Let us work through a concrete example to make the scaling rule tangible. Suppose you have determined through experimentation that batch size $B = 256$ trains well with learning rate $\eta = 0.1$ using SGD.

You now want to scale to a 64-GPU cluster where, to keep all GPUs busy, you need to use at least 16,384 examples per step.

The scaling factor is:

$$k = \frac{16384}{256} = 64$$

The linearly scaled learning rate is:

$$\eta' = k \cdot \eta = 64 \times 0.1 = 6.4$$

For warmup, a common practice is to linearly increase the learning rate from a small initial value to the target $\eta' = 6.4$ over the first few epochs, often 5 epochs for ImageNet-scale tasks.

At the end of training, the model has processed the same number of gradient updates relative to data seen. If the original run did 100 epochs of 1,000 steps with batch size 256, the scaled run does 100 epochs of roughly 16 steps with batch size 16,384. The same amount of data is processed overall, but in far fewer, more accurate steps.

In practice, this specific example is near the edge of where the rule holds. For batch sizes beyond roughly 8,192 on ImageNet, researchers have found that extra techniques (gradient noise injection, LARS/LAMB, more aggressive warmup) are needed to maintain quality. This example sits at 16,384, so you would likely need at least LAMB or careful warmup tuning to close the generalization gap.