Fix NCCL test hang (#9367)

src/collective/nccl_device_communicator.cuh

@@ -12,7 +12,20 @@ namespace collective {
 
 class NcclDeviceCommunicator : public DeviceCommunicator {
  public:
-  explicit NcclDeviceCommunicator(int device_ordinal);
+  /**
+   * @brief Construct a new NCCL communicator.
+   * @param device_ordinal The GPU device id.
+   * @param needs_sync Whether extra CUDA stream synchronization is needed.
+   *
+   * In multi-GPU tests when multiple NCCL communicators are created in the same process, sometimes
+   * a deadlock happens because NCCL kernels are blocking. The extra CUDA stream synchronization
+   * makes sure that the NCCL kernels are caught up, thus avoiding the deadlock.
+   *
+   * The Rabit communicator runs with one process per GPU, so the additional synchronization is not
+   * needed. The in-memory communicator is used in tests with multiple threads, each thread
+   * representing a rank/worker, so the additional synchronization is needed to avoid deadlocks.
+   */
+  explicit NcclDeviceCommunicator(int device_ordinal, bool needs_sync);
   ~NcclDeviceCommunicator() override;
   void AllReduce(void *send_receive_buffer, std::size_t count, DataType data_type,
                  Operation op) override;
@@ -60,6 +73,7 @@ class NcclDeviceCommunicator : public DeviceCommunicator {
                         Operation op);
 
   int const device_ordinal_;
+  bool const needs_sync_;
   int const world_size_;
   int const rank_;
   ncclComm_t nccl_comm_{};