Is it expected for DistributedDataParallel to use more memory on 1 GPU in a 1GPU:1process setup?

jramapuram: <p>Is it expected for <code>nn.parallel.DistributedDataParallel</code> in a 1GPU:1Process setup to use a little extra memory on one of the GPUs? The use isn’t exorbitant ( 3614 MiB vs. 4189 MiB). If so, what is this extra memory used for? Is it the <code>all_reduce</code> call on the gradients? If not what would this be attributed to?</p>
<ul>
<li>1 Gpu per 1 process spun up with: <code>mp.spawn(run, nprocs=args.num_replicas, args=(args.num_replicas,))</code>
</li>
<li>Entire module wrapped with <code>nn.parallel.DistributedDataParallel</code>
</li>
<li>loss function is a member of above module.</li>
<li>pytorch 1.4</li>
</ul>
<p>Multiprocessing init created via:</p>
<pre><code class="lang-python">def handle_multiprocessing_logic(rank, num_replicas):
    """Sets the appropriate flags for multi-process jobs."""
    args.gpu = rank  # Set the GPU device to use

    if num_replicas &gt; 1:
        torch.distributed.init_process_group(
            backend='nccl', init_method='env://',
            world_size=args.num_replicas, rank=rank
        )

        # Update batch size appropriately
        args.batch_size = args.batch_size // num_replicas

        # Set the cuda device
        torch.cuda.set_device(rank)
        print("Replica {} / {} using GPU: {}".format(
            rank + 1, num_replicas, torch.cuda.get_device_name(rank)))
</code></pre>

mrshenli: <p>Yes, this is expected currently, because DDP creates buckets to consolidate gradient communication. Checkout <a href="https://github.com/pytorch/pytorch/blob/1f8267931176cc9ecbf00493e5a359a57baca3df/torch/nn/parallel/distributed.py#L173" rel="nofollow noopener">this</a> and  <a href="https://github.com/pytorch/pytorch/blob/1f8267931176cc9ecbf00493e5a359a57baca3df/torch/csrc/distributed/c10d/reducer.cpp#L108" rel="nofollow noopener">this</a>. We could potentially mitigate this problem by setting <code>param.grad</code> to point to different offsets in the bucket so that we don’t need two copies of grads.</p>

mrshenli: <p>tracking issue: <a href="https://github.com/pytorch/pytorch/issues/37030" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/37030</a></p>

jramapuram: <p>Thanks for the response; just to clarify: you mean it is expected that 1 of the GPUs in say a 2 GPU (single-process-single-gpu) DDP setup will use more memory because of bucketing? Wouldn’t the buckets be of the same size on both devices?</p>

mrshenli: <p>oh, sorry, I misread the question. I thought you mean DDP uses a little more memory than local model.</p>
<blockquote>
<p>you mean it is expected that 1 of the GPUs in say a 2 GPU (single-process-single-gpu) DDP setup will use more memory because of bucketing?</p>
</blockquote>
<p>no, they should be the same I think</p>

jramapuram: <p>The general problem of more memory makes sense in the current way things are setup (from your src links), let me see if I can come up with a minimum viable example for this effect (more mem on one of the GPUs).</p>

mrshenli: <p>Is the DDP process the only process using that GPU? The extra size ~500MB looks like an extra cuda context. Does this behavior still persist if you set <code>CUDA_VISIBLE_DEVICES </code> env var properly (instead of using <code>torch.cuda.set_device(rank)</code>) before launching each process?</p>

jramapuram: <aside class="quote no-group" data-username="mrshenli" data-post="7" data-topic="77748" data-full="true">
<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/mrshenli/48/12220_2.png" class="avatar"> mrshenli:</div>
<blockquote>
<p>Is the DDP process the only process using that GPU? The extra size ~500MB looks like an extra cuda context. Does this behavior still persist if you set <code>CUDA_VISIBLE_DEVICES </code> env var properly (instead of using <code>torch.cuda.set_device(rank)</code> ) before launching each process?</p>
</blockquote>
</aside>
<p>Yup, I don’t have X or anything else running: <img src="https://discuss.pytorch.org/uploads/default/original/3X/8/4/8401a5dd6daf338c37be2fc57f980e556cf21203.png" alt="image" data-base62-sha1="iPMyw4nnlFB3nZtkJ9NIj9REHnl" width="562" height="340"></p>
<p>I don’t set <code>CUDA_VISIBLE_DEVICES</code> but I do call <code>torch.cuda.set_device(rank)</code> before instantiating the model / optimizers, etc (see function in initial post).</p>
<p><strong>Edit</strong>: not sure how I would be able to set the ENV var when using <code>mp.spawn</code> such that it doesn’t apply to both processes.</p>
<p><strong>Edit2</strong>: didn’t realize children can modify their ENV independently of the parent and other processes: <a href="https://stackoverflow.com/questions/24642811/set-env-var-in-python-multiprocessing-process" class="inline-onebox" rel="noopener nofollow ugc">environment variables - set env var in Python multiprocessing.Process - Stack Overflow</a></p>

jramapuram: <p>Using pytotch 1.5, setting <code>CUDA_VISIBLE_DEVICES</code> appropriately per process, doing any CUDA related stuff before the ENV var is set and changing the DDP constructor with:</p>
<pre data-code-wrap="python"><code class="lang-python">device_ids=[0],   # set w/cuda environ var
output_device=0,  # set w/cuda environ var
</code></pre>
<p>it seems to be working as expected:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/2/0/209b87f0f083c3adc9af2ca8a8b8378862f9f4a3.png" alt="image" data-base62-sha1="4EsvaueFI45gYp1s1MbH6tsLUcP" width="571" height="220"></p>

mrshenli: <p>It might be relevant to <a href="https://discuss.pytorch.org/t/distributed-training-creates-multiple-processes-in-gpu0/77881/4">this post</a>. If <code>CUDA_VISIBLE_DEVICES</code> is not set to one device per process, and the application program calls <code>clear_cache</code> somewhere without a device context, it could try to initialize the CUDA context on device 0.</p>