perf: fix cuda-aware mpi in v3 (#4977)

This pull request updates the MPI CUDA-awareness detection and handling
logic in the `Border` autograd function, simplifying how CUDA support is
determined and removing some legacy checks. The changes ensure that
CUDA-aware MPI support is queried more directly, and some unnecessary
device synchronization calls are removed.

* The logic for checking CUDA-aware MPI support has been simplified:
version checks and redundant branches have been removed, and the code
now directly queries `MPIX_Query_cuda_support()` unless `NO_CUDA_AWARE`
is defined.
[[1]](diffhunk://#diff-7b7590fd4222d9c50f1dd7dde5ce7ed4b27695fbe591b536787db7575c35e32cL102-L112)
[[2]](diffhunk://#diff-7b7590fd4222d9c50f1dd7dde5ce7ed4b27695fbe591b536787db7575c35e32cL227-L237)

* Removed explicit `gpuDeviceSynchronize()` calls from both the forward
and backward paths, relying on PyTorch's internal synchronization
mechanisms instead.
[[1]](diffhunk://#diff-7b7590fd4222d9c50f1dd7dde5ce7ed4b27695fbe591b536787db7575c35e32cL196-L198)
[[2]](diffhunk://#diff-7b7590fd4222d9c50f1dd7dde5ce7ed4b27695fbe591b536787db7575c35e32cL332-L334)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- Performance
- Reduced explicit GPU synchronization, potentially improving throughput
during distributed forward/backward operations.
- Compatibility
- Safer default when CUDA-aware MPI isn’t present: automatically falls
back to CPU-based transfers unless support is detected, improving
stability across varied clusters.
- Reliability
- Simplified CUDA-aware detection reduces edge-case misconfigurations in
mixed MPI environments.
- No API Changes
- Public interfaces remain unchanged; existing workflows continue to
work.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: caic99

source/op/pt/comm.cc

@@ -86,7 +86,7 @@ class Border : public torch::autograd::Function<Border> {
 #ifdef USE_MPI
     int mpi_init = 0;
     MPI_Initialized(&mpi_init);
-    int cuda_aware = 1;
+    int cuda_aware = 0;
     int me = 0;
     MPI_Comm world;
     int world_size = 0;
@@ -99,17 +99,9 @@ class Border : public torch::autograd::Function<Border> {
     MPI_Request request;
 #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
     if (world_size >= 1) {
-      int version, subversion;
-      MPI_Get_version(&version, &subversion);
-      if (version >= 4) {
-#ifdef NO_CUDA_AWARE
-        cuda_aware = 0;
-#else
-        cuda_aware = MPIX_Query_cuda_support();
+#ifndef NO_CUDA_AWARE
+      cuda_aware = MPIX_Query_cuda_support();
 #endif
-      } else {
-        cuda_aware = 0;
-      }
       if (cuda_aware == 0) {
         recv_g1_tensor = torch::empty_like(g1).to(torch::kCPU);
         recv_g1_tensor.copy_(g1);
@@ -193,10 +185,6 @@ class Border : public torch::autograd::Function<Border> {
   static torch::autograd::variable_list backward_t(
       torch::autograd::AutogradContext* ctx,
       torch::autograd::variable_list grad_output) {
-#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
-    gpuDeviceSynchronize();
-#endif
-
     torch::autograd::variable_list saved_variables = ctx->get_saved_variables();
     torch::Tensor sendlist_tensor = saved_variables[0];
     torch::Tensor sendproc_tensor = saved_variables[1];
@@ -212,7 +200,7 @@ class Border : public torch::autograd::Function<Border> {
     int mpi_init = 0;
     MPI_Initialized(&mpi_init);
     int world_size = 0;
-    int cuda_aware = 1;
+    int cuda_aware = 0;
     int me = 0;
     MPI_Comm world;
     if (mpi_init) {
@@ -224,17 +212,9 @@ class Border : public torch::autograd::Function<Border> {
     MPI_Request request;
 #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
     if (world_size >= 1) {
-      int version, subversion;
-      MPI_Get_version(&version, &subversion);
-      if (version >= 4) {
-#ifdef NO_CUDA_AWARE
-        cuda_aware = 0;
-#else
-        cuda_aware = MPIX_Query_cuda_support();
+#ifndef NO_CUDA_AWARE
+      cuda_aware = MPIX_Query_cuda_support();
 #endif
-      } else {
-        cuda_aware = 0;
-      }
       if (cuda_aware == 0) {
         d_local_g1_tensor = torch::empty_like(grad_output[0]).to(torch::kCPU);
         d_local_g1_tensor.copy_(grad_output[0]);
@@ -329,9 +309,6 @@ class Border : public torch::autograd::Function<Border> {
                                      recv_g1_tensor.slice(0, 0, nrecv));
       }
     }
-#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
-    gpuDeviceSynchronize();
-#endif
 #ifdef USE_MPI
 #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
     if (cuda_aware == 0) {