From ae7d19bf9e0a228317b575e854d5dc2a349d4af1 Mon Sep 17 00:00:00 2001 From: Jonathan Zdziarski Date: Sat, 25 Jul 2026 19:57:00 -0400 Subject: [PATCH] Fix crashes in the ring and jaccl distributed backends Three unrelated stability issues in the CPU-side distributed backends, all of which show up as SIGSEGVs on long-running multi-node jobs. Use-after-free in the collectives. `Encoder::dispatch` only enqueues the lambda, and `set_input_array`/`set_output_array` are no-ops on the CPU backend, so nothing keeps the arrays' buffers alive until the task runs. `cpu::eval` queues a keep-alive holding the inputs, but excludes the output, which the collectives allocate themselves. Under memory pressure the allocator can reclaim and reuse the output block before the collective runs, and the stream thread then writes through a stale pointer. Capture `data_shared_ptr()` alongside the raw pointers in the jaccl and ring collectives so the buffers are pinned for the task's lifetime. Unvalidated work completions in the jaccl ring. A failed or spurious completion carries an undefined `wr_id`, and the wire/buff decoded from it was used unchecked to index the per-wire counters, offsets, limits and the buffer pools. Out-of-range values produced a wild pointer in `reduce_op`/`std::copy`. Check `wc.status` and range-check the decoded indices, then drop the completion and keep draining. The in-flight counter is decremented before the check so the loop still terminates. 32-bit offset overflow. `read_offset`/`write_offset` in the mesh all_gather and `n_steps` in the ring all_gather were `int` while indexing payloads sized by `int64_t`, wrapping past 2GB and ~4GB per rank respectively. Widen them to `int64_t`. --- mlx/distributed/jaccl/jaccl.cpp | 59 ++++++++++---- mlx/distributed/jaccl/lib/jaccl/mesh_impl.h | 6 +- mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 76 +++++++++++++++--- mlx/distributed/ring/ring.cpp | 86 ++++++++++++--------- 4 files changed, 163 insertions(+), 64 deletions(-) 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);