Why torch.cuda.allocated_memory reports that GPU Memory usage was decreasing during forwarding?

raining_day513: <p>Hi, I’m trying to record the CUDA GPU memory usage using the API <code>torch.cuda.memory_allocated</code>. The target I want to achieve is that I want to draw a diagram of GPU memory usage(in MB) during forwarding.</p>
<p>This is the <code>nn.Module</code> class I’m using that makes use of the class method <code>register_forward_hook</code> of <code>nn.Module</code> to get the memory usage before the forward method being called:</p>
<pre><code class="lang-auto">class Segment(torch.nn.Module):
    def __init__(self,
                 specs_truncated,
                 in_channels,
                 use_checkpoint=False,
                 use_batch_norm=True,
                 ):
        super().__init__()
        self.subspecs = specs_truncated
        self.in_channels = in_channels
        self.use_checkpoint = use_checkpoint
        self.use_batch_norm = use_batch_norm
        self.seq_of_nodes = self._make_sequential_by_specs_truncated()
        self.handles_forward_hooks = self._register_forward_hooks(
            [torch.nn.Conv2d, torch.nn.BatchNorm2d, torch.nn.ReLU, torch.nn.MaxPool2d],
            lambda d: lambda m, i:
                print("Memory allocated before execute {}-{}: {} MB".format(
                    d['name'], d['class'],
                    torch.cuda.memory_allocated()/1e6)),
            lambda d: lambda m, i, o:
                print("Memory allocated after execute {}-{}: {} MB".format(
                    d['name'], d['class'],
                    torch.cuda.memory_allocated()/1e6)),
        )

    def forward(self, x):
        if self.use_checkpoint:
            return checkpoint(self.seq_of_nodes, x, use_reentrant=False)
        return self.seq_of_nodes(x)

    def _make_sequential_by_specs_truncated(self):
        layers = []
        in_channels = self.in_channels
        for node_spec in self.subspecs:
            if node_spec == 'M':
                layers += [
                        torch.nn.MaxPool2d(kernel_size=2),
                        ]
            else:
                layers += [
                        torch.nn.Conv2d(in_channels, node_spec, kernel_size=3, padding='same'),
                        ]
                layers += [
                        torch.nn.BatchNorm2d(node_spec),
                        torch.nn.ReLU(),
                        ][not self.use_batch_norm:]
                in_channels = node_spec

        return torch.nn.Sequential(*layers)

    def _register_forward_hooks(self, list_targets, pre_hook, post_hook):
        handles = {}
        for name, module in self.seq_of_nodes.named_modules():
            if isinstance(module, tuple(list_targets)):
                # handles[name + '.pre'] = module.register_forward_pre_hook(
                #    pre_hook({ 'name': name, 'class': module.__class__.__name__ }))
                handles[name + '.post'] = module.register_forward_hook(
                    post_hook({ 'name': name, 'class': module.__class__.__name__ }))
        return handles

    def deregister_forward_hooks(self):
        for handle in self.handles_forward_hooks.values():
            handle.remove()
</code></pre>
<p>But I got a very unintuitive result as I have described by the title:</p>
<pre><code class="lang-auto">Memory allocated after execute 0-Conv2d: 609.58208 MB
Memory allocated after execute 1-BatchNorm2d: 643.136512 MB
Memory allocated after execute 2-ReLU: 643.136512 MB
Memory allocated after execute 3-Conv2d: 643.136512 MB
Memory allocated after execute 4-BatchNorm2d: 643.136512 MB
Memory allocated after execute 5-ReLU: 643.136512 MB
Memory allocated after execute 6-MaxPool2d: 617.970688 MB
Memory allocated after execute 0-Conv2d: 601.193472 MB
Memory allocated after execute 1-BatchNorm2d: 617.970688 MB
Memory allocated after execute 2-ReLU: 617.970688 MB
Memory allocated after execute 3-Conv2d: 617.970688 MB
Memory allocated after execute 4-BatchNorm2d: 617.970688 MB
Memory allocated after execute 5-ReLU: 617.970688 MB
Memory allocated after execute 6-MaxPool2d: 605.912064 MB
Memory allocated after execute 0-Conv2d: 597.523456 MB
Memory allocated after execute 1-BatchNorm2d: 606.96064 MB
Memory allocated after execute 2-ReLU: 606.96064 MB
Memory allocated after execute 3-Conv2d: 606.96064 MB
Memory allocated after execute 4-BatchNorm2d: 606.96064 MB
Memory allocated after execute 5-ReLU: 606.96064 MB
Memory allocated after execute 6-Conv2d: 606.96064 MB
Memory allocated after execute 7-BatchNorm2d: 606.96064 MB
Memory allocated after execute 8-ReLU: 606.96064 MB
Memory allocated after execute 9-Conv2d: 606.96064 MB
Memory allocated after execute 10-BatchNorm2d: 606.96064 MB
Memory allocated after execute 11-ReLU: 606.96064 MB
Memory allocated after execute 12-MaxPool2d: 600.669184 MB
Memory allocated after execute 13-Conv2d: 595.426304 MB
Memory allocated after execute 14-BatchNorm2d: 597.523456 MB
Memory allocated after execute 15-ReLU: 597.523456 MB
Memory allocated after execute 16-Conv2d: 597.523456 MB
Memory allocated after execute 17-BatchNorm2d: 597.523456 MB
Memory allocated after execute 18-ReLU: 597.523456 MB
Memory allocated after execute 19-Conv2d: 597.523456 MB
Memory allocated after execute 20-BatchNorm2d: 597.523456 MB
Memory allocated after execute 21-ReLU: 597.523456 MB
Memory allocated after execute 22-Conv2d: 597.523456 MB
Memory allocated after execute 23-BatchNorm2d: 597.523456 MB
Memory allocated after execute 24-ReLU: 597.523456 MB
Memory allocated after execute 25-MaxPool2d: 594.377728 MB
Memory allocated after execute 26-Conv2d: 591.232 MB
Memory allocated after execute 27-BatchNorm2d: 591.232 MB
Memory allocated after execute 28-ReLU: 591.232 MB
Memory allocated after execute 29-Conv2d: 591.232 MB
Memory allocated after execute 30-BatchNorm2d: 591.232 MB
Memory allocated after execute 31-ReLU: 591.232 MB
Memory allocated after execute 32-Conv2d: 591.232 MB
Memory allocated after execute 33-BatchNorm2d: 591.232 MB
Memory allocated after execute 34-ReLU: 591.232 MB
Memory allocated after execute 35-Conv2d: 591.232 MB
Memory allocated after execute 36-BatchNorm2d: 591.232 MB
Memory allocated after execute 37-ReLU: 591.232 MB
</code></pre>
<p>As you can see, the memory usage reported by <code>allocated_memory</code> shows that GPU usage was decreasing!</p>
<p>Could anyone help me identify whether or not I might misuse any API in my code? Thanks</p>

ptrblck: <p>Autograd will store intermediate tensors needed for the gradient computation and will free other unused temporary tensors. Depending on the model architecture and the used modules larger tensors might not be needed and freeing them would thus reduce the memory usage.<br>
To check your model you could try to verify the memory usage of each submodule/layer and compare it to your output.</p>

raining_day513: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Sorry, but I still don’t know how to prevent this to happen. Did you mean that this is a normal behavior specific to PyTorch’s autograd implementation?</p>

eqy: <p>I would also check what happens if you add <code>torch.cuda.synchronize()</code> calls before each print just to rule out any possibility that temporary workspaces for convolutions are being counted in the memory usage. Additionally, if you are using checkpointing, then indeed some temporary activations would be garbage collected with their memory being reused which would free memory as <a class="mention" href="/u/ptrblck">@ptrblck</a> suggested.</p>

raining_day513: <blockquote>
<p>I would also check what happens if you add <code>torch.cuda.synchronize()</code> calls before each print […]</p>
</blockquote>
<p><a class="mention" href="/u/eqy">@eqy</a> First of all, many thanks for reading my post <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20">!</p>
<p>I did, and I got the same result. I even tried adding <code>torch.cuda.empty_cache</code> before&amp;after the <code>synchronize()</code> call, which is not included in the image below since no effect. This is the callback for the <code>register_forward_hook</code>.</p>
<p><a href="/uploads/short-url/ipBv5wvzcmBd9tNroGzTIIwepRi.png"><img src="https://discuss.pytorch.org/uploads/default/original/3X/8/1/810c17a484eeb5db9f9d38ea2c89fee4b10ca0dc.png" data-base62-sha1="ipBv5wvzcmBd9tNroGzTIIwepRi" width="300" height="132"></a></p>
<blockquote>
<p>Additionally, if you are using checkpointing, then […]</p>
</blockquote>
<p>I’m happy that you mentioned checkpointing. Could you shed some light on how could I check this: Say I’ve divided my model into many <code>Segment</code>s (are <code>nn.Module</code>), I want to know whether <code>torch.utils.checkpoint</code> will use some allocator(so far I only know there are two kinds: sync&amp;async one) to reserve GPU memory for all nodes (e.g. build-in classes like <code>nn.Conv2d</code>, etc) of a <code>Segment</code> at once before it will <code>forward</code> them one-by-one. Otherwise, PyTorch might allocate the GPU memory before running each node, which seems unlikely as GPU memory allocation is expensive.</p>
<p>(Sorry for asking an additional question after the reply. My intention is to make my reply compact, and any insight into the question itself will help me a lot)</p>

eqy: <p>You can read the implementation of <code>torch.utils.checkpoint</code> directly here: <a href="https://pytorch.org/docs/stable/_modules/torch/utils/checkpoint.html#checkpoint" class="inline-onebox" rel="noopener nofollow ugc">torch.utils.checkpoint — PyTorch 2.0 documentation</a><br>
(The sequential version will be in the same source file in case you are using that instead)<br>
You will see that there is mostly nothing special about the checkpointing use case in the forward pass (especially with regard to calling an allocator explicitly or switching allocators), as the main difference is that it is running each checkpointed segment under a <code>with torch.no_grad():</code> context. This context allows the intermediate forward activation(s) to be discarded <em>within</em> each segment rather than being saved for backward, which explains why the memory usage can increase and then decrease within a <em>segment</em>. If you check the memory usage at the beginning and end of each <em>segment</em>, rather than each <em>layer</em>, I believe you should see nondecreasing memory usage.</p>
<p>Simply put, if you have a model with layers A-&gt;B-&gt;C-&gt;D-&gt;E-&gt;F, and segments A-&gt;B-&gt;C, D-&gt;E-&gt;F, all checkpointing will do is discard the inputs and outputs of B and E when they are no longer needed after the forward pass (this is automatically done by putting A-&gt;B-&gt;C and D-&gt;E-&gt;F in the no-grad context), and recompute them by rerunning A-&gt;B and D-&gt;E when they are needed again in the backward pass.</p>
<p>Finally, as you noted, GPU memory allocation via raw <code>cudaMalloc</code> calls is expensive, but the caching allocator continuously recycles allocations by maintaining its own pool(s) of GPU memory. <a href="https://pytorch.org/docs/stable/notes/cuda.html#memory-management" class="inline-onebox" rel="noopener nofollow ugc">CUDA semantics — PyTorch 2.0 documentation</a> In practice once a training loop is warmed-up, we would not expect to see any more raw <code>cudaMalloc</code> calls.</p>

raining_day513: <aside class="quote no-group" data-username="eqy" data-post="7" data-topic="177448">
<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/eqy/48/46233_2.png" class="avatar"> eqy:</div>
<blockquote>
<p>If you check the memory usage at the beginning and end of each <em>segment</em>, rather than each <em>layer</em>, I believe you should see nondecreasing memory usage.</p>
</blockquote>
</aside>
<p><a class="mention" href="/u/eqy">@eqy</a> I just confirmed this, thank you! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"> So, to conclude: in fact, the <code>torch.utils.checkpoint</code> is not related to GPU memory allocation, which is the job of the caching allocator. But being recycled or not, which can be controlled by <code>with torch.no_grad():</code>, is indeed related to <code>torch.utils.checkpoint</code> as per implementation. (please correct me if I did it wrong.)</p>
<p>TIL these two facts, appreciate your explanations:</p>
<ul>
<li>This context allows the intermediate forward activation(s) to be discarded <em>within</em> each segment rather than being saved for backward, which explains why the memory usage can increase and then decrease within a <em>segment</em> .</li>
<li>In practice, once a training loop is warmed-up, we would not expect to see any more raw <code>cudaMalloc</code> calls</li>
</ul>
<hr>
<p>If you’re still interested in this topic, one last question please:</p>
<p>Now, for me, the advantage of API <code>torch.utils.checkpoint</code> is that it can reduce the GPU memory usage(i.e. <code>torch.cuda.memory_allocated</code>) <em>within</em> a fixed pre-allocated amount by the caching allocator. Would it be practical to rewrite the (or maybe one of those) caching allocator so that the pre-allocation of GPU memory by the allocator won’t be too much larger than the memory actually being used, i.e. the value of <code>torch.cuda.memory_allocated()</code>?</p>
<p>My intention behind this is simple: For me, <code>torch.utils.checkpoint</code> is used to achieve “training a large model using a small GPU”. The pre-allocation of caching allocator seems to negate this effect. (I haven’t investigated how PyTorch implementing those caching allocators, so I’m probably wrong with this conclusion.) Could you also share some insight regarding this concern?</p>

eqy: <p>If you are concerned that the caching allocator is “wasting” memory by preallocating more than what is actually used, let me clarify that this should not be the case unless you are concerned with fragmentation (which I don’t think is the intent of your question). When there is more GPU memory <em>available</em> than is actually being <em>used</em>, the caching allocator may reserve more memory than is needed, but at the limit (your case of training a large model on a small GPU), the caching allocator can only reserve as much memory as the GPU has, and all of this memory would be available to model training (minus some hopefully small fraction due to fragmentation and GPU kernels). From this perspective, the caching allocator is meant to provide speedy allocations with minimal actual wasted memory usage.</p>
<p>With that said, the main limitation of <code>torch.utils.checkpoint</code> when it comes to training large models with the least amount of memory possible is that it can only recompute intermediate activations at most <em>once</em>. In theory, a more sophisticated checkpointing algorithm could recompute intermediate activations many times. Consider the A-&gt;B-&gt;C-&gt;D-&gt;E-&gt;F example shown previously. A more sophisticated checkpointing algorithm could store just the input to A and the output of F following the forward pass, and e.g.,<br>
recompute B, recompute C, throw out B, recompute D, throw out C, recompute E, throw out D, compute F’s backward, throw out F, recompute B, recompute C, throw out B, recompute D, throw out C, compute E’s backward, throw out E, … and so on. This approach is more formally described and implemented in this paper: <a href="https://arxiv.org/abs/2006.09616" class="inline-onebox" rel="noopener nofollow ugc">[2006.09616] Dynamic Tensor Rematerialization</a> which you may be interested in.</p>

raining_day513: <blockquote>
<p>[…] the caching allocator can only reserve as much memory as the GPU has, and all of this memory would be available to model training […] From this perspective, the caching allocator is meant to provide speedy allocations with minimal actual wasted memory usage.</p>
</blockquote>
<p><a class="mention" href="/u/eqy">@eqy</a> Could you elaborate more on this part? I think I haven’t been able to understand it fully. When a cache allocator reserves say 20GiB of GPU memory, and only 12GiB were used during the training, isn’t this a kind of <em>wasting</em> memory? From my current intuition (since I haven’t fully read the implementation details), the part being reserved would not be able to be used by other programs. In short, I think I might misunderstand your words.</p>
<p>I’m fine with the other part of your last comment, and I will read the paper you kindly linked.</p>
<p>Again, many thanks for your reading!</p>

eqy: <p>Regarding the first point, PyTorch is usually expected to be the only process running on a GPU. If you know the maximum amount of memory PyTorch should use and expect to use memory for other process, you can limit the amount of memory the caching allocator can reserve with e.g., <a href="https://pytorch.org/docs/stable/generated/torch.cuda.set_per_process_memory_fraction.html" class="inline-onebox" rel="noopener nofollow ugc">torch.cuda.set_per_process_memory_fraction — PyTorch 2.0 documentation</a></p>

raining_day513: <blockquote>
<p>[…] the caching allocator can only reserve as much memory as the GPU has, and all of this memory would be available to model training (minus some hopefully small fraction due to fragmentation and GPU kernels).<br>
<a href="https://discuss.pytorch.org/t/why-torch-cuda-allocated-memory-reports-that-gpu-memory-usage-was-decreasing-during-forwarding/177448/9">scroll to the reply</a></p>
</blockquote>
<p><a class="mention" href="/u/eqy">@eqy</a> So am I correct to say that there are some algorithms that PyTorch has implemented for the caching allocator to predict how much GPU will be reserved before the training? Since you only mentioned that it will “at most” use all the memory of the current single GPU. (Am I correct to say that this is the reason for doing the warm-up round in the first few batches?)</p>
<blockquote>
<p>If you know the maximum amount of memory PyTorch should use and expect to use memory for other process […] you can limit the amount of memory the caching allocator can reserve with e.g., <a href="https://pytorch.org/docs/stable/generated/torch.cuda.set_per_process_memory_fraction.html" rel="noopener nofollow ugc">torch.cuda.set_per_process_memory_fraction</a> […]</p>
</blockquote>
<p>Yes, I’m working on an algorithm to make the prediction. So I will assume that this API will be useful only if the default algorithm of caching allocator might not set the optimal fraction(a memory top-bound that is very close to what will actually be used), right? I think I need to know where are these files located in the source, could you help? (I selected your reply above due to the title. So, this thread basically ends. But I’m more interested in our further discuss here.)</p>
<p>Again, it’s amazing that you replied all these with valuable details and references, very thank you!</p>

eqy: <p>PyTorch by default should not make strong assumptions about how much memory is used during training as it is meant to support dynamism during training. The caching allocator’s reservation of additional memory is based on heuristics rather than trying to “predict” how much memory should be reserved. I think the code for allocation size calculation/rounding might be relevant: <a href="https://github.com/pytorch/pytorch/blob/d5f15d351510d1cb5e2294deb07569d343eea51f/c10/cuda/CUDACachingAllocator.cpp#L1307" class="inline-onebox" rel="noopener nofollow ugc">pytorch/CUDACachingAllocator.cpp at d5f15d351510d1cb5e2294deb07569d343eea51f · pytorch/pytorch · GitHub</a></p>
<p>As the caching allocator has no foresight into the maximum memory required and will need to reserve it on-the-fly at least once, warmup iteration(s) will observe raw <code>cudaMalloc</code> calls while later iterations should not.</p>
<p>As for setting the maximum memory via the fraction environment variable, that code should be in the same file at <a href="https://github.com/pytorch/pytorch/blob/main/c10/cuda/CUDACachingAllocator.cpp#L1633" rel="noopener nofollow ugc">pytorch/CUDACachingAllocator.cpp at main · pytorch/pytorch (github.com)</a></p>