What algorithm does NCCL use to perform distributed training?

username: watalou
cooked: <h2><a name="p-462001-i-have-two-gpu-servers-a-docker-is-deployed-on-each-gpu-server-and-each-docker-uses-one-gpu-card-to-conduct-ddp-and-the-ip-addresses-of-the-two-dockers-are-1010102-and-1010103-respectively-the-script-is-as-follows-1" class="anchor" href="#p-462001-i-have-two-gpu-servers-a-docker-is-deployed-on-each-gpu-server-and-each-docker-uses-one-gpu-card-to-conduct-ddp-and-the-ip-addresses-of-the-two-dockers-are-1010102-and-1010103-respectively-the-script-is-as-follows-1"></a>I have two GPU servers,  a docker is deployed on each GPU server, and each docker uses one GPU card to conduct DDP. And the IP addresses of the two dockers are 10.10.10.2 and 10.10.10.3, respectively. The script is as follows:</h2>
<p>import torch<br>
import torch.distributed as dist<br>
import torch.nn as nn<br>
import torch.optim as optim<br>
from torch.utils.data import DataLoader, DistributedSampler<br>
from torchvision import datasets, transforms<br>
import os</p>
<p>def setup_distributed(local_rank, world_size, backend=‘nccl’):<br>
rank = int(os.environ[‘RANK’])<br>
dist.init_process_group(backend=backend, rank=rank, world_size=world_size)<br>
torch.cuda.set_device(local_rank)</p>
<p>def cleanup():<br>
dist.destroy_process_group()</p>
<p>class SimpleModel(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(SimpleModel, self).<strong>init</strong>()<br>
self.fc = nn.Linear(28 * 28, 10)</p>
<pre><code>def forward(self, x):
    x = x.view(-1, 28 * 28)
    return self.fc(x)
</code></pre>
<p>def train(rank, world_size):<br>
setup_distributed(rank, world_size)</p>
<pre><code># 设置数据集和分布式采样器
transform = transforms.Compose([transforms.ToTensor()])
dataset = datasets.MNIST(root='/workspace/data', train=True, download=False, transform=transform)
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
dataloader = DataLoader(dataset, batch_size=64, sampler=sampler)

model = SimpleModel().cuda(rank)
model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])
criterion = nn.CrossEntropyLoss().cuda(rank)
optimizer = optim.SGD(model.parameters(), lr=0.01)
for epoch in range(1): 
    model.train()
    sampler.set_epoch(epoch)
    for batch_idx, (data, target) in enumerate(dataloader):
        data, target = data.cuda(rank), target.cuda(rank)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

cleanup()
</code></pre>
<h2><a name="p-462001-if-__name__-__main__-local_rank-intosenvironlocal_rank-world_size-intosenvironworld_size-trainlocal_rank-world_size-2" class="anchor" href="#p-462001-if-__name__-__main__-local_rank-intosenvironlocal_rank-world_size-intosenvironworld_size-trainlocal_rank-world_size-2"></a>if <strong>name</strong> == “<strong>main</strong>”:<br>
local_rank = int(os.environ[‘LOCAL_RANK’])<br>
world_size = int(os.environ[‘WORLD_SIZE’])<br>
train(local_rank, world_size)</h2>
<p>On docker A (10.10.10.2), we run:<br>
RANK=0 WORLD_SIZE=2 MASTER_ADDR=10.10.10.2 MASTER_PORT=12345 LOCAL_RANK=0 python script.py &amp;</p>

username: ptrblck
cooked: <p>You can find the internal design of <code>DDP</code> <a href="https://pytorch.org/docs/main/notes/ddp.html#internal-design">here</a> in case you are looking for such docs.</p>