I followed the tutorial to create custom c++/cuda extensions https://pytorch.org/tutorials/advanced/cpp_extension.html 11.
What I have works locally (only 1 pytorch capable GPU), but I have problems running it on our cluster with 4 GPUs per node:

when I start learning, I see in nvidia-smi, that 4 python processes use GPU 0. The rest of the GPUs have one python process. I used to see only one process on each GPU before I implemented the extension.
after some time I get “RuntimeError: CUDA error: an illegal memory access was encountered” errors on GPUs 1, 2 and 3 (but not 0)
I started the python script using CUDA_LAUNCH_BLOCKING=1 and could trace the error to my own CUDA extension.
I have the feeling that wrong data is used, because the system is not learning. But I have no proof for that, it might be an unrelated issue.

I set the CUDA device by putting all pytorch tensors onto the right device using “cuda:2” etc device names.
The extension is a computation including a gradient computation for the backward pass. I have to create some new matrices in the c++ part and I do that by
torch::Tensor sum = torch::zeros({n.batch, n.layers, n.xes}, torch::dtype(xes.dtype()).device(xes.device()));

I hope, that I put them onto the right GPU by doing that, and I think that I’m creating all tensors using such code (the code is short, I have one in the forward and 2 in the backward pass).
My interpretation of the issue is, that the kernel runs on GPU0 anyway, and that there might be some synchronisation issues or similar later.
I would like to use cudaSetDevice from the cuda driver, but I wouldn’t know where to get the correct gpu id from pytorch. Besides there might be a nicer way to do it. Is there anything that sticks into your eyes? If not, I’ll create a minimal example.
Thanks, Adam