Skip to content

Fix crashes in the ring and jaccl distributed backends #3933

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

Open
jzdziarski wants to merge 1 commit into
base: main
Choose a base branch
from
+163 −64
Open

Fix crashes in the ring and jaccl distributed backends #3933

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
59 changes: 44 additions & 15 deletions mlx/distributed/jaccl/jaccl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,9 +83,16 @@ class JACCLGroup : public GroupImpl {
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, n_bytes, dtype, this]() {
group_->all_sum(in_ptr, out_ptr, n_bytes, dtype);
});
// Capture the buffers, not just the raw pointers, so they outlive the
// queued task. See all_gather below.
encoder.dispatch(
[in_ptr,
out_ptr,
n_bytes,
dtype,
in_buf = input.data_shared_ptr(),
out_buf = output.data_shared_ptr(),
this]() { group_->all_sum(in_ptr, out_ptr, n_bytes, dtype); });
}

void all_max(const array& input, array& output, Stream stream) override {
Expand All @@ -96,9 +103,14 @@ class JACCLGroup : public GroupImpl {
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, n_bytes, dtype, this]() {
group_->all_max(in_ptr, out_ptr, n_bytes, dtype);
});
encoder.dispatch(
[in_ptr,
out_ptr,
n_bytes,
dtype,
in_buf = input.data_shared_ptr(),
out_buf = output.data_shared_ptr(),
this]() { group_->all_max(in_ptr, out_ptr, n_bytes, dtype); });
}

void all_min(const array& input, array& output, Stream stream) override {
Expand All @@ -109,9 +121,14 @@ class JACCLGroup : public GroupImpl {
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, n_bytes, dtype, this]() {
group_->all_min(in_ptr, out_ptr, n_bytes, dtype);
});
encoder.dispatch(
[in_ptr,
out_ptr,
n_bytes,
dtype,
in_buf = input.data_shared_ptr(),
out_buf = output.data_shared_ptr(),
this]() { group_->all_min(in_ptr, out_ptr, n_bytes, dtype); });
}

void all_gather(const array& input, array& output, Stream stream) override {
Expand All @@ -121,9 +138,18 @@ class JACCLGroup : public GroupImpl {
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, n_bytes, this]() {
group_->all_gather(in_ptr, out_ptr, n_bytes);
});
// dispatch only enqueues the lambda, and nothing else keeps the buffers
// alive until it runs: set_input_array/set_output_array are no-ops on the
// CPU backend, and cpu::eval's keep-alive task holds the inputs but
// excludes the output, which the collective allocates here. Capturing the
// shared_ptr for both sides pins them for the task's lifetime.
encoder.dispatch(
[in_ptr,
out_ptr,
n_bytes,
in_buf = input.data_shared_ptr(),
out_buf = output.data_shared_ptr(),
this]() { group_->all_gather(in_ptr, out_ptr, n_bytes); });
}

void send(const array& input, int dst, Stream stream) override {
Expand All @@ -132,16 +158,19 @@ class JACCLGroup : public GroupImpl {
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.dispatch(
[data, n_bytes, dst, this]() { group_->send(data, n_bytes, dst); });
[data, n_bytes, dst, buf = input.data_shared_ptr(), this]() {
group_->send(data, n_bytes, dst);
});
}

void recv(array& out, int src, Stream stream) override {
auto data = out.data<char>();
size_t n_bytes = out.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_output_array(out);
encoder.dispatch(
[data, n_bytes, src, this]() { group_->recv(data, n_bytes, src); });
encoder.dispatch([data, n_bytes, src, buf = out.data_shared_ptr(), this]() {
group_->recv(data, n_bytes, src);
});
}

void sum_scatter(const array& input, array& output, Stream stream) override {
Expand Down
6 changes: 4 additions & 2 deletions mlx/distributed/jaccl/lib/jaccl/mesh_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,9 +214,11 @@ class MeshImpl {

// Counters to maintain the state of transfers
int in_flight = 0;
int read_offset = 0;
// Byte offsets into a payload sized by the int64_t `total`, so they have
// to be 64-bit as well or they wrap past 2GB per rank.
int64_t read_offset = 0;
int completed_send_count[PIPELINE] = {0};
int write_offset[MESH_MAX_PEERS] = {0};
int64_t write_offset[MESH_MAX_PEERS] = {0};

// Prefill the pipeline
int buff = 0;
Expand Down
76 changes: 65 additions & 11 deletions mlx/distributed/jaccl/lib/jaccl/ring_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

#pragma once

#include <iostream>
#include <span>

#include "jaccl/rdma.h"
Expand Down Expand Up @@ -130,13 +131,25 @@ class RingImpl {
ibv_wc wc[WC_NUM];
int n = poll(left_, right_, WC_NUM, wc);
for (int i = 0; i < n; i++) {
in_flight--;

// A failed or spurious completion carries an undefined wr_id, so
// the decoded wire and buff would index the counters and buffer
// pools out of bounds. Drop it and keep draining; the in-flight
// slot is already retired above.
int work_type = wc[i].wr_id >> 16;
int buff = (wc[i].wr_id >> 8) & 0xff;
int wire = wc[i].wr_id & 0xff;
int lr = wire / RING_MAX_CONNS;
int lw = wire % RING_MAX_CONNS;

in_flight--;
if (wc[i].status != IBV_WC_SUCCESS || lr >= MAX_DIR ||
lw >= n_wires || buff >= PIPELINE) {
std::cerr << IBV_TAG << " rank " << rank_
<< ": dropped bad all_reduce completion (status="
<< wc[i].status << ", wr_id=0x" << std::hex << wc[i].wr_id
<< std::dec << ")" << std::endl;
continue;
}

if (work_type == SEND_WR && send_count[wire] < n_steps) {
int64_t offset = lw * N + send_count[wire] * n_wires * N +
Expand Down Expand Up @@ -229,13 +242,25 @@ class RingImpl {
ibv_wc wc[WC_NUM];
int n = poll(left_, right_, WC_NUM, wc);
for (int i = 0; i < n; i++) {
in_flight--;

// A failed or spurious completion carries an undefined wr_id, so
// the decoded wire and buff would index the counters and buffer
// pools out of bounds. Drop it and keep draining; the in-flight
// slot is already retired above.
int work_type = wc[i].wr_id >> 16;
int buff = (wc[i].wr_id >> 8) & 0xff;
int wire = wc[i].wr_id & 0xff;
int lr = wire / RING_MAX_CONNS;
int lw = wire % RING_MAX_CONNS;

in_flight--;
if (wc[i].status != IBV_WC_SUCCESS || lr >= MAX_DIR ||
lw >= n_wires || buff >= PIPELINE) {
std::cerr << IBV_TAG << " rank " << rank_
<< ": dropped bad all_reduce completion (status="
<< wc[i].status << ", wr_id=0x" << std::hex << wc[i].wr_id
<< std::dec << ")" << std::endl;
continue;
}

if (work_type == SEND_WR && send_count[wire] < n_steps) {
int64_t offset = lw * N + send_count[wire] * n_wires * N +
Expand Down Expand Up @@ -306,7 +331,9 @@ class RingImpl {
size_t n_bytes_per_wire = (n_bytes + (2 * n_wires) - 1) / (2 * n_wires);
size_t out_bytes = n_bytes * size_;
auto [sz, N] = buffer_size_from_message(n_bytes_per_wire);
int n_steps = (n_bytes_per_wire + N - 1) / N;
// 64-bit to match all_reduce above: the offsets derived from the step
// count wrap an int once the per-rank payload reaches ~4GB.
int64_t n_steps = (n_bytes_per_wire + N - 1) / N;

// Counters to maintain the state of transfers
int in_flight = 0;
Expand Down Expand Up @@ -354,13 +381,22 @@ class RingImpl {
ibv_wc wc[WC_NUM];
int n = poll(left_, right_, WC_NUM, wc);
for (int i = 0; i < n; i++) {
in_flight--;

// Same guard as all_reduce. all_gather is always 2-directional.
int work_type = wc[i].wr_id >> 16;
int buff = (wc[i].wr_id >> 8) & 0xff;
int wire = wc[i].wr_id & 0xff;
int lr = wire / RING_MAX_CONNS;
int lw = wire % RING_MAX_CONNS;

in_flight--;
if (wc[i].status != IBV_WC_SUCCESS || lr >= 2 || lw >= n_wires ||
buff >= PIPELINE) {
std::cerr << IBV_TAG << " rank " << rank_
<< ": dropped bad all_gather completion (status="
<< wc[i].status << ", wr_id=0x" << std::hex << wc[i].wr_id
<< std::dec << ")" << std::endl;
continue;
}

if (work_type == SEND_WR && send_count[wire] < n_steps) {
int64_t offset = lw * N + send_count[wire] * n_wires * N +
Expand Down Expand Up @@ -450,11 +486,20 @@ class RingImpl {
ibv_wc wc[WC_NUM];
int n = poll(conns, WC_NUM, wc);
for (int i = 0; i < n; i++) {
in_flight--;

// Same guard as all_reduce.
int buff = (wc[i].wr_id >> 8) & 0xff;
int wire = wc[i].wr_id & 0xff;
int lw = wire % RING_MAX_CONNS;

in_flight--;
if (wc[i].status != IBV_WC_SUCCESS || lw >= n_wires ||
buff >= PIPELINE) {
std::cerr << IBV_TAG << " rank " << rank_
<< ": dropped bad send completion (status=" << wc[i].status
<< ", wr_id=0x" << std::hex << wc[i].wr_id << std::dec
<< ")" << std::endl;
continue;
}

if (read_offset[lw] < limits[lw]) {
std::copy(
Expand Down Expand Up @@ -513,11 +558,20 @@ class RingImpl {
ibv_wc wc[WC_NUM];
int n = poll(conns, WC_NUM, wc);
for (int i = 0; i < n; i++) {
in_flight--;

// Same guard as all_reduce.
int buff = (wc[i].wr_id >> 8) & 0xff;
int wire = wc[i].wr_id & 0xff;
int lw = wire % RING_MAX_CONNS;

in_flight--;
if (wc[i].status != IBV_WC_SUCCESS || lw >= n_wires ||
buff >= PIPELINE) {
std::cerr << IBV_TAG << " rank " << rank_
<< ": dropped bad recv completion (status=" << wc[i].status
<< ", wr_id=0x" << std::hex << wc[i].wr_id << std::dec
<< ")" << std::endl;
continue;
}

std::copy(
recv_buffer(sz, buff, dir, lw).begin<char>(),
Expand Down
86 changes: 50 additions & 36 deletions mlx/distributed/ring/ring.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -497,9 +497,13 @@ class RingGroup : public GroupImpl {
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
// Retain both buffers for the lifetime of the queued task. See
// mlx/distributed/jaccl/jaccl.cpp for the rationale.
encoder.dispatch([input_ptr = input.data<char>(),
nbytes = input.nbytes(),
output_ptr = output.data<char>(),
in_buf = input.data_shared_ptr(),
out_buf = output.data_shared_ptr(),
this]() {
constexpr size_t min_send_size = 262144;
size_t n_gathers = std::max(
Expand Down Expand Up @@ -533,46 +537,50 @@ class RingGroup : public GroupImpl {
void send(const array& input, int dst, Stream stream) override {
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.dispatch(
[input_ptr = input.data<char>(), nbytes = input.nbytes(), dst, this]() {
int right = (rank_ + 1) % size_;
int left = (rank_ + size_ - 1) % size_;
if (dst == right) {
send(sockets_right_, input_ptr, nbytes);
} else if (dst == left) {
send(sockets_left_, input_ptr, nbytes);
} else {
std::ostringstream msg;
msg << "[ring] Send only supported to direct neighbors "
<< "but tried to send to " << dst << " from " << rank_
<< std::endl;
throw std::runtime_error(msg.str());
}
});
encoder.dispatch([input_ptr = input.data<char>(),
nbytes = input.nbytes(),
dst,
in_buf = input.data_shared_ptr(),
this]() {
int right = (rank_ + 1) % size_;
int left = (rank_ + size_ - 1) % size_;
if (dst == right) {
send(sockets_right_, input_ptr, nbytes);
} else if (dst == left) {
send(sockets_left_, input_ptr, nbytes);
} else {
std::ostringstream msg;
msg << "[ring] Send only supported to direct neighbors "
<< "but tried to send to " << dst << " from " << rank_ << std::endl;
throw std::runtime_error(msg.str());
}
});
}

void recv(array& out, int src, Stream stream) override {
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_output_array(out);
encoder.dispatch(
[out_ptr = out.data<char>(), nbytes = out.nbytes(), src, this]() {
// NOTE: We 'll check the sockets with the opposite order of send so
// that they work even with 2 nodes where left and right is the same
// neighbor.
int right = (rank_ + 1) % size_;
int left = (rank_ + size_ - 1) % size_;
if (src == left) {
recv(sockets_left_, out_ptr, nbytes);
} else if (src == right) {
recv(sockets_right_, out_ptr, nbytes);
} else {
std::ostringstream msg;
msg << "[ring] Recv only supported from direct neighbors "
<< "but tried to recv from " << src << " to " << rank_
<< std::endl;
throw std::runtime_error(msg.str());
}
});
encoder.dispatch([out_ptr = out.data<char>(),
nbytes = out.nbytes(),
src,
out_buf = out.data_shared_ptr(),
this]() {
// NOTE: We 'll check the sockets with the opposite order of send so
// that they work even with 2 nodes where left and right is the same
// neighbor.
int right = (rank_ + 1) % size_;
int left = (rank_ + size_ - 1) % size_;
if (src == left) {
recv(sockets_left_, out_ptr, nbytes);
} else if (src == right) {
recv(sockets_right_, out_ptr, nbytes);
} else {
std::ostringstream msg;
msg << "[ring] Recv only supported from direct neighbors "
<< "but tried to recv from " << src << " to " << rank_ << std::endl;
throw std::runtime_error(msg.str());
}
});
}

void sum_scatter(const array& input, array& output, Stream stream) override {
Expand All @@ -590,7 +598,13 @@ class RingGroup : public GroupImpl {
auto out_ptr = output.data<char>();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, size = input.size(), this, reduce_op]() {
encoder.dispatch([in_ptr,
out_ptr,
size = input.size(),
in_buf = input.data_shared_ptr(),
out_buf = output.data_shared_ptr(),
this,
reduce_op]() {
// If the input data cannot be split into size_ segments then copy it and
// all reduce a local buffer prefilled with 0s.
size_t nbytes = size * sizeof(T);
Expand Down