I don’t know what’s wrong, but I do know it’s likely going wrong somewhere in the MPI code.
PyTorch performs a runtime check for CUDA awareness of the MPI distribution before running any collective (including send/recv) with a CUDA tensor. You’re passing CUDA tensors, and not seeing this error, so the runtime check must be successful. What happens beyond that, in the MPI code, is beyond my purview. To be sure there is nothing wrong with your code you can try running it with CPU side tensors. If that passes, and it only fails with CUDA tensors, I’d try and run MPI in some kind of debug mode. Good luck!