Pytorch NCCL DDP freezes but Gloo Works

username: nyxynyx
cooked: <p>I am trying to get NCCL backend working on my Ubuntu 20.04 system that has two Nvidia 2070S GPUs and runs Pytorch 1.8.</p>
<p>My test script is based on <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html#basic-use-case" rel="noopener nofollow ugc">the Pytorch docs</a>, but with the backend changed from <code>"gloo"</code> to <code>"nccl"</code>.</p>
<p>When the backend is <code>"gloo"</code>, the script finishes running in less than a minute.</p>
<pre><code class="lang-auto">$ time python test_ddp.py 
Running basic DDP example on rank 0.
Running basic DDP example on rank 1.

real    0m4.839s
user    0m4.980s
sys     0m1.942s
</code></pre>
<p>However, when the backend is set to <code>"nccl"</code>, the script gets stuck with the below output and never returns to the bash prompt.</p>
<pre><code class="lang-auto">$ python test_ddp.py 
Running basic DDP example on rank 1.
Running basic DDP example on rank 0.
</code></pre>
<p>Same problem when disabling IB</p>
<pre><code class="lang-auto">$ NCCL_IB_DISABLE=1 python test_ddp.py
Running basic DDP example on rank 1.
Running basic DDP example on rank 0.
</code></pre>
<p>I’m using the packages:</p>
<ul>
<li>pytorch 1.8.1</li>
<li>cudatoolkit 11.1.1</li>
<li>python 3.8.8</li>
</ul>
<p>How can we fix the problem when using NCCL? Thank you!</p>
<p><strong>Python code used for testing NCCL:</strong></p>
<pre><code class="lang-auto">import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

from torch.nn.parallel import DistributedDataParallel as DDP


def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"

    # gloo: works
    # dist.init_process_group("gloo", rank=rank, world_size=world_size)

    # nccl: hangs forever
    dist.init_process_group(
        "nccl", init_method="tcp://10.1.1.20:23456", rank=rank, world_size=world_size
    )


def cleanup():
    dist.destroy_process_group()


class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def demo_basic(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)

    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(rank)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup()


def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn, args=(world_size,), nprocs=world_size, join=True)


if __name__ == "__main__":
    run_demo(demo_basic, 2)
</code></pre>

username: H-Huang
cooked: <p>Thanks for posting your code and details!</p>
<p>I see that when you init the process group with NCCL you specify the init_method as “tcp” and provide a local IP and port. Can you ensure that those are reachable? Alternatively, since you are already setting the address and port as environment variables and it works for gloo, you can remove the “init_method” parameter from init_process_group and it will default to use “env://” and that should work as well.</p>
<pre><code class="lang-auto">dist.init_process_group("nccl", rank=rank, world_size=world_size)
</code></pre>
<p>Here is the documentation for init_process_group: <a href="https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group" class="inline-onebox" rel="noopener nofollow ugc">Distributed communication package - torch.distributed — PyTorch 1.8.1 documentation</a></p>
<p>Please let me know if this works.</p>

username: nyxynyx
cooked: <p>Yes, you are right! I’ve got it running with NCCL by changing <code>setup</code> function as suggested</p>
<pre><code class="lang-auto">def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"

    dist.init_process_group("nccl", rank=rank, world_size=world_size)
</code></pre>

username: nyxynyx
cooked: <p>Yes, you are right! I’ve got it running with NCCL by changing <code>setup</code> function as suggested</p>
<pre><code class="lang-auto">def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"

    dist.init_process_group("nccl", rank=rank, world_size=world_size)
</code></pre>
<p>However, is this now using TCP initialization?</p>

username: H-Huang
cooked: <p>Yes, the underlying communication uses TCP. During initialization it uses the variables you defined for addr, port, rank, and world_size to create TCPStore instances on all workers.<br>
<a href="https://pytorch.org/docs/master/distributed.html#torch.distributed.Store" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/master/distributed.html#torch.distributed.Store</a></p>