High GPU memory usage problem

[stas (Stas Bekman)]

@albanD

The backward is going though the computational graph and computing gradients for every Tensors that is a leaf (in nn, usually nn.Parameters). The max memory increases here because halfway through the backprop, we had to allocate more buffers for the computed gradients but we could not clear the memory of the unecessary buffers yet. It is expected that the max memory usage is attained during the backward phase. Note as well that at the end of this phase, the allocated memory is greatly reduced again as we have freed all the intermediary buffers. You don’t go all the way down as before the forward call because the gradients are still there (and some internal buffers that some ops have like batchnorm).

@albanD, your explanation was very helpful and timely. Thank you!

Is there anything that can be done to trade off that inefficient setup phase that leads to a huge at times peak memory usage, but not affect the normal training (i.e. not reducing bs, model sizes, etc.)?

For example, MNIST (28x28) w/ bs=512 and resnet50 peaks at 6GB on the first pass, and then goes down to a steady 1GB peak for subsequent epochs - 6 vs 1 is an insanely huge overhead!

Can the computational graph/gradients’ setup be done in stages and then combined to keep the peak memory usage at a smaller multiple?