diff --git a/mlx/distributed/jaccl/jaccl.cpp b/mlx/distributed/jaccl/jaccl.cpp index 01cc415a00..35c80c554f 100644 --- a/mlx/distributed/jaccl/jaccl.cpp +++ b/mlx/distributed/jaccl/jaccl.cpp @@ -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 { @@ -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 { @@ -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 { @@ -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 { @@ -132,7 +158,9 @@ 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 { @@ -140,8 +168,9 @@ class JACCLGroup : public GroupImpl { 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 { diff --git a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h index 6327f41d35..4f59bbe501 100644 --- a/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h @@ -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; diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index 766f55386d..a45795a982 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -2,6 +2,7 @@ #pragma once +#include #include #include "jaccl/rdma.h" @@ -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 + @@ -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 + @@ -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; @@ -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 + @@ -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( @@ -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(), diff --git a/mlx/distributed/ring/ring.cpp b/mlx/distributed/ring/ring.cpp index ea40042844..5d5161bf6d 100644 --- a/mlx/distributed/ring/ring.cpp +++ b/mlx/distributed/ring/ring.cpp @@ -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(), nbytes = input.nbytes(), output_ptr = output.data(), + 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( @@ -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(), 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(), + 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(), 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(), + 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 { @@ -590,7 +598,13 @@ class RingGroup : public GroupImpl { auto out_ptr = output.data(); 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);