0
$\begingroup$

I am currently training some models using gradient accumulation since the model batches do not fit in GPU memory. Since I am using gradient accumulation, I had to tweak the training configuration a bit. There are two parameters that I tweaked: the batch size and the gradient accumulation steps. However I am not sure about the effects of this modification so I would like to fully understand what is the relationship between the gradient accumulation steps parameter and the batch size.

I know that when you accumulate the gradient you are just adding the gradient contributions for some steps before updating the weights. Normally, you would update the weights every time you compute the gradients (traditional approach):

$w_{t+1} = w_t - \alpha \cdot \nabla_{w_t}loss$

But when accumulating gradients you compute the gradients several times before updating the weights (being $N$ the number of gradient accumulation steps):

$w_{t+1} = w_t - \alpha \cdot \sum_{0}^{N-1} \nabla_{w_t}loss$

My question is: What is the relationship between the batch size $B$ and the gradient accumulation steps $N$?

By example: are the following configurations equivalent?

  • $B=8, N=1$: No gradient accumluation (accumulating every step), batch size of 8 since it fits in memory.
  • $B=2, N=4$: Gradient accumuluation (accumulating every 4 steps), reduced batch size to 2 so it fits in memory.

My intuition is that they are but I am not sure. I am not sure either if I would have to modify the learning rate $\alpha$.

$\endgroup$
1
$\begingroup$

There isn't any explicit relation between the batch size and the gradient accumulation steps, except for the fact that gradient accumulation helps one to fit models with relatively larger batch sizes (typically in single-GPU setups) by cleverly avoiding memory issues. The core idea of gradient accumulation is to perform multiple backward passes using the same model parameters before updating them all at once for multiple batches. This is unlike the conventional manner, where the model parameters are updated once every batch-size number of samples. Therefore, finding the correct batch-size and accumulation steps is a design trade-off that has to be made based on two things: (i) how much increase in the batch-size can the GPU handle, and (ii) whether the gradient accumulation steps result in at least as much better performance than without accumulation.

As for your example configurations, there are the same in theory. But, there are a few important caveats that need to be addressed before proceeding with this intuition.

  1. Using Batch Normalization with gradient accumulation generally does not work well, simply because BatchNorm statistics cannot be accumulated. A better solution would be to use Group Normalization instead of BatchNorm.
  2. When performing a combined update in gradient accumulation, it must be ensured that the optimizer is not initialized to zero (i.e. optimizer.zero_grad()) for every backward update (i.e. loss.backward()). It is easy to include both statements in the same for loop while training which defeats the purpose of gradient accumulation.

Here are some interesting resources to find out more about it detail.

  1. Thomas Wolf's article on different ways of combating memory issues.
  2. Kaggle discussions on the effect on learning rate: here and here.
  3. A comprehensive answer on gradient accumulation in PyTorch.
| improve this answer | |
$\endgroup$
  • $\begingroup$ Thank you for your answer. The caveats section was specially helpful. I was already aware of point (2) but point (1) was very ilustrative. As you hint, it seems it is better to move to multi GPU... I'll add it to the TODO list. $\endgroup$ – JVGD Jun 18 at 6:56

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.