Memory bloat in gradient descent model

username: IsaacBoyd
cooked: <p>I’m running into memory bloat issues when backpropagating through this loop-heavy simulation. Even though the code looks simple, the backward pass consumes a <em>lot</em> of memory. I’d like to understand if there is a way that I can trim the tree at all or if because I am using ODEs where the current state depends on the previous state this is the most efficient implementation. For context, this model currently utilizes 14gb. When I calculate how much it should be using by hand, accounting for the auto grad nodes, it shouldn’t even take 1gb.</p>
<p>Any insight into whether the loop is causing unnecessary graph storage, or how to refactor it to reduce memory overhead, would be much appreciated!</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import gc
from torch.cuda.amp import autocast


class LIF_ODE(nn.Module):
    def __init__(self):
        super().__init__()
        self.dt = 0.0001
        self.tau = 0.02
        self.V_th = 1.0
        self.V_reset = 0.0
        self.V_rest = 0.0
        self.input_current = nn.Parameter(torch.tensor(1.5, dtype=torch.float32))

    def forward(self):
        T = 3500000 
        V = torch.zeros(T, dtype=torch.float32)
        spikes = torch.zeros(T, dtype=torch.float32)
        v_t = self.V_rest

        for t in range(1, T):
            dv = (-v_t + self.input_current) * (self.dt / self.tau)
            v_t = v_t + dv

            if v_t &gt;= self.V_th:
                spikes[t] = 1.0
                v_t = self.V_reset

            V[t] = v_t

        return spikes

def main():
    model = LIF_ODE()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01, betas=(0.0, 0.999))
    num_epochs = 1
    target_spikes = torch.tensor(50.0, dtype=torch.float32)

    for epoch in range(num_epochs):
        optimizer.zero_grad()

        with autocast():
            output = model()

        fr = output.sum() / 10 / 3
        loss = (fr - target_spikes) ** 2
        loss.backward()
        optimizer.step()

        gc.collect()
        print(f"Epoch {epoch}: Loss = {loss.item()}", flush=True)

if __name__ == "__main__":
    main()
</code></pre>

username: soulitzer
cooked: <p>Hey,<br>
This tool might be helpful for you for this types of investigation if you haven’t seen it already</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://pytorch.org/blog/understanding-gpu-memory-1/">
  <header class="source">
      <img src="https://pytorch.org/wp-content/uploads/2024/10/cropped-favicon-32x32.webp" class="site-icon" width="32" height="32">

      <a href="https://pytorch.org/blog/understanding-gpu-memory-1/" target="_blank" rel="noopener nofollow ugc">pytorch.org</a>
  </header>

  <article class="onebox-body">
    

<h3><a href="https://pytorch.org/blog/understanding-gpu-memory-1/" target="_blank" rel="noopener nofollow ugc">Understanding GPU Memory 1: Visualizing All Allocations over Time – PyTorch</a></h3>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


username: IsaacBoyd
cooked: <p>Thank you,</p>
<p>Perhaps I can try moving everything over to the GPU and then visualizing things. Currently I am storing my model in RAM and just running everything on the CPU.</p>

username: KFrank
cooked: <p>Hi Isaac!</p>
<aside class="quote no-group" data-username="IsaacBoyd" data-post="1" data-topic="220931" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/i/3da27b/48.png" class="avatar"> IsaacBoyd:</div>
<blockquote>
<p>this model currently utilizes 14gb.</p>
</blockquote>
</aside>
<p>I don’t have an opinion about whether this memory usage is to be expected or not (or whether it can be reduced), but I do have a couple of comments.</p>
<aside class="quote no-group">
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">    def forward(self):
        T = 3500000 
        ...
        for t in range(1, T):
            dv = (-v_t + self.input_current) * (self.dt / self.tau)
            v_t = v_t + dv
</code></pre>
</blockquote>
</aside>
<p>With each iteration of your loop – 3.5 million of them – you grow your computation graph<br>
by a small number of nodes.  With the memory usage you quote (presumably not all used<br>
by your growing computation graph), I get 4000 bytes per loop iteration.  This doesn’t seems<br>
outlandish to me (and I don’t have an opinion about whether it can be reduced  significantly).</p>
<p>The core use cases that pytorch is designed for don’t involve networks that are millions of<br>
layers deep, and typically involve pumping batches of larger objects (e.g, images) through<br>
those smaller numbers of layers.  So a computation-graph node might need to store a rather<br>
large tensor for backpropagation for a smaller number of nodes, rather than a single scalar<br>
for a much larger number of nodes.</p>
<p>As such, I could well believe that requiring a few hundred or a few thousand bytes of<br>
bookkeeping / overhead per node could be considered an acceptable design choice.</p>
<aside class="quote no-group">
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">            if v_t &gt;= self.V_th:
                spikes[t] = 1.0
                v_t = self.V_reset

            V[t] = v_t

        return spikes
</code></pre>
</blockquote>
</aside>
<p>As written, the return value of your <code>forward()</code> function, <code>spikes</code>, doesn’t carry<br>
<code>requires_grad = True</code> and is not part of any computation graph.  Furthermore, because<br>
<code>spikes</code> depends on <code>v_t &gt;= self.V_th</code>, it is not usefully differentiable, so (regardless of<br>
how much memory your unused computation graph consumes) <code>spikes</code> and your <code>forward()</code><br>
function don’t give you anything you can backpropagate through.</p>
<p>Right now, each entry in <code>spikes</code> is a (not-usefully-differentiable) step function in <code>v_t</code>.<br>
I’m not saying it would work (I don’t understand your use case), but you could potentially<br>
use <code>sigmoid()</code> as a (usefully) differentiable approximation to your step function.  (Of<br>
course, this wouldn’t address the issue with your 3.5-million-iteration computation graph.)</p>
<p>Best.</p>
<p>K. Frank</p>

username: IsaacBoyd
cooked: <p>Hi KFrank,</p>
<p>Thank you so much for your response. I have tried reimplementing what I have done and just taking the gradients manually. This seems to have gotten rid of the memory bloat. It was able to reduce the model from 14gb to &lt;100mb. The functions that I am using in my larger model are pretty differentiable, so it shouldn’t be too bad to convert out of PyTorch into a manual implementation.</p>
<p>Thanks!</p>