Stochastic Illegal Memory Access error mid-epoch on AWS p4d instances
Labels: module: cuda, triaged, module: cublas, matrix multiplication

### 🐛 Describe the bug

While training on a large cluster of AWS p4d instances with PyTorch, PyTorch Lightning, and DeepSpeed, we observe an IllegalMemoryAccess error that happens stochastically at different places and different steps mid-epoch.

The error is stochastic in the sense that when training a 1000-sized epoch on 120GPUs, the error sometimes happens at step ~400, sometimes at ~600, or sometimes at ~700.
We did some memory tracking on our end and found out that there was no memory leak happening and that the peak memory usage remains the same from start of epoch (after some stabilization) till the end and to the next epochs. It is unclear whether a specific location in the code is faulty since we have only been training on the large compute instances (and we have limited budget for that). Upon trying the same version of code in a comparatively very small cluster of 18 GPUs, with same specifications and same datasets, we have not observed the issue so far.

The traceback pasted at the end is with `CUDA_LAUNCH_BLOCKING=1` and hence should reflect the real exception being thrown at the point of the error.

Selected List of our package versions:
```
cuda                      11.6.2                        0    nvidia
cuda-cccl                 11.6.55              hf6102b2_0    nvidia
cuda-command-line-tools   11.6.2                        0    nvidia
cuda-compiler             11.6.2                        0    nvidia
cuda-cudart               11.6.55              he381448_0    nvidia
cuda-cudart-dev           11.6.55              h42ad0f4_0    nvidia
cuda-cuobjdump            11.6.124             h2eeebcb_0    nvidia
cuda-cupti                11.6.124             h86345e5_0    nvidia
cuda-cuxxfilt             11.6.124             hecbf4f6_0    nvidia
cuda-driver-dev           11.6.55                       0    nvidia
cuda-gdb                  11.8.86                       0    nvidia
cuda-libraries            11.6.2                        0    nvidia
cuda-libraries-dev        11.6.2                        0    nvidia
cuda-memcheck             11.8.86                       0    nvidia
cuda-nsight               11.8.86                       0    nvidia
cuda-nsight-compute       11.8.0                        0    nvidia
cuda-nvcc                 11.6.124             hbba6d2d_0    nvidia
cuda-nvdisasm             11.8.86                       0    nvidia
cuda-nvml-dev             11.6.55              haa9ef22_0    nvidia
cuda-nvprof               11.8.87                       0    nvidia
cuda-nvprune              11.6.124             he22ec0a_0    nvidia
cuda-nvrtc                11.6.124             h020bade_0    nvidia
cuda-nvrtc-dev            11.6.124             h249d397_0    nvidia
cuda-nvtx                 11.6.124             h0630a44_0    nvidia
cuda-nvvp                 11.8.87                       0    nvidia
cuda-runtime              11.6.2                        0    nvidia
cuda-samples              11.6.101             h8efea70_0    nvidia
cuda-sanitizer-api        11.8.86                       0    nvidia
cuda-toolkit              11.6.2                        0    nvidia
cuda-tools                11.6.2                        0    nvidia
cuda-visual-tools         11.6.2                        0    nvidia
cudatoolkit               11.6.0              hecad31d_10    conda-forge
deepspeed                 0.5.10                   pypi_0    pypi

python                    3.7.12          hf930737_100_cpython    conda-forge
python_abi                3.7                     2_cp37m    conda-forge
pytorch                   1.13.0          py3.7_cuda11.6_cudnn8.3.2_0    pytorch
pytorch-cuda              11.6                 h867d48c_0    pytorch
pytorch-lightning         1.5.10                   pypi_0    pypi
pytorch-mutex             1.0                        cuda    pytorch

```

Sample Code
All of our code is in the repository: https://github.com/aqlaboratory/openfold/
The snapshot of the code that I am using is uploaded here: https://github.com/aqlaboratory/openfold/tree/aws-finetuning-code

Traceback
```
Epoch 0:  46%|████▌     | 4/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py:430: UserWarning: torch.distributed.distributed_c10d._get_global_rank is deprecated please use torch.distributed.distributed_c10d.get_global_rank instead
  "torch.distributed.distributed_c10d._get_global_rank is deprecated "
Traceback (most recent call last):
  File "/shared/openfold-release/train_openfold.py", line 585, in <module>
    main(args)
  File "/shared/openfold-release/train_openfold.py", line 369, in main
    ckpt_path=ckpt_path,
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 741, in fit
    self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 685, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 777, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1199, in _run
    self._dispatch()
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1279, in _dispatch
    self.training_type_plugin.start_training(self)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 202, in start_training
    self._results = trainer.run_stage()
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1289, in run_stage
    return self._run_train()
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1319, in _run_train
    self.fit_loop.run()
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/loops/base.py", line 145, in run
    self.advance(*args, **kwargs)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/loops/fit_loop.py", line 234, in advance
    self.epoch_loop.run(data_fetcher)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/loops/base.py", line 145, in run
    self.advance(*args, **kwargs)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 193, in advance
    batch_output = self.batch_loop.run(batch, batch_idx)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/loops/base.py", line 145, in run
    self.advance(*args, **kwargs)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 88, in advance
    outputs = self.optimizer_loop.run(split_batch, optimizers, batch_idx)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/loops/base.py", line 145, in run
    self.advance(*args, **kwargs)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 219, in advance
    self.optimizer_idx,
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 266, in _run_optimization
    self._optimizer_step(optimizer, opt_idx, batch_idx, closure)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 386, in _optimizer_step
    using_lbfgs=is_lbfgs,
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/core/lightning.py", line 1652, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/core/optimizer.py", line 164, in step
    trainer.accelerator.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/accelerators/accelerator.py", line 339, in optimizer_step
    self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/plugins/precision/deepspeed_precision.py", line 65, in optimizer_step
    closure_result = closure()
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 160, in __call__
    self._result = self.closure(*args, **kwargs)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 155, in closure
    self._backward_fn(step_output.closure_loss)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 327, in backward_fn
    self.trainer.accelerator.backward(loss, optimizer, opt_idx)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/accelerators/accelerator.py", line 314, in backward
    self.precision_plugin.backward(self.lightning_module, closure_loss, *args, **kwargs)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/pytorch_lightning/plugins/precision/deepspeed_precision.py", line 48, in backward
    deepspeed_engine.backward(closure_loss, *args, **kwargs)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/deepspeed/runtime/engine.py", line 1697, in backward
    self.optimizer.backward(loss)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1910, in backward
    self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 53, in backward
    scaled_loss.backward(retain_graph=retain_graph)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/_tensor.py", line 488, in backward
    self, gradient, retain_graph, create_graph, inputs=inputs
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/autograd/__init__.py", line 199, in backward
    allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to run the backward_fnard pass
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/autograd/function.py", line 267, in apply
    return user_fn(self, *args)
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/deepspeed/runtime/activation_checkpointing/checkpointing.py", line 693, in backward
    outputs = ctx.run_function(*detached_inputs)
  File "/shared/openfold-release/openfold/utils/checkpointing.py", line 69, in exec_sliced
    return exec(blocks[s:e], a)
  File "/shared/openfold-release/openfold/utils/checkpointing.py", line 64, in exec
    a = wrap(block(*a))
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/shared/openfold-release/openfold/model/template.py", line 229, in forward
    inplace_safe=inplace_safe,
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/shared/openfold-release/openfold/model/triangular_attention.py", line 142, in forward
    use_lma=use_lma
  File "/home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/shared/openfold-release/openfold/model/primitives.py", line 469, in forward
    o = _attention(q, k, v, biases)
  File "/shared/openfold-release/openfold/model/primitives.py", line 237, in _attention
    a = torch.matmul(a, value)
RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasSgemmStridedBatched( handle, opa, opb, m, n, k, &alpha, a, lda, stridea, b, ldb, strideb, &beta, c, ldc, stridec, num_batches)`
terminate called after throwing an instance of 'c10::Error'
  what():  CUDA error: an illegal memory access was encountered
Exception raised from c10_cuda_check_implementation at /opt/conda/conda-bld/pytorch_1666643004612/work/c10/cuda/CUDAException.cpp:31 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x57 (0x7f4db1cce457 in /home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) + 0x64 (0x7f4db1c983ec in /home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/lib/libc10.so)
frame #2: c10::cuda::c10_cuda_check_implementation(std::string const&, std::string const&, int, bool) + 0xb4 (0x7f4db1d68044 in /home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/lib/libc10_cuda.so)
frame #3: <unknown function> + 0x164bc (0x7f4db1d3f4bc in /home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/lib/libc10_cuda.so)
frame #4: c10::cuda::CUDACachingAllocator::raw_delete(void*) + 0x244 (0x7f4db1d42434 in /home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/lib/libc10_cuda.so)
frame #5: <unknown function> + 0x4edcc3 (0x7f4d88ab1cc3 in /home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/lib/libtorch_python.so)
frame #6: c10::TensorImpl::~TensorImpl() + 0x1a0 (0x7f4db1cae9e0 in /home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/lib/libc10.so)
frame #7: c10::TensorImpl::~TensorImpl() + 0x9 (0x7f4db1caeaf9 in /home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/lib/libc10.so)
frame #8: <unknown function> + 0x75e5a8 (0x7f4d88d225a8 in /home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/lib/libtorch_python.so)
frame #9: THPVariable_subclass_dealloc(_object*) + 0x2f8 (0x7f4d88d229a8 in /home/ec2-user/miniconda3/envs/openfold_train_env/lib/python3.7/site-packages/torch/lib/libtorch_python.so)
frame #10: <unknown function> + 0xe59d2 (0x5584083189d2 in /home/ec2-user/miniconda3/envs/openfold_train_env/bin/python)
frame #11: <unknown function> + 0xe707f (0x55840831a07f in /home/ec2-user/miniconda3/envs/openfold_train_env/bin/python)
frame #12: <unknown function> + 0x17fcc1 (0x5584083b2cc1 in /home/ec2-user/miniconda3/envs/openfold_train_env/bin/python)
frame #13: <unknown function> + 0xe707f (0x55840831a07f in /home/ec2-user/miniconda3/envs/openfold_train_env/bin/python)
frame #14: <unknown function> + 0xe7782 (0x55840831a782 in /home/ec2-user/miniconda3/envs/openfold_train_env/bin/python)
frame #15: <unknown function> + 0xfc8c7 (0x55840832f8c7 in /home/ec2-user/miniconda3/envs/openfold_train_env/bin/python)
frame #16: _PyGC_CollectNoFail + 0x2b (0x5584084542db in /home/ec2-user/miniconda3/envs/openfold_train_env/bin/python)
frame #17: PyImport_Cleanup + 0x367 (0x558408466307 in /home/ec2-user/miniconda3/envs/openfold_train_env/bin/python)
frame #18: Py_FinalizeEx + 0x67 (0x558408466407 in /home/ec2-user/miniconda3/envs/openfold_train_env/bin/python)
frame #19: <unknown function> + 0x248b1b (0x55840847bb1b in /home/ec2-user/miniconda3/envs/openfold_train_env/bin/python)
frame #20: _Py_UnixMain + 0x3c (0x55840847be9c in /home/ec2-user/miniconda3/envs/openfold_train_env/bin/python)
frame #21: __libc_start_main + 0xea (0x7f4dc323813a in /lib64/libc.so.6)
frame #22: <unknown function> + 0x1c761d (0x5584083fa61d in /home/ec2-user/miniconda3/envs/openfold_train_env/bin/python)

srun: error: train-p4d-st-train-p4d-15: task 118: Aborted
srun: error: Node failure on train-p4d-st-train-p4d-1
slurmstepd: error: *** JOB 44 ON train-p4d-st-train-p4d-1 CANCELLED AT 2022-11-18T13:21:41 DUE TO NODE FAILURE, SEE SLURMCTLD LOG FOR DETAILS ***
srun: Job step aborted: Waiting up to 32 seconds for job step to finish.

```

### Versions

Collecting environment information...
PyTorch version: 1.13.0
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Amazon Linux release 2 (Karoo) (x86_64)
GCC version: (GCC) 7.3.1 20180712 (Red Hat 7.3.1-15)
Clang version: Could not collect
CMake version: version 2.8.12.2
Libc version: glibc-2.10

Python version: 3.7.12 | packaged by conda-forge | (default, Oct 26 2021, 06:08:21)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.10.149-133.644.amzn2.x86_64-x86_64-with-glibc2.10
Is CUDA available: False
CUDA runtime version: 11.6.124
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.2
[pip3] pytorch-lightning==1.5.10
[pip3] torch==1.13.0
[pip3] torchmetrics==0.10.2
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.6.0              hecad31d_10    conda-forge
[conda] mkl                       2021.4.0           h8d4b97c_729    conda-forge
[conda] mkl-service               2.4.0            py37h402132d_0    conda-forge
[conda] mkl_fft                   1.3.1            py37h3e078e5_1    conda-forge
[conda] mkl_random                1.2.2            py37h219a48f_0    conda-forge
[conda] numpy                     1.21.2                   pypi_0    pypi
[conda] pytorch                   1.13.0          py3.7_cuda11.6_cudnn8.3.2_0    pytorch
[conda] pytorch-cuda              11.6                 h867d48c_0    pytorch
[conda] pytorch-lightning         1.5.10                   pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchmetrics              0.10.2                   pypi_0    pypi


cc @ngimel @csarofeen @ptrblck @xwang233