Merge branch 'develop_v2.3' into 'main'

[DCU] remove cudaStreamSynchronize for tp overlap

See merge request dcutoolkit/deeplearing/TransformerEngine!13

transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp

@@ -306,7 +306,6 @@ void CommOverlapBase::bulk_overlap(const TensorWrapper &A, bool transa, const Te
                                    bool accumulate, bool use_split_accumulator,
                                    CommOverlapType comm_type, TensorWrapper &rs_output,
                                    cudaStream_t stream_main) {
-  NVTE_CHECK_CUDA(cudaStreamSynchronize(stream_main));
   int ori_sms = _ub_comm->sms;
   _ub_comm->use_ce = _use_ce;
   _ub_comm->sms = _num_comm_sm;
@@ -356,7 +355,6 @@ void CommOverlapBase::bulk_overlap(const TensorWrapper &A, bool transa, const Te
   NVTE_CHECK_CUDA(cudaStreamWaitEvent(stream_main, _stop_comm, 0));
   NVTE_CHECK_CUDA(cudaEventRecord(_stop_comm, _stream_compute[0]));
   NVTE_CHECK_CUDA(cudaStreamWaitEvent(stream_main, _stop_comm, 0));
-  NVTE_CHECK_CUDA(cudaStreamSynchronize(stream_main));
 
 }  // CommOverlapBase::bulk_overlap
 
@@ -465,7 +463,6 @@ void CommOverlapBase::split_overlap_rs(const TensorWrapper &A, bool transa, cons
                                        TensorWrapper &pre_gelu_out, TensorWrapper &workspace,
                                        bool grad, bool accumulate, bool use_split_accumulator,
                                        TensorWrapper &rs_output, cudaStream_t stream_main) {
-  NVTE_CHECK_CUDA(cudaStreamSynchronize(stream_main));
   // Get GEMM dimensions
   int ori_sms = _ub_comm->sms;
   _ub_comm->use_ce = _use_ce;
@@ -620,7 +617,6 @@ void CommOverlapBase::split_overlap_rs(const TensorWrapper &A, bool transa, cons
   }
   NVTE_CHECK_CUDA(cudaEventRecord(_stop_comm, _stream_comm));
   NVTE_CHECK_CUDA(cudaStreamWaitEvent(stream_main, _stop_comm, 0));
-  NVTE_CHECK_CUDA(cudaStreamSynchronize(stream_main));
 }  // CommOverlapBase::split_overlap_rs
 
 /***************************************************************************************************
@@ -843,7 +839,6 @@ void CommOverlapP2PBase::split_overlap_ag(const TensorWrapper &A, bool transa,
                                           TensorWrapper &workspace, bool grad, bool accumulate,
                                           bool use_split_accumulator, TensorWrapper &B_copy,
                                           cudaStream_t stream_main) {
-  NVTE_CHECK_CUDA(cudaStreamSynchronize(stream_main));
   int ori_sms = _ub_comm->sms;
   _ub_comm->use_ce = _use_ce;
   _ub_comm->sms = _num_comm_sm;
@@ -1012,7 +1007,6 @@ void CommOverlapP2PBase::split_overlap_ag(const TensorWrapper &A, bool transa,
   NVTE_CHECK_CUDA(cudaStreamWaitEvent(stream_main, _stop_send, 0));
   NVTE_CHECK_CUDA(cudaEventRecord(_stop_recv, _stream_recv));
   NVTE_CHECK_CUDA(cudaStreamWaitEvent(stream_main, _stop_recv, 0));
-  NVTE_CHECK_CUDA(cudaStreamSynchronize(stream_main));
 }  // CommOverlapP2PBase::split_overlap_ag
 
 /*
@@ -1087,7 +1081,6 @@ void CommOverlapP2PBase::split_overlap_rs(const TensorWrapper &A, bool transa,
                                           TensorWrapper &workspace, bool grad, bool accumulate,
                                           bool use_split_accumulator, TensorWrapper &rs_output,
                                           cudaStream_t stream_main) {
-  NVTE_CHECK_CUDA(cudaStreamSynchronize(stream_main));
   int ori_sms = _ub_comm->sms;
   _ub_comm->use_ce = _use_ce;
   _ub_comm->sms = _num_comm_sm;
@@ -1163,7 +1156,6 @@ void CommOverlapP2PBase::split_overlap_rs(const TensorWrapper &A, bool transa,
   }
   NVTE_CHECK_CUDA(cudaEventRecord(_stop_recv, _stream_recv));
   NVTE_CHECK_CUDA(cudaStreamWaitEvent(stream_main, _stop_recv, 0));
-  NVTE_CHECK_CUDA(cudaStreamSynchronize(stream_main));
 
   // Reduce GEMM output chunks
   char *reduce_buf_ptr = reinterpret_cast<char *>(_ubufs[_tp_size - 1].dptr());