Increased GPU memory demand when loading from checkpoint

TJKlein: <p>Hi,</p>
<p>I have been observing some memory issues for quite a while. Whenever I resume training from a checkpoint more GPU memory is required than when running from scratch. I have observed this independent of architecture/application (NLP, Vision). It’s hard to quantify because it depends on the model, but I’d guess it is around 5-10%.  Consequentially, I have to reduce the batch size in order to continue training from a checkpoint, which is a bit painful.<br>
Any idea?</p>

ptrblck: <p>How are you measuring the memory usage?<br>
Note that PyTorch uses a custom memory allocation, which caches the device memory.<br>
<code>nvidia-smi</code> will this give you the memory usage of the CUDA context, the allocated as well as the cached memory.<br>
Since you have to reduce the batch size, it seems that the memory might indeed to growing.<br>
Since you are observing this behavior in a lot of use cases, could you post a simple code snippet to reproduce this behavior, please?</p>

TJKlein: <pre><code class="lang-auto">model = VAE(...)
pretrained_dict = torch.load(os.path.join(FLAGS.savedir,FLAGS.loadmodel))
model.load_state_dict(pretrained_dict)
</code></pre>
<p>I observed this also using the Huggingface BERT implementations. Yes, nvidia-smi.</p>
<p>Here is the memory footprint when training from scratch:</p>
<pre><code class="lang-auto">+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0     32630      C   python                                     11985MiB |
+-----------------------------------------------------------------------------+
</code></pre>
<p>Memory footprint when continuing from checkpoint:</p>
<pre><code class="lang-auto">+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0     32707      C   python                                     15045MiB |
+-----------------------------------------------------------------------------+
</code></pre>
<pre><code class="lang-auto"> File "/home/ubuntu/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/torch/tensor.py", line 195, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/ubuntu/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/torch/autograd/__init__.py", line 99, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: CUDA out of memory. Tried to allocate 576.00 MiB (GPU 0; 15.75 GiB total capacity; 12.94 GiB already allocated; 392.88 MiB free; 14.38 GiB reserved in total by PyTorch)
</code></pre>

ptrblck: <p>Could you try to delete the <code>pretrained_dict</code> after loading it?<br>
The difference in the memory usage seems to be quite large.<br>
Also, could you post the definition of <code>VAE</code>?</p>

TJKlein: <p>Yes, in this case, the difference is quite large. Deleting the dictionary does not change the memory footprint unfortunately.</p>
<p>The  VAE code is a bit more complex. It consists of a Conv2D encoder stack, with a bunch of FCs, and the Deconv decoder stack. Maybe at some point, it would be easier to look at the NLP problem, as it is more accessible.</p>
<pre><code class="lang-auto">class VAE2(nn.Module):
    '''
    Variational Autoencoder
    '''
    def __init__(self, img_channels, img_dim, latent_dim, filters,
                 kernel_sizes, strides, activation=nn.LeakyReLU,
                 out_activation=nn.Tanh, batch_norm=True, no_samples=10, sp_activation=None, public_stream=None, private_stream=None):
        '''
        in_dim (int): number of pixels on each row / column of the images
                      (assumes the images are square).
        in_channels (int): number of channels of the images (e.g.: 1 for
                           grayscale, 3 for color images).
        latent_dim (int): dimension of the latent space.
        filters (list of length n_conv): number of filters for each conv.
                                         layer.
        kernel_sizes (list of length n_conv): kernel size for each conv.
        strides (list of length n_conv): strides for each conv. layer.
        activation (nn.Module): activation used in all layers (default:
                                LeakyReLU).
        out_activation (subclass of nn.Module): activation used in the output
                                                layer (default: Tanh).
        batch_norm (boolean): if True, batch normalization is applied in every
                              layer before the activation (default: True).
        '''
        super(VAE2, self).__init__()

        self.img_dim = img_dim
        self.img_channels = img_channels
        self.latent_dim = latent_dim
        self.filters = filters
        self.kernel_sizes = kernel_sizes
        self.strides = strides
        self.activation = activation
        self.out_activation = out_activation
        self.batch_norm = batch_norm
        self.no_samples = no_samples

        n_conv = len(self.filters)

        # compute the paddings and the flattened dimension at the output of the
        # last conv.
        paddings = []
        dims = [self.img_dim]
        for i in range(n_conv):
            if (dims[i] - self.kernel_sizes[i]) % strides[i] == 0:
                paddings.append((self.kernel_sizes[i] - 1)//2)
            else:
                paddings.append((self.kernel_sizes[i] - strides[i] + 1)//2)

            dims.append((dims[i] + 2*paddings[i] - self.kernel_sizes[i])
                        // self.strides[i] + 1)
        flat_dim = self.filters[-1] * (dims[-1]**2)

        self.encoder = Encoder(self.img_channels, self.img_dim,
                               self.latent_dim, self.filters,
                               self.kernel_sizes, self.strides,
                               paddings, flat_dim,
                               activation=self.activation,
                               batch_norm=self.batch_norm)

        # the decoder architecture will be the transposed of the encoder's
        filters_dec = (list(reversed(self.filters[0:n_conv-1]))
                       + [img_channels])
        kernel_sizes_dec = list(reversed(self.kernel_sizes))
        strides_dec = list(reversed(self.strides))
        paddings = list(reversed(paddings))
        dims = list(reversed(dims))

        # compute the output paddings
        out_paddings = []
        for i in range(n_conv):
            out_dim = ((dims[i] - 1)*strides_dec[i] - 2*paddings[i] +
                       kernel_sizes_dec[i])
            out_paddings.append(dims[i+1] - out_dim)

        self.decoder = Decoder(self.latent_dim, self.filters[-1], dims[0],
                               filters_dec, kernel_sizes_dec, strides_dec,
                               paddings=paddings, out_paddings=out_paddings,
                               activation=self.activation,
                               out_activation=self.out_activation,
                               batch_norm=self.batch_norm)
       
            
        self.bottleneck_fc = SpLinear(in_features=flat_dim, out_features=int(np.ceil(flat_dim/4)), bias=True, activation=sp_activation)
       .............
        self.fc_model_public = nn.Sequential(stream2dict(public_stream))
        self.fc_model_private = nn.Sequential(stream2dict(private_stream))
</code></pre>

ptrblck: <p>Could you post the arguments as well as the input shapes?<br>
If it’s easier for you to post to another repository, it would be fine as well.<br>
I would just want to make sure that you are seeing the memory increase with a specific setup, which I’m able to run locally.</p>

TJKlein: <p>The input ist are RGB images:<br>
<code>3 x 128 128</code></p>
<p>The 2d filters are:</p>
<pre><code class="lang-auto">filters: '64, 128, 256, 384'
kernel-sizes: '3, 3, 3, 3, 3'
strides: '2, 2, 2, 2, 2'
</code></pre>
<p>The FC that maps from flattened convs to a compressed rep:</p>
<pre><code class="lang-auto">int: 12288, out: 6144, 
</code></pre>
<p>The FC stream (I am using a variant of linear layers, so I need more memory), each has dimensions:</p>
<pre><code class="lang-auto">'2x in: 6044, out: 2048; in: 2048, out: 512'
</code></pre>
<p>And another FC:</p>
<pre><code class="lang-auto">In: 1024, Out: 8400
</code></pre>

zaptrem: <p><a class="mention" href="/u/tjklein">@TJKlein</a> Did you ever figure out the cause of this?</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> I’m running into the same problem. Loading from checkpoint takes a couple more gigabytes of VRAM than starting fresh and so I need to reduce batch size to avoid OOMing. I originally blamed PyTorch Lightning but the issue persists after reimplementing my entire training loop in vanilla PyTorch. <a href="https://gist.github.com/zaptrem/7be75641411937e61b35dc05a928682d" rel="noopener nofollow ugc">Here’s my entire unedited training loop if specifics are needed</a>.</p>

Jeremy_Cowles: <p>I’m hitting the same problem - training on 4 nodes, with 4 gpus per node, the training runs fine, but when resuming from checkpoint I get a Cuda out of memory error</p>