Understanding why memory allocation occurs during inference, backpropagation, and model update
In the process of tracking down a GPU OOM error, I made the following checkpoints in my Pytorch code (running on Google Colab P100):
learning_rate = 0.001
num_epochs = 50

device = torch.device(&quot;cuda:0&quot; if torch.cuda.is_available() else &quot;cpu&quot;)

print('check 1')
!nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

model = MyModel()

print('check 2')
!nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

model = model.to(device)

print('check 3')
!nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

print('check 4')
!nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

for epoch in range(num_epochs):
    train_running_loss = 0.0
    train_accuracy = 0.0

    model = model.train()

    print('check 5')
    !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

    ## training step
    for i, (name, output_array, input) in enumerate(trainloader):
        
        output_array = output_array.to(device)
        input = input.to(device)
        comb = torch.zeros(1,1,100,1632).to(device)

        print('check 6')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        ## forward + backprop + loss
        output = model(input, comb)

        print('check 7')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        loss = my_loss(output, output_array)

        print('check 8')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        optimizer.zero_grad()

        print('check 9')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        loss.backward()

        print('check 10')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        ## update model params
        optimizer.step()

        print('check 11')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        train_running_loss += loss.detach().item()

        print('check 12')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        temp = get_accuracy(output, output_array)

        print('check 13')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        train_accuracy += temp     

with the following output:
check 1
2MiB/16160MiB
check 2
2MiB/16160MiB
check 3
3769MiB/16160MiB
check 4
3769MiB/16160MiB
check 5
3769MiB/16160MiB
check 6
3847MiB/16160MiB
check 7
6725MiB/16160MiB
check 8
6725MiB/16160MiB
check 9
6725MiB/16160MiB
check 10
9761MiB/16160MiB
check 11
16053MiB/16160MiB
check 12
16053MiB/16160MiB
check 13
16053MiB/16160MiB
check 6
16053MiB/16160MiB
check 7
16071MiB/16160MiB
check 8
16071MiB/16160MiB
check 9
16071MiB/16160MiB
check 10
16071MiB/16160MiB
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-11-f566d09448f9&gt; in &lt;module&gt;()
     65 
     66         ## update model params
---&gt; 67         optimizer.step()
     68 
     69         print('check 11')

3 frames
/usr/local/lib/python3.7/dist-packages/torch/optim/optimizer.py in wrapper(*args, **kwargs)
     86                 profile_name = &quot;Optimizer.step#{}.step&quot;.format(obj.__class__.__name__)
     87                 with torch.autograd.profiler.record_function(profile_name):
---&gt; 88                     return func(*args, **kwargs)
     89             return wrapper
     90 

/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     26         def decorate_context(*args, **kwargs):
     27             with self.__class__():
---&gt; 28                 return func(*args, **kwargs)
     29         return cast(F, decorate_context)
     30 

/usr/local/lib/python3.7/dist-packages/torch/optim/adam.py in step(self, closure)
    116                    lr=group['lr'],
    117                    weight_decay=group['weight_decay'],
--&gt; 118                    eps=group['eps'])
    119         return loss

/usr/local/lib/python3.7/dist-packages/torch/optim/_functional.py in adam(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, amsgrad, beta1, beta2, lr, weight_decay, eps)
     92             denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps)
     93         else:
---&gt; 94             denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
     95 
     96         step_size = lr / bias_correction1

RuntimeError: CUDA out of memory. Tried to allocate 2.32 GiB (GPU 0; 15.78 GiB total capacity; 11.91 GiB already allocated; 182.75 MiB free; 14.26 GiB reserved in total by PyTorch)

It makes sense to me that model = model.to(device) creates 3.7G of memory.
But why does running the model output = model(input, comb) create another 3G of memory?
And then loss.backward() creates another 3G of memory?
And then optimizer.step() creates another 6.3G of memory?
I would appreciate it if someone could explain how the PyTorch GPU memory allocation model is working in this example.


Inference
By default, an inference on your model will allocate memory to store the activations of each layer (activation as in intermediate layer inputs). This is needed for backpropagation where those tensors are used to compute the gradients. A simple but effective example is a function defined by f: x -&gt; x². Here, df/dx = 2x, i.e. in order to compute df/dx you are required to keep x in memory.
If you use the torch.no_grad() context manager, you will allow PyTorch to not save those values thus saving memory. This is particularly useful when evaluating or testing your model, i.e. when backpropagation is performed. Of course, you won't be able to use this during training!

Backward propagation
The backward pass call will allocate additional memory on the device to store each parameter's gradient value. Only leaf tensor nodes (model parameters and inputs) get their gradient stored in the grad attribute. This is why the memory usage is only increasing between the inference and backward calls.

Model parameter update
Since you are using a stateful optimizer (Adam), some additional memory is required to save some parameters. Read related PyTorch forum post on that. If you try with a stateless optimizer (for instance SGD) you should not have any memory overhead on the step call.



All three steps can have memory needs. In summary, the memory allocated on your device will effectively depend on three elements:

The size of your neural network: the bigger the model, the more layer activations and gradients will be saved in memory.

Whether you are under the torch.no_grad context: in this case, only the state of your model needs to be in memory (no activations or gradients necessary).

The type of optimizer used: whether it is stateful (saves some running estimates during parameter update, or stateless (doesn't require to).


whether you require to do back
