PyTorch Distributed is going out of CPU RAM

username: maralm
cooked: <p>I have two separate models in my algorithm. A large model that resides on CPU and a small model that goes to GPU. I am using the DDP to train the small model on multiple GPUs while the large model remains on CPU. I have 4 GPUs and I observe that the CPU model also loads 4 times and cause OOM in CPU. Is there anyway to keep a sigle CPU model and multiple GPU models in DDP?</p>

username: mrshenli
cooked: <p>Do you mean you want to share the same CPU model across 4 DDP processes? If so, you can use <a href="https://pytorch.org/docs/stable/notes/multiprocessing.html#reuse-buffers-passed-through-a-queue" rel="nofollow noopener">torch.multiprocessing.Queue</a> to share the model. Does the CPU model participate in training or just inference?</p>

username: maralm
cooked: <p>sorry for late reply, missed your comment. Yes, I want to use DDP but have a single copy in the CPU. my model runs forward and backward on GPU and optimizer on CPU.</p>

username: mrshenli
cooked: <blockquote>
<p>Yes, I want to use DDP but have a single copy in the CPU.</p>
</blockquote>
<p>In this case, you might need to use <a href="https://pytorch.org/docs/stable/multiprocessing.html#sharing-cuda-tensors" rel="noopener nofollow ugc">multiprocess communication</a> to share tensors.</p>
<p>Sorry that I still don’t fully understand the use case. Some pseudo code would be helpful.</p>

username: maralm
cooked: <p>Thank you. I am trying to load the cpu model in one thread and broadcast in gpu. However, I have an issue with the barrier. Here is my code:</p>
<pre><code class="lang-auto">if args.local_rank in [-1,0]:
    for net1,net2 in zip(self.encoder.layer[0].named_parameters(), model_cpu.bert.encoder.layer[i].named_parameters()):
         net1[1].data.copy_(net2[1].data.clone(), non_blocking=True)
 
torch.distributed.barrier()

if args.local_rank not in [-1,0]:
    for name, p in self.encoder.layer[0].named_parameters():
        torch.distributed.broadcast(p, src=0)
</code></pre>
<p>my code stucks at the barrier. Any idea what could be wrong with my code? Should I even use barrier?</p>

username: mrshenli
cooked: <aside class="quote no-group" data-username="maralm" data-post="5" data-topic="76261">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/e36b37/48.png" class="avatar"> maralm:</div>
<blockquote>
<p>I am trying to load the cpu model in one thread and broadcast in gpu.</p>
</blockquote>
</aside>
<p>The default process group is per-process object, instead of per-thread. Is this just a typo and you actually mean “process”?</p>
<pre data-code-wrap="python"><code class="lang-python">if args.local_rank not in [-1,0]:
    for name, p in self.encoder.layer[0].named_parameters():
        torch.distributed.broadcast(p, src=0)
</code></pre>
<p>For collective communications, it requires all ranks to make the same number of <code>c10d</code> API invocations in the same order. It seems that, with the above code, rank 0 is not participating in the broadcast? If you need broadcast in a subgroup, you will need to first create a subgroup using the <a href="https://pytorch.org/docs/stable/distributed.html#torch.distributed.new_group" rel="noopener nofollow ugc"><code>new_group</code></a> API and then call <code>boradcast</code> in that group.</p>

username: maralm
cooked: <p>Sorry, yes I mean process. I was thinking to load the model on cpu in one process, transfer it to the gpu 0 (which is in the same process) and from gpu 0, copy weights to other gpus in other processes. Broadcasting is the correct way to do so?</p>

username: mrshenli
cooked: <blockquote>
<p>I was thinking to load the model on cpu in one process, transfer it to the gpu 0 (which is in the same process) and from gpu 0, copy weights to other gpus in other processes. Broadcasting is the correct way to do so?</p>
</blockquote>
<p>Yes, this looks correct to me. <a href="https://github.com/pytorch/pytorch/blob/96b512be07ce9789cfe3daa0b1c97e5641861715/torch/nn/parallel/distributed.py#L280-L285" rel="noopener nofollow ugc"><code>DistributedDataParallel</code></a> is actually doing the similar thing in its constructor.</p>

username: maralm
cooked: <p>Thank you, it worked!</p>