Repository: pytorch
Issue: dataloader hangs with fork and pin memory=True
Labels: module: dataloader,triaged

### 🚀 The feature, motivation and pitch

I train my model on DDP when num_workers>0, pin_memory=True and multiple dataloaders(train, valid, test). There is a small chance to hang that dataloader can not load data.  I set multiprocessing_context='spawn' or close pin memory or just use one dataloader can solve the hang. So I guess fork and pin_memory may cause potential deadlock.

I have spent some time to debug it.
I use this network, and you can just see how it construct three dataloader iterator.
https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/training.py#L1512 

I attach the hanged num_worker subprocess, it hangs on torch.Generator in _worker_loop. It waits for a lock, and
it is in test_dataloader_iterator's num_worker start function. 
```
#0  futex_wait (private=0, expected=2, futex_word=0x7f70128d6ee8 <__exit_funcs_lock>) at ../sysdeps/nptl/futex-internal.h:146
#1  __GI___lll_lock_wait_private (futex=futex@entry=0x7f70128d6ee8 <__exit_funcs_lock>) at ./nptl/lowlevellock.c:34
#2  0x00007f7012700958 in __internal_atexit (listp=0x7f70128d5838 <__exit_funcs>, d=<optimized out>, arg=<optimized out>, func=0x7f7010a35710 <torch::PythonArgParser::~PythonArgParser()>)
    at ./stdlib/cxa_atexit.c:43
#3  __GI___cxa_atexit (func=0x7f7010a35710 <torch::PythonArgParser::~PythonArgParser()>, arg=<optimized out>, d=<optimized out>) at ./stdlib/cxa_atexit.c:70
#4  0x00007f7010c211ec in THPGenerator_pynew(_typeobject*, _object*, _object*) () from /torch/venv3/pytorch/lib/python3.10/site-packages/torch/lib/libtorch_python.so
#5  0x00007f7012a524c5 in tp_new_wrapper (self=<type at remote 0x7f7011561e00>, args=(None,), kwds=0x0) at Objects/typeobject.c:7017
#6  0x00007f7012a1c952 in cfunction_call (func=<built-in method __new__ of type object at remote 0x7f7011561e00>, args=<optimized out>, kwargs=<optimized out>) at Objects/methodobject.c:543
#7  0x00007f7012a19b48 in _PyObject_MakeTpCall (tstate=0x5584c0f5e560, callable=<built-in method __new__ of type object at remote 0x7f7011561e00>, args=0x7f6fc82386d8, nargs=2, keywords=0x0)
    at Objects/call.c:215
```
```
Traceback (most recent call first):
  File "/torch/venv3/pytorch/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 235, in _worker_loop
    shared_rng = torch.Generator()
  File "/opt/py3.10/lib/python3.10/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/opt/py3.10/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/opt/py3.10/lib/python3.10/multiprocessing/popen_fork.py", line 71, in _launch
    code = process_obj._bootstrap(parent_sentinel=child_r)
  File "/opt/py3.10/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/opt/py3.10/lib/python3.10/multiprocessing/context.py", line 281, in _Popen
    return Popen(process_obj)
  File "/opt/py3.10/lib/python3.10/multiprocessing/context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "/opt/py3.10/lib/python3.10/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/torch/venv3/pytorch/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1047, in __init__
    w.start()
  File "/torch/venv3/pytorch/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 386, in _get_iterator
    return _MultiProcessingDataLoaderIter(self)
  File "/torch/venv3/pytorch/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 433, in __iter__
    self._iterator = self._get_iterator()
  <built-in method iter of module object at remote 0x7f7011f80900>
  File "/workspace/Megatron-LM/megatron/training/training.py", line 1401, in _get_iterator
    return iter(dataloader)
  File "/workspace/Megatron-LM/megatron/training/training.py", line 1424, in build_train_valid_test_data_iterators
    test_data_iterator = _get_iterator(dl_type, test_dataloader)
  File "/workspace/Megatron-LM/megatron/training/training.py", line 250, in pretrain
    = build_train_valid_test_data_iterators(
  File "/workspace/Megatron-LM/megatron/addons/training/training.py", line 97, in pretrain
    orig_func(train_valid_test_dataset_provider,
  File "/workspace/Megatron-LM/pretrain_gpt.py", line 220, in <module>
    pretrain(train_valid_test_datasets_provider,
```

Then I found that from train_dataloader start num_worker to test_dataloader start num_worker,  only pin_memory thread will enter __internal_atexit and acquire lock. If it locks and another dataloader fork num_worker process, it may cause potential deadlock.

I can provide some backtraces
```
#0  __internal_atexit (listp=0x7f4fd8fd8838 <__exit_funcs>, d=0x7f4fd6882f00, arg=0x7f4fd6bf6d60 <caffe2::pthreadpool()::threadpool>, 
    func=0x7f4fc500bda0 <std::unique_ptr<caffe2::PThreadPool, std::default_delete<caffe2::PThreadPool> >::~unique_ptr()>) at ./stdlib/cxa_atexit.c:41
#1  __GI___cxa_atexit (func=0x7f4fc500bda0 <std::unique_ptr<caffe2::PThreadPool, std::default_delete<caffe2::PThreadPool> >::~unique_ptr()>, arg=0x7f4fd6bf6d60 <caffe2::pthreadpool()::threadpool>, 
    d=0x7f4fd6882f00) at ./stdlib/cxa_atexit.c:70
#2  0x00007f4fc500bc6e in caffe2::pthreadpool() () from /torch/venv3/pytorch/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
#3  0x00007f4fc2b5c7f7 in at::set_num_threads(int) () from /torch/venv3/pytorch/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so
#4  0x00007f4fd7327ada in THPModule_setNumThreads(_object*, _object*) () from /torch/venv3/pytorch/lib/python3.10/site-packages/torch/lib/libtorch_python.so
```

```
#0  __internal_atexit (listp=0x7f4fd8fd8838 <__exit_funcs>, d=0x7f4fd7bfff00, arg=0x7f4fd7ca8b00 <torch::autograd::THPVariable_set_(_object*, _object*, _object*)::parser>, 
    func=0x7f4fd7138710 <torch::PythonArgParser::~PythonArgParser()>) at ./stdlib/cxa_atexit.c:41
#1  __GI___cxa_atexit (func=0x7f4fd7138710 <torch::PythonArgParser::~PythonArgParser()>, arg=0x7f4fd7ca8b00 <torch::autograd::THPVariable_set_(_object*, _object*, _object*)::parser>, d=0x7f4fd7bfff00)
    at ./stdlib/cxa_atexit.c:70
#2  0x00007f4fd7133b62 in torch::autograd::THPVariable_set_(_object*, _object*, _object*) () from /torch/venv3/pytorch/lib/python3.10/site-packages/torch/lib/libtorch_python.so
#3  0x00007f4fd91203fd in method_vectorcall_VARARGS_KEYWORDS (func=<method_descriptor at remote 0x7f4fd821d6c0>, args=0x7f4f7c230598, nargsf=<optimized out>, kwnames=<optimized out>)
    at Objects/descrobject.c:344

Traceback (most recent call first):
  File "/torch/venv3/pytorch/lib/python3.10/site-packages/torch/_utils.py", line 184, in _rebuild_tensor
    return t.set_(storage._untyped_storage, storage_offset, size, stride)
  File "/torch/venv3/pytorch/lib/python3.10/site-packages/torch/multiprocessing/reductions.py", line 110, in rebuild_tensor
    t = torch._utils._rebuild_tensor(storage, storage_offset, size, stride)
  <built-in method loads of module object at remote 0x7f4fbe33ae30>
  File "/opt/py3.10/lib/python3.10/multiprocessing/queues.py", line 122, in get
    return _ForkingPickler.loads(res)
  File "/torch/venv3/pytorch/lib/python3.10/site-packages/torch/utils/data/_utils/pin_memory.py", line 31, in do_one_step
    r = in_queue.get(timeout=MP_STATUS_CHECK_INTERVAL)
  File "/torch/venv3/pytorch/lib/python3.10/site-packages/torch/utils/data/_utils/pin_memory.py", line 54, in _pin_memory_loop
    do_one_step()
  File "/opt/py3.10/lib/python3.10/threading.py", line 953, in run
    self._target(*self._args, **self._kwargs)
  File "/opt/py3.10/lib/python3.10/threading.py", line 1016, in _bootstrap_inner
    self.run()
  File "/opt/py3.10/lib/python3.10/threading.py", line 973, in _bootstrap
    self._bootstrap_inner()
```

```
#0  __internal_atexit (listp=0x7f4fd8fd8838 <__exit_funcs>, d=0x7f4fd7bfff00, arg=0x7f4fd7caa600 <torch::autograd::THPVariable_pin_memory(_object*, _object*, _object*)::parser>, 
    func=0x7f4fd7138710 <torch::PythonArgParser::~PythonArgParser()>) at ./stdlib/cxa_atexit.c:41
#1  __GI___cxa_atexit (func=0x7f4fd7138710 <torch::PythonArgParser::~PythonArgParser()>, arg=0x7f4fd7caa600 <torch::autograd::THPVariable_pin_memory(_object*, _object*, _object*)::parser>, d=0x7f4fd7bfff00)
    at ./stdlib/cxa_atexit.c:70
#2  0x00007f4fd70efaae in torch::autograd::THPVariable_pin_memory(_object*, _object*, _object*) () from /torch/venv3/pytorch/lib/python3.10/site-packages/torch/lib/libtorch_python.so
#3  0x00007f4fd91203fd in method_vectorcall_VARARGS_KEYWORDS (func=<method_descriptor at remote 0x7f4fd82295d0>, args=0x7f4f7c20d718, nargsf=<optimized out>, kwnames=<optimized out>)
    at Objects/descrobject.c:344

Traceback (most recent call first):
  File "/torch/venv3/pytorch/lib/python3.10/site-packages/torch/utils/data/_utils/pin_memory.py", line 58, in pin_memory
    return data.pin_memory(device)
  File "/torch/venv3/pytorch/lib/python3.10/site-packages/torch/utils/data/_utils/pin_memory.py", line 63, in <dictcomp>
    return type(data)({k: pin_memory(sample, device) for k, sample in data.items()})  # type: ignore[call-arg]
  File "/torch/venv3/pytorch/lib/python3.10/site-packages/torch/utils/data/_utils/pin_memory.py", line 63, in pin_memory
    return type(data)({k: pin_memory(sample, device) for k, sample in data.items()})  # type: ignore[call-arg]
  File "/torch/venv3/pytorch/lib/python3.10/site-packages/torch/utils/data/_utils/pin_memory.py", line 37, in do_one_step
    data = pin_memory(data, device)
  File "/torch/venv3/pytorch/lib/python3.10/site-packages/torch/utils/data/_utils/pin_memory.py", line 54, in _pin_memory_loop
    do_one_step()
  File "/opt/py3.10/lib/python3.10/threading.py", line 953, in run
    self._target(*self._args, **self._kwargs)
  File "/opt/py3.10/lib/python3.10/threading.py", line 1016, in _bootstrap_inner
    self.run()
  File "/opt/py3.10/lib/python3.10/threading.py", line 973, in _bootstrap
    self._bootstrap_inner()

```



So do you think whether there is a potential deadlock here and have some solutions to avoid the deadlock when use fork and pin memory? Can we change default value of multiprocessing_context to spawn?

### Alternatives

_No response_

### Additional context

_No response_

cc @andrewkho @gokulavasan @SsnL @VitalyFedyunin @dzhulgakov