Capture CUDAGraph without using double the memory

minzhao.liu: <p>Hi!</p>
<p>I can successfully capture the CUDAGraph and replay. I took the API example from <a href="https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/" rel="noopener nofollow ugc">this blog</a> and modified it for my own model. Basically, I can forward and backward run my model normally with a certain batch size (25). But when I need to capture the graph, the reserved memory suddenly doubles, making me only able to do batch size 8.</p>
<pre><code class="lang-python">...
device = 'cpu'
if torch.cuda.is_available():
    device = 'cuda'

train_x, train_y = data_generator(T, seq_len, n_train)
model = TCN(1, n_classes, channel_sizes, kernel_size, dropout=dropout).to(device)
criterion = torch.nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

print('Before stream s, allocated: ', torch.cuda.memory_allocated('cuda'), '; reserved: ', torch.cuda.memory_reserved('cuda'))
# warmup
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for i in range(3):
        inputs = train_x[batch_size*i:batch_size*(i+1)].unsqueeze(1).contiguous().to(device)
        labels = train_y[batch_size*i:batch_size*(i+1)].to(device)
        optimizer.zero_grad(set_to_none=True)
        y_pred = model(inputs)
        loss = criterion(y_pred.view(-1, n_classes), labels.view(-1))
        loss.backward()
        optimizer.step()
torch.cuda.current_stream().wait_stream(s)

print('After stream s before graph g, allocated: ', torch.cuda.memory_allocated('cuda'), '; reserved: ', torch.cuda.memory_reserved('cuda'))
print('Start capturing CUDA graph.')

# capture
g = torch.cuda.CUDAGraph()
# Sets grads to None before capture, so backward() will create
# .grad attributes with allocations from the graph's private pool
optimizer.zero_grad(set_to_none=True)
static_inputs = train_x[:batch_size].unsqueeze(1).contiguous().to(device)
static_labels = train_y[:batch_size].to(device)
with torch.cuda.graph(g):
    static_y_pred = model(static_inputs)
    static_loss = criterion(static_y_pred.view(-1, n_classes), static_labels.view(-1))
    static_loss.backward()
    optimizer.step()
    
print('After graph g before replay, allocated: ', torch.cuda.memory_allocated('cuda'), '; reserved: ', torch.cuda.memory_reserved('cuda'))

for i in range(30):
    inputs = train_x[batch_size*i:batch_size*(i+1)].unsqueeze(1).contiguous().to(device)
    labels = train_y[batch_size*i:batch_size*(i+1)].to(device)
    static_inputs.copy_(inputs)
    static_labels.copy_(labels)
    g.replay()

print('After replay, allocated: ', torch.cuda.memory_allocated('cuda'), '; reserved: ', torch.cuda.memory_reserved('cuda'))
</code></pre>
<p>The output is:</p>
<pre><code class="lang-bash">Before stream s, allocated:  2978816 ; reserved:  4194304
After stream s before graph g, allocated:  13397522944 ; reserved:  15101591552
Start capturing CUDA graph.
After graph g before replay, allocated:  13400187392 ; reserved:  30045896704
After replay, allocated:  13400187392 ; reserved:  30045896704
</code></pre>
<p>If I try a batch size larger (9, 25, etc), it works up to capturing the graph. The error output is:</p>
<pre><code class="lang-bash">Traceback (most recent call last):
  File "/gpfs/fs1/home/minzhao.liu/TCN/cuda_graph.py", line 123, in &lt;module&gt;
    optimizer.step()
  File "/home/minzhao.liu/.conda/envs/qtensor-torch/lib/python3.9/site-packages/torch/cuda/graphs.py", line 149, in __exit__
    self.cuda_graph.capture_end()
  File "/home/minzhao.liu/.conda/envs/qtensor-torch/lib/python3.9/site-packages/torch/cuda/graphs.py", line 71, in capture_end
    super(CUDAGraph, self).capture_end()
RuntimeError: CUDA error: out of memory
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
</code></pre>
<p>Also, can I save the graph and run it on a different machine? If anyone can recommend some resources and more examples of CUDA graphs with PyTorch, I would really appreciate since I have only found a few minimal examples.</p>

mcarilli: <p>What jumps out at me is</p>
<pre><code class="lang-auto">Before stream s, allocated:  2978816...
After stream s before graph g, allocated:  13397522944
</code></pre>
<p>The model evidently isn’t that big (about 3 MB) but somehow after warmup the <em>allocated</em> memory shoots up to 13 GB.<br>
I’m not surprised the <em>reserved</em> memory gets pushed up to ~15 GB (it increases to accommodate the high-water mark of the built graph). But after backward(), the graph should be torn down, and the <em>allocated</em> memory should settle back to whatever the model params and optimizer states need which should be several times the model’s param set size, in other words 3 MB x N where N is a small integer.</p>
<p>Can you add</p>
<pre><code class="lang-auto">gc.collect()
torch.cuda.empty_cache()
</code></pre>
<p>just above</p>
<pre><code class="lang-auto">print('After stream s before graph g, allocated: ', torch.cuda.memory_allocated('cuda'), '; reserved: ', torch.cuda.memory_reserved('cuda'))
</code></pre>
<p>and report the values?</p>
<blockquote>
<p>can I save the graph and run it on a different machine</p>
</blockquote>
<p>No, there’s currently no way to serialize and deserialize graphs.</p>
<blockquote>
<p>more examples of CUDA graphs with PyTorch</p>
</blockquote>
<p>The master docs are the best resource.<br>
<a href="https://pytorch.org/docs/master/notes/cuda.html#cuda-graphs" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/master/notes/cuda.html#cuda-graphs</a><br>
Some internal people are adding cuda graphs to <a href="https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch" rel="noopener nofollow ugc">https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch</a>, but I don’t know how far that effort has progressed.</p>

minzhao.liu: <p>Thank you so much for the suggestions! After adding your two lines of code,  the allocated memory is still:</p>
<pre><code class="lang-bash">After stream s before graph g, allocated:  13397522944 ; reserved:  14986248192.
</code></pre>
<p>I think I need to figure out why the memory is not clearing. Is this the reason why the reserved memory doubles during capture graph?</p>
<p>The model is a quantum simulation model, which generates enormous tensors and uses torch.einsum everywhere with complicated python logic. This is why it gets so large. The simulator itself is a black box and I gotta dig deeper.</p>

minzhao.liu: <p>I solved the problem. It was because there are tensors in lists and dictionaries. After del list, it was able to release the memory after backward passes.</p>