Repository: pytorch
Issue: Evaluating on single GPU (DDP)
Labels: triaged,module: ddp

I am currently trying to get used to DistributedDataParallel.

I trained the network with 4 gpus using DDP, and tried to evaluate with a single gpu, but got a following error:

```
Traceback (most recent call last):
  File "/home/lthilnklover/.conda/envs/torch_1.8/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/lthilnklover/.conda/envs/torch_1.8/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/lthilnklover/sam_torch/parallel_test.py", line 119, in <module>
    mp.spawn(main, args=(world_size, args), nprocs=world_size, join=True)
  File "/home/lthilnklover/.conda/envs/torch_1.8/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 230, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/home/lthilnklover/.conda/envs/torch_1.8/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 188, in start_processes
    while not context.join():
  File "/home/lthilnklover/.conda/envs/torch_1.8/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 150, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException: 

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/home/lthilnklover/.conda/envs/torch_1.8/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 59, in _wrap
    fn(i, *args)
  File "/home/lthilnklover/sam_torch/parallel_test.py", line 101, in main
    accuracy = evaluate(test_loader, model, rank)
  File "/home/lthilnklover/sam_torch/parallel_test.py", line 32, in evaluate
    preds = model(inputs)
  File "/home/lthilnklover/.conda/envs/torch_1.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lthilnklover/.conda/envs/torch_1.8/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 791, in forward
    self._sync_params()
  File "/home/lthilnklover/.conda/envs/torch_1.8/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1318, in _sync_params
    self._distributed_broadcast_coalesced(
  File "/home/lthilnklover/.conda/envs/torch_1.8/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 1278, in _distributed_broadcast_coalesced
    dist._broadcast_coalesced(
RuntimeError: [/pytorch/third_party/gloo/gloo/transport/tcp/pair.cc:598] Connection closed by peer
```

I works fine if I evaluate on all 4 gpus. 


The code I ran is

```
import argparse
import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import CIFAR10
import torchvision
import torch.distributed as dist
import torch.multiprocessing as mp


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12345'

    dist.init_process_group("gloo", rank=rank, world_size=world_size)


def clean_up():
    dist.destroy_process_group()


def evaluate(test_loader, model, rank):
    model.eval()
    accuracy = 0.
    cnt = 0.

    with torch.no_grad():
        for data in test_loader:
            inputs, labels = data[0].to(rank), data[1].to(rank)
            preds = model(inputs)

            accuracy += (torch.argmax(preds, 1) == labels).sum().item()
            cnt += len(labels)
    accuracy *= 100 / cnt
    return accuracy


def main(rank, world_size, args):
    setup(rank, world_size)

    mean = (0.4914, 0.4822, 0.4465)
    std = (0.2023, 0.1994, 0.2010)


    # Transforms
    train_transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4, fill=128),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean, std)])

    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean, std)])

    # DataLoader
    try:
        train_set = CIFAR10(root='./data', train=True, download=False, transform=train_transform)
        test_set = CIFAR10(root='./data', train=False, download=False, transform=test_transform)
    except:
        train_set = CIFAR10(root='./data', train=True, download=True, transform=train_transform)
        test_set = CIFAR10(root='./data', train=False, download=True, transform=test_transform)

    train_sampler = torch.utils.data.distributed.DistributedSampler(train_set)

    train_loader = DataLoader(
        train_set,
        batch_size=args.batch_size, shuffle=(train_sampler is None), num_workers=args.num_workers,
        sampler=train_sampler)
    test_loader = DataLoader(test_set, batch_size=world_size * args.batch_size, shuffle=False,
                             num_workers=args.num_workers)

    model = torchvision.models.resnet18(pretrained=False)
    model = model.to(rank)

    model = nn.parallel.DistributedDataParallel(
        model, device_ids=[rank]
    )

    criterion = torch.nn.CrossEntropyLoss()

    optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate)

    for epoch in range(args.epochs):
        train_sampler.set_epoch(epoch)
        print(f"Rank {rank} Epoch {epoch}")

        model.train()

        for data in train_loader:
            inputs, labels = data[0].to(rank), data[1].to(rank)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

        if rank == 0:
            accuracy = evaluate(test_loader, model, rank)
            print(f"Epoch {epoch} Accuracy {accuracy}")

        dist.barrier()
    clean_up()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--learning_rate", default=0.1, type=float, help="Initial learning rate.")
    parser.add_argument("--batch_size", default=32, type=int, help="Batch size")
    parser.add_argument("--epochs", default=200, type=int, help="Number of epochs.")
    parser.add_argument("--num_workers", default=2, type=int, help="Number of Workers")
    args = parser.parse_args()

    n_gpus = torch.cuda.device_count()
    world_size = n_gpus

    mp.spawn(main, args=(world_size, args), nprocs=world_size, join=True)

```