Repository: pytorch
Issue: MPI causing job to hang --- unresponsive to external (termination) signals
Labels: oncall: distributed

## Issue description
When running `torch.distributed` with the MPI backend in a cluster, the job fails to exit and hangs unresponsively --- essentially any node in the cluster that the runs the script becomes unresponsive when you try to terminate the job. This varies from run to run, but it happens often enough that it is problematic.

Here is a minimal reproduction script. High level steps:

1.  initialize multiple DataLoader workers
2. communicate model parameters
3. update model parameters
4. run forward-backward pass

## Code example

For a single machine with `$world_size` GPUs, run

`mpirun -n $world_size python issues.py`

issues.py:

```
import os
import signal

import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.utils.data
import torch.utils.data.distributed
import torch.multiprocessing as mp
import torchvision.datasets as datasets
import torchvision.models as models
import torchvision.transforms as transforms

# path to dataset (REPLACE WITH YOUR OWN LOCAL PATH TO DATASET)
TRAIN_DIRECTORY = '/datasets/imagenet_full_size/train'


def main():

    signal.signal(signal.SIGTERM, SIGTERMHandler)

    # initialize torch distributed mpi backend
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '40101'
    dist.init_process_group(backend='mpi')

    # run each task on a single GPU
    torch.cuda.set_device(dist.get_rank())

    # seed for reproducibility
    torch.manual_seed(37)
    torch.cuda.manual_seed(37)
    torch.backends.cudnn.deterministic = True

    model = models.resnet50()
    model.cuda()
    model.train()
    criterion = torch.nn.CrossEntropyLoss().cuda()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9,
                                weight_decay=0.0001, nesterov=True)

    cudnn.benchmark = True

    # initialize a dataloader with 4 worker processes
    loader, sampler = make_dataloader(TRAIN_DIRECTORY)

    for i, (input, target) in enumerate(loader):

        print('itr {}'.format(i))

        input = input.cuda(non_blocking=True)
        target = target.cuda(non_blocking=True)

        # communicate model parameters
        comm_buffer = _flatten_tensors(list(model.parameters())).detach_()
        send_recieve(comm_buffer)
        comm_buffer = _unflatten_tensors(comm_buffer, list(model.parameters()))
        with torch.no_grad():
            for p, e in zip(model.parameters(), comm_buffer):
                p.data.copy_(e)

        # forward/backward pass
        output = model(input)
        loss = criterion(output, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        prec1, prec5 = accuracy(output, target, topk=(1, 5))
        print('* @Prec1 {prec1:.3f}\t@Prec5 {prec5:.3f}'
              .format(prec1=prec1.item(), prec5=prec5.item()))


def send_recieve(buffer):

    buffer.mul_(0.5)

    # ring communication
    destination = (dist.get_rank() + 1) % dist.get_world_size()
    source = dist.get_rank() - 1 if dist.get_rank() != 0 else dist.get_world_size() - 1

    # non-blocking send message
    send_buffer = buffer.clone()
    req = dist.isend(tensor=send_buffer, dst=destination)
    out_msg = (req, send_buffer)  # keep in scope

    # blocking receive
    receive = buffer.clone()
    dist.recv(receive, src=source)

    # update buffer
    buffer.add_(receive)

    # wait for neighbours to receive message
    out_msg[0].wait()


def make_dataloader(train_directory):
""" Create distributed dataloader """
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    train_dataset = datasets.ImageFolder(train_directory, transforms.Compose([
                        transforms.RandomResizedCrop(224),
                        transforms.RandomHorizontalFlip(),
                        transforms.ToTensor(),
                        normalize]))

    # sampler produces indices used to assign data samples to each agent
    train_sampler = torch.utils.data.distributed.DistributedSampler(
                        dataset=train_dataset,
                        num_replicas=dist.get_world_size(),
                        rank=dist.get_rank())

    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=32,
        shuffle=False,
        num_workers=4,
        pin_memory=True, sampler=train_sampler)

    return train_loader, train_sampler


def accuracy(output, target, topk=(1,)):
    """ Computes the precision@k for the specified values of k """
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res


def _flatten_tensors(tensors):
    if len(tensors) == 1:
        return tensors[0].view(-1)
    flat = torch.cat([t.view(-1) for t in tensors], dim=0)
    return flat


def _unflatten_tensors(flat, tensors):
    outputs = []
    offset = 0
    for tensor in tensors:
        numel = tensor.numel()
        outputs.append(flat.narrow(0, offset, numel).view_as(tensor))
        offset += numel
    return tuple(outputs)


def SIGTERMHandler(signum, frame):
    """
    Ignore SIGTERM preemption signal (doesn't stop preemption);
    instead SIGUSR1 will be moved up and handled accordingly
    """
    print('Received SIGTERM')


if __name__ == '__main__':
    mp.set_start_method('forkserver')
    main()

```

## System Info

PyTorch (built from source) version: 0.5.0a0+f8c18e0
Is debug build: No
CUDA used to build PyTorch: 9.0.176

OS: Ubuntu 16.04.4 LTS
GCC version: (Ubuntu 5.5.0-12ubuntu1~16.04) 5.5.0 20171010
CMake version: version 3.11.1

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 9.0.176
GPU models and configuration:
GPU[0-8]: V100 (NVIDIA voltas)

Nvidia driver version: 384.81
cuDNN version: cudnn/v7.0-cuda.9.0

Versions of relevant libraries:

openmpi/3.0.0/gcc.5.4.0
NCCL/2.2.12-1-cuda.9.0

[pip] numpy (1.14.3)
[pip] torch (0.5.0a0+f8c18e0)
[pip] torchvision (0.2.1)
[conda] magma-cuda90   2.3.0       1         pytorch
[conda] torch                     0.5.0   <pip>
[conda] torchvision            0.2.1   <pip>

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd