Extra 10GB memory on GPU 0 in DDP tutorial

jastern33: <p>When running the basic DDP (distributed data parallel) example from the <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" rel="noopener nofollow ugc">tutorial here</a>, GPU 0 gets an extra 10 GB of memory on this line:</p>
<pre><code class="lang-auto">ddp_model = DDP(model, device_ids=[rank])
</code></pre>
<p>What I’ve tried:</p>
<ul>
<li>Setting the ‘CUDA_VISIBLE_DEVICES’ environment variable so that each subprocess can only see the GPU of its rank. Then I set rank = 0.</li>
</ul>
<pre><code class="lang-auto">os.environ['CUDA_VISIBLE_DEVICES'] = str(rank)
rank = 0
setup(rank, world_size)
# etc.
</code></pre>
<p>This results in the error</p>
<pre><code class="lang-auto">File "/usr/local/lib/python3.8/dist-packages/torch/multiprocessing/spawn.py", line 59, in _wrap
    fn(i, *args)
  File "/code/test.py", line 49, in demo_basic
    setup(rank, world_size)
  File "/code/test.py", line 29, in setup
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
  File "/usr/local/lib/python3.8/dist-packages/torch/distributed/distributed_c10d.py", line 500, in init_process_group
    store, rank, world_size = next(rendezvous_iterator)
  File "/usr/local/lib/python3.8/dist-packages/torch/distributed/rendezvous.py", line 190, in _env_rendezvous_handler
    store = TCPStore(master_addr, master_port, world_size, start_daemon, timeout)
RuntimeError: Address already in use
</code></pre>
<p>I also tried rewriting this to use the <a href="https://gist.github.com/sgraaf/5b0caa3a320f28c27c12b5efeb35aa4c#file-ddp_example-sh-L5" rel="noopener nofollow ugc">scripting version of DDP</a>, and saw the same problem.</p>
<p>There is also nothing here involving a profiler, as described <a href="https://github.com/pytorch/pytorch/issues/37313" rel="noopener nofollow ugc">here</a>, loading a model from disk, as described <a href="https://github.com/pytorch/pytorch/issues/23138" rel="noopener nofollow ugc">here</a>, or calls to torch.cuda.empty_cache, as described <a href="https://discuss.pytorch.org/t/ddp-taking-up-too-much-memory-on-rank-0/62443/4">here</a>. My machine is a DGX-1 with 8 Tesla V100 GPUs. I tried this on another DGX-1 with the same result.</p>
<p>Any ideas what is consuming so much memory on GPU 0, and how to resolve it?</p>
<p>Screenshot of the problem:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/0/e090597cc28384c394b21486cfd6cc4f2874e083.png" data-download-href="https://discuss.pytorch.org/uploads/default/e090597cc28384c394b21486cfd6cc4f2874e083" title="Screenshot from 2021-04-14 00-01-09"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/0/e090597cc28384c394b21486cfd6cc4f2874e083_2_379x500.png" alt="Screenshot from 2021-04-14 00-01-09" data-base62-sha1="w2AdH5pX3cNkjJrL4GaBf6PjFZ1" width="379" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/0/e090597cc28384c394b21486cfd6cc4f2874e083_2_379x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/0/e090597cc28384c394b21486cfd6cc4f2874e083_2_568x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/0/e090597cc28384c394b21486cfd6cc4f2874e083.png 2x" data-dominant-color="3D1A32"><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">Screenshot from 2021-04-14 00-01-09</span><span class="informations">731×964 117 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>Below is the full MWE, copied from the tutorial:</p>
<pre><code class="lang-auto">import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))

def demo_basic(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)

    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(rank)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup()

def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             join=True)

if __name__ == "__main__":
    run_demo(demo_basic, torch.cuda.device_count())
</code></pre>

jastern33: <p>Adding these two lines above the model initialization line solved the problem for me:</p>
<pre><code>torch.cuda.set_device(rank)
torch.cuda.empty_cache()
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd340f86fa9f40e2746921eb46c7d9336436a4f5.png" data-download-href="https://discuss.pytorch.org/uploads/default/bd340f86fa9f40e2746921eb46c7d9336436a4f5" title="Screenshot from 2021-04-14 10-56-30"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd340f86fa9f40e2746921eb46c7d9336436a4f5_2_388x500.png" alt="Screenshot from 2021-04-14 10-56-30" data-base62-sha1="qZLMrvEsVpBrfme88dyWPJLK9hj" width="388" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd340f86fa9f40e2746921eb46c7d9336436a4f5_2_388x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd340f86fa9f40e2746921eb46c7d9336436a4f5_2_582x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd340f86fa9f40e2746921eb46c7d9336436a4f5.png 2x" data-dominant-color="3C1A31"><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">Screenshot from 2021-04-14 10-56-30</span><span class="informations">759×977 110 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>

Zhang_Kin: <aside class="quote no-group" data-username="jastern33" data-post="2" data-topic="118113">
<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/jastern33/48/11917_2.png" class="avatar"> jastern33:</div>
<blockquote>
<pre><code class="lang-auto">torch.cuda.set_device(rank)
torch.cuda.empty_cache()
</code></pre>
</blockquote>
</aside>
<p>! This is the solution!!! Finally! Thanks!!!</p>

windj007: <p>Hi <a class="mention" href="/u/jastern33">@jastern33</a> <a class="mention" href="/u/zhang_kin">@Zhang_Kin</a> ! What version of pytorch did you use? I’m facing the same problem with pytorch==1.12.1+cu116 and the fix with set_device and empty_cache does not help.</p>
<p>Also, do I understand it correctly that these two lines should be inserted just before <code>model = ToyModel().to(rank)</code> ?</p>

Hongdou_Liu: <p>What a hack! Solve the same problem for me. I wonder where did you get this solution.</p>