Freeing gradients memory after optimizer step

username: omarfoq
cooked: <p>I am training multiple models in a sequential way on the same GPU, and I need them to share the parameters after a given number of iterations. For GPU sonsumption optimization I need to free the gradients of each model at the end of each optimizer iteration. A simple solution is to set all gradients to None manually, i.e.,</p>
<pre><code class="lang-auto">for param in model.parameters():
    param.grad = None
</code></pre>
<p>Is this is a good practice? If not what are the alternatives</p>
<p>Thanks in advance</p>

username: albanD
cooked: <p>Hi,</p>
<p>Depending on the particular model and training loop, it may improve perf and not.<br>
Note that a simpler way to do this is via the regular zero grad: <code>model.zero_grad(set_to_none=True)</code>.</p>

username: omarfoq
cooked: <p>I think in the current scenario I don’t have enough memory if I keep all gradients. I think, however that setting gradients to None well decrease the training speed, because we will need to allocate new memory for gradients after each time we set them to None, Am I correct? In that case, how much do you think this will cost in term of time.</p>
<p>Also her the training loop is very standard, and looks like</p>
<pre><code class="lang-auto">        for x, y, indices in iterator:
            x = x.to(self.device).type(torch.float32)
            y = y.to(self.device)

            if self.is_binary_classification:
                y = y.type(torch.float32)

            self.optimizer.zero_grad()

            y_pred = self.model(x).squeeze()

            loss_vec = self.criterion(y_pred, y)
            if weights is not None:
                weights = weights.to(self.device)
                loss = (loss_vec.T @ weights[indices]) / loss_vec.size(0)
            else:
                loss = loss_vec.mean()
            loss.backward()
</code></pre>
<p><strong>Rq</strong>: criterion is initialized with <code>reduction=None</code>.</p>

username: albanD
cooked: <p>Setting gradients to None will not necessarily slow down training as we have some optimizations in place to avoid re-allocating the gradient buffer and just re-use the intermediary buffers from the backward.<br>
But that does not work all the time, depending on many factors.<br>
So you will have to try out for your model to know the exact impact.</p>

username: omarfoq
cooked: <p>Hello <a class="mention" href="/u/alband">@albanD</a>,</p>
<p>I have an other different, yet closely related, question. Normally when freeing gradients, I expect the used memory to be freed, and go back to the initial state before running backpropagation, however this is not what I observe, so I imagine memory is still allocate, what do you thin is the reason behind this?</p>

username: albanD
cooked: <p>How do you check the memory?<br>
PyTorch uses a custom allocator to speed up GPU allocations. So it is expected that the memory in nvidia-smi doesn’t go down. You can use <code>torch.cuda.memory_allocated()</code> to see the memory that is actually used by Tensors.</p>

username: omarfoq
cooked: <p>Indeed I use nvidia-smi to get memory consumption. My guess is that the memory consumption shown by nvidia-smi is the one that matters, in the sense I don’t want it to exceed memory capacity, what happens now is that when I free gradients by setting them to None, it doesn’t solve the issue I had, since the memory kept allocated, is there a way to force freeing memory?</p>

username: omarfoq
cooked: <p>Just an update on this, I guess the proper way is to free optimizer instead of the model, so it should be<br>
<code>optimizer.zero_grad(set_to_none=True)</code>, otherwisze the optimizer will keep the reference to the gradients, and thus they will keep place in memory.</p>

username: albanD
cooked: <p>Well that memory is available to allocate more Tensors (even though other processes can’t use it). So if you’re only using the GPU for pytorch, then its doing what you want.</p>
<p>zeroing from the optimizer or model does the same thing.</p>

username: omarfoq
cooked: <p>Hello <a class="mention" href="/u/alband">@albanD</a>,</p>
<p>Do you know how to free gradients in the case of LSTM, for some reason when I use LSTM and I free gradients setting them to 0, a part of memory is still used. What do you think is the reason behind this?</p>

username: albanD
cooked: <aside class="quote no-group" data-username="omarfoq" data-post="10" data-topic="116346">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/omarfoq/48/27473_2.png" class="avatar"> omarfoq:</div>
<blockquote>
<p>gradients setting them to 0,</p>
</blockquote>
</aside>
<p>You mean setting them to None?</p>
<p>Not sure why LSTM would be different.</p>

username: omarfoq
cooked: <p>yeah I mean stetting them to None, The following outputs</p>
<pre><code class="lang-auto">905216
2793472
1888256
</code></pre>
<p>Do you think it has something to due with <a href="https://github.com/pytorch/pytorch/issues/1536" rel="noopener nofollow ugc">LSTM memory leak</a></p>
<p>Here is the code,</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.optim as optim 
import string


class MyLSTM(nn.Module):
    def __init__(self, input_size, embed_size, hidden_size, output_size, n_layers):
        super(MyLSTM, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.embed_size = embed_size
        self.output_size = output_size
        self.n_layers = n_layers

        self.encoder = nn.Embedding(input_size, embed_size)
        self.rnn = nn.LSTM(input_size=embed_size,
                                       hidden_size=hidden_size,
                                      num_layers=n_layers,
                                     batch_first=True,
                                    dropout=0.5)
        self.decoder = nn.Linear(hidden_size, output_size)

    def forward(self, input_):
        encoded = self.encoder(input_)
        output = self.rnn(encoded)
        output = self.decoder(output)
        output = output.permute(0, 2, 1)  # change dimension to (B, C, T)
        return output


device = torch.device("cuda")

model =\
    NextCharacterLSTM(
        input_size=100,
        embed_size=8,
        hidden_size=256,
        output_size=100,
        n_layers=1
    ).to(device)

criterion = nn.CrossEntropyLoss().to(device)

def fit_epoch(model, weights=None):

    model.train()

    for _ in range(10):

        x = torch.zeros(32, 80).type(torch.long)
        y = torch.zeros(32, 100, 80)

        x = x.to(device)
        y = y.to(device)

        model.zero_grad()

        y_pred = model(x).squeeze()

        loss = criterion(y_pred, y)

        loss.backward()


print(torch.cuda.memory_allocated())
fit_epoch(model)
print(torch.cuda.memory_allocated())
model.zero_grad(set_to_none=True)
print(torch.cuda.memory_allocated())

</code></pre>

username: albanD
cooked: <p>I am not sure your code is correct, I had to modify a couple things to make it work:</p>
<pre><code class="lang-python">import torch
import torch.nn as nn
import torch.optim as optim 
import string


class MyLSTM(nn.Module):
    def __init__(self, input_size, embed_size, hidden_size, output_size, n_layers):
        super(MyLSTM, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.embed_size = embed_size
        self.output_size = output_size
        self.n_layers = n_layers

        self.encoder = nn.Embedding(input_size, embed_size)
        self.rnn = nn.LSTM(input_size=embed_size,
                                       hidden_size=hidden_size,
                                      num_layers=n_layers,
                                     batch_first=True,
                                    dropout=0.5)
        self.decoder = nn.Linear(hidden_size, output_size)

    def forward(self, input_):
        encoded = self.encoder(input_)
        output, _ = self.rnn(encoded)
        output = self.decoder(output)
        output = output.permute(0, 2, 1)  # change dimension to (B, C, T)
        return output


device = torch.device("cuda")

model =\
    MyLSTM(
        input_size=100,
        embed_size=8,
        hidden_size=256,
        output_size=100,
        n_layers=1
    ).to(device)

criterion = nn.CrossEntropyLoss().to(device)

def fit_epoch(model, weights=None):

    model.train()

    for _ in range(10):

        x = torch.zeros(32, 80, dtype=torch.long)
        y = torch.zeros(32, 80, dtype=torch.long)

        x = x.to(device)
        y = y.to(device)

        model.zero_grad()

        y_pred = model(x).squeeze()

        loss = criterion(y_pred, y)

        loss.backward()


print(torch.cuda.memory_allocated())
fit_epoch(model)
print(torch.cuda.memory_allocated())
model.zero_grad(set_to_none=True)
print(torch.cuda.memory_allocated())
</code></pre>
<p>But then running this on colab gives me<br>
108730368<br>
110368768<br>
108730368</p>
<p>Which is what we expect right?</p>

username: omarfoq
cooked: <aside class="quote no-group" data-username="omarfoq" data-post="12" data-topic="116346">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/omarfoq/48/27473_2.png" class="avatar"> omarfoq:</div>
<blockquote>
<p><code>output = self.rnn(encoded)</code></p>
</blockquote>
</aside>
<p>You are right, I had some typos in the previous code (I edited when writing the answer, I am sorry for that). But still the problem is that if you run this code in the first time it will give:</p>
<pre><code class="lang-auto">1196032
3858432
2179072
</code></pre>
<p>Then if you un it a second time it will give:</p>
<pre><code class="lang-auto">2179072
3817472
2179072
</code></pre>
<p>But I think this is not related to gradients at all, I observe the same thing when running with <code>torch.no_grad()</code>. Apparently there is some part of the memory that is allocated when doing the forward pass, I guess that this is part of memory allocated by the hidden state in the case of RNN, what do you think?</p>
<p>Here is a <a href="https://colab.research.google.com/drive/1MYGbYzQUcTPrMo7TK_mOSZQcU9Uzl5ex?usp=sharing" rel="noopener nofollow ugc">link</a> if you want to try it directly.</p>

username: albanD
cooked: <p>I am not super familiar with how LSTM work in details. But it is indeed possible that the default hidden state is lazily initialized the first time it is used.</p>