-
Notifications
You must be signed in to change notification settings - Fork 174
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
MultiSteps OOM on TPU #320
Comments
Hi! I think the problem is that you are increasing the training batch size - this isn't necessary. The wrapper counts batches as you pass them in and returns updates=0 whenever the count isn't k and then the accumulated gradients at the kth step. This means that there shouldn't be an extra memory requirement. Is this something that's confusing in the documentation? If so, is there something we could change to make it clearer? Thanks a lot! |
Thanks @mkunesch for your reply. Please, correct me if I am mistaken. From my experience in Pytorch + GPU + gradient accumulation, I can increase the accumulation steps without causing the GPU memory to increase and gives OOM. In the case of MultuSteps, if the actual maximum TPU batch size is 16, then I can only make the MultiStep work with 1 gradient accumulation step, if I want to increase the gradient accumulation step, then I have to decrease the actual TPU batch size. For example, if I want to make the gradient accumulation step = 2, then I must decrease the TPU batch size down to 8. This makes the virtual global batch size is always the same. My understanding is that I can increase the gradient accumulation steps in MultiSteps without having to decrease the actual TPU batch size. Please, correct me if my understanding is incorrect, or if my understanding is correct but there is a bug hat causing the TPU to give OOM. |
Hi! Yes, you should totally be able to increase the accumulation steps in So the following line in your code shouldn't be there:
Does that make sense or am I misunderstanding the problem? |
Yes, it makes a lot of sense now. So the train_batch_size should always be fixed and independent of the gradient_accumulation_steps. The correct code should be :
The MultiSteps and optimizer are expecting to receive the actual physical batch size, and then it will not calculate these batches until it hits the gradient_accumulation_steps. It is very clear now. It will be great if you could explain that in the documentation. Thanks a lot for your explanation. |
My suggestion is :
Maybe something like that. |
My initial assumption is that the MultiSteps will take the virtual batch size and perform a loop over the mini-batches. Thanks again for your excellent explanation. |
Hello,
I am testing "MultiSteps" for gradient accumulation on Colab TPU V2, but every time I use anything above 1, I get an OOM.
It seems that it increases the memory requirements equivalently to increase the batch size, which should not be the case.
My understanding is that all I need to do is to use it, is to call it after the optimizer and increase the training batch size by multiplying it by the gradient accumulation steps.
My understanding is also that I should be able to increase the batch size by using gradient accumulation.
In my case, if the per_device_train_batch_size is 8, then I can only set gradient_accumulation_steps to 1, and if the per_device_train_batch_size is 2, then I can then increase the gradient_accumulation_steps to 4.
The rest of the code should be the same.
Is my understanding is correct or there is something else we should take care of while using the MultiSteps function?
The text was updated successfully, but these errors were encountered: