Repository: pytorch
Issue: Native code assert or hang on process exit with a torch.jit.script model on multi-GPU
Labels: triage review,oncall: jit

## 🐛 Bug

Running a multi-GPU module with TorchScript on multiple GPUs can hang on process exit or trigger a native side assertion.

## To Reproduce

Steps to reproduce the behavior:

Run the below code sample on a V100 based DGX-1 system with multiple GPUs.  This may trigger also on other types of hardware, but a repro does seem to require multiple GPUs.

Note: the code has comments WORKS/HANGS/CRASHES.  The repro can be modified in a few different ways to make the bug go away.  Please see the code comments.  The currently copy&pasted version always either hung or crashed for me.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

```
import torch
#----------------------------------------------------------------------------
class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.img_mean = torch.nn.Parameter(torch.zeros([3, 1, 1]))
        self.img_std = torch.nn.Parameter(torch.zeros([3, 1, 1]))

    def forward(self, x):
        batch_size, channels, height, width = x.shape # [NCHW]
        # Emulate tf.image.resize_bilinear(x, [224, 224]), including the funky alignment.
        new_width, new_height = 224, 224
        #theta = [[1.0, 0.0, 1.0 / width - 1.0 / new_width], [0.0, 1.0, 1.0 / height - 1.0 / new_height]]
        #theta = torch.as_tensor(theta, dtype=x.dtype, device=x.device).unsqueeze(0).repeat([batch_size, 1, 1])
        theta = torch.eye(2, 3, device=x.device)
        theta[0, 2] += theta[0, 0] / width - theta[0, 0] / new_width
        theta[1, 2] += theta[1, 1] / height - theta[1, 1] / new_height
        theta = theta.to(x.dtype).unsqueeze(0).repeat([batch_size, 1, 1])
        grid = torch.nn.functional.affine_grid(theta, [batch_size, channels, new_height, new_width], align_corners=False)
        x = torch.nn.functional.grid_sample(x, grid, mode='bilinear', padding_mode='border', align_corners=False)

        # Scale dynamic range from [0,255] to [0,1].
        x /= 255

        # Normalize according to training data statistics.
        x = x - self.img_mean; x = x / self.img_std # HANGS/CRASHES
        #x -= self.img_mean; x /= self.img_std # WORKS
        #x = x.contiguous() - self.img_mean; x = x.contiguous() / self.img_std # WORKS
        return x

#----------------------------------------------------------------------------

def process_fn(rank):
    #module = MyModule() # WORKS
    module = torch.jit.script(MyModule()) # HANGS/CRASHES
    device = torch.device('cuda', rank)
    module = module.to(device)
    #if rank == 0: print(module.code)
    if rank == 0: print('rank 0 start')
    for _ in range(1000): module(torch.zeros([64, 3, 224, 224], device=device))
    if rank == 0: print('rank 0 end')

#----------------------------------------------------------------------------

def run():
    torch.multiprocessing.set_start_method('spawn')
    # nprocs 2, 4, 8 fail, no failure with nprocs=1
    torch.multiprocessing.spawn(fn=process_fn, args=(), nprocs=8)
```

Here's a couple of different stack traces that trigger with the repro case:

```
python: ../nptl/pthread_mutex_lock.c:81: __pthread_mutex_lock: Assertion `mutex->__data.__owner == 0' failed.
Traceback (most recent call last):
  File "/mnt/nvr_results/jhellsten/gan/debug/00027-torch_jit_multiprocess_hang-hel-8gpu/src/dnnlib/submission/submit.py", line 338, in run_wrapper
    run_func_obj(**submit_config.run_func_kwargs)
  File "/mnt/nvr_results/jhellsten/gan/debug/00027-torch_jit_multiprocess_hang-hel-8gpu/src/torch_jit_multiprocess_hang.py", line 54, in run
    torch.multiprocessing.spawn(fn=process_fn, args=(), nprocs=dnnlib.submit_config.num_gpus)
  File "/opt/conda/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 247, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/opt/conda/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 205, in start_processes
    while not context.join():
  File "/opt/conda/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 146, in join
    raise ProcessExitedException(
torch.multiprocessing.spawn.ProcessExitedException: process 7 terminated with signal SIGABRT
```

Another:

```
python: tpp.c:82: __pthread_tpp_change_priority: Assertion `new_prio == -1 || (new_prio >= fifo_min_prio && new_prio <= fifo_max_prio)' failed.
Traceback (most recent call last):
  File "/mnt/debug/00028-torch_jit_multiprocess_hang-hel-8gpu/src/dnnlib/submission/submit.py", line 338, in run_wrapper
    run_func_obj(**submit_config.run_func_kwargs)
  File "/mnt/debug/00028-torch_jit_multiprocess_hang-hel-8gpu/src/torch_jit_multiprocess_hang.py", line 54, in run
    torch.multiprocessing.spawn(fn=process_fn, args=(), nprocs=dnnlib.submit_config.num_gpus)
  File "/opt/conda/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 247, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/opt/conda/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 205, in start_processes
    while not context.join():
  File "/opt/conda/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 146, in join
    raise ProcessExitedException(
torch.multiprocessing.spawn.ProcessExitedException: process 2 terminated with signal SIGABRT
```


## Expected behavior

I'd expect the multi-GPU case to run to completion and exit (including a clean exit of all the subprocesses).

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```

```
WARNING: Your kernel does not support swap limit capabilities or the cgroup is not mounted. Memory limited without swap.

=============
== PyTorch ==
=============

NVIDIA Release 20.12 (build 17950526)
PyTorch Version 1.8.0a0+1606899

Container image Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.

Copyright (c) 2014-2020 Facebook Inc.
Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
Copyright (c) 2011-2013 NYU                      (Clement Farabet)
Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
Copyright (c) 2015      Google Inc.
Copyright (c) 2015      Yangqing Jia
Copyright (c) 2013-2016 The Caffe contributors
All rights reserved.

NVIDIA Deep Learning Profiler (dlprof) Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.

Various files include modifications (c) NVIDIA CORPORATION.  All rights reserved.
NVIDIA modifications are covered by the license terms that apply to the underlying project or file.

NOTE: Legacy NVIDIA Driver detected.  Compatibility mode ENABLED.

dnnlib: Running collect_env.farmmain() on fi-dgx01...
Collecting environment information...
PyTorch version: 1.8.0a0+1606899
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.1 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.18.2

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-16GB
GPU 1: Tesla V100-SXM2-16GB
GPU 2: Tesla V100-SXM2-16GB
GPU 3: Tesla V100-SXM2-16GB
GPU 4: Tesla V100-SXM2-16GB
GPU 5: Tesla V100-SXM2-16GB
GPU 6: Tesla V100-SXM2-16GB
GPU 7: Tesla V100-SXM2-16GB

Nvidia driver version: 450.80.02
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] pytorch-transformers==1.1.0
[pip3] torch==1.8.0a0+1606899
[pip3] torchtext==0.9.0a0
[pip3] torchvision==0.9.0a0
[conda] magma-cuda110             2.5.2                         5    local
[conda] mkl                       2019.4                      243  
[conda] mkl-include               2019.4                      243  
[conda] nomkl                     3.0                           0  
[conda] numpy                     1.19.2           py38h6163131_0  
[conda] numpy-base                1.19.2           py38h75fe3a5_0  
[conda] pytorch-transformers      1.1.0                    pypi_0    pypi
[conda] torch                     1.8.0a0+1606899          pypi_0    pypi
[conda] torchtext                 0.9.0a0                  pypi_0    pypi
[conda] torchvision               0.9.0a0                  pypi_0    pypi
dnnlib: Finished collect_env.farmmain() in 5s.
```

## Additional context

I tried on multiple pytorch container images.  The one used in this bug report should be available publicly.  This was with a few different pytorch versions but can confirm that I got the repro with also the standard pytorch 1.7.1 version when basing the install on cuda-runtime container images directly.

<!-- Add any other context about the problem here. -->


cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @gmagogsfm