Thank you so much for the amazingly clear explanation on the GPU usage.
I myself is confused with the backward step of second step, I accept your theory of allocating buffer to update weights. But I come across with one case which is that the allocating buffer size is very huge (not a little) leading to OOM. Below is the memory usage log. Do you think it’s reasonable? The main memory user in my model is embedding layer with size 12 million item X 64 dimension, which is just around 3.07G. I believe it’s not a coincidence. Why another huge memory has to allocate to the embedding weights when the embedding.weight and embedding.weight.grad have already been allocated memory in the first step?
-- beginning
allocated: 3144M, max allocated: 3144M, cached: 3148M, max cached: 3148M
-- zero grads
allocated: 3144M, max allocated: 3144M, cached: 3148M, max cached: 3148M
-- forward
allocated: 3157M, max allocated: 3249M, cached: 3164M, max cached: 3254M
-- compute loss
allocated: 3158M, max allocated: 3249M, cached: 3164M, max cached: 3254M
-- backward
allocated: 6291M, max allocated: 12590M, cached: 6328M, max cached: 12624M
-- optimizer step
allocated: 6291M, max allocated: 12590M, cached: 6328M, max cached: 12624M
***************** finish train batch: 1

-- zero grads
allocated: 6291M, max allocated: 12590M, cached: 6328M, max cached: 12624M
-- forward
allocated: 6302M, max allocated: 12590M, cached: 6332M, max cached: 12624M
-- compute loss
allocated: 6303M, max allocated: 12590M, cached: 6332M, max cached: 12624M
....
....
RuntimeError: CUDA out of memory. Tried to allocate 3.07 GiB (GPU 0; 15.75 GiB total capacity; 12.30 GiB already allocated; 2.46 GiB free; 11.35 MiB cached)

I also made a comment in the issue https://github.com/pytorch/pytorch/issues/16655#issuecomment-498951709 22