Profiling: occasional slow cudaMalloc calls

ArchieGertsman: <p>I am profiling my policy gradient RL code using Tensorboard. I call an agent at each step of an episode, and the observation’s batch dimension varies. Throughout the episode, the batch size gradually increases until it hits a peak, then decreases from then on. I notice that during some calls to the agent, there is an expensive call to <code>cudaMalloc</code>:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/8/785c0d3281c9aa85c0c240b615aa8337445803c0.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/785c0d3281c9aa85c0c240b615aa8337445803c0" title="Screen Shot 2022-12-23 at 6.41.15 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/8/785c0d3281c9aa85c0c240b615aa8337445803c0_2_503x500.jpeg" alt="Screen Shot 2022-12-23 at 6.41.15 PM" data-base62-sha1="haKvZtTrsMb4zxgUCZaGDy80myI" width="503" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/8/785c0d3281c9aa85c0c240b615aa8337445803c0_2_503x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/7/8/785c0d3281c9aa85c0c240b615aa8337445803c0_2_754x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/8/785c0d3281c9aa85c0c240b615aa8337445803c0_2_1006x1000.jpeg 2x" data-dominant-color="DFDEE3"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2022-12-23 at 6.41.15 PM</span><span class="informations">1130×1122 229 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>In step 12 of this this episode, the malloc occurs in  <code>atten::relu -&gt; atten::clamp_min -&gt; atten::resize_</code>, which is often the case. I’ve also seen a slow <code>cudaMalloc</code> call come from <code>atten::mul</code>. The malloc calls alone account for 30-40% of my program’s runtime.</p>
<p>I thought that maybe this happens when the batch size increases significantly from the previous observation, but this is not the case. E.g. I saw two observations in a row with batch size 400, but only the second observation resulted in a <code>cudaMalloc</code>. Any other ideas on what could be going on? These occasional calls are bottlenecking my performance.</p>
<p>Edit: I’ve changed <code>ReLU()</code> to <code>ReLU(inplace=True)</code>, so this no longer happens with relu. Now, I mostly see the malloc stem from <code>repeat_interleave</code>, <code>cat</code>, <code>mm</code>, <code>mul</code>, etc. This makes more sense, but I still wonder if there’s anything I can do to make the malloc faster. Maybe it has to do with memory fragmentation? I tried allocate a bunch of memory ahead of time instead of having these slow on-the-fly calls, but that didn’t help–the <code>resize</code> calls still take long.</p>

ptrblck: <p><code>cudaMalloc</code> calls are synchronizing and expensive which is why PyTorch uses a caching allocator to reuse already allocated memory. In case you are clearing the cache via <code>torch.cuda.empty_cache()</code>, remove these calls. PyTorch does not allocate all available memory from the GPU and will allocate new memory if needed (e.g. if the memory in the cache cannot be reused), which is usually visible during a few initial warmup iterations.<br>
Changing the input size might trigger this behavior and you might be able to avoid it by starting with the largest size of your input.</p>

ArchieGertsman: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, I appreciate your reply. Do you think that starting with the largest size of my input will result in less calls to <code>cudaMalloc</code> than initially allocating a large chunk of memory (e.g. 1 GB) like in <a href="https://discuss.pytorch.org/t/reserving-gpu-memory/25297/2">this thread</a>?</p>
<p>How exactly does the caching allocator reuse memory? If I have a call like <code>t_repeat = torch.repeat_interleave(t, n, dim=0)</code> in my <code>forward</code> method, then does the caching allocator try to reuse the same memory for the output of <code>repeat_interleave</code> each time? I.e. if the dimensions of <code>t</code> stay constant and <code>n</code> is largest on the very first call, then does each subsequent call necessarily reuse that same chunk of memory?</p>

ptrblck: <aside class="quote no-group" data-username="ArchieGertsman" data-post="3" data-topic="168934">
<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/archiegertsman/48/51153_2.png" class="avatar"> ArchieGertsman:</div>
<blockquote>
<p>Do you think that starting with the largest size of my input will result in less calls to <code>cudaMalloc</code> than initially allocating a large chunk of memory (e.g. 1 GB) like in <a href="https://discuss.pytorch.org/t/reserving-gpu-memory/25297/2">this thread </a>?</p>
</blockquote>
</aside>
<p>Yes, my suggestion was to use the largest size first to allocate enough memory which could be reused.</p>
<aside class="quote no-group" data-username="ArchieGertsman" data-post="3" data-topic="168934">
<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/archiegertsman/48/51153_2.png" class="avatar"> ArchieGertsman:</div>
<blockquote>
<p>I.e. if the dimensions of <code>t</code> stay constant and <code>n</code> is largest on the very first call, then does each subsequent call necessarily reuse that same chunk of memory?</p>
</blockquote>
</aside>
<p>No, not necessarily since other tensors could reuse the cached memory before.<br>
The caching allocator will move the memory of deleted tensors into its cache without freeing it via (expensive) <code>cudaFree</code> calls. If a new tensors needs to be allocated (e.g. as a result tensor from an operation), the caching allocator will try to reuse memory from the cache before trying to allocate new memory via <code>cudaMalloc</code> calls.</p>

ArchieGertsman: <p>Okay, I understand your second answer. But can you please help me understand why reserving a very large chunk of memory ahead of time won’t work? E.g. if I do</p>
<pre><code class="lang-auto">t = torch.empty(num_bytes, dtype=torch.int8)
del t
</code></pre>
<p>then won’t <code>num_bytes</code> bytes of memory now be in the cache? So, as long as I am using less than <code>num_bytes</code> of memory at any given time, I can just reuse this memory without needing any more <code>cudaMalloc</code> calls. I know my logic is flawed somewhere, as I tried this and ended up with a bunch of “reserved” memory, but more memory kept being allocated nevertheless. Why isn’t the memory reusable?</p>

ptrblck: <p>Generally, your workflow should work as seen here:</p>
<pre><code class="lang-python">import torch
from torch.profiler import profile, record_function, ProfilerActivity

# initial check
print(torch.cuda.memory_allocated() / 1024**2)
# 0.0
print(torch.cuda.memory_reserved() / 1024**2)
# 0.0

# allocate new memory
with profile(activities=[
        ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
    with record_function("allocation"):
        x = torch.empty(1024**3, dtype=torch.uint8, device='cuda')
# note the cudaMalloc call!
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
# -------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
#                      Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
# -------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
#                allocation         0.44%     750.000us       100.00%     171.706ms     171.706ms             1  
#               aten::empty         1.24%       2.127ms        99.56%     170.956ms     170.956ms             1  
#     cudaStreamIsCapturing        97.67%     167.717ms        97.67%     167.717ms     167.717ms             1  
#                cudaMalloc         0.65%       1.112ms         0.65%       1.112ms       1.112ms             1  
#     cudaDeviceSynchronize         0.00%       5.000us         0.00%       5.000us       5.000us             1  
# -------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
# Self CPU time total: 171.711ms

# Verify the memory is allocated and reserved
print(torch.cuda.memory_allocated() / 1024**2)
# 1024.0
print(torch.cuda.memory_reserved() / 1024**2)
# 1024.0

# delete x - the memory should still be reserved
del x
print(torch.cuda.memory_allocated() / 1024**2)
# 0.0
print(torch.cuda.memory_reserved() / 1024**2)
# 1024.0

# allocate the same tensor again
with profile(activities=[
        ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
    with record_function("allocation"):
        x = torch.empty(1024**3, dtype=torch.uint8, device='cuda')
        
# no cudaMalloc calls!
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
# -------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
#                      Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
# -------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
#                allocation        62.10%     480.000us        98.71%     763.000us     763.000us             1  
#               aten::empty        36.61%     283.000us        36.61%     283.000us     283.000us             1  
#     cudaDeviceSynchronize         1.29%      10.000us         1.29%      10.000us      10.000us             1  
# -------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
# Self CPU time total: 773.000us
</code></pre>
<p>As you can see no <code>cudaMalloc</code> call is visible and I can directly reuse the memory.<br>
However, note that the cache is using a <code>large</code> and <code>small</code> segment size. Once your newly created tensor is too small the caching allocator will create a new pool for small segments and call into <code>cudaMalloc</code> again. Take a look at the <a href="https://pytorch.org/docs/stable/notes/cuda.html#memory-management">Memory management</a> section of the docs and also check <code>torch.cuda.memory_stats()</code> and <code>torch.cuda.memory_snapshot()</code> as debugging tools for your use case.</p>

ArchieGertsman: <p>Thank you for your explanation! I read some more about the small and large allocators, and I have a better understanding now.</p>
<p>I tried doing a “warmup” episode, in which the environment makes a few observations of maximum size, but this didn’t help. I did some more profiling using <code>torch.cuda.memory_stats()</code>, and I have a suspicion that the computational graph is just massive and keeps on growing, which is why a few warmup iterations don’t really help:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/7/a787dc6474d8a3ea5c6dfcc979f7f66d9e52bf4e.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/a787dc6474d8a3ea5c6dfcc979f7f66d9e52bf4e" title="Screen Shot 2022-12-25 at 12.52.15 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/7/a787dc6474d8a3ea5c6dfcc979f7f66d9e52bf4e_2_500x500.jpeg" alt="Screen Shot 2022-12-25 at 12.52.15 AM" data-base62-sha1="nU2OJea8gJLavatnwab5jQWA6No" width="500" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/7/a787dc6474d8a3ea5c6dfcc979f7f66d9e52bf4e_2_500x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/a/7/a787dc6474d8a3ea5c6dfcc979f7f66d9e52bf4e_2_750x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/7/a787dc6474d8a3ea5c6dfcc979f7f66d9e52bf4e_2_1000x1000.jpeg 2x" data-dominant-color="F8F8F8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2022-12-25 at 12.52.15 AM</span><span class="informations">1152×1150 121 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Note that stats from the same episode are displayed in both columns; I’ve just split the stats by allocator type (small and large). I’m not sure if these plots reveal much, but it seems like memory that can be freed is getting freed. I believe the only memory that can’t be freed is the computation graph. Also, <code>cudaMalloc</code> reserves an increasing number segments as the episode progresses.</p>
<p>I’m pretty stuck right now and don’t know how to move forward. Can I profile the computational graph somehow? I’ll mention that I have a lot of computing power: A100 GPU with 80 GB RAM.</p>