Is it possible to keep memory for only one row in a batch

username: Kexin_Wang
cooked: <p>Imagine I am doing a forward function for an input x (batch size = 32) and get the result y. Is it possible to keep only one row of x and release all other memory? The example code is as follows:</p>
<pre><code class="lang-auto">import torch
from torch.nn import functional as F

base = 4000256
params = torch.nn.Parameter(torch.randn(1000, 1000)).to("cuda")


def f(params: torch.nn.Parameter, x, keep, n=10):
    # x: (bsz, 1000)
    for i in range(n):
        x = F.linear(x, params)
        print(i, torch.cuda.memory_allocated(0) - base)
    result = x[keep]
    return result


print(torch.cuda.memory_allocated(0) - base)
x = torch.randn(32, 1000).to("cuda")  # 1408000
# x = torch.randn(1, 1000).to("cuda")  # 45056
print(torch.cuda.memory_allocated(0) - base)
y = f(params, x, 0)
torch.cuda.empty_cache()
print("e", torch.cuda.memory_allocated(0) - base)
</code></pre>
<p>It prints:</p>
<pre><code class="lang-auto">0
128000
0 256000
1 384000
2 512000
3 640000
4 768000
5 896000
6 1024000
7 1152000
8 1280000
9 1408000
e 1408000
</code></pre>
<p>One can find even if I need only one row, it keeps the memory (1408000) for the whole batch. What I expect is: it retains only the memory for one row (as if there is only one row processed, i.e. 45056 in this case).</p>
<p>Is there any possible way to achieve what I want? Thanks in advance</p>

username: soulitzer
cooked: <p>Maybe you could do the following?</p>
<pre><code class="lang-auto">y = x[i,:].clone()
del x
</code></pre>

username: Kexin_Wang
cooked: <p>Thanks! I tried your idea</p>
<pre><code class="lang-auto">import torch
from torch.nn import functional as F

base = 4000256
params = torch.nn.Parameter(torch.randn(1000, 1000)).to("cuda")


def f(params: torch.nn.Parameter, x, keep, n=10):
    # x: (bsz, 1000)
    for i in range(n):
        x = F.linear(x, params)
        print(i, torch.cuda.memory_allocated(0) - base)
    result = x[keep, :].clone()
    del x
    return result


print(torch.cuda.memory_allocated(0) - base)
x = torch.randn(32, 1000).to("cuda")  # 1408000
# x = torch.randn(1, 1000).to("cuda")  # 45056
print(torch.cuda.memory_allocated(0) - base)
y = f(params, x, 0)
torch.cuda.empty_cache()
print("e", torch.cuda.memory_allocated(0) - base)
</code></pre>
<p>and it prints:</p>
<pre><code class="lang-auto">0
128000
0 256000
1 384000
2 512000
3 640000
4 768000
5 896000
6 1024000
7 1152000
8 1280000
9 1408000
e 1284096
</code></pre>
<p>1284096 is a bit lower than 1408000, but ideally, it should reduce itself to 45056 (i.e. when batch size = 1 since the very beginning)</p>

username: soulitzer
cooked: <p>The memory is being occupied by the tensors that are saved for backward. You can either run forward with gradients disabled <code>with torch.no_grad(): ... </code> or running <code>.backward()</code> would have the same effect of clearing the saved tensors. After calling <code>.backward()</code>, I see the following on colab:</p>
<pre><code class="lang-auto">0
128000
0 256000
1 384000
2 512000
3 640000
4 768000
5 896000
6 1024000
7 1152000
8 1280000
9 1408000
e 132096
</code></pre>

username: Kexin_Wang
cooked: <p>Thanks!</p>
<p>I need to keep the gradient and backward indeed can clear the saved tensor. However, I want the redundant memory to be cleared before calling backward, since I can only compute the loss and do backward after encoding the whole sequence and the redundant vectors take a lot of memory.</p>
<p>Do you have any idea how? Or can this be achieved?</p>