Hanging with NeMo

Nobody has claimed this yet.

### Bug description

The pretraining job of a GPT model with [NeMo](https://github.com/NVIDIA/NeMo/tree/r1.21.0) hangs after 24 hours of training. Config is based on https://github.com/NVIDIA/NeMo-Megatron-Launcher/blob/master/launcher_scripts/conf/training/gpt3/175b.yaml with smaller model size. The training job uses tensor parallelism of `4`. It does not look like the hanging is related to NCCL ops since all the communication-computation overlapping has been turned off. Most processes hang at the following place:

`py-spy` log:
```
__to_tensor (pytorch_lightning/core/module.py:619)
apply_to_collection (lightning_utilities/core/apply_func.py:51)
log (pytorch_lightning/core/module.py:447)
training_step (nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py:653)
wrap_training_step (nemo/utils/model_utils.py:381)
forward (pytorch_lightning/overrides/base.py:90)
_call_impl (torch/nn/modules/module.py:1514)
_wrapped_call_impl (torch/nn/modules/module.py:1505)
training_step (pytorch_lightning/strategies/ddp.py:330)
_call_strategy_hook (pytorch_lightning/trainer/call.py:293)
_training_step (pytorch_lightning/loops/optimization/automatic.py:315)
closure (pytorch_lightning/loops/optimization/automatic.py:128)
decorate_context (torch/utils/_contextlib.py:115)
__call__ (pytorch_lightning/loops/optimization/automatic.py:142)
_wrap_closure (pytorch_lightning/plugins/precision/precision_plugin.py:103)
step (apex/contrib/optimizers/distributed_fused_adam.py:1960)
wrapper (torch/optim/optimizer.py:316)
wrapper (torch/optim/lr_scheduler.py:69)
optimizer_step (pytorch_lightning/plugins/precision/precision_plugin.py:116)
optimizer_step (pytorch_lightning/plugins/precision/amp.py:73)
optimizer_step (pytorch_lightning/strategies/strategy.py:231)
optimizer_step (pytorch_lightning/strategies/ddp.py:257)
step (pytorch_lightning/core/optimizer.py:161)
optimizer_step (pytorch_lightning/core/module.py:1270)
_call_lightning_module_hook (pytorch_lightning/trainer/call.py:145)
_optimizer_step (pytorch_lightning/loops/optimization/automatic.py:266)
run (pytorch_lightning/loops/optimization/automatic.py:188)
advance (pytorch_lightning/loops/training_epoch_loop.py:219)
run (pytorch_lightning/loops/training_epoch_loop.py:133)
advance (pytorch_lightning/loops/fit_loop.py:355)
run (pytorch_lightning/loops/fit_loop.py:202)
_run_stage (pytorch_lightning/trainer/trainer.py:1023)
_run (pytorch_lightning/trainer/trainer.py:980)
_fit_impl (pytorch_lightning/trainer/trainer.py:571)
launch (pytorch_lightning/strategies/launchers/subprocess_script.py:93)
_call_and_handle_interrupt (pytorch_lightning/trainer/call.py:42)
fit (pytorch_lightning/trainer/trainer.py:532)
main (megatron_gpt_pretraining.py:46)
run_job (hydra/core/utils.py:186)
run (hydra/_internal/hydra.py:119)
(hydra/_internal/utils.py:453)
run_and_report (hydra/_internal/utils.py:213)
_run_app (hydra/_internal/utils.py:452)
_run_hydra (hydra/_internal/utils.py:389)
wrapper (nemo/core/config/hydra_runner.py:126)
(megatron_gpt_pretraining.py:50)
Thread 0x151A1251B640 (idle): "Thread-1 (_pin_memory_loop)"
select (selectors.py:416)
wait (multiprocessing/connection.py:931)
_poll (multiprocessing/connection.py:424)
poll (multiprocessing/connection.py:257)
get (multiprocessing/queues.py:113)
do_one_step (torch/utils/data/_utils/pin_memory.py:31)
_pin_memory_loop (torch/utils/data/_utils/pin_memory.py:54)
run (threading.py:953)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)
Thread 0xB85D76FF640 (idle): "QueueFeederThread"
wait (threading.py:320)
_feed (multiprocessing/queues.py:231)
run (threading.py:953)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)
Thread 0xB85D74FE640 (idle): "QueueFeederThread"
wait (threading.py:320)
_feed (multiprocessing/queues.py:231)
run (threading.py:953)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)
Thread 0xB85D72FD640 (idle): "Thread-2 (_pin_memory_loop)"
select (selectors.py:416)
wait (multiprocessing/connection.py:931)
_poll (multiprocessing/connection.py:424)
poll (multiprocessing/connection.py:257)
get (multiprocessing/queues.py:113)
do_one_step (torch/utils/data/_utils/pin_memory.py:31)
_pin_memory_loop (torch/utils/data/_utils/pin_memory.py:54)
run (threading.py:953)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)
Thread 0xB85D6DFC640 (idle): "QueueFeederThread"
wait (threading.py:320)
_feed (multiprocessing/queues.py:231)
run (threading.py:953)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)
Thread 0xB85D6BFB640 (idle): "QueueFeederThread"
wait (threading.py:320)
_feed (multiprocessing/queues.py:231)
run (threading.py:953)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)
Thread 0xB85D69FA640 (idle): "Thread-3"
wait (threading.py:324)
wait (threading.py:607)
run (tqdm/_monitor.py:60)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)

```
`gdb` log:
```
(gdb) where
#0 0x000014631bbb4494 in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#1 0x000014631b90909f in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#2 0x000014631bc95d2f in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#3 0x000014631bc9691e in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#4 0x000014631b90dbbd in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#5 0x000014631ba279eb in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#6 0x000014631bc773e5 in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#7 0x000014631bac05bd in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#8 0x000014631da138b5 in ?? () from target:/usr/local/cuda/lib64/libcudart.so.12
#9 0x000014631da70898 in cudaStreamSynchronize () from target:/usr/local/cuda/lib64/libcudart.so.12
#10 0x00001462d52d5ca6 in at::native::copy_kernel_cuda(at::TensorIterator&, bool) () from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cuda.so
#11 0x0000146311259f64 in at::native::copy_impl(at::Tensor&, at::Tensor const&, bool) [clone .isra.0] () from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so
#12 0x000014631125b1f0 in at::native::copy_(at::Tensor&, at::Tensor const&, bool) () from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so
#13 0x0000146311f213bc in at::_ops::copy_::call(at::Tensor&, at::Tensor const&, bool) () from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so
#14 0x0000146311507867 in at::native::_to_copy(at::Tensor const&, c10::optional, c10::optional, c10::optional, c10::optional, bool, c10::optional) ()
from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so
#15 0x00001463122a8841 in c10::impl::wrap_kernel_functor_unboxed_, c10::optional, c10::optional, c10::optional, bool, c10::optional), &at::(anonymous namespace)::(anonymous namespace)::wrapper_CompositeExplicitAutograd___to_copy>, at::Tensor, c10::guts::typelist::typelist, c10::optional, c10::optional, c10::optional, bool, c10::optional > >, at::Tensor (at::Tensor const&, c10::optional, c10::optional, c10::optional, c10::optional, bool, c10::optional)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, c10::optional, c10::optional, c10::optional, c10::optional, bool, c10::optional) ()
from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so
#16 0x00001463119da78c in at::_ops::_to_copy::redispatch(c10::DispatchKeySet, at::Tensor const&, c10::optional, c10::optional, c10::optional, c10::optional, bool, c10::optional) ()
from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so
#17 0x00001463120ed018 in c10::impl::wrap_kernel_functor_unboxed_, c10::optional, c10::optional, c10::optional, bool, c10::optional), &at::(anonymous namespace)::_to_copy>, at::Tensor, c10::guts::typelist::typelist, c10::optional, c10::optional, c10::optional, bool, c10::optional > >, at::Tensor (at::Tensor const&, c10::optional, c10::optional, c10::optional, c10::optional, bool, c10::optional)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, c10::optional, c10::optional, c10::optional, c10::optional, bool, c10::optional) () from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so
#18 0x0000146311a8b338 in at::_ops::_to_copy::call(at::Tensor const&, c10::optional, c10::optional, c10::optional, c10::optional, bool, c10::optional) ()
from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so
#19 0x00001463114faf88 in at::native::to(at::Tensor const&, c10::Device, c10::ScalarType, bool, bool, c10::optional) () from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so
#20 0x0000146312460c2d in c10::impl::wrap_kernel_functor_unboxed_), &at::(anonymous namespace)::(anonymous namespace)::wrapper_CompositeImplicitAutograd_device_to>, at::Tensor, c10::guts::typelist::typelist > >, at::Tensor (at::Tensor const&, c10::Device, c10::ScalarType, bool, bool, c10::optional)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, c10::Device, c10::ScalarType, bool, bool, c10::optional) ()
from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so
#21 0x0000146311c28eaf in at::_ops::to_device::call(at::Tensor const&, c10::Device, c10::ScalarType, bool, bool, c10::optional) () from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so
#22 0x000014631b27fb5a in torch::utils::(anonymous namespace)::internal_new_from_data(c10::TensorOptions, c10::ScalarType, c10::optional, _object*, bool, bool, bool, bool) () from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_python.so
#23 0x000014631b285570 in torch::utils::tensor_ctor(c10::DispatchKey, c10::ScalarType, torch::PythonArgs&) () from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_python.so
#24 0x000014631ae9a262 in torch::autograd::THPVariable_tensor(_object*, _object*, _object*) () from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_python.so
#25 0x000055d630e52e0e in ?? ()
#26 0x000055d630e495eb in _PyObject_MakeTpCall ()
#27 0x000055d630e42908 in _PyEval_EvalFrameDefault ()
#28 0x000055d630e6162e in ?? ()
#29 0x000055d630e3e2c1 in _PyEval_EvalFrameDefault ()
#30 0x000055d630e5370c in _PyFunction_Vectorcall ()
#31 0x000055d630e3be0d in _PyEval_EvalFrameDefault ()
#32 0x000055d630e614e1 in ?? ()
#33 0x000055d630e3d0d1 in _PyEval_EvalFrameDefault ()
#34 0x000055d630e6162e in ?? ()
#35 0x000055d630e3e2c1 in _PyEval_EvalFrameDefault ()
#36 0x000055d630e5370c in _PyFunction_Vectorcall ()
```
Some other processes stuck at
`py-spy` log:
```
synchronize (torch/cuda/__init__.py:721)
stop (nemo/utils/timers.py:115)
_on_batch_end (nemo/utils/exp_manager.py:195)
on_train_batch_end (nemo/utils/exp_manager.py:210)
_call_callback_hooks (pytorch_lightning/trainer/call.py:194)
advance (pytorch_lightning/loops/training_epoch_loop.py:231)
run (pytorch_lightning/loops/training_epoch_loop.py:133)
advance (pytorch_lightning/loops/fit_loop.py:355)
run (pytorch_lightning/loops/fit_loop.py:202)
_run_stage (pytorch_lightning/trainer/trainer.py:1023)
_run (pytorch_lightning/trainer/trainer.py:980)
_fit_impl (pytorch_lightning/trainer/trainer.py:571)
launch (pytorch_lightning/strategies/launchers/subprocess_script.py:93)
_call_and_handle_interrupt (pytorch_lightning/trainer/call.py:42)
fit (pytorch_lightning/trainer/trainer.py:532)
main (megatron_gpt_pretraining.py:46)
run_job (hydra/core/utils.py:186)
run (hydra/_internal/hydra.py:119)
(hydra/_internal/utils.py:453)
run_and_report (hydra/_internal/utils.py:213)
_run_app (hydra/_internal/utils.py:452)
_run_hydra (hydra/_internal/utils.py:389)
wrapper (nemo/core/config/hydra_runner.py:126)
(megatron_gpt_pretraining.py:50)
Thread 0x1502C1907640 (idle): "Thread-2 (_pin_memory_loop)"
select (selectors.py:416)
wait (multiprocessing/connection.py:931)
_poll (multiprocessing/connection.py:424)
poll (multiprocessing/connection.py:257)
get (multiprocessing/queues.py:113)
do_one_step (torch/utils/data/_utils/pin_memory.py:31)
_pin_memory_loop (torch/utils/data/_utils/pin_memory.py:54)
run (threading.py:953)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)
Thread 0xB6CD09FF640 (idle): "QueueFeederThread"
wait (threading.py:320)
_feed (multiprocessing/queues.py:231)
run (threading.py:953)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)
Thread 0xB6CD07FE640 (idle): "QueueFeederThread"
wait (threading.py:320)
_feed (multiprocessing/queues.py:231)
run (threading.py:953)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)
Thread 0xB6CD05FD640 (idle): "Thread-3 (_pin_memory_loop)"
select (selectors.py:416)
wait (multiprocessing/connection.py:931)
_poll (multiprocessing/connection.py:424)
poll (multiprocessing/connection.py:257)
get (multiprocessing/queues.py:113)
do_one_step (torch/utils/data/_utils/pin_memory.py:31)
_pin_memory_loop (torch/utils/data/_utils/pin_memory.py:54)
run (threading.py:953)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)
Thread 0xB6CCBFFF640 (idle): "QueueFeederThread"
wait (threading.py:320)
_feed (multiprocessing/queues.py:231)
run (threading.py:953)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)
Thread 0xB6CCBDFE640 (idle): "QueueFeederThread"
wait (threading.py:320)
_feed (multiprocessing/queues.py:231)
run (threading.py:953)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)
Thread 0xB6CCBBFD640 (idle): "Thread-4"
wait (threading.py:324)
wait (threading.py:607)
run (tqdm/_monitor.py:60)
_bootstrap_inner (threading.py:1016)
_bootstrap (threading.py:973)

```
and `gdb` log:
```
#0 0x00007ffe2d1cf991 in clock_gettime ()
#1 0x000015074c42e6dd in clock_gettime () from target:/usr/lib/x86_64-linux-gnu/libc.so.6
#2 0x0000150704043944 in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#3 0x0000150703f0de2e in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#4 0x0000150703fa388c in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#5 0x0000150703fa5796 in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#6 0x0000150703ea264f in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#7 0x0000150704061375 in ?? () from target:/usr/local/cuda/compat/lib.real/libcuda.so.1
#8 0x000015070601116d in ?? () from target:/usr/local/cuda/lib64/libcudart.so.12
#9 0x00001507060489e5 in cudaDeviceSynchronize () from target:/usr/local/cuda/lib64/libcudart.so.12
#10 0x000015074adccfc8 in c10::cuda::device_synchronize() () from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libc10_cuda.so
#11 0x00001507038a44ac in THCPModule_cudaSynchronize(_object*, _object*) () from target:/usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_python.so
#12 0x000055abc1fe564e in ?? ()
#13 0x000055abc1fd68a2 in _PyEval_EvalFrameDefault ()
#14 0x000055abc1fe870c in _PyFunction_Vectorcall ()
#15 0x000055abc1fd68a2 in _PyEval_EvalFrameDefault ()
#16 0x000055abc1fe870c in _PyFunction_Vectorcall ()
#17 0x000055abc1fd0f52 in _PyEval_EvalFrameDefault ()
#18 0x000055abc1fe870c in _PyFunction_Vectorcall ()
#19 0x000055abc1fd0f52 in _PyEval_EvalFrameDefault ()

```

The Lightning version is
```
>>> pytorch_lightning.__version__
'2.0.7'
```

Any idea is appreciated.

### What version are you seeing the problem on?

v2.0

### How to reproduce the bug

_No response_

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):
```

### More info

_No response_

Start by reproducing the GPT pretraining run from the linked NeMo-Megatron-Launcher configuration with tensor parallelism set to 4. Inspect the reported paths in nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py, pytorch_lightning/core/module.py, and apex/contrib/optimizers/distributed_fused_adam.py, then compare the py-spy and gdb locations. Done means identifying the cause of the long-running hang and verifying a fix in a comparable training run.

Written by the indexing model from the issue text.