From f011c62f02f96777f1bc28cf5785d234a0fcb697 Mon Sep 17 00:00:00 2001 From: "Ma, Guokai" Date: Wed, 26 Jul 2023 11:41:54 -0400 Subject: [PATCH 1/2] Fix deadlock when allreduce spin too fast --- csrc/cpu/comm/ccl.cpp | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/csrc/cpu/comm/ccl.cpp b/csrc/cpu/comm/ccl.cpp index 69859cc0c0d2..2b4a22acd32d 100644 --- a/csrc/cpu/comm/ccl.cpp +++ b/csrc/cpu/comm/ccl.cpp @@ -75,6 +75,14 @@ void wait_buffer_state_until(int index, int state) ; } +void wait_buffer_state_until_not(int index, int state) +{ + volatile int* state_ptr = &(workspace[index].state); + + while (*state_ptr == state) + ; +} + __m512 cvt_bf16_to_fp32(const __m256i src) __attribute__((target("avx512bw"))); inline __m512 cvt_bf16_to_fp32(const __m256i src) { @@ -526,7 +534,10 @@ void inference_all_reduce(torch::Tensor& data, py::object op, py::object group, workspace[world_rank].state = 0; } if (world_rank != 0) { - wait_buffer_state_until(0, 0); + // if rank 0 spin too fast it could be in state 1 of next allreduce + // in this case wait_buffer_state_until(0, 0) may cause deadlock + // what we are certain is when rank 0 finishes the state won't be 2 + wait_buffer_state_until_not(0, 2); workspace[world_rank].state = 0; } } From 6361ebf2b53cf4b8bb47f5d18e4c13278925792e Mon Sep 17 00:00:00 2001 From: "Ma, Guokai" Date: Wed, 26 Jul 2023 21:29:19 -0400 Subject: [PATCH 2/2] Change state to enum to increase readability --- csrc/cpu/comm/ccl.cpp | 41 ++++++++++++++++++++++++++--------------- 1 file changed, 26 insertions(+), 15 deletions(-) diff --git a/csrc/cpu/comm/ccl.cpp b/csrc/cpu/comm/ccl.cpp index 2b4a22acd32d..11c28415d6ce 100644 --- a/csrc/cpu/comm/ccl.cpp +++ b/csrc/cpu/comm/ccl.cpp @@ -16,6 +16,15 @@ #include #include +// states for collectives +enum coll_state { + coll_begin = 0, + // coll states for naive allreduce + coll_allreduce_naive__copy_in_done, // this state is for rank != 0 + coll_allreduce_naive__reduce_done, // this state is for rank == 0 + coll_allreduce_naive__copy_out_done, // this state is for rank != 0 +}; + // SHM building blocks struct SharedData { const char* name; @@ -62,22 +71,22 @@ void shared_close(SharedData* data) #define SHM_BUFFER_NAME "deepspeed_allreduce_buffer" SharedData allreduce_buffer; struct allreduce_workspace { - int state; + enum coll_state state; char buffer[MAX_BUF_SIZE]; }; struct allreduce_workspace* workspace; -void wait_buffer_state_until(int index, int state) +void wait_buffer_state_until(int index, enum coll_state state) { - volatile int* state_ptr = &(workspace[index].state); + volatile enum coll_state* state_ptr = &(workspace[index].state); while (*state_ptr != state) ; } -void wait_buffer_state_until_not(int index, int state) +void wait_buffer_state_until_not(int index, enum coll_state state) { - volatile int* state_ptr = &(workspace[index].state); + volatile enum coll_state* state_ptr = &(workspace[index].state); while (*state_ptr == state) ; @@ -316,7 +325,7 @@ void initialize(int size, int rank, torch::Tensor& kvs_data) workspace, size * sizeof(struct allreduce_workspace)); workspace = (struct allreduce_workspace*)allreduce_buffer.bytes; - for (int i = 0; i < size; i++) { workspace[i].state = 0; } + for (int i = 0; i < size; i++) { workspace[i].state = coll_begin; } } CCLCHECK(ccl::barrier(_get_comm_from_group()).wait()); if (rank != 0) { @@ -509,36 +518,38 @@ void inference_all_reduce(torch::Tensor& data, py::object op, py::object group, memcpy(workspace[world_rank].buffer, data_ptr, data_size); std::atomic_thread_fence(std::memory_order_release); - workspace[world_rank].state = 1; + workspace[world_rank].state = coll_allreduce_naive__copy_in_done; if (world_rank == 0) { // compute allreduce result on rank 0 for (int i = 1; i < world_size; i++) { // wait until the other rank copy the buffer - wait_buffer_state_until(i, 1); + wait_buffer_state_until(i, coll_allreduce_naive__copy_in_done); } reduce_all_buffers(workspace, numel, data.scalar_type(), world_size); std::atomic_thread_fence(std::memory_order_release); - workspace[world_rank].state = 2; + workspace[world_rank].state = coll_allreduce_naive__reduce_done; memcpy(data_ptr, workspace[0].buffer, data_size); } if (world_rank != 0) { - wait_buffer_state_until(0, 2); + wait_buffer_state_until(0, coll_allreduce_naive__reduce_done); memcpy(data_ptr, workspace[0].buffer, data_size); std::atomic_thread_fence(std::memory_order_release); - workspace[world_rank].state = 2; + workspace[world_rank].state = coll_allreduce_naive__copy_out_done; } if (world_rank == 0) { - for (int i = 1; i < world_size; i++) { wait_buffer_state_until(i, 2); } + for (int i = 1; i < world_size; i++) { + wait_buffer_state_until(i, coll_allreduce_naive__copy_out_done); + } std::atomic_thread_fence(std::memory_order_release); - workspace[world_rank].state = 0; + workspace[world_rank].state = coll_begin; } if (world_rank != 0) { // if rank 0 spin too fast it could be in state 1 of next allreduce // in this case wait_buffer_state_until(0, 0) may cause deadlock // what we are certain is when rank 0 finishes the state won't be 2 - wait_buffer_state_until_not(0, 2); - workspace[world_rank].state = 0; + wait_buffer_state_until_not(0, coll_allreduce_naive__reduce_done); + workspace[world_rank].state = coll_begin; } }