ensure all GPUs having same allreduces to avoid hanging

rapidsai · Nov 22, 2024 · 6c56fa1 · 6c56fa1
1 parent 06958c4
commit 6c56fa1
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/cpp/src/glm/qn/mg/standardization.cuh b/cpp/src/glm/qn/mg/standardization.cuh
@@ -169,6 +169,13 @@ void mean_stddev(const raft::handle_t& handle,
 
     SimpleVec<T> stddevVec(stddev_vector, D);
     stddevVec.fill(0., stream);
+
+    // call allreduces on zeroes to sync with other GPUs to avoid hanging
+    auto& comm = handle.get_comms();
+    comm.allreduce(mean_vector, mean_vector, D, raft::comms::op_t::SUM, stream);
+    comm.sync_stream(stream);
+    comm.allreduce(stddev_vector, stddev_vector, D, raft::comms::op_t::SUM, stream);
+    comm.sync_stream(stream);
     return;
   }