RuntimeError during train_npe_model

[Clownshift]

Hello Devs,

I have been trying to get some toy example of cryoSBI to run, but I am experiencing some crashes that I cannot make sense of. I am not sure, if this is really a bug, but I am hoping that you can help me understand why the program is crashing. The training actually runs for many epochs (i.e., many hours) but eventually crashes. The whole output looks like that:

Training neural netowrk:
 65%|██████████████████████████████████████████████████████████████████████████▍                                       | 98/150 [18:47:34<9:58:18, 690.35s/epoch, loss=-1.01]
Traceback (most recent call last):
  File "/home/paq/programs/micromamba/envs/cryoSBI/bin/train_npe_model", line 8, in <module>
    sys.exit(cl_npe_train_no_saving())
  File "/home/paq/programs/micromamba/envs/cryoSBI/lib/python3.10/site-packages/cryo_sbi/inference/command_line_tools.py", line 50, in cl_npe_train_no_saving
    npe_train_no_saving(
  File "/home/paq/programs/micromamba/envs/cryoSBI/lib/python3.10/site-packages/cryo_sbi/inference/train_npe_model.py", line 159, in npe_train_no_saving
    step(
  File "/home/paq/programs/micromamba/envs/cryoSBI/lib/python3.10/site-packages/lampe/utils.py", line 36, in __call__
    if loss.isfinite().all():
RuntimeError: CUDA error: unspecified launch failure
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Have you encountered this error before? Can you give me any hints what went wrong?

---

My input was:

export CUDA_VISIBLE_DEVICES=1
train_npe_model \
    --image_config_file 2_simulation/_sim_config.json \
    --train_config_file 3_train/_train_config.json\
    --epochs 150 \
    --estimator_file 3_train/posterior.estimator \
    --loss_file 3_train/posterior.loss \
    --n_workers 4 \
    --simulation_batch_size 5120 \
    --train_device cuda

I am running the program on an NVIDIA GeForce RTX 4080, with CUDA 12.3.

---

If you need any additional information, or if I can be of help for any eventual debugging, please let me know.

Best Regards,
Patrick

[Clownshift]

I think the error might be connected to the fact that I requested 4 workers on a single GPU, but I am not sure. Furthermore, I think that the error comes from torch, therefore, not actually from your code.
I am still investigating, if this definitely does not occur anymore, if I only request a single worker.

[Clownshift]

So, I tried with a single walker and this error still occurs every now and then.
I cannot really pinpoint where it comes from, because I cannot really reproduce this error reliably with a specific script.

However, this problem seems to be related to cryoSBI somehow.

[DSilva27]

Found a discussion of this bug in this post of the Torch forum: https://discuss.pytorch.org/t/cuda-error-unspecified-launch-failure/153122/5.

It seems this might be a hardware issue, which makes sense considering that this is failing after a while. Could you check the commands they mention in that post and see if you are facing a similar issue?

[Clownshift]

Thanks. Yes, I also found this issue in other online forums. It seems to have something to do with torch not being compatible with the GPU driver.
So the problem is not solved, yet, but I will just try it on another machine and see if it works there. (I do not have the user rights to update the driver on the first machine, I tested it on).
If I encounter the same issue on other machines, too, I will keep you posted.