Hi, all
I was trying out a very simple example to use DistributedDataParallel but the code got stuck at data loading for some reason. The code I used is pasted below in its entirety.
import os
import time
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torchvision import datasets, transforms
from torch import nn

class Model(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.fc = nn.Linear(7*7*32, num_classes)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.reshape(x.size(0), -1)
        x = self.fc(x)

        return x

def main(rank, world_size):
    # Initialisation
    dist.init_process_group(
        backend="nccl",
        init_method="env://",
        world_size=world_size,
        rank=rank
    )
    # Fix random seed
    torch.manual_seed(0)
    # Initialize network
    net = Model()
    net.cuda(rank)
    # Initialize loss function
    criterion = torch.nn.CrossEntropyLoss().to(rank)
    optimizer = torch.optim.SGD(net.parameters(), 1e-4) 

    net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[rank])
    # Prepare dataset
    trainset = datasets.MNIST('./data', train=True, download=True,
        transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])
    )
    # Prepare sampler
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        trainset, num_replicas=world_size, rank=rank
    )
    # Prepare dataloader
    train_loader = torch.utils.data.DataLoader(
        trainset, batch_size=100, shuffle=False,
        num_workers=0, pin_memory=True, sampler=train_sampler)
 
    epoch = 0
    iteration = 0

    for _ in range(5):
        epoch += 1
        train_loader.sampler.set_epoch(epoch)

        timestamp = time.time()
        print("Rank: {}. Before dataloader".format(rank))
        for batch in train_loader:
            print("Rank: {}. Batch loaded".format(rank))
            inputs = batch[0]
            targets = batch[1]

            iteration += 1
            inputs = inputs.cuda(rank, non_blocking=True)
            targets = targets.cuda(rank, non_blocking=True)

            output = net(inputs)
            loss = criterion(output, targets)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

if __name__ == '__main__':

    # Number of GPUs to run the experiment with
    WORLD_SIZE = 2

    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "8888"
    mp.spawn(main, nprocs=WORLD_SIZE, args=(WORLD_SIZE,))

As I ran the program, the print message I got looks like this
Rank: 1. Before dataloader
Rank: 0. Before dataloader
Rank: 0. Batch loaded
Rank: 0. Batch loaded
Rank: 0. Batch loaded
Rank: 0. Batch loaded
Rank: 0. Batch loaded
Rank: 0. Batch loaded
Rank: 0. Batch loaded
Rank: 0. Batch loaded
Rank: 0. Batch loaded
Rank: 0. Batch loaded
Rank: 0. Batch loaded
Rank: 0. Batch loaded

And from the GPU utilisation, I noticed that the first GPU (corresponding to subprocess with rank 0) is at its full capacity (100%) while the second one is at 0%. And more interestingly, the second subprocess (with rank 1) also occupied a small amount of memory in the first GPU. I can’t seem to figure out the problem. Please let me know if you spotted anything that might help.
Many thanks,
Fred