Repository: pytorch
Issue: Segmentation fault in dataloader worker sub-process
Labels: oncall: distributed,module: crash,module: dataloader

### 🐛 Describe the bug

I have been coming across a segmentation fault coming from the threadpool implementation leveraged by the dataloader for some time (it occurred in version 2.0.1 and it also does in 2.2.0).  In my environment this requires a distributed run to trigger the issue (what I show bellow uses 16 ranks, one per GPU). A simple MNIST training example is sufficient to see this:
```
import os
from datetime import datetime
import argparse
import torch.multiprocessing as mp
import torchvision
import torchvision.transforms as transforms
import torch
import torch.nn as nn
import torch.distributed as dist


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-lrank', '--local_rank', default=0, type=int,
                        help='ranking within the nodes')
    parser.add_argument('--epochs', default=2, type=int, metavar='N',
                        help='number of total epochs to run')
                        
    args = parser.parse_args()
    args.world_size = int(os.environ['WORLD_SIZE'])
    args.rank = int(os.environ['RANK'])
    train(args)


class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.layer1 = 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.layer2 = 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):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out


def train( args):
    rank = args.rank
    gpu = args.local_rank
    dist.init_process_group(backend='nccl', init_method='env://', world_size=args.world_size, rank=rank)
    torch.manual_seed(0)
    model = ConvNet()
    torch.cuda.set_device(gpu)
    model.cuda(gpu)
    batch_size = 8
    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss().cuda(gpu)
    optimizer = torch.optim.SGD(model.parameters(), 1e-4)
    # Wrap the model
    model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu])
    # Data loading code
    train_dataset = torchvision.datasets.MNIST(root=os.getenv('MNIST_DESTINATION_FOLDER','/myrun/datasets'),
                                               train=True,
                                               transform=transforms.ToTensor(),
                                               download=False)
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,
                                                                    num_replicas=args.world_size,
                                                                    rank=rank)
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=batch_size,
                                               shuffle=False,
                                               num_workers=1,
                                               pin_memory=True,
                                               sampler=train_sampler)

    start = datetime.now()
    total_step = len(train_loader)
    for epoch in range(args.epochs):
        for i, (images, labels) in enumerate(train_loader):
            images = images.cuda(non_blocking=True)
            labels = labels.cuda(non_blocking=True)
            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if (i + 1) % 10 == 0 and rank == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, args.epochs, i + 1, total_step,
                                                                         loss.item()))
    if rank == 0:
        print("Training complete in: " + str(datetime.now() - start))


if __name__ == '__main__':
    main()
```
Some threads report:
```
RuntimeError: DataLoader worker (pid(s) 92712) exited unexpectedly
    raise RuntimeError(f'DataLoader worker (pid(s) {pids_str}) exited unexpectedly') from e
```
Looking at the core files this is the stack trace (Pytorch 2.2.0):
```
Core file information:                                                                                   
state: D zombie: True niceness: 0                                                                        
pid: 1577 ppid: 1356 sid: 964                     
uid: 10015212 gid: 10015212 pgrp: 970               
executable: python arguments: python -u ../mnist.py --local_rank 0 --epochs 5                            
                                                                                                                                                                                                                  
The process died due receiving signal SIGSEGV
Traceback for thread 1577 [Has the GIL] (most recent call last):                                                                                                                                                  
    (C) File "???", line 0, in _start (python)                                                                                                                                                                    
    (C) File "???", line 0, in __libc_start_main (libc.so.6)  
    (C) File "???", line 0, in _start (python)                                                                                                                                                            [8/1973]
    (C) File "???", line 0, in __libc_start_main (libc.so.6)                                                                                                                                                      
    (C) File "/usr/local/src/conda/python-3.10.13/Modules/main.c", line 1090, in Py_BytesMain (python)   
    (C) File "/usr/local/src/conda/python-3.10.13/Modules/main.c", line 670, in Py_RunMain (python)                                                                                                               
    (C) File "/usr/local/src/conda/python-3.10.13/Modules/main.c", line 591, in pymain_run_python (inlined) (python)                                                                                              
    (C) File "/usr/local/src/conda/python-3.10.13/Modules/main.c", line 376, in pymain_run_file (inlined) (python)        
    (C) File "/usr/local/src/conda/python-3.10.13/Modules/main.c", line 357, in pymain_run_file_obj (inlined) (python)                                                                                            
    (C) File "/usr/local/src/conda/python-3.10.13/Python/pythonrun.c", line 1208, in pyrun_file.cold (python)                                       
    (C) File "/usr/local/src/conda/python-3.10.13/Python/pythonrun.c", line 1312, in run_mod (python)                                                                                                             
    (C) File "/usr/local/src/conda/python-3.10.13/Python/pythonrun.c", line 1291, in run_eval_code_obj (python)           
    (Python) File "/myrun/folder-0/../mnist.py", line 103, in <module>                                   
        main()                                                                                                                                                                                                    
    (Python) File "/myrun/folder-0/../mnist.py", line 24, in main                                        
        train(args)                                                                                                                                                                                               
    (Python) File "/myrun/folder-0/../mnist.py", line 82, in train                                       
        for i, (images, labels) in enumerate(train_loader):                                                                                                                                                       
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/clinic/enumobject.c.h", line 46, in enum_new (python)                                                                                                   
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/enumobject.c", line 70, in enum_new_impl (inlined) (python)         
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/abstract.c", line 2813, in PyObject_GetIter (python)                                                                                                    
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/typeobject.c", line 7650, in slot_tp_iter (python)                                                                                                      
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/typeobject.c", line 1636, in call_unbound_noarg (inlined) (python)                                                                                      
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 439, in __iter__                                                                               
        return self._get_iterator()                                                                      
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 387, in _get_iterator                                                                          
        return _MultiProcessingDataLoaderIter(self)                                                      
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/typeobject.c", line 7734, in slot_tp_init (python)                                                                                                      
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1040, in __init__                                                                              
        w.start()                                                                                        
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/process.py", line 121, in start                                                                                                    
        self._popen = self._Popen(self)          
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/context.py", line 224, in _Popen                                                                                                   
        return _default_context.get_context().Process._Popen(process_obj)                                
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/context.py", line 281, in _Popen                                                                                                   
        return Popen(process_obj)                                                                                                                                                                                 
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/typeobject.c", line 7734, in slot_tp_init (python)                                                                                                      
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__                                                                                               
        self._launch(process_obj)
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/popen_fork.py", line 71, in _launch                                                                                                
        code = process_obj._bootstrap(parent_sentinel=child_r)                                           
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap                                                                                               
        self.run()                                                                                                                                                                                                
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/process.py", line 108, in run                                 
        self._target(*self._args, **self._kwargs)                                                        
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 222, in _worker_loop                                      
        torch.set_num_threads(1)                                                                         
    (C) File "???", line 0, in THPModule_setNumThreads(_object*, _object*) (libtorch_python.so)
    (C) File "???", line 0, in at::set_num_threads(int) (libtorch_cpu.so)                                
    (C) File "???", line 0, in caffe2::PThreadPool::set_thread_count(unsigned long) (libtorch_cpu.so)    
    (C) File "???", line 0, in pthreadpool_destroy (libtorch_cpu.so)                                     
    (C) File "???", line 0, in syscall (libc.so.6)
```
I happen to have a version 2.0.1 built with debug symbols handy and this is the stack trace for that one:
```
Using executable found in the core file: /opt/miniconda3/envs/pytorch/bin/python

Core file information:
state: D zombie: True niceness: 0
pid: 65559 ppid: 65276 sid: 64903
uid: 10015212 gid: 10015212 pgrp: 64909
executable: python arguments: python -u ../mnist.py --local_rank 0 --epochs 5 

The process died due receiving signal SIGSEGV
Traceback for thread 65559 (systemd) [Has the GIL] (most recent call last):
    (C) File "???", line 0, in _start (python)
    (C) File "???", line 0, in __libc_start_main (libc.so.6)
    (C) File "/usr/local/src/conda/python-3.10.13/Modules/main.c", line 1090, in Py_BytesMain (python)
    (C) File "/usr/local/src/conda/python-3.10.13/Modules/main.c", line 670, in Py_RunMain (python)
    (C) File "/usr/local/src/conda/python-3.10.13/Modules/main.c", line 591, in pymain_run_python (inlined) (python)
    (C) File "/usr/local/src/conda/python-3.10.13/Modules/main.c", line 376, in pymain_run_file (inlined) (python)
    (C) File "/usr/local/src/conda/python-3.10.13/Modules/main.c", line 357, in pymain_run_file_obj (inlined) (python)
    (C) File "/usr/local/src/conda/python-3.10.13/Python/pythonrun.c", line 1208, in pyrun_file.cold (python)
    (C) File "/usr/local/src/conda/python-3.10.13/Python/pythonrun.c", line 1312, in run_mod (python)
    (C) File "/usr/local/src/conda/python-3.10.13/Python/pythonrun.c", line 1291, in run_eval_code_obj (python)
    (Python) File "/myrun/folder-0/../mnist.py", line 102, in <module>
        main()
    (Python) File "/myrun/folder-0/../mnist.py", line 24, in main
        train(args)
    (Python) File "/myrun/folder-0/../mnist.py", line 82, in train
        for i, (images, labels) in enumerate(train_loader):
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/clinic/enumobject.c.h", line 46, in enum_new (python)
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/enumobject.c", line 70, in enum_new_impl (inlined) (python)
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/abstract.c", line 2813, in PyObject_GetIter (python)
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/typeobject.c", line 7650, in slot_tp_iter (python)
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/typeobject.c", line 1636, in call_unbound_noarg (inlined) (python)
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 441, in __iter__
        return self._get_iterator()
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 388, in _get_iterator
        return _MultiProcessingDataLoaderIter(self)
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/typeobject.c", line 7734, in slot_tp_init (python)
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1042, in __init__
        w.start()
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/process.py", line 121, in start
        self._popen = self._Popen(self)
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/context.py", line 224, in _Popen
        return _default_context.get_context().Process._Popen(process_obj)
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/context.py", line 281, in _Popen
        return Popen(process_obj)
    (C) File "/usr/local/src/conda/python-3.10.13/Objects/typeobject.c", line 7734, in slot_tp_init (python)
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__
        self._launch(process_obj)
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/popen_fork.py", line 71, in _launch
        code = process_obj._bootstrap(parent_sentinel=child_r)
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap
        self.run()
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/multiprocessing/process.py", line 108, in run
        self._target(*self._args, **self._kwargs)
    (Python) File "/opt/miniconda3/envs/pytorch/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 222, in _worker_loop
        torch.set_num_threads(1)
    (C) File "/opt/pytorch-build/torch/csrc/Module.cpp", line 230, in THPModule_setNumThreads(_object*, _object*) (libtorch_python.so)
    (C) File "/opt/pytorch-build/aten/src/ATen/ParallelOpenMP.cpp", line 66, in at::set_num_threads(int) (libtorch_cpu.so)
    (C) File "/opt/pytorch-build/caffe2/utils/threadpool/pthreadpool-cpp.cc", line 46, in caffe2::PThreadPool::set_thread_count(unsigned long) (libtorch_cpu.so)
    (C) File "/usr/include/c++/10/bits/unique_ptr.h", line 456, in std::unique_ptr<pthreadpool, void (*)(pthreadpool*)>::reset(pthreadpool*) (inlined) (libtorch_cpu.so)
    (C) File "/usr/include/c++/10/bits/unique_ptr.h", line 182, in std::__uniq_ptr_impl<pthreadpool, void (*)(pthreadpool*)>::reset(pthreadpool*) (inlined) (libtorch_cpu.so)
    (C) File "/opt/pytorch-build/third_party/pthreadpool/src/pthreads.c", line 444, in pthreadpool_destroy (libtorch_cpu.so)
    (C) File "???", line 0, in __pthread_clockjoin_ex (libpthread.so.0)
```
It looks like the threadpool becomes rubbish after the fork or maybe a race. This has some resemblance with #54752 so not sure if the thread pool could in some cases be used prior to being leaked.

Changing the multiprocessing method to `mp.set_start_method('spawn')` has been my workaround. Setting `num_workers=0` also result in a successful run.

### Versions

```
PyTorch version: 2.2.0+rocm5.6                  
Is debug build: False                            
CUDA used to build PyTorch: N/A                 
ROCM used to build PyTorch: 5.6.31061-8c743ae5d 
                                                    
OS: SUSE Linux Enterprise Server 15 SP4 (x86_64)
GCC version: Could not collect                  
Clang version: 16.0.0 (https://github.com/RadeonOpenCompute/llvm-project roc-5.6.1 23332 4f9bb99d78a4d8d9770be38b91ebd004ea4d2a3a)
CMake version: version 3.26.4                   
Libc version: glibc-2.31                                                                                 
                                                                                                         
Python version: 3.10.13 (main, Sep 11 2023, 13:44:35) [GCC 11.2.0] (64-bit runtime)                     
Python platform: Linux-5.14.21-150400.24.81_12.0.75-cray_shasta_c-x86_64-with-glibc2.31                                                                                                                           
Is CUDA available: True                         
CUDA runtime version: Could not collect         
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: AMD Instinct MI250X (gfx90a:sramecc+:xnack-)
Nvidia driver version: Could not collect
cuDNN version: Could not collect 
HIP runtime version: 5.6.31061
MIOpen runtime version: 2.20.0  
Is XNNPACK available: True 
                                                    
CPU:                              
Architecture:                       x86_64                                                               
CPU op-mode(s):                     32-bit, 64-bit                                                       
Address sizes:                      48 bits physical, 48 bits virtual    
Byte Order:                         Little Endian                                                        
CPU(s):                             128                                                                  
On-line CPU(s) list:                0-127                                                                
Vendor ID:                          AuthenticAMD                                                         
Model name:                         AMD EPYC 7A53 64-Core Processor
CPU family:                         25
Model:                              48
Thread(s) per core:                 2
Core(s) per socket:                 64
Socket(s):                          1
Stepping:                           1
Frequency boost:                    enabled
CPU max MHz:                        2000.0000
CPU min MHz:                        1500.0000
BogoMIPS:                           3992.25
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl n
onstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse
 3dnowprefetch osvw ibs wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm
 rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save t
sc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca fsrm
Virtualization:                     AMD-V
L1d cache:                          2 MiB (64 instances)
L1i cache:                          2 MiB (64 instances)
L2 cache:                           32 MiB (64 instances)
L3 cache:                           256 MiB (8 instances)
NUMA node(s):                       4
NUMA node0 CPU(s):                  0-15,64-79
NUMA node1 CPU(s):                  16-31,80-95
NUMA node2 CPU(s):                  32-47,96-111
NUMA node3 CPU(s):                  48-63,112-127
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Mitigation; safe RET, no microcode
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.3
[pip3] pytorch-triton-rocm==2.2.0
[pip3] torch==2.2.0+rocm5.6
[pip3] torchaudio==2.2.0+rocm5.6
[pip3] torchdata==0.7.1+cpu
[pip3] torchtext==0.17.0+cpu
[pip3] torchvision==0.17.0+rocm5.6
[conda] numpy                     1.26.3                   pypi_0    pypi
[conda] pytorch-triton-rocm       2.2.0                    pypi_0    pypi
[conda] torch                     2.2.0+rocm5.6            pypi_0    pypi
[conda] torchaudio                2.2.0+rocm5.6            pypi_0    pypi
[conda] torchdata                 0.7.1+cpu                pypi_0    pypi
[conda] torchtext                 0.17.0+cpu               pypi_0    pypi
[conda] torchvision               0.17.0+rocm5.6           pypi_0    pypi
```

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @SsnL @VitalyFedyunin @ejguan @dzhulgakov