Gradient checkpointing and its effect on memory and runtime

Dxyk: <p>Hello,</p>
<p>I am trying to understand how the number of checkpoints in gradient checkpointing affects the memory and runtime for computing gradients.</p>
<p>I found this <a href="https://github.com/rasbt/deeplearning-models/blob/18e046926551378cd691fd871dda0f21dcd272ab/pytorch_ipynb/mechanics/gradient-checkpointing-nin.ipynb" rel="noopener nofollow ugc">notebook</a> that explains how gradient checkpointing works. I modified it a bit and ran a couple of experiments to see how the <code>use_reentrant</code> and the <code>segments</code> arguments affect the memory and runtime.</p>
<p>I was surprised to find the following results</p>
<ul>
<li>For memory, the data shows almost no improvement when checkpointing is enabled compared to no checkpointing, no matter the number of checkpoints and whether we are using reentrant. From <a href="https://discuss.pytorch.org/t/gradient-checkpointing-more-checkpoints-higher-memory-usage/97677/2">this thread</a>, I understand that the optimal number of checkpoints depends on the specific scenario. However, I still expected at least some improvement in the memory consumption. I’m not quite sure why this is?</li>
<li>For runtime, the experiment shows the runtime increases as the number of checkpoints increases. However, from the Chen2016 [1] paper, my understanding is the runtime should roughly increase by 30% when checkpointing is enabled due to the extra forward pass, and should remain roughly the same regardless the number of checkpoints. Is the increasing runtime due to the extra overhead required for keeping more checkpoints?</li>
</ul>
<p>Bellow are my results and their corresponding graphs. Note segments=0 means no checkpointing is enabled.</p>
<p>Any insight will help, and thanks in advance <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>
<div class="md-table">
<table>
<thead>
<tr>
<th style="text-align:center">use_reentrant</th>
<th style="text-align:center">segments</th>
<th style="text-align:center">Runtime (s)</th>
<th style="text-align:center">Current Memory (B)</th>
<th style="text-align:center">Peak Memory (B)</th>
</tr>
</thead>
<tbody>
<tr>
<td style="text-align:center">FALSE</td>
<td style="text-align:center">0</td>
<td style="text-align:center">109.1</td>
<td style="text-align:center">413527</td>
<td style="text-align:center">154106538</td>
</tr>
<tr>
<td style="text-align:center">TRUE</td>
<td style="text-align:center">0</td>
<td style="text-align:center">109.1</td>
<td style="text-align:center">413527</td>
<td style="text-align:center">154106538</td>
</tr>
<tr>
<td style="text-align:center">FALSE</td>
<td style="text-align:center">1</td>
<td style="text-align:center">109.12</td>
<td style="text-align:center">402938</td>
<td style="text-align:center">154105009</td>
</tr>
<tr>
<td style="text-align:center">TRUE</td>
<td style="text-align:center">1</td>
<td style="text-align:center">110.14</td>
<td style="text-align:center">394474</td>
<td style="text-align:center">154121963</td>
</tr>
<tr>
<td style="text-align:center">FALSE</td>
<td style="text-align:center">2</td>
<td style="text-align:center">116.23</td>
<td style="text-align:center">382560</td>
<td style="text-align:center">154127403</td>
</tr>
<tr>
<td style="text-align:center">TRUE</td>
<td style="text-align:center">2</td>
<td style="text-align:center">114.86</td>
<td style="text-align:center">378951</td>
<td style="text-align:center">154162031</td>
</tr>
<tr>
<td style="text-align:center">FALSE</td>
<td style="text-align:center">5</td>
<td style="text-align:center">116.82</td>
<td style="text-align:center">428681</td>
<td style="text-align:center">154214720</td>
</tr>
<tr>
<td style="text-align:center">TRUE</td>
<td style="text-align:center">5</td>
<td style="text-align:center">117.23</td>
<td style="text-align:center">507778</td>
<td style="text-align:center">154334766</td>
</tr>
<tr>
<td style="text-align:center">FALSE</td>
<td style="text-align:center">10</td>
<td style="text-align:center">119.44</td>
<td style="text-align:center">384482</td>
<td style="text-align:center">154244424</td>
</tr>
<tr>
<td style="text-align:center">TRUE</td>
<td style="text-align:center">10</td>
<td style="text-align:center">118.84</td>
<td style="text-align:center">515236</td>
<td style="text-align:center">154438154</td>
</tr>
</tbody>
</table>
</div><p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/d/6d219c2661ebbc1400920f9c9b2b711a8ad79c5c.png" data-download-href="https://discuss.pytorch.org/uploads/default/6d219c2661ebbc1400920f9c9b2b711a8ad79c5c" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/d/6d219c2661ebbc1400920f9c9b2b711a8ad79c5c.png" alt="image" data-base62-sha1="fzq3iPVcPuDpB0lOk1J5R5EIfRa" width="416" height="500" data-dominant-color="F9F8F8"><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">image</span><span class="informations">479×575 18.7 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>[1] Chen, Tianqi, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. “Training deep nets with sublinear memory cost.” <em>arXiv preprint arXiv:1604.06174</em> (2016).</p>

soulitzer: <p>Activation checkpoint produces more pronounced reductions in peak memory when your activations are larger relative to model size. What do those numbers look like for your example?</p>
<blockquote>
<p>Is the increasing runtime due to the extra overhead required for keeping more checkpoints?</p>
</blockquote>
<p>In the checkpoint sequential the last segment is not checkpointed. The fewer segments you have, the larger that non-checkpointed region, and the smaller the runtime.</p>

Dxyk: <p>Thanks for your reply!</p>
<p>Yes there are dropout layers in the model so the neuron units are not always active if that’s what you mean. I’ve removed the dropout layers and tried again, but the memory trend is not much different than before. The peak memory I have before training is 146.98MB, with checkpointing enabled and <code>segments=1</code>, the memory consumption is 146.94MB, with <code>segments=2</code> the memory consumption is 146.96MB. It doesn’t seem to improve the performance much, but I might be misunderstanding something here.</p>
<blockquote>
<p>In the checkpoint sequential the last segment is not checkpointed.</p>
</blockquote>
<p>That makes a lot more sense. Do you know if <code>checkpoint</code> and <code>checkpoint_sequential </code> behave the same way in this aspect? I.e. if I use <code>checkpoint</code> with fewer segments, should I also expect a smaller runtime because the last segment is not checkpointed?</p>
<p>I’ve also had experiments where checkpointing a function actually decreased the runtime compared to the uncheckpointed version. Do you have any idea why that could be?</p>