Hi,
Thanks for the detailed question and measures !
1- when you compute the loss, you allocate memory for the new output Tensors and the intermediary results within the loss function itself. So it is expected to see an increased memory usage during that step. In particular 400M is not too large.
2- The backward does not corresponds to the Adam step. 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).
3- The extra 500M you see during the first optimizer step is exactly what you computed in question 2: the extra buffers needed by the adam optimizer. model.zero_grad() will zero out the gradients Tensors but not free them. Since you will reuse them, it would be bad performance-wise to free them and re-allocate them everytime. I guess the memory allocation here is: 500M gradients Tensors, 500M adam buffers, 300M others (I guess stats from batchnorm or similar layer).
4- The max memory increases here again because of the allocator most likely. It has some buffer that he uses and cannot reuse exactly the same ones and so have to re-allocate a bit more. This can happen for the first two steps but should be stable after that. As you can see this only changes the cached values and not the actually used one. If you’re familiar with memory fragmentation, that’s what is happening here.
5- That depends why you want to reduce memory usage.

If it’s to be able to fit larger batch size or larger net, you don’t need to do anything. Even though more memory is allocated on the device, all the cached but not used memory is still available to put Tensors. So it should not be a problem. If you still OOM, you can use the checkpoint 109 tool to trade memory for compute.
If you want to give that memory back to the OS to use it with another program, that’s going to be a bit trickier. There is an empty_cache() function to return all cached but not used memory back to the OS but you should be very careful when using it. It will slow down significantly the next allocations (as we need to re-allocate memory again from cuda) and won’t necessarily reduce the peak memory usage as it may not decrease fragmentation.

Hope this helps, let me know if you have further questions !