Skip to content

Fix deadlock when SHM based allreduce spin too fast #4048

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tjruwase merged 3 commits into from
Jul 28, 2023
Merged

Fix deadlock when SHM based allreduce spin too fast #4048

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f011c62
Fix deadlock when allreduce spin too fast
delock Jul 26, 2023
6361ebf
Change state to enum to increase readability
delock Jul 27, 2023
679bfd2
Merge branch 'master' into gma/fix_shm_allreduce
tjruwase Jul 28, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 35 additions & 13 deletions csrc/cpu/comm/ccl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,15 @@
#include <iostream>
#include <oneapi/ccl.hpp>

// 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;
Expand Down Expand Up @@ -63,19 +72,27 @@ 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, enum coll_state state)
{
volatile enum coll_state* 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)
{
Expand Down Expand Up @@ -313,7 +330,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) {
Expand Down Expand Up @@ -506,33 +523,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) {
wait_buffer_state_until(0, 0);
workspace[world_rank].state = 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, coll_allreduce_naive__reduce_done);
workspace[world_rank].state = coll_begin;
}
}

Expand Down