Resolving CUDA Being Out of Memory With Gradient Accumulation and AMP
By: Rishik C. Mourya
Tags: Pytorch, Optimization

A rticle Overview

Do you remember that delightful time, when you were so happily preprocessing your data, building your model with so many efforts, and then right at the moment when you expect to enjoy that soda while your model trains, an error pops up, telling you to get your butt down here again because you don’t have enough CUDA memory to enjoy? Well, I certainly can’t forget that, and this seriously is a really common issue among ML engineers. Alright then, let’s figure out if there’s a solution to our soda problem.

Explaining The Error

RuntimeError: CUDA error: out of memory

There’s nothing to explain actually, I mean the error message is already self-explanatory, but still, let’s have a quick brush up. The issue is, to train the model using GPU, you need the error between the labels and predictions, and for the error, you need to make predictions, and for making the predictions, you need both the model and the input data to be allocated in the CUDA memory. So when you try to execute the training, and you don’t have enough free CUDA memory available, then the framework you’re using throws this out of memory error.

Causes Of This Error

So keeping that in mind, here are the common reasons why this error is so prone to occur:

When you’re model is big, by big I mean lot’s of parameters to train.

When you’re using such a model architecture that performs a lot of concatenation operations (in DenseNet). Or huge channels operations like additions (in ResNet)

When the size of your input data is too high. I don’t mean to say the size of data you have available, but the input size of single data instance, like in case of an image, it could be 3x224x224.

And finally, when you’re batch size is too high with respect to the dimensions of a data instance.

Now the issue is not that you can’t solve the problems above, I mean switching to a smaller model might take not more than 20 to 30 minutes, and decreasing the batch size is, of course, a piece of cake, and such solutions do solve issue for most of the cases, but at the cost of time and model’s accuracy.

For instance, if you set your batch size too low, then it might just take forever to make the loss converge because your model would be just oscillating too much, not only that, some model architectures don’t even perform well when batch size is not up to the requirements. And, if you decide to switch to a smaller model or reduce the dimensions of the input data, then you might just lose quite a lot of accuracy for obvious reasons.

So as you can see solving this error is not that hard to tackle, but its the cost of the solution which is too high for a good ML engineer.

Coming up with a solution

Alright, so now it is pretty obvious that just one solution wouldn’t be enough to solve all of the above problems, in fact, a limited solution would only create even more problems. So after quite a lot of headaches I’ve finally came to these solutions, when applied together, do wonders :)

Gradient Accumulation.

Training with AMP(Automatic Mixed Precision).

Few of my own tricks.

Let’s have a look at them one by one.

Gradient Accumulation

This solution has been a de facto for solving out of memory issues. In fact, this is the very only reason why this technique exists in the first place. This technique solves the issue of model oscillation when the batch size is set too low. The nub here is actually very easy to understand.

The reason why model oscillates when batch size is too low is that the gradients are averages for less number of data instances. So after each training iteration (training a single batch), the model gets too specific updates, and thus parameters are changed too much towards a single direction. And reducing the learning rate would not be helpful here, in fact, it would just make the training even slower and just horrible because not only the updates are less generalized but with low magnitude too.

So the idea behind gradient accumulation is, keep adding gradients of the parameters for n number of batches, where batch size is small, and then after n batch iterations, apply the updates using the average of all those gradients accumulated over n iterations, that’s why its called gradient accumulation.

Below is the sample procedure for Pytorch implementation.

In the above example, note that we are dividing the loss by gradient_accumulations for keeping the scale of gradients same as if were training with 64 batch size. For an effective batch size of 64, ideally, we want to average over 64 gradients to apply the updates, so if we don’t divide by gradient_accumulations then we would be applying updates using an average of gradients over the batch size of 4 only, because we just kept adding the gradients for gradient_accumulations number of batch iterations before calling optimizer.step() .

So, this technique effectively solves the problem of limited gradient averaging. Now mathematically we are totally good here, but the problem is the training time. As our batch size is really low, so definitely we are doing drastically less parallel computations as compared to high batch size.

And if you are accumulating the gradients for k batch iterations then it would approximately take k times the time for training with the batch size that you actually want. Let’s fix this.

Automatic Mixed Precision

Almost all of the deep learning frameworks operate on 32-bit floating-point or float32 data type by default. Though there are many operations the does not need to be this much precisely accurate.

3 years ago Nvidia researchers created a new methodology, that in a nutshell combines single-precision with the half-precision floating-point for training deep learning models, that achieves the same level of accuracy as float32.

It enables automatic conversion of certain GPU operations from float32 precision to mixed-precision, while improving performance and maintaining same accuracy. For an exact performance and accuracy comparisons make sure to check out Nvidia’s official benchmarks.

Main plus points of this method are:

Less training time.

Enabling larger batch sizes.

Larger models and large inputs.

Lower memory requirements.

So now you can see why I decided to combine AMP with gradient accumulation. Anyways, below is its Pytorch implementation.

Why mixed, why not simply use float16?

I’ll let this paper Multi-Node BERT-Pretraining explain this to you

Typically in a computation graph, not all FP16 operators are numerically safe. This means operators that are considered numerically dangerous will have its calculation in full precision. For example, a plus operator is marked as safe while a 5 power or a log operator is considered numerically dangerous in half precision. Automated mixed precision handles the categorization of the numerical safety level through the rewriting of computation graph — Authors

For more info regarding different precisions, make sure to check out this brief article by non other than Nvidia.

Finishing up the issue

So now, since we are set to use both of the techniques I’ve explained above, let’s put the together to come up with an awesome combo.

Yuppah, that’s pretty much it :)

Depending on the model and input sizes, and how much gradients you are accumulating, you could easily get somewhere around 1.2x to 1.5x (and sometimes even more) speed ups with AMP, and of course your loss would converge thanks to gradient accumulation.

Few of my own tricks

Now for 95% of the times, you wouldn’t need to use techniques that I’m sharing here because AMP and gradient accumulation applied together works juuuuust fine. Though recently I’ve started working for a startup, where we are trying to detect violence in CCTV footages. This task basically is a video classification, where in a batch we would have to deal with one more dimension as compared to image classification, making total 5 dimensions of a single batch. So this definitely means massive input size.

I wanted to further speed up the training, cuz we had over 16GB of training data so speeding the things was definitely a requirement.

Here’s what I did.

Reducing the input data features . In my case I reduced the temporal dimensions to reasonable 28, rather than 36 as before. This not sound much of a decrease, but keep that in mind that a single step in temporal dimension in a video means covering 3 * 224 * 244 feature values… so in total I decreased over 1 million features :), not only that I also reduced the image size from 224 to reasonable 196.

. In my case I reduced the temporal dimensions to reasonable 28, rather than 36 as before. This not sound much of a decrease, but keep that in mind that a single step in temporal dimension in a video means covering 3 * 224 * 244 feature values… so in total I decreased over 1 million features :), not only that I also reduced the image size from 224 to reasonable 196. Making classes distribution uniform over a batch. This is very simple actually, the idea is to make sure that every batch contains the same number of input data for every class. So in my case each batch always contained 4 non violence videos and 4 violence videos for batch size of 8. This would help to minimize the data imbalance (sort of, not 100%) that I had and finally improve the testing accuracy within same training time.

Yes, these ideas are not necessarily for solving the out of CUDA memory issue, but while applying these techniques, there was a well noticeable amount decrease in time for training, and helped me to get ahead by 3 training epochs where each epoch was approximately taking over 25 minutes.

Conclusion

So we saw how to use gradient accumulation and automatic mixed precision together to come out limited memory hell.

Gradient accumulation workouts the math and helps to solve the loss oscillation issue, but small batch size causes slower training

And this is where automatic mixed precision and few of my tweaks helps to speed up the whole training process.

Okie then, I hope after reading this article you’d be able to tackle out this really common issue and yeah, enjoy the soda.