Send/Recv is slower in NCCL than in Gloo

username: Yuki_Takezawa
cooked: <p>Hi!<br>
I tried to replace gloo with NCCL in my code because PyTorch 1.11 starts to support send/recv in NCCL.</p>
<p>Because the <a href="https://pytorch.org/docs/stable/distributed.html" rel="noopener nofollow ugc">official document</a> recommends using NCCL for GPU training, I expected that NCCL is faster than Gloo. However, I found that send/recv is slower in NCCL than in Gloo.</p>
<p>Is this result to be expected?<br>
Moreover, is there any way to make send/recv for NCCL faster?</p>
<p>Code:</p>
<pre><code class="lang-auto">import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import time

def run(rank, size, backend):
    total_time = 0.
    
    for _ in range(10):
        if rank == 0:
            tensor = torch.randn((100,100), device="cuda:1")
            time.sleep(5) # wait untill process 1 initializes the variable 'tensor'.
            
            t1_time = time.time()
            if backend == "nccl":
                dist.send(tensor=tensor, dst=1)
            else:
                dist.send(tensor=tensor.to("cpu"), dst=1)
                
            total_time += time.time() - t1_time
        else:
            if backend == "nccl":
                tensor = torch.randn((100,100), device="cuda:2")
                dist.recv(tensor=tensor, src=0)
            else:
                tensor = torch.randn((100,100), device="cpu")
                dist.recv(tensor=tensor, src=0)
                
    if rank == 0:
        print(f"{backend}: {total_time} sec")
        
def init_process(rank, size, fn, backend='nccl'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29501'
    dist.init_process_group(backend, rank=rank, world_size=size)
    fn(rank, size, backend)


if __name__ == "__main__":
    mp.set_start_method("spawn")
    
    for backend in ["nccl", "gloo"]:
        size = 2
        processes = []
        for rank in range(size):
            p = mp.Process(target=init_process, args=(rank, size, run, backend))
            p.start()
            processes.append(p)

        for p in processes:
            p.join()

</code></pre>
<p>Results:</p>
<pre><code class="lang-auto">nccl: 3.5865402221679688 sec
gloo: 0.00415349006652832 sec
</code></pre>
<p>Enviroment:</p>
<ul>
<li>Python 3.10.4</li>
<li>cuda 11.4</li>
<li>PyTorch 1.11.0</li>
<li>GPUs: NVIDIA Quadro A6000 * 2</li>
</ul>
<p>Best Regard!</p>

username: kumpera
cooked: <p>NCCL has an extremely slow warmup compared to GLOO. It does lazily so you will notice the slowness on the first call.</p>
<p>I can reproduce your results locally and the vast majority of the time is spent on the first call to send, after that performance between the two is what you’d expected.</p>

username: Maxwell_Albert
cooked: <p>Also, he did not perform <code>torch.synchronize</code> when testing the time. cuda <code>dist.send</code> and <code>dist.recv</code> should use that.</p>

username: kumpera
cooked: <p>Here’s an updated example that does warmup and incorporates Maxwell’s point on waiting for cuda ops.</p>
<p>Note that while waiting for CUDA ops makes sense when benchmarking, it’s usually not necessary otherwise.</p>
<pre><code class="lang-python">import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import time

def run(rank, size, backend):

    def run_once():
        if rank == 0:
            tensor = torch.randn((100,100), device="cuda:0")
            
            t1_time = time.time()
            if backend == "nccl":
                dist.send(tensor=tensor, dst=1)
                # Wait for all data to be sent
                torch.cuda.synchronize()
            else:
                dist.send(tensor=tensor.to("cpu"), dst=1)
                
            return time.time() - t1_time
        else:
            if backend == "nccl":
                tensor = torch.randn((100,100), device="cuda:1")
                dist.recv(tensor=tensor, src=0)
            else:
                tensor = torch.randn((100,100), device="cpu")
                dist.recv(tensor=tensor, src=0)
            return 0

    # execute a few rounds of warmup
    warmup_time = 0.
    for _ in range(2):
        warmup_time += run_once()
    # measure runtime
    benchmark_time = 0.
    for _ in range(10):
        benchmark_time += run_once()
 
    if rank == 0:
        print(f"{backend}: warmup: {warmup_time} sec benchmark time: {benchmark_time} sec")
        
def init_process(rank, size, fn, backend='nccl'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29501'
    dist.init_process_group(backend, rank=rank, world_size=size)
    fn(rank, size, backend)


if __name__ == "__main__":
    mp.set_start_method("spawn")
    
    for backend in ["nccl", "gloo"]:
        size = 2
        processes = []
        for rank in range(size):
            p = mp.Process(target=init_process, args=(rank, size, run, backend))
            p.start()
            processes.append(p)

        for p in processes:
            p.join()
</code></pre>
<p>On my system I get the following output:</p>
<pre><code class="lang-auto">nccl: warmup: 1.2995800971984863 sec benchmark time: 0.00042319297790527344 sec
gloo: warmup: 0.0008029937744140625 sec benchmark time: 0.0012202262878417969 sec
</code></pre>

username: Yuki_Takezawa
cooked: <p>Thank you for your replies.<br>
I can reproduce the results in my environment and confirm that send/recv in NCCL is faster than in GLOO except for the first communication.</p>