DDP taking up too much memory on rank 0

PCerles: <p><strong>Edit: Mistaken!</strong><br>
This was my issue:</p><aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/23138">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/23138" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/23138" target="_blank" rel="noopener nofollow ugc">DistributedDataParallel: resume training from a checkpoint results in additional processes on GPU 0</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2019-07-21" data-time="01:23:03" data-timezone="UTC">01:23AM - 21 Jul 19 UTC</span>
      </div>

        <div class="date">
          closed <span class="discourse-local-date" data-format="ll" data-date="2019-07-21" data-time="03:00:34" data-timezone="UTC">03:00AM - 21 Jul 19 UTC</span>
        </div>

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

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          oncall: distributed
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">Hi,

When I was trying the[ imagenet example](https://github.com/pytorch/examp<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">les/tree/master/imagenet) with DistributedDataParallel, using single node with 4 gpus, I found that when I add `--resume /path/to/checkpoint `to the command, gpu 0 has additional processes like below, the pid of each process is exactly the same as running on every other gpu, and each process consumes ~725M. These additional processes keep lingering on gpu 0 instead of disapear when loading finished. Is there any way to solve this? If there are 8 gpus, that means ~5G on gpu 0 are wasted.

&gt; +-----------------------------------------------------------------------------+
&gt; | Processes:                                                       GPU Memory |
&gt; |  GPU       PID   Type   Process name                             Usage      |
&gt; |=============================================================================|
&gt; |    0     62250      C   /usr/bin/python3.6                          4903MiB |
&gt; |    0     62251      C   /usr/bin/python3.6                           725MiB |
&gt; |    0     62252      C   /usr/bin/python3.6                           725MiB |
&gt; |    0     62253      C   /usr/bin/python3.6                           725MiB |
&gt; |    1     62251      C   /usr/bin/python3.6                          4693MiB |
&gt; |    2     62252      C   /usr/bin/python3.6                          4693MiB |
&gt; |    3     62253      C   /usr/bin/python3.6                          4693MiB |
&gt; 
&gt; 

If I train from scratch, there is no such issue.

&gt; +-----------------------------------------------------------------------------+
&gt; | Processes:                                                       GPU Memory |
&gt; |  GPU       PID   Type   Process name                             Usage      |
&gt; |=============================================================================|
&gt; |    0      2080      C   /usr/bin/python3.6                          4687MiB |
&gt; |    1      2081      C   /usr/bin/python3.6                          4687MiB |
&gt; |    2      2082      C   /usr/bin/python3.6                          4687MiB |
&gt; |    3      2083      C   /usr/bin/python3.6                          4687MiB |
&gt;</span></p>
  </div>

  </article>

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

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

<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/b/cb63a779ce29d8a7ddb168539aad24eaf9810f45.png" data-download-href="https://discuss.pytorch.org/uploads/default/cb63a779ce29d8a7ddb168539aad24eaf9810f45" title="13%20PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/b/cb63a779ce29d8a7ddb168539aad24eaf9810f45_2_690x223.png" alt="13%20PM" data-base62-sha1="t1gqWM4FyQM3kkfYlQXJEylr0vb" width="690" height="223" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/b/cb63a779ce29d8a7ddb168539aad24eaf9810f45_2_690x223.png, https://discuss.pytorch.org/uploads/default/optimized/3X/c/b/cb63a779ce29d8a7ddb168539aad24eaf9810f45_2_1035x334.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/c/b/cb63a779ce29d8a7ddb168539aad24eaf9810f45_2_1380x446.png 2x" data-dominant-color="0C0C0C"><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">13%20PM</span><span class="informations">1416×458 32.4 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><br>
Here’s a screenshot of distributed training in Pytorch when I call the train function like:<br>
<code>CUDA_VISIBLE_DEVICES=1,2,3,4 python -m torch.distributed.launch --nproc_per_node=4 train_new.py</code>. You can see that the first rank has also initted 3 separate processes for each other GPU. When I use 10 GPUs on a box this severely limits the batch size, since the 0th dimension node has so much less capacity. What is it storing? I thought gradients in DDP were all-reduced. I’ve also tried turning broadcast_buffers to False to no avail.<br>
Model is stacked modules of 1D-conv, relu, batch norm, LSTM, followed by a large softmax layer and CTC loss. Backend is NCCL<br>
Pytorch 1.3.0, Cuda 10.1, Titan RTX, Ubuntu 18.04. Can provide more code upon request.</p>

songzw: <p>any solution?<br>
github issue solution does not work for me</p>

mrshenli: <p>Discussion <a href="https://discuss.pytorch.org/t/torch-not-able-to-utilize-gpu-ram-properly/83245/8">here</a> might be helpful.</p>
<p>This is likely due to some tensors/context is unintentionally created on the 1st GPU, e.g., when calling <code>torch.cuda.empty_cache()</code> without a device guard. Solutions would be either 1) carefully walking though libs/codes to make sure no states leaks to <code>cuda:0</code>, or 2) set <code>CUDA_VISIBLE_DEVICES</code> to let each process only see one GPU.The second approach might be easier.</p>

Felix_Kreuk: <p><a class="mention" href="/u/pcerles">@PCerles</a> I’m having a similar issue. Were you able to resolve your problem? Thanks.</p>

seungjun: <p><a class="mention" href="/u/pcerles">@PCerles</a> <a class="mention" href="/u/felix_kreuk">@Felix_Kreuk</a></p>
<p>What <a class="mention" href="/u/mrshenli">@mrshenli</a> mentioned could seamlessly happen when you load saved parameters without specifying <code>map_location</code>.<br>
<code>torch.load</code> by default loads parameters to the device where they were, usually the rank 0 device.<br>
<code>load_state_dict</code> then copies the loaded value from that device to the target device.<br>
After the intermediate use, torch still occupies the GPU memory as cached memory.<br>
I had a similar issue and solved it by directly loading parameters to the target device.</p>
<p>For example:</p>
<pre><code class="lang-python">state_dict = torch.load(model_name, map_location=self.args.device)
self.load_state_dict(state_dict)
</code></pre>
<p>Full code <a href="https://github.com/SeungjunNah/DeepDeblur-PyTorch/blob/master/src/model/__init__.py#L102" rel="nofollow noopener">here</a>.</p>