Using DDP for a large computation graph

hos-b: <p>Hi,<br>
I’ve been trying to train a GNN with pytorch. My sample graphs can have 1-8 nodes.</p>
<p>Unfortunately, the computation graph is too large to fit inside the resources I have. I split up the computation for each node in the graph network using a for loop that computes/propagates the messages, calculates the loss &amp; calls backward. Inside this loop, I call a couple of functions from my model, not <code>forward()</code>. This helps me fit the whole thing inside a Titan with 12GBs of memory but I have 2 and the second one has remained unused.</p>
<p>I couldn’t use <code>DataParallel</code> or <code>DistributedDataParallel</code> because the wrapped models only have the forward function. I could technically fit everything in a single forward function and call it multiple times instead of once, but sample graphs do not have the same number of nodes. e.g. I can’t call <code>losses.backward()</code> 3 times in one process and 5 times in the other.</p>
<p>Is there any other way for me to utilize both GPUs?<br>
I appreciate any tips or advice.</p>

mrshenli: <p>Hey <a class="mention" href="/u/hos-b">@hos-b</a> does model parallelism work in your case? i.e., split your model into two shards and put each shard into one GPU.</p>
<p>If the two GPUs are on the same machine, here is a single-machine model parallel tutorial:</p>
<ul>
<li><a href="https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">Single-Machine Model Parallel Best Practices — PyTorch Tutorials 1.8.0 documentation</a></li>
</ul>
<p>If the two GPUs are on different machines, here are some tutorials on RPC:</p>
<ul>
<li>RPC: <a href="https://pytorch.org/docs/master/rpc.html" class="inline-onebox" rel="noopener nofollow ugc">Distributed RPC Framework — PyTorch master documentation</a>
</li>
<li>tutorials: <a href="https://pytorch.org/tutorials/beginner/dist_overview.html#general-distributed-training" class="inline-onebox" rel="noopener nofollow ugc">PyTorch Distributed Overview — PyTorch Tutorials 1.8.0 documentation</a>
</li>
</ul>

hos-b: <p>Hello <a class="mention" href="/u/mrshenli">@mrshenli</a>. Thank you for your reply.</p>
<p>Model parallelism will probably also not work because even though the torch calls are async, I have to call <code>backward()</code> at the end of the for loop, creating a blocking bottleneck. The computation graph would otherwise get too large for a single GPU.</p>
<p>Your reply made me realize I had silently forgotten about <code>DataParallel</code>. Since my batches contain whole graphs, I somehow forgot I could further divide the graph into batches of size 2. I had stopped thinking about <code>DataParallel</code> when I found out about the advantages of <code>DistributedDataParallel</code>.</p>

mrshenli: <aside class="quote no-group" data-username="hos-b" data-post="3" data-topic="114986">
<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/h/e47774/48.png" class="avatar"> hos-b:</div>
<blockquote>
<p>Model parallelism will probably also not work because even though the torch calls are async, I have to call <code>backward()</code> at the end of the for loop, creating a blocking bottleneck. The computation graph would otherwise get too large for a single GPU.</p>
</blockquote>
</aside>
<p>I see. Have you tried checkpointing some part of the autograd graph? This would drop the the activations and autograd graph in the forward and recompute them again in the backward. It’s like paying more compute to reduce memory footprint.</p>
<p><a href="https://pytorch.org/docs/stable/checkpoint.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/checkpoint.html</a></p>
<blockquote>
<p>I have to call <code>backward()</code> at the end of the for loop</p>
</blockquote>
<p>Curious, is this because you need to compute the global loss across all iterations, so that you cannot run multiple fwd-bwd-fwd-bwd to accumulate grads (DDP can support this with <code>no_sync()</code> context manager) and then run <code>optimizer.step()</code> once to update parameters?</p>

hos-b: <p>I did not know about checkpointing. It looks promising. Thank you</p>
<aside class="quote no-group" data-username="mrshenli" data-post="4" data-topic="114986">
<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>Curious, is this because you need to compute the global loss across all iterations, so that you cannot run multiple fwd-bwd-fwd-bwd to accumulate grads (DDP can support this with <code>no_sync()</code> context manager) and then run <code>optimizer.step()</code> once to update parameters?</p>
</blockquote>
</aside>
<p>Sorry, I didn’t explain my training well. Each graph has up to 8 nodes. A full forward pass for a single node + the detached messages from the other nodes takes about 8GBs of memory. I go through the nodes using a for loop. I have to call <code>backward()</code> at the end of each iteration to calculate the loss, so it does get accumulated at each iteration. I do a single <code>optimizer.step()</code> after the loop to update the parameters. I technically do fwd-bwd x [1 - 8], into <code>step()</code>.</p>
<p>From my understanding of DDP, the gradient reduction across multiple processes starts with the <code>backward()</code> call. That’s why I thought it wouldn’t work if different processes called it different amount of times. I did not know about <code>no_sync()</code> either. I’ll have to read about it. Thanks.</p>

mrshenli: <p>here is the link: <a href="https://pytorch.org/docs/master/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync" class="inline-onebox" rel="noopener nofollow ugc">DistributedDataParallel — PyTorch master documentation</a></p>
<p><code>no_sync</code> basically disables DDP comm within the context.</p>

hos-b: <p>Thanks. It took me some time to adapt the code to distributed training but <code>no_sync()</code> did the trick.</p>