GPU running out of memory in the middle of validation

mattaq31: <p>Hi all, I’m working on a super-resolution CNN model and for some reason or another I’m running into GPU memory issues.  I’m using the following training and validation loops in separate functions, and I am taking care to detach tensor data as appropriate, to prevent the computational graph from being replicated needlessly (as discussed in many other issues flagged in this forum):<br>
Training Function:</p>
<pre><code class="lang-auto">    def run_train(self, x, y, *args, **kwargs):
        if self.eval_mode:
            raise RuntimeError('Model initialized in eval mode, training not possible.')
        self.net.train()  # sets model to training mode (activates appropriate procedures for certain layers)
        x, y = x.to(device=self.device), y.to(device=self.device)
        out = self.run_model(x, **kwargs)  # run data through model

        loss = self.criterion(out, y)   # compute loss

        self.optimizer.zero_grad()  # set all weight grads from previous training iters to 0
        loss.backward()  # backpropagate to compute gradients for current iter loss
        if self.grad_clip is not None:  # gradient clipping
            nn.utils.clip_grad_norm_(self.net.parameters(), self.grad_clip)
        self.optimizer.step()  # update network parameters

        if self.learning_rate_scheduler is not None:
            self.learning_rate_scheduler.step()

        return loss.detach().cpu().numpy()
</code></pre>
<p>Validation Function:</p>
<pre><code class="lang-auto">    def run_eval(self, x, y=None, request_loss=False, tag=None, *args, **kwargs):
        self.net.eval()  # sets the system to validation mode
        with torch.no_grad():
            x = x.to(device=self.device)
            out = self.run_model(x, image_names=tag, **kwargs)  # forward the data in the model
            if request_loss:
                y = y.to(device=self.device)
                loss = self.criterion(out, y).detach().cpu().numpy()  # compute loss
            else:
                loss = None
        return out.detach().cpu(), loss
</code></pre>
<p>For some reason, the GPU runs out of memory only in the middle of either the training run or in the middle of a validation run (i.e. after a number of images have already been tested/fed into the model without issue).  This seems to be due to memory building up throughout validation/training.  I have attempted to probe the issue by clearing the pytorch cache and deleting variables before exiting the function, but nothing seems to help this problem.  This GPU buildup reaches a certain limit before stopping, and seems to be dependant on the training batch size (validation batch size is always 1).  I.e. if I, for example, set the train batch size to 16, the GPU memory builds up to ~8GB during validation, then stops there for the remainder of the training run.  If I set the batch size to 8, the GPU memory buildup stops at ~4GB, and sticks there (these are both hypothetical examples).  This means that I have to severely limit my batch size in order to allow training to occur, which is too much of a tradeoff for me.</p>
<p>Do you have any further insight into what could be going wrong?  Thanks for your help!</p>

ptrblck: <p>I’m not sure I understand the issue completely, since you were mentioning a memory increase in the first part, but also explain that the memory usage reaches a certain level for a smaller batch size.<br>
Is this memory peak stable for the smaller batch sizes or does it also increase for multiple training and validation epochs?</p>

Megh_Bhalerao: <p>Does the GPU memory build-up remain at those particular points i.e. 8 and 4 GB once they have reached upto those points, or do they increase further?</p>
<p>If you’re not already - you can try using <code>nvidia-smi -l 1</code> on the terminal to continuously monitor the GPU memory usage</p>

mattaq31: <p>I apologize, I might have not been so clear in my explanation.  The situation is this:</p>
<ul>
<li>When I run a super-res model for a particular training batch size (say 4) with a constant validation batch size (always 1), the model slowly accumulates GPU memory through each batch of training and validation inputs until a certain stable limit (say 4GB), where it remains until the end of my training.  This limit is normally achieved by the end of my first epoch (a single run-through all training input images and validation images).</li>
<li>When I increase the batch size (to say 8), this memory limit increases to, say, 8GB, which is also fine for an 11GB card.  This is of course normal to consume more memory with a higher batch size.</li>
<li>However, when I increase the training batch size to 16, the memory limit error does not trigger during training, but actually during validation, where the batch size should not have changed at all.  This means that technically I could achieve a batch size of 16 with my GPU, but is somehow being held back during the validation pass.  I tried to flush the GPU memory using <code>torch.cuda.empty_cache()</code> after my training pass (right before my validation pass) but this seemingly wasn’t enough to prevent the error from occurring.</li>
<li>I do observe the GPU memory increasing when I send my model to the GPU, but I also observe that, after training, the GPU memory does not drop back to the value it had before training (with just the model loaded in).  This is measured using <code>nvidia-smi</code>, after I run the empty cache command.</li>
</ul>

Robke: <p>I observe the same thing and am surprised that no-one has responded to this thread in the meanwhile.<br>
Apparently, Pytorch by default does start the validation run already before the training run is finished. This unnecessarily increases used GPU memory.<br>
Is this also what you are experiencing <a class="mention" href="/u/mattaq31">@mattaq31</a> ?<br>
Since all limitations on GPUs are related to memory for training large models, you would expect Pytorch to take specific care to not keep using memory without an actual use.<br>
At least I would expect that all memory will be cleaned completely when switching between training and validation runs, as this only happens once every few minuts/hours/days (depending on the model and dataset size).</p>
<p>Is there a way now to do this more carefully?<br>
Model sizes are increasing way more rapidly than GPU memory sizes are increasing, so it would be nice to think about this a bit more…</p>

Robke: <p>Found that when using Pytorch Lightning, there are two parameters that control the early start of the validation (which obviously causes the GPU memory to load both train/val data):</p>
<ul>
<li>val_check_interval</li>
<li>check_val_every_n_epoch</li>
</ul>
<p>see: <a href="https://lightning.ai/docs/pytorch/stable/advanced/speed.html" class="inline-onebox" rel="noopener nofollow ugc">Speed Up Model Training — PyTorch Lightning 2.2.1 documentation</a></p>

ptrblck: <aside class="quote no-group" data-username="Robke" data-post="5" data-topic="121316">
<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/robke/48/68912_2.png" class="avatar"> Robke:</div>
<blockquote>
<p>Apparently, Pytorch by default does start the validation run already before the training run is finished.</p>
</blockquote>
</aside>
<p>That’s not the case since by default all CUDA operations will be enqueued and executed in the default stream. PyTorch does not start the validation run before the Python script reaches it.</p>
<aside class="quote no-group" data-username="Robke" data-post="5" data-topic="121316">
<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/robke/48/68912_2.png" class="avatar"> Robke:</div>
<blockquote>
<p>At least I would expect that all memory will be cleaned completely when switching between training and validation runs, as this only happens once every few minuts/hours/days (depending on the model and dataset size).</p>
</blockquote>
</aside>
<p>PyTorch uses an internal CUDA caching mechanism and is able to reuse the GPU memory directly after it is released. I guess you see a different behavior on the host using Python’s garbage collector?</p>

Robke: <p>The training and validation runs are automated in Pytorch Lightning, but I guess I could probably try it using by implementing one of the callbacks for start of validation run.<br>
I did not notice any slow increase of GPU memory though, that was reported by <a class="mention" href="/u/mattaq31">@mattaq31</a>.<br>
My practical issue was resolved as mentioned by starting the validation run completely after each training epoch (Pytorch Lightning configuration thingy).<br>
What do you refer to explictly with “Python’s garbage collector”?<br>
I am probably not able to delete objects (they are all managed by Pytorch Lightning), but I was hoping there was a kind of memory-free function in Pytorch/Cuda that enables all gradient information of training epochs to be removed as to free GPU memory for the validation run. But as said, not running into practical issues anymore now.</p>

ptrblck: <aside class="quote no-group" data-username="Robke" data-post="8" data-topic="121316">
<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/robke/48/68912_2.png" class="avatar"> Robke:</div>
<blockquote>
<p>I was hoping there was a kind of memory-free function in Pytorch/Cuda that enables all gradient information of training epochs to be removed as to free GPU memory for the validation run.</p>
</blockquote>
</aside>
<p>Which is already the case since the internal caching allocator will move GPU memory to its cache once all references are freed of the corresponding tensor.</p>

mattaq31: <p>Hi both, I posted this query a few years ago now - I do not know if the issue is still valid today.  I ended up not being able to resolve the issue and I had to circumvent it by keeping the batch size low when I was using 11GB memory cards and using higher memory cards for higher batch size runs.</p>

sadmankiba: <p>Faced this in evaluation. Using <code>torch.no_grad()</code> kept the memory usage low.</p>