Reducing RAM usage with DDP

neel_g: <p>Hi, I’m working on <code>16</code>x nodes, each with <code>8x A100</code> gpus and when using DDP, it appears that there are certain spiked in the used <strong>RAM</strong> which reach the limit of the <code>1.1TB</code> available on <code>p4d</code> instances.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/2/7226f9f91d9d0bf9a15dfc432ea78e2eac1462a2.png" alt="image" data-base62-sha1="ghPWaQD3v7qNlw77vGUmaxr0Znc" width="624" height="322"></p>
<ul>
<li>This arises on the host node (above graph), but I can’t confirm whether it arises on other nodes as well. Presumably, that doesn’t happen but if the host node goes down - the entire training does too.</li>
</ul>
<pre data-code-wrap="shell"><code class="lang-shell">============================================================
scripts/ddp_convnext.py FAILED
------------------------------------------------------------
Failures:
[1]:
  time      : 2022-08-26_22:41:31
  host      : gpu-st-p4d-24xlarge-56.hpc-1click-prod450.pcluster
  rank      : 77 (local_rank: 5)
  exitcode  : -6 (pid: 28176)
  error_file: &lt;N/A&gt;
  traceback : Signal 6 (SIGABRT) received by PID 28176
------------------------------------------------------------

.......
.... [Multiple other nodes's stdout] ...
.......
srun: error: gpu-st-p4d-24xlarge-57: task 13: Exited with exit code 1
srun: error: gpu-st-p4d-24xlarge-46: task 2: Exited with exit code 1
slurmstepd: error: Detected 1142 oom-kill event(s) in StepId=4236.0. Some of your processes may have been killed by the cgroup out-of-memory handler.
srun: error: gpu-st-p4d-24xlarge-51: task 7: Exited with exit code 1
srun: error: gpu-st-p4d-24xlarge-55: task 11: Out Of Memory

</code></pre>
<hr>
<p>Code: <a href="https://rentry.co/683t9" class="inline-onebox" rel="noopener nofollow ugc">Main script</a></p>
<p>How can we reduce the memory footprint when doing DDP multi-node? I feel <code>1.1 TiB</code> per node should be enough - but that doesn’t seem to be.</p>
<p>Interestingly, this problem arises only when I scale <em>processes</em>/GPUs - in a certain limit, stuff works quite well enough. I also had to disable some GPUs to save memory, so now I can only run <code>8x nodes, each with 6xA100</code> instead of the <code>8xA100s</code> available per node. Increasing processes also makes training <em>slower</em>. I’m unsure where the problem here - some guidance would be very, very welcome <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><strong>EDIT:-</strong> Here’s also a truncated SLURM log: <a href="https://gist.githubusercontent.com/neel04/22052c61512657aec4b364888ed74a95/raw/4caff5d0b5b6864991d572e909995e44f38bb810/slurm_nccl_oom.log" rel="noopener nofollow ugc">Gist</a> - It’s quite a few thousand lines, so It’s the raw link. I’ve only removed argparse, some <code>echo</code>s and stuff like model summaries.</p>

ptrblck: <p>Did you try to profile your code and check where the spikes might be coming from?<br>
Based on the posted code snippet it seems you are pre-loading the dataset:</p>
<pre><code class="lang-python">self.dataset = tfds.as_numpy(self.ds)
</code></pre>
<p>and might have a large batch size:</p>
<pre><code class="lang-python">return 45491349 // args.batch_size
</code></pre>
<p>which might increase the memory usage significantly depending on the number of workers in the <code>DataLoader</code> in each DDP process.<br>
Also, it seems you are mixing TensorFlow with PyTorch and I don’t know how TF would behave in this setup.</p>

neel_g: <p>I wasn’t using a dataloader - I profiled without it - <a href="https://rentry.co/c9kaf" rel="noopener nofollow ugc">rentry.co link</a></p>
<p>The usage seems nominal, and not too much on the terabyte of RAM available.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="160145">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Also, it seems you are mixing TensorFlow with PyTorch and I don’t know how TF would behave in this setup.</p>
</blockquote>
</aside>
<p>True, but I am converting it to a standard numpy iteratore, wrapping it into <code>IterableDataset</code> so the it should be like any standard PyTorch <code>Dataset</code> - no?</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="160145">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>and might have a large batch size:</p>
</blockquote>
</aside>
<p>The batch size <code>48</code> <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title=":thinking:" class="emoji" alt=":thinking:" loading="lazy" width="20" height="20"></p>

ptrblck: <aside class="quote no-group" data-username="neel_g" data-post="3" data-topic="160145">
<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/neel_g/48/42759_2.png" class="avatar"> neel_g:</div>
<blockquote>
<p>True, but I am converting it to a standard numpy iteratore, wrapping it into <code>IterableDataset</code> so the it should be like any standard PyTorch <code>Dataset</code> - no?</p>
</blockquote>
</aside>
<p>Yes, this would behave like the “standard” eager <code>Dataset</code>, is pre-loading the large dataset to the host in each DDP process (and each worker if multi-processing is used).</p>
<aside class="quote no-group" data-username="neel_g" data-post="3" data-topic="160145">
<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/neel_g/48/42759_2.png" class="avatar"> neel_g:</div>
<blockquote>
<p>The batch size <code>48</code> <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title=":thinking:" class="emoji" alt=":thinking:" loading="lazy" width="20" height="20"></p>
</blockquote>
</aside>
<p>Ah, sorry. I was staring at other C++ code and saw the <code>// args.batch_size</code> as a comment, not a integer division.</p>

neel_g: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="160145">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Yes, this would behave like the “standard” eager <code>Dataset</code>, is pre-loading the large dataset to the host in each DDP process (and each worker if multi-processing is used).</p>
</blockquote>
</aside>
<p><img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"> That’s unfortunate, it’s supposed to be an iterator because the dataset is humongous (~20TB size). Would you happen to have any general direction I can take to debug this? would reducing workers for the <code>tf.dataset</code> itself help?</p>

ptrblck: <p>You could try to use <code>np.memmap</code> to avoid loading the data. However, given the size of the dataset (20TB) and the RAM of your system (1TB), I don’t think you are indeed pre-loading the data and would need to explain what a</p>
<aside class="quote no-group" data-username="neel_g" data-post="3" data-topic="160145">
<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/neel_g/48/42759_2.png" class="avatar"> neel_g:</div>
<blockquote>
<p>but I am converting it to a standard numpy iteratore</p>
</blockquote>
</aside>
<p>is.</p>

neel_g: <p>Pretty much this:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://www.tensorflow.org/datasets/api_docs/python/tfds/as_numpy">
  <header class="source">
      <img src="https://www.gstatic.com/devrel-devsite/prod/v6bd99c7d8fa6220aea6020b29137b212e2eacdba13242535487f52ab3557b0cb/tensorflow/images/favicon.png" class="site-icon" width="32" height="32">

      <a href="https://www.tensorflow.org/datasets/api_docs/python/tfds/as_numpy" target="_blank" rel="noopener nofollow ugc">TensorFlow</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/388;"><img src="https://www.tensorflow.org/images/tf_logo_social.png" class="thumbnail" width="690" height="388"></div>

<h3><a href="https://www.tensorflow.org/datasets/api_docs/python/tfds/as_numpy" target="_blank" rel="noopener nofollow ugc">tfds.as_numpy  |  TensorFlow Datasets</a></h3>

  <p>Converts a tf.data.Dataset to an iterable of NumPy arrays.</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Converting it to a generator of <code>numpy</code> arrays</p>
<p><strong>EDIT:</strong> It appears <code>del</code>-ing a few variables helps to reduce memory usage. I can use <code>6x A100s</code> on <code>16</code> nodes, but still can’t use all 8 GPUs per node due to memory errors. Any ideas <a class="mention" href="/u/ptrblck">@ptrblck</a></p>