Mitigating CUDA GPU memory fragmentation and OOM issues

eric.mitchell: <p>I’m currently playing around with some transformers with variable batch sizes, and I’m running into pretty severe memory fragmentation issues, with CUDA OOM occurring at less than 70% GPU memory utilization. For example (see the GitHub link below for more extreme cases, of failure at &lt;50% GPU memory):</p>
<pre><code class="lang-auto">RuntimeError: CUDA out of memory. Tried to allocate 1.48 GiB (GPU 0; 23.65 GiB total capacity; 16.22 GiB already allocated; 111.12 MiB free; 22.52 GiB reserved in total by PyTorch)
</code></pre>
<p>This has been discussed before on the PyTorch forums [<a>1</a>, <a href="https://discuss.pytorch.org/t/gpu-ram-fragmentation-diagnostics/34073/9">2</a>] and <a href="https://github.com/pytorch/pytorch/issues/31252" rel="noopener nofollow ugc">GitHub</a>. Fragmentation is also mentioned briefly in the <a href="https://pytorch.org/docs/stable/cuda.html" rel="noopener nofollow ugc">docs</a> for <code>torch.cuda.empty_cache()</code>:</p>
<blockquote>
<p><a href="https://pytorch.org/docs/stable/cuda.html#torch.cuda.empty_cache" rel="noopener nofollow ugc"><code>empty_cache()</code></a> doesn’t increase the amount of GPU memory available for PyTorch. However, it may help reduce fragmentation of GPU memory in certain cases. See <a href="https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management" rel="noopener nofollow ugc">Memory management</a> for more details about GPU memory management.</p>
</blockquote>
<p>According to <a href="https://discuss.pytorch.org/t/about-torch-cuda-empty-cache/34232">this thread</a>, we shouldn’t be relying on <code>torch.cuda.empty_cache()</code>. However (in my use case, at least), <code>torch.cuda.empty_cache()</code> is the difference between running code with GiBs of GPU memory to spare and CUDA OOM errors. None of these past threads (as far as I can tell) really led to conclusive mitigation strategies for this problem.</p>
<p>My questions are:</p>
<ol>
<li>Are there <strong>“official” best practices</strong> for handling cases when GPU memory fragmentation is severely reducing effective usable GPU memory?</li>
<li>Can we get more documentation about <strong>when <code>torch.cuda.empty_cache()</code> is the right tool</strong> for addressing OOM/fragmentation issues?</li>
</ol>
<p>As far as I can tell, moving all allocated tensors into main memory, emptying the cache, then moving them back to GPU memory a la <a href="https://github.com/pytorch/pytorch/issues/31252" rel="noopener nofollow ugc">this post</a> seems a reasonable (if potentially expensive) strategy. Is this recommended?</p>

stas: <p>In no way to steal from the focus of this super important issue you raise, I am curious how your code would fair with the DeepSpeed integration - since it manages temp memory allocations itself. And in my few experiments it needs 3-5x less gpu RAM with its ZeRO algorithms. that is if you’re using the HF trainer. If you don’t use the latter, you can see from the code how to activate deepspeed in your own trainer.</p>
<p>We are getting close to completing the integration but you can already try my branch:</p><aside class="onebox githubpullrequest">
  <header class="source">
      <a href="https://github.com/huggingface/transformers/pull/9211" target="_blank" rel="noopener nofollow ugc">github.com/huggingface/transformers</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Pull Request">
    <svg width="60" height="60" class="github-icon" viewBox="0 0 12 16" aria-hidden="true"><path d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/huggingface/transformers/pull/9211" target="_blank" rel="noopener nofollow ugc">[trainer] deepspeed integration</a>
    </h4>

    <div class="branches">
      <code>huggingface:master</code> ← <code>stas00:ds</code>
    </div>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-12-19" data-time="07:13:32" data-timezone="UTC">07:13AM - 19 Dec 20 UTC</span>
      </div>

      <div class="user">
        <a href="https://github.com/stas00" target="_blank" rel="noopener nofollow ugc">
          <img alt="stas00" src="https://avatars2.githubusercontent.com/u/10676103?v=4" class="onebox-avatar-inline" width="20" height="20">
          stas00
        </a>
      </div>

      <div class="lines" title="38 commits changed 7 files with 384 additions and 23 deletions">
        <a href="https://github.com/huggingface/transformers/pull/9211/files" target="_blank" rel="noopener nofollow ugc">
          <span class="added">+384</span>
          <span class="removed">-23</span>
        </a>
      </div>
    </div>

  </div>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
<p>
There is a doc in it that explains how to activate it.<br>
You’d need to install DeepSpeed from its master.</p>
<p>If you have any follow up questions please ask in that PR thread so that we don’t derail this thread.</p>
<p>We also recently added <a href="https://huggingface.co/transformers/master/training.html#trainer-integrations" rel="noopener nofollow ugc">support for Sharded DDP via fairscale</a>. if you use HF trainer, just install fairscale and add <code>--sharded_ddp</code> to the training args and you should also see a huge improvement in memory utilization. Again if you have questions let’s continue this discussion on HF forums.</p>
<p>Back to the topic of this thread.</p>

eric.mitchell: <p>Very interesting, thanks for sharing! I’m not using the HuggingFace trainer, but I went ahead and ran with deepspeed with my current code.</p>
<p>I normally run:</p>
<pre><code>ENV=var python -m train --arg1 val1 --arg2 val2
</code></pre>
<p>I ran:</p>
<pre><code>ENV=var deepspeed ./train.py --arg1 val1 --arg2 val2
</code></pre>
<p>I don’t see any difference in runtime/memory usage, though. Is there anything else I need to do? My <code>ds_report</code> looks like this:</p>
<pre><code class="lang-auto">JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires one of the following commands '['llvm-config', 'llvm-config-9']', but it does not exist!
sparse_attn ............ [NO] ....... [NO]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
--------------------------------------------------
</code></pre>
<p>Did you pre-build the extension ops? I’ve got a CUDA Version mismatch that’s preventing me from pre-building them with <code>DS_BUILD_OPS=1 pip install deepspeed</code>.</p>

stas: <p>Based on previous experience - continuing the ZeRO solution discussion here will derail your attempt to find core solutions. I highly recommend removing these follow ups and opening a separate thread on HF forums (and tag <a class="mention" href="/u/stas">@stas</a>) and we can continue there.</p>
<p>But the quick answer is that based on your reply - you’re not using deepspeed, you’re just using its launcher. You need to:</p>
<ol>
<li>create a configuration file which sets up its ZeRO magic.See: <a href="https://www.deepspeed.ai/docs/config-json/" class="inline-onebox" rel="noopener nofollow ugc">DeepSpeed Configuration JSON - DeepSpeed</a>
</li>
<li>activate deepspeed <a href="https://www.deepspeed.ai/getting-started/#writing-deepspeed-models" class="inline-onebox" rel="noopener nofollow ugc">Getting Started - DeepSpeed</a>
</li>
<li>use its wrapped model in training <a href="https://www.deepspeed.ai/getting-started/#training" class="inline-onebox" rel="noopener nofollow ugc">Getting Started - DeepSpeed</a>
</li>
<li>launch it as <a href="https://www.deepspeed.ai/getting-started/#launching-deepspeed-training" class="inline-onebox" rel="noopener nofollow ugc">Getting Started - DeepSpeed</a>
</li>
</ol>
<p>That’s when you will see huge improvements.</p>

zyeric: <p>Stas, thanks for your information very much. I am trying to train a GPT3-350M with a large batch size by recompute. Fragmentation is a very serious problem under this setting. For example, if the self-attention is selected for recompute, the size of the attention matrix (temporary variable) cab be 1024 MB. After recomputation, this large memory block may be occupied by other tensors which don’t need recompute. If another self-attention needs recompute, torch’s allocator has to allocate another 1024 MB, leading to the fragmentation.</p>
<p>You mentioned that deepspeed ‘manages temp memory allocations itself’. I spent some time in reading deepspeed’s memory management mechanism but did not find related implementations. Could you kindly point out the source code?</p>

stas: <p>Oh, I see, you’re doing some custom tweaking of the forward path.</p>
<p>Deepspeed avoids fragmentation by pre-allocating all the tensors it needs ahead of time and then just working directly with <code>.data</code> when it needs to move/update something. So unless user’s code does some unusual allocations past model’s inits the memory remains quite steady through the whole training.</p>
<p>Perhaps you could adopt a somewhat similar approach, by pre-allocating your large tensors and then re-using those and thus avoiding fragmentation.</p>