Improving inter-process communication costs for `pytorch.distributed`

kernelmethod: <p>Hi there –</p>
<p>TL;DR: I’m running into some performance issues where the inter-process communication cost increases linearly in the number of processes that I spawn, and wanted to know what options are available to improve communication performance.</p>
<hr>
<p>I’m using PyTorch for a distributed learning application that requires me to create a number of process groups that is ~linear in the number of processes that I’ve created, and I have to perform a decent amount of communication within those groups during the forward and backward passes through my model. Unfortunately, it seems like the cost of communication is increasing linearly in the number of processes I’m using for training, which is negating any performance benefits I’d otherwise be getting by scaling out training across multiple CPUs / GPUs.</p>
<p>Here’s a MWE showing how the communication costs are scaling for me:</p>
<pre><code class="lang-python">import datetime
import os
import sys
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

def target(process_group, n_iters = 1024):
    for _ in range(n_iters):
        x = torch.randn(10, 10)
        dist.all_reduce(x, group=process_group)

def _launch(rank, world_size):
    # Initialize process groups for every pair of consecutive processes

    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

    if world_size % 2 != 0:
        raise ValueError("The world size must be a multiple of 2")

    for i in range(0, world_size, 2):
        new_group = dist.new_group([i, i+1])
        if rank == i or rank == i + 1:
            process_group = new_group

    start = datetime.datetime.now()
    target(process_group)
    end = datetime.datetime.now()

    if rank == 0:
        print(f"Time to run target() = {end - start}")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        raise RuntimeError("Usage: ./test.py &lt;n_process_groups&gt;")

    n_process_groups = int(sys.argv[1])
    world_size = 2 * n_process_groups

    mp.spawn(
        _launch,
        args=(world_size,),
        nprocs=world_size,
        join=True,
    )
</code></pre>
<p>Here’s my results when I run the script:</p>
<pre><code class="lang-nohighlight">$ for p in {1..8}; do echo -n "p = $p; " &amp;&amp; python3 test.py $p; done
p = 1; Time to run target() = 0:00:00.386286
p = 2; Time to run target() = 0:00:00.478979
p = 3; Time to run target() = 0:00:00.728055
p = 4; Time to run target() = 0:00:00.920900
p = 5; Time to run target() = 0:00:01.150142
p = 6; Time to run target() = 0:00:01.505277
p = 7; Time to run target() = 0:00:01.937280
p = 8; Time to run target() = 0:00:02.068987
</code></pre>
<p>I’m not surprised that communication performance decreases as the number of processes increases, but I don’t understand why the cost is growing as quickly as it is.</p>
<p>So my question is – are the growing communication costs I’m seeing here just something that’s expected out of the <code>gloo</code> backend or other elements of the PyTorch distributed communication internals? Or is there a way I could improve the script I have written above so that the <code>target()</code> function runs faster?</p>
<p>Any help would be greatly appreciated. Thank you!</p>

kernelmethod: <p>(Incidentally, I’m currently doing all of my training on CPU. Here’s some additional info about the machine I’m currently running this on, if it helps:)</p>
<pre><code class="lang-auto">$ lscpu | grep -E "^(CPU\(s\):|Model name:)"
CPU(s):                48
Model name:            Intel(R) Xeon(R) Silver 4214 CPU @ 2.20GHz
</code></pre>

wayi: <p>Hi! Gloo backend is much slower than NCCL if can do trainings on GPU. You may want to switch to NCCL if possible.</p>
<p>A few relatively easy options on my mind:</p>
<ol>
<li>It will be much more efficient if you use <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" rel="noopener nofollow ugc">DDP</a> instead of implementing your own communication by allreduce primitive.</li>
<li>Try <code>no_sync</code> context manager on the tutorial, which can reduce the sync frequency by accumulating gradients locally.</li>
<li>If you can use NCCL, it will be much faster. You can also try <a href="https://pytorch.org/docs/master/ddp_comm_hooks.html" rel="noopener nofollow ugc">DDP communication hooks</a> for gradient compression.</li>
</ol>

kernelmethod: <p>Thanks for the reply!</p>
<aside class="quote no-group" data-username="wayi" data-post="3" data-topic="115991">
<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/wayi/48/32374_2.png" class="avatar"> wayi:</div>
<blockquote>
<ol>
<li>It will be much more efficient if you use DDP instead of implementing your own communication by allreduce primitive.</li>
<li>Try <code>no_sync</code> context manager on the tutorial, which can reduce the sync frequency by accumulating gradients locally.</li>
</ol>
</blockquote>
</aside>
<p>Unfortunately I can’t currently wrap my models with DDP (I’m integrating with another library that doesn’t play very well with DDP), but my plan is to do so as soon as possible.</p>
<p>That said, the application I’m developing for requires inter-process communication in both the forwards and backwards passes. The total communication is quite a bit higher than just the <a href="https://pytorch.org/docs/stable/notes/ddp.html#internal-design" rel="noopener nofollow ugc">setup communication / gradient synchronization</a> performed by DDP, so I suspect that wrapping my model with DDP won’t do very much to reduce communication costs.</p>
<aside class="quote no-group" data-username="wayi" data-post="3" data-topic="115991">
<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/wayi/48/32374_2.png" class="avatar"> wayi:</div>
<blockquote>
<ol start="3">
<li>If you can use NCCL, it will be much faster. You can also try DDP communication hooks for gradient compression.</li>
</ol>
</blockquote>
</aside>
<p>I’d like to use NCCL, but it doesn’t directly support some of the communication primitives I need. In any case, for now I’m less concerned about raw communication speed and more concerned with how it scales as the number of processes increases.</p>
<p>If I used the NCCL backend, would the communication speed decrease at the same rate as Gloo (so that e.g. doubling the number of processes / process groups causes <code>all_reduce</code> to run at half speed)?</p>

wayi: <p>A few other thoughts:</p>
<p>Do you have a chance to set <code>async_op=True</code> in <code>dist.all_reduce</code>, so that you can overlap some computation with allreduce?</p>
<p>Additionally, probably you can try to tune some Gloo environment variables. I am not familiar with Gloo, but I see tuning some NCCL parameters can help communication performance a lot.</p>
<blockquote>
<p>I’m integrating with another library that doesn’t play very well with DDP</p>
</blockquote>
<p>Also interested in knowing the reasons that cause DDP not feasible for this case.</p>
<blockquote>
<p>If I used the NCCL backend, would the communication speed decrease at the same rate as Gloo (so that e.g. doubling the number of processes / process groups causes <code>all_reduce</code> to run at half speed)?</p>
</blockquote>
<p>NCCL will have both a better scaling efficiency and higher raw speed than Gloo.</p>

kernelmethod: <aside class="quote no-group" data-username="wayi" data-post="5" data-topic="115991">
<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/wayi/48/32374_2.png" class="avatar"> wayi:</div>
<blockquote>
<p>Do you have a chance to set <code>async_op=True</code> in <code>dist.all_reduce</code>, so that you can overlap some computation with allreduce?</p>
</blockquote>
</aside>
<p>I have tried this, yes! It gives a pretty good speedup for the test script (and a smaller speedup for my actual code), but not enough to counteract the growth of communications costs. There may be some other comms operations I could make async, though; I’ll definitely look into it.</p>
<aside class="quote no-group" data-username="wayi" data-post="5" data-topic="115991">
<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/wayi/48/32374_2.png" class="avatar"> wayi:</div>
<blockquote>
<p>Additionally, probably you can try to tune some Gloo environment variables. I am not familiar with Gloo, but I see tuning some NCCL parameters can help communication performance a lot.</p>
</blockquote>
</aside>
<p>I’d definitely like to know if there are any environmental variables I can tune for Gloo. I think I’ll need to dig in some more and see if anything like that exists. <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=12" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:" loading="lazy" width="20" height="20"></p>
<aside class="quote no-group" data-username="wayi" data-post="5" data-topic="115991">
<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/wayi/48/32374_2.png" class="avatar"> wayi:</div>
<blockquote>
<p>Also interested in knowing the reasons that cause DDP not feasible for this case.</p>
</blockquote>
</aside>
<p>The library I’m using on top of PyTorch wraps a model and mostly tries to imitate PyTorch’s API for training. Unfortunately, the only API it currently exposes for updating parameters after backprop is a single function that basically just loops over the model parameters and applies a single step of SGD on them.</p>
<p>I decided after some experimentation that the amount of effort required to either (a) modify the library to be able to wrap DDP or (b) create a modified version of DDP that could wrap the library was nontrivial. I might end up having a flash of inspiration that’ll help me figure out how to do one or the other, but for now it’s not possible to use DDP for my purposes.</p>
<aside class="quote no-group" data-username="wayi" data-post="5" data-topic="115991">
<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/wayi/48/32374_2.png" class="avatar"> wayi:</div>
<blockquote>
<p>NCCL will have both a better scaling efficiency and higher raw speed than Gloo.</p>
</blockquote>
</aside>
<p>Thank you, that’s good to know. <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=12" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:" loading="lazy" width="20" height="20"> If I can’t figure out how to speed up Gloo, I’ll try to see if I can modify my code to use NCCL.</p>