L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

JACCL update (#3094)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2026-02-05 15:16:07 -08:00
committed by GitHub
parent 99ca62c4d3
commit ceea571490
11 changed files with 2455 additions and 1256 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/distributed/jaccl/CMakeLists.txt
+6 -1
View File
@@ -1,7 +1,12 @@
if(MLX_BUILD_CPU
AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin"
AND MACOS_SDK_VERSION GREATER_EQUAL 26.2)
target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jaccl.cpp)
target_sources(
mlx
PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jaccl.cpp
${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
${CMAKE_CURRENT_SOURCE_DIR}/mesh.cpp
${CMAKE_CURRENT_SOURCE_DIR}/ring.cpp)
else()
target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_jaccl.cpp)
endif()
mlx/distributed/jaccl/jaccl.cpp
+109 -1129
View File
File diff suppressed because it is too large Load Diff
mlx/distributed/jaccl/mesh.cpp
+451
View File
@@ -0,0 +1,451 @@
// Copyright © 2026 Apple Inc.
#include "mlx/distributed/jaccl/mesh.h"
#include "mlx/backend/cpu/encoder.h"
#include "mlx/distributed/reduction_ops.h"
#include "mlx/dtype_utils.h"
constexpr int MAX_PEERS = 8;
namespace mlx::core::distributed::jaccl {
MeshGroup::MeshGroup(
int rank,
const std::vector<std::string>& device_names,
const char* coordinator_addr)
: rank_(rank),
size_(device_names.size()),
side_channel_(rank_, size_, coordinator_addr),
connections_(create_connections(device_names)) {
if (size_ > MAX_PEERS) {
std::ostringstream msg;
msg << "[jaccl] The JACCL mesh supports up to " << MAX_PEERS
<< " peers but " << size_ << " were provided.";
throw std::runtime_error(msg.str());
}
// Initialize all the connections and allocate buffers
initialize();
// Make sure every node has reached here before continuing
side_channel_.all_gather<int>(0);
}
void MeshGroup::initialize() {
// Create the queue pairs
for (auto& conn : connections_) {
if (conn.ctx == nullptr) {
continue;
}
conn.allocate_protection_domain();
conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR);
conn.create_queue_pair();
}
allocate_buffers();
// First init all connections
for (int peer = 0; peer < size_; peer++) {
if (peer == rank_) {
continue;
}
connections_[peer].queue_pair_init();
}
// Gather the information to be exchanged, this also serves as a barrier so
// that all peers have initialized their connections before attempting to
// transition to RTS.
std::vector<Destination> info;
for (auto& conn : connections_) {
info.emplace_back(conn.info());
}
auto all_infos = side_channel_.all_gather(info);
// Transition queue pairs to RTS
for (int peer = 0; peer < size_; peer++) {
if (peer == rank_) {
continue;
}
auto peer_info = all_infos[peer][rank_];
connections_[peer].queue_pair_rtr(peer_info);
connections_[peer].queue_pair_rts();
}
}
void MeshGroup::allocate_buffers() {
// Deregister any buffers and free the memory
buffers_.clear();
// Allocate the memory
for (int k = 0; k < BUFFER_SIZES; k++) {
for (int i = 0; i < NUM_BUFFERS; i++) {
for (int j = 0; j < size_; j++) {
buffers_.emplace_back(FRAME_SIZE * (1 << k));
}
}
}
for (int k = 0; k < BUFFER_SIZES; k++) {
for (int i = 0; i < NUM_BUFFERS; i++) {
for (int j = 0; j < size_; j++) {
// This is our send buffer so register it with all pds so we can send
// it to all connected devices.
if (j == rank_) {
for (auto& conn : connections_) {
if (conn.ctx != nullptr) {
buffers_[k * NUM_BUFFERS * size_ + i * size_ + j]
.register_to_protection_domain(conn.protection_domain);
}
}
}
// This is the recv buffer from rank j so register it to rank j's
// protection domain.
else {
buffers_[k * NUM_BUFFERS * size_ + i * size_ + j]
.register_to_protection_domain(connections_[j].protection_domain);
}
}
}
}
}
void MeshGroup::all_sum(const array& input, array& output, Stream stream) {
dispatch_all_types(output.dtype(), [&](auto type_tag) {
using T = MLX_GET_TYPE(type_tag);
all_reduce<T>(input, output, stream, detail::SumOp<T>{});
});
}
void MeshGroup::all_max(const array& input, array& output, Stream stream) {
dispatch_all_types(output.dtype(), [&](auto type_tag) {
using T = MLX_GET_TYPE(type_tag);
all_reduce<T>(input, output, stream, detail::MaxOp<T>{});
});
}
void MeshGroup::all_min(const array& input, array& output, Stream stream) {
dispatch_all_types(output.dtype(), [&](auto type_tag) {
using T = MLX_GET_TYPE(type_tag);
all_reduce<T>(input, output, stream, detail::MinOp<T>{});
});
}
void MeshGroup::all_gather(const array& input, array& output, Stream stream) {
auto in_ptr = input.data<char>();
auto out_ptr = output.data<char>();
size_t n_bytes = input.nbytes();
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]() {
// Copy our data to the appropriate place
std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes);
// Fully connected all gather
char* data = out_ptr;
char* our_data = out_ptr + rank_ * n_bytes;
auto [sz, N] = buffer_size_from_message(n_bytes);
constexpr int PIPELINE = 2;
constexpr int WC_NUM = PIPELINE * MAX_PEERS * 2;
int64_t total = static_cast<int64_t>(n_bytes);
int num_peers = size_ - 1;
// Counters to maintain the state of transfers
int in_flight = 0;
int read_offset = 0;
int completed_send_count[PIPELINE] = {0};
int write_offset[MAX_PEERS] = {0};
// Prefill the pipeline
int buff = 0;
while (read_offset < total && buff < PIPELINE) {
post_recv_all(sz, buff);
std::copy(
our_data + read_offset,
our_data + std::min(read_offset + N, total),
send_buffer(sz, buff).begin<char>());
post_send_all(sz, buff);
buff++;
in_flight += 2 * num_peers;
read_offset += N;
}
// Main loop
//
// Keep going until we have no longer data in flight.
while (in_flight > 0) {
ibv_wc wc[WC_NUM];
int n = poll(connections_, WC_NUM, wc);
for (int i = 0; i < n; i++) {
int work_type = wc[i].wr_id >> 16;
int buff = (wc[i].wr_id >> 8) & 0xff;
int rank = wc[i].wr_id & 0xff;
in_flight--;
// Send completed. If all sends completed then send the next chunk.
if (work_type == SEND_WR && read_offset < total) {
completed_send_count[buff]++;
if (completed_send_count[buff] == num_peers) {
std::copy(
our_data + read_offset,
our_data + std::min(read_offset + N, total),
send_buffer(sz, buff).begin<char>());
post_send_all(sz, buff);
completed_send_count[buff] = 0;
in_flight += num_peers;
read_offset += N;
}
}
// Recv completed. If we have more chunks then post another recv.
else if (work_type == RECV_WR) {
std::copy(
recv_buffer(sz, buff, rank).begin<char>(),
recv_buffer(sz, buff, rank).begin<char>() +
std::min(N, total - write_offset[rank]),
data + rank * n_bytes + write_offset[rank]);
write_offset[rank] += N;
if (write_offset[rank] + N * (PIPELINE - 1) < total) {
recv_from(sz, rank, buff);
in_flight++;
}
}
}
}
});
}
void MeshGroup::send(const array& input, int dst, Stream stream) {
auto data = input.data<char>();
int64_t n_bytes = input.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.dispatch([data, n_bytes, dst, this]() {
constexpr int PIPELINE = 2;
constexpr int WC_NUM = PIPELINE;
auto [sz, N] = buffer_size_from_message(n_bytes);
int in_flight = 0;
int64_t read_offset = 0;
// Prefill the pipeline
int buff = 0;
while (read_offset < n_bytes && buff < PIPELINE) {
std::copy(
data + read_offset,
data + std::min(read_offset + N, n_bytes),
send_buffer(sz, buff).begin<char>());
send_to(sz, dst, buff);
buff++;
read_offset += N;
in_flight++;
}
// Main loop
while (in_flight > 0) {
// Poll the hardware for completions.
//
// If a send was completed and we have more data to send then go ahead
// and send them.
ibv_wc wc[WC_NUM];
int n = connections_[dst].poll(WC_NUM, wc);
for (int i = 0; i < n; i++) {
int buff = (wc[i].wr_id >> 8) & 0xff;
int rank = wc[i].wr_id & 0xff;
in_flight--;
if (read_offset < n_bytes) {
std::copy(
data + read_offset,
data + std::min(read_offset + N, n_bytes),
send_buffer(sz, buff).begin<char>());
send_to(sz, dst, buff);
read_offset += N;
in_flight++;
}
}
}
});
}
void MeshGroup::recv(array& out, int src, Stream stream) {
auto data = out.data<char>();
int64_t n_bytes = out.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_output_array(out);
encoder.dispatch([data, n_bytes, src, this]() {
constexpr int PIPELINE = 2;
constexpr int WC_NUM = PIPELINE;
auto [sz, N] = buffer_size_from_message(n_bytes);
int in_flight = 0;
int64_t write_offset = 0;
// Prefill the pipeline
int buff = 0;
while (N * buff < n_bytes && buff < PIPELINE) {
recv_from(sz, src, buff);
in_flight++;
buff++;
}
// Main loop
while (in_flight > 0) {
// Poll the hardware for completions.
//
// If a recv was completed copy it to the output and if we have more
// data to fetch post another recv.
ibv_wc wc[WC_NUM];
int n = connections_[src].poll(WC_NUM, wc);
for (int i = 0; i < n; i++) {
int buff = (wc[i].wr_id >> 8) & 0xff;
int rank = wc[i].wr_id & 0xff;
in_flight--;
std::copy(
recv_buffer(sz, buff, src).begin<char>(),
recv_buffer(sz, buff, src).begin<char>() +
std::min(n_bytes - write_offset, static_cast<int64_t>(N)),
data + write_offset);
write_offset += N;
if (write_offset + (PIPELINE - 1) * N < n_bytes) {
recv_from(sz, src, buff);
in_flight++;
}
}
}
});
}
template <typename T, typename ReduceOp>
void MeshGroup::all_reduce(
const array& input,
array& output,
Stream stream,
ReduceOp reduce_op) {
auto in_ptr = input.data<T>();
auto out_ptr = output.data<T>();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, size = input.size(), this, reduce_op]() {
// If not inplace all reduce then copy the input to the output first
if (in_ptr != out_ptr) {
std::memcpy(out_ptr, in_ptr, size * sizeof(T));
}
// Fully connected all reduce
T* data = out_ptr;
auto [sz, buffer_size] = buffer_size_from_message(size * sizeof(T));
int64_t N = buffer_size / sizeof(T);
constexpr int PIPELINE = 2;
constexpr int WC_NUM = PIPELINE * MAX_PEERS * 2;
int64_t total = static_cast<int64_t>(size);
int num_peers = size_ - 1;
// Counters to maintain the state of transfers
int in_flight = 0;
int64_t read_offset = 0;
int completed_send_count[PIPELINE] = {0};
int completed_recv_begin[MAX_PEERS] = {0};
int completed_recv_end[MAX_PEERS] = {0};
// Prefill the pipeline
int buff = 0;
while (read_offset < total && buff < PIPELINE) {
post_recv_all(sz, buff);
std::copy(
data + read_offset,
data + std::min(read_offset + N, total),
send_buffer(sz, buff).begin<T>());
post_send_all(sz, buff);
buff++;
in_flight += 2 * num_peers;
read_offset += N;
}
// Main loop
//
// Keep going until we have no longer data in flight.
while (in_flight > 0) {
// Poll the hardware for completions.
//
// If a send was completed mark how many completions we have received
// for that buffer. If we have sent the buffer to all peers we can
// reuse the buffer so copy the next chunk of data and send it to all.
//
// If a receive is completed then advance the pointer of completed
// receives.
ibv_wc wc[WC_NUM];
int n = poll(connections_, WC_NUM, wc);
for (int i = 0; i < n; i++) {
int work_type = wc[i].wr_id >> 16;
int buff = (wc[i].wr_id >> 8) & 0xff;
int rank = wc[i].wr_id & 0xff;
in_flight--;
if (work_type == SEND_WR && read_offset < total) {
completed_send_count[buff]++;
if (completed_send_count[buff] == num_peers) {
std::copy(
data + read_offset,
data + std::min(read_offset + N, total),
send_buffer(sz, buff).begin<T>());
post_send_all(sz, buff);
completed_send_count[buff] = 0;
in_flight += num_peers;
read_offset += N;
}
}
else if (work_type == RECV_WR) {
completed_recv_end[rank]++;
}
}
// Process the completed recv
//
// For each rank we have a range of completed recv defined by a begin
// and end inclusive and exlusive in standard C++ fashion.
//
// When there is an unprocessed receive we first check if we have
// finished sending the write location. If so then we reduce in-place
// and then check if there is more to be received and post a recv.
for (int r = 0; r < size_; r++) {
int s = completed_recv_begin[r];
int e = completed_recv_end[r];
int w = s * N;
while (w < read_offset && e - s > 0) {
int buff = s % PIPELINE;
reduce_op(
recv_buffer(sz, buff, r).begin<T>(),
data + w,
std::min(N, total - w));
w += N;
s++;
if (w + (PIPELINE - 1) * N < total) {
recv_from(sz, r, buff);
in_flight++;
}
}
completed_recv_begin[r] = s;
}
}
});
}
} // namespace mlx::core::distributed::jaccl
mlx/distributed/jaccl/mesh.h
+122
View File
@@ -0,0 +1,122 @@
// Copyright © 2026 Apple Inc.
#pragma once
#include "mlx/distributed/distributed_impl.h"
#include "mlx/distributed/jaccl/utils.h"
using GroupImpl = mlx::core::distributed::detail::GroupImpl;
namespace mlx::core::distributed::jaccl {
/**
* The JACCL communication group for a fully connected mesh. We expect one
* connection per peer and it should be the lowest latency communication group
* for small to medium size messages.
*
* Like all JACCL groups it uses a side channel to exchange the necessary
* information and then configure the connections to be ready for RDMA
* operations.
*/
class MeshGroup : public GroupImpl {
public:
MeshGroup(
int rank,
const std::vector<std::string>& device_names,
const char* coordinator_addr);
Stream communication_stream(StreamOrDevice s) override {
return to_stream(s, Device::cpu);
}
int rank() override {
return rank_;
}
int size() override {
return size_;
}
void all_sum(const array& input, array& output, Stream stream) override;
void all_max(const array& input, array& output, Stream stream) override;
void all_min(const array& input, array& output, Stream stream) override;
void all_gather(const array& input, array& output, Stream stream) override;
void send(const array& input, int dst, Stream stream) override;
void recv(array& out, int src, Stream stream) override;
void sum_scatter(const array& input, array& output, Stream stream) override {
throw std::runtime_error("[jaccl] sum_scatter not supported.");
}
std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
throw std::runtime_error("[jaccl] Group split not supported.");
}
private:
template <typename T, typename ReduceOp>
void all_reduce(
const array& input,
array& output,
Stream stream,
ReduceOp reduce_op);
/**
* Performs the connection initialization. Namely, after this call all
* Connection objects should have a queue pair in RTS state and all buffers
* should have been allocated.
*/
void initialize();
/**
* Allocate all the buffers that we will use in the communication group.
*/
void allocate_buffers();
void send_to(int sz, int rank, int buff) {
connections_[rank].post_send(
send_buffer(sz, buff), SEND_WR << 16 | buff << 8 | rank);
}
void recv_from(int sz, int rank, int buff) {
connections_[rank].post_recv(
recv_buffer(sz, buff, rank), RECV_WR << 16 | buff << 8 | rank);
}
SharedBuffer& send_buffer(int sz, int buff) {
return buffers_[sz * NUM_BUFFERS * size_ + buff * size_ + rank_];
}
SharedBuffer& recv_buffer(int sz, int buff, int rank) {
return buffers_[sz * NUM_BUFFERS * size_ + buff * size_ + rank];
}
void post_send_all(int sz, int buff) {
auto& b = send_buffer(sz, buff);
int wr_id = SEND_WR << 16 | buff << 8;
for (int i = 0; i < size_; i++) {
if (i == rank_) {
continue;
}
connections_[i].post_send(b, wr_id | i);
}
}
void post_recv_all(int sz, int buff) {
int b = sz * NUM_BUFFERS * size_ + buff * size_;
int wr_id = RECV_WR << 16 | buff << 8;
for (int i = 0; i < size_; i++) {
if (i == rank_) {
continue;
}
connections_[i].post_recv(buffers_[b + i], wr_id | i);
}
}
int rank_;
int size_;
SideChannel side_channel_;
std::vector<Connection> connections_;
std::vector<SharedBuffer> buffers_;
};
} // namespace mlx::core::distributed::jaccl
mlx/distributed/jaccl/ring.cpp
+692
View File
@@ -0,0 +1,692 @@
// Copyright © 2026 Apple Inc.
#include "mlx/distributed/jaccl/ring.h"
#include "mlx/backend/cpu/encoder.h"
#include "mlx/distributed/reduction_ops.h"
#include "mlx/dtype_utils.h"
namespace mlx::core::distributed::jaccl {
RingGroup::RingGroup(
int rank,
int size,
const std::vector<std::string>& left_devices,
const std::vector<std::string>& right_devices,
const char* coordinator_addr)
: rank_(rank),
size_(size),
side_channel_(rank_, size_, coordinator_addr),
left_(create_connections(left_devices)),
right_(create_connections(right_devices)) {
if (left_.size() > MAX_CONNS || right_.size() > MAX_CONNS) {
std::ostringstream msg;
msg << "[jaccl] Up to " << MAX_CONNS << " per direction supported but "
<< left_.size() << " were provided.";
throw std::runtime_error(msg.str());
}
// Initialize all the connections and allocate buffers
initialize();
// Make sure every node has reached here before continuing
side_channel_.all_gather<int>(0);
}
void RingGroup::initialize() {
// Create the queue pairs
for (auto& conn : left_) {
conn.allocate_protection_domain();
conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR);
conn.create_queue_pair();
}
for (auto& conn : right_) {
conn.allocate_protection_domain();
conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR);
conn.create_queue_pair();
}
// Allocate the buffers
allocate_buffers();
// Initialize the conections
for (auto& conn : left_) {
conn.queue_pair_init();
}
for (auto& conn : right_) {
conn.queue_pair_init();
}
// Gather the information to be exchanged, this also serves as a barrier so
// that all peers have initialized their connections before attempting to
// transition to RTS.
std::vector<Destination> left_info;
for (auto& conn : left_) {
left_info.emplace_back(conn.info());
}
std::vector<Destination> right_info;
for (auto& conn : right_) {
right_info.emplace_back(conn.info());
}
auto all_left_infos = side_channel_.all_gather(left_info);
auto all_right_infos = side_channel_.all_gather(right_info);
// Transition queue pairs to RTS
int left_peer = (rank_ + size_ - 1) % size_;
for (int i = 0; i < left_.size(); i++) {
auto peer_info = all_right_infos[left_peer][i];
left_[i].queue_pair_rtr(peer_info);
left_[i].queue_pair_rts();
}
int right_peer = (rank_ + 1) % size_;
for (int i = 0; i < right_.size(); i++) {
auto peer_info = all_left_infos[right_peer][i];
right_[i].queue_pair_rtr(peer_info);
right_[i].queue_pair_rts();
}
}
void RingGroup::allocate_buffers() {
// Deregister any buffers and free the memory
send_buffers_.clear();
recv_buffers_.clear();
// Allocate the memory
for (int k = 0; k < BUFFER_SIZES; k++) {
for (int i = 0; i < NUM_BUFFERS; i++) {
for (int j = 0; j < MAX_CONNS * 2; j++) {
send_buffers_.emplace_back(FRAME_SIZE * (1 << k));
recv_buffers_.emplace_back(FRAME_SIZE * (1 << k));
}
}
}
// Register the buffers with the corresponding connections
for (int k = 0; k < BUFFER_SIZES; k++) {
for (int i = 0; i < NUM_BUFFERS; i++) {
for (int j = 0; j < MAX_CONNS * 2; j++) {
int wire = j % MAX_CONNS;
int lr = j / MAX_CONNS;
if (wire >= left_.size()) {
continue;
}
if (lr) {
send_buffers_[k * NUM_BUFFERS * MAX_CONNS * 2 + i * MAX_CONNS * 2 + j]
.register_to_protection_domain(left_[wire].protection_domain);
recv_buffers_[k * NUM_BUFFERS * MAX_CONNS * 2 + i * MAX_CONNS * 2 + j]
.register_to_protection_domain(right_[wire].protection_domain);
} else {
send_buffers_[k * NUM_BUFFERS * MAX_CONNS * 2 + i * MAX_CONNS * 2 + j]
.register_to_protection_domain(right_[wire].protection_domain);
recv_buffers_[k * NUM_BUFFERS * MAX_CONNS * 2 + i * MAX_CONNS * 2 + j]
.register_to_protection_domain(left_[wire].protection_domain);
}
}
}
}
}
void RingGroup::all_sum(const array& input, array& output, Stream stream) {
dispatch_all_types(output.dtype(), [&](auto type_tag) {
using T = MLX_GET_TYPE(type_tag);
all_reduce<T>(input, output, stream, detail::SumOp<T>{});
});
}
void RingGroup::all_max(const array& input, array& output, Stream stream) {
dispatch_all_types(output.dtype(), [&](auto type_tag) {
using T = MLX_GET_TYPE(type_tag);
all_reduce<T>(input, output, stream, detail::MaxOp<T>{});
});
}
void RingGroup::all_min(const array& input, array& output, Stream stream) {
dispatch_all_types(output.dtype(), [&](auto type_tag) {
using T = MLX_GET_TYPE(type_tag);
all_reduce<T>(input, output, stream, detail::MinOp<T>{});
});
}
void RingGroup::all_gather(const array& input, array& output, Stream stream) {
auto in_ptr = input.data<char>();
auto out_ptr = output.data<char>();
size_t n_bytes = input.nbytes();
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]() {
// Copy our data to the appropriate place
std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes);
constexpr int PIPELINE = 2;
constexpr int WC_NUM = PIPELINE * MAX_CONNS * 2 * 2;
int n_wires = left_.size();
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;
// Counters to maintain the state of transfers
int in_flight = 0;
int64_t send_offset[2];
int64_t recv_offset[2];
int64_t limits[2];
int send_count[2 * MAX_CONNS] = {0};
int recv_count[2 * MAX_CONNS] = {0};
send_offset[0] = send_offset[1] = rank_ * n_bytes;
recv_offset[0] = ((rank_ + size_ - 1) % size_) * n_bytes;
recv_offset[1] = ((rank_ + 1) % size_) * n_bytes;
limits[0] = n_wires * n_bytes_per_wire;
limits[1] = n_bytes;
// Possible perf improvement by not syncing at every step but running ahead
// as needed.
for (int k = 0; k < size_ - 1; k++) {
// Prefill the pipeline
int buff = 0;
while (buff < n_steps && buff < PIPELINE) {
post_recv_all(sz, buff);
for (int lr = 0; lr < 2; lr++) {
for (int lw = 0; lw < n_wires; lw++) {
int64_t offset = lw * N +
send_count[lr * MAX_CONNS + lw] * n_wires * N +
lr * n_wires * n_bytes_per_wire;
std::copy(
out_ptr + send_offset[lr] + offset,
out_ptr + send_offset[lr] +
std::max(offset, std::min(offset + N, limits[lr])),
send_buffer(sz, buff, lr, lw).begin<char>());
send_count[lr * MAX_CONNS + lw]++;
}
}
post_send_all(sz, buff);
buff++;
in_flight += 2 * 2 * n_wires;
}
// Main loop
//
// Keep going until we have no longer data in flight.
while (in_flight > 0) {
ibv_wc wc[WC_NUM];
int n = poll(left_, right_, WC_NUM, wc);
for (int i = 0; i < n; i++) {
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 / MAX_CONNS;
int lw = wire % MAX_CONNS;
in_flight--;
if (work_type == SEND_WR && send_count[wire] < n_steps) {
int64_t offset = lw * N + send_count[wire] * n_wires * N +
lr * n_wires * n_bytes_per_wire;
std::copy(
out_ptr + send_offset[lr] + offset,
out_ptr + send_offset[lr] +
std::max(offset, std::min(offset + N, limits[lr])),
send_buffer(sz, buff, lr, lw).begin<char>());
send_to(sz, buff, lr, lw);
in_flight++;
send_count[wire]++;
}
else if (work_type == RECV_WR) {
int64_t offset = lw * N + recv_count[wire] * n_wires * N +
lr * n_wires * n_bytes_per_wire;
std::copy(
recv_buffer(sz, buff, lr, lw).begin<char>(),
recv_buffer(sz, buff, lr, lw).begin<char>() +
std::max<int64_t>(0, std::min(N, limits[lr] - offset)),
out_ptr + recv_offset[lr] + offset);
recv_count[wire]++;
if (recv_count[wire] + (PIPELINE - 1) < n_steps) {
recv_from(sz, buff, lr, lw);
in_flight++;
}
}
}
}
send_offset[0] = (send_offset[0] + out_bytes - n_bytes) % out_bytes;
recv_offset[0] = (recv_offset[0] + out_bytes - n_bytes) % out_bytes;
send_offset[1] = (send_offset[1] + n_bytes) % out_bytes;
recv_offset[1] = (recv_offset[1] + n_bytes) % out_bytes;
for (int i = 0; i < 2 * MAX_CONNS; i++) {
send_count[i] = recv_count[i] = 0;
}
}
});
}
void RingGroup::send(const array& input, int dst, Stream stream) {
int right = (rank_ + 1) % size_;
int left = (rank_ + size_ - 1) % size_;
if (dst != right && dst != left) {
std::ostringstream msg;
msg << "[jaccl] In ring mode send is only supported to direct neighbors "
<< "but tried to send to " << dst << " from " << rank_ << std::endl;
throw std::runtime_error(msg.str());
}
auto data = input.data<char>();
int64_t n_bytes = input.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.dispatch([data, n_bytes, dst, left, this]() {
// In the case that size_ == 2 then left == right so we bias send towards
// left and recv towards right so that the selections will be correct for
// the 2 node case.
auto& conns = (dst == left) ? left_ : right_;
int dir = dst == left;
constexpr int PIPELINE = 2;
constexpr int WC_NUM = PIPELINE * MAX_CONNS;
int n_wires = conns.size();
int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires;
auto [sz, N] = buffer_size_from_message(bytes_per_wire);
int in_flight = 0;
int64_t read_offset[MAX_CONNS];
int64_t limits[MAX_CONNS];
for (int lw = 0; lw < n_wires; lw++) {
read_offset[lw] = std::min(lw * bytes_per_wire, n_bytes);
limits[lw] = std::min((lw + 1) * bytes_per_wire, n_bytes);
}
// Prefill the pipeline
for (int lw = 0; lw < n_wires; lw++) {
int buff = 0;
while (read_offset[lw] < limits[lw] && buff < PIPELINE) {
std::copy(
data + read_offset[lw],
data + std::min(read_offset[lw] + N, limits[lw]),
send_buffer(sz, buff, dir, lw).begin<char>());
send_to(sz, buff, dir, lw);
buff++;
read_offset[lw] += N;
in_flight++;
}
}
// Main loop
while (in_flight > 0) {
// Poll the hardware for completions.
//
// If a send was completed and we have more data to send then go ahead
// and send them.
ibv_wc wc[WC_NUM];
int n = poll(conns, WC_NUM, wc);
for (int i = 0; i < n; i++) {
int buff = (wc[i].wr_id >> 8) & 0xff;
int wire = wc[i].wr_id & 0xff;
int lw = wire % MAX_CONNS;
in_flight--;
if (read_offset[lw] < limits[lw]) {
std::copy(
data + read_offset[lw],
data + std::min(read_offset[lw] + N, limits[lw]),
send_buffer(sz, buff, dir, lw).begin<char>());
send_to(sz, buff, dir, lw);
read_offset[lw] += N;
in_flight++;
}
}
}
});
}
void RingGroup::recv(array& out, int src, Stream stream) {
int right = (rank_ + 1) % size_;
int left = (rank_ + size_ - 1) % size_;
if (src != right && src != left) {
std::ostringstream msg;
msg << "[jaccl] In ring mode recv is only supported to direct neighbors "
<< "but tried to recv from " << src << " to " << rank_ << std::endl;
throw std::runtime_error(msg.str());
}
auto data = out.data<char>();
int64_t n_bytes = out.nbytes();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_output_array(out);
encoder.dispatch([data, n_bytes, src, right, this]() {
// In the case that size_ == 2 then left == right so we bias send towards
// left and recv towards right so that the selections will be correct for
// the 2 node case.
auto& conns = (src == right) ? right_ : left_;
int dir = src == right;
constexpr int PIPELINE = 2;
constexpr int WC_NUM = PIPELINE * MAX_CONNS;
int n_wires = conns.size();
int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires;
auto [sz, N] = buffer_size_from_message(bytes_per_wire);
int in_flight = 0;
int64_t write_offset[MAX_CONNS];
int64_t limits[MAX_CONNS];
for (int lw = 0; lw < n_wires; lw++) {
write_offset[lw] = std::min(lw * bytes_per_wire, n_bytes);
limits[lw] = std::min((lw + 1) * bytes_per_wire, n_bytes);
}
// Prefill the pipeline
for (int lw = 0; lw < n_wires; lw++) {
int buff = 0;
while (N * buff < limits[lw] && buff < PIPELINE) {
recv_from(sz, buff, dir, lw);
buff++;
in_flight++;
}
}
// Main loop
while (in_flight > 0) {
// Poll the hardware for completions.
//
// If a recv was completed copy it to the output and if we have more
// data to fetch post another recv.
ibv_wc wc[WC_NUM];
int n = poll(conns, WC_NUM, wc);
for (int i = 0; i < n; i++) {
int buff = (wc[i].wr_id >> 8) & 0xff;
int wire = wc[i].wr_id & 0xff;
int lw = wire % MAX_CONNS;
in_flight--;
std::copy(
recv_buffer(sz, buff, dir, lw).begin<char>(),
recv_buffer(sz, buff, dir, lw).begin<char>() +
std::max<int64_t>(
0, std::min<int64_t>(limits[lw] - write_offset[lw], N)),
data + write_offset[lw]);
write_offset[lw] += N;
if (write_offset[lw] + (PIPELINE - 1) * N < limits[lw]) {
recv_from(sz, buff, dir, lw);
in_flight++;
}
}
}
});
}
template <typename T, typename ReduceOp>
void RingGroup::all_reduce(
const array& input,
array& output,
Stream stream,
ReduceOp reduce_op) {
auto in_ptr = input.data<T>();
auto out_ptr = output.data<T>();
auto& encoder = cpu::get_command_encoder(stream);
encoder.set_input_array(input);
encoder.set_output_array(output);
encoder.dispatch([in_ptr, out_ptr, size = input.size(), this, reduce_op]() {
if (size < size_ * 2 * left_.size()) {
all_reduce_impl<1, T, ReduceOp>(in_ptr, out_ptr, size, 1, reduce_op);
return;
}
all_reduce_impl<2, T, ReduceOp>(
in_ptr, out_ptr, size, left_.size(), reduce_op);
});
}
template <int MAX_DIR, typename T, typename ReduceOp>
void RingGroup::all_reduce_impl(
const T* in_ptr,
T* out_ptr,
int64_t size,
int n_wires,
ReduceOp reduce_op) {
// If not inplace all reduce then copy the input to the output first
if (in_ptr != out_ptr) {
std::memcpy(out_ptr, in_ptr, size * sizeof(T));
}
constexpr int PIPELINE = 2;
constexpr int WC_NUM = PIPELINE * MAX_CONNS * 2 * MAX_DIR;
int64_t chunk_size = (size + size_ - 1) / size_;
int64_t size_per_wire =
(chunk_size + (MAX_DIR * n_wires) - 1) / (MAX_DIR * n_wires);
auto [sz, N] = buffer_size_from_message(size_per_wire * sizeof(T));
N /= sizeof(T);
int64_t n_steps = (size_per_wire + N - 1) / N;
// Counters to maintain the state of transfers
int in_flight = 0;
int64_t chunk_multiple_size = size_ * chunk_size;
int64_t send_offset[MAX_DIR];
int64_t recv_offset[MAX_DIR];
int64_t send_limits[MAX_DIR];
int64_t recv_limits[MAX_DIR];
int send_count[MAX_DIR * MAX_CONNS] = {0};
int recv_count[MAX_DIR * MAX_CONNS] = {0};
send_offset[0] = rank_ * chunk_size;
recv_offset[0] = ((rank_ + size_ - 1) % size_) * chunk_size;
if constexpr (MAX_DIR == 2) {
send_offset[1] = rank_ * chunk_size;
recv_offset[1] = ((rank_ + 1) % size_) * chunk_size;
send_limits[0] = std::min(
n_wires * size_per_wire, std::max<int64_t>(0, size - send_offset[0]));
send_limits[1] =
std::min(chunk_size, std::max<int64_t>(0, size - send_offset[1]));
recv_limits[0] = std::min(
n_wires * size_per_wire, std::max<int64_t>(0, size - recv_offset[0]));
recv_limits[1] =
std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[1]));
} else {
send_limits[0] =
std::min(chunk_size, std::max<int64_t>(0, size - send_offset[0]));
recv_limits[0] =
std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[0]));
}
// First reduce scatter
//
// Possible perf improvement by not syncing at every step but running ahead
// as needed.
for (int k = 0; k < size_ - 1; k++) {
// Prefill the pipeline
int buff = 0;
while (buff < n_steps && buff < PIPELINE) {
post_recv_all<MAX_DIR>(sz, buff, n_wires);
for (int lr = 0; lr < MAX_DIR; lr++) {
for (int lw = 0; lw < n_wires; lw++) {
int64_t offset = lw * N +
send_count[lr * MAX_CONNS + lw] * n_wires * N +
lr * n_wires * size_per_wire;
std::copy(
out_ptr + send_offset[lr] + offset,
out_ptr + send_offset[lr] +
std::max(offset, std::min(offset + N, send_limits[lr])),
send_buffer(sz, buff, lr, lw).begin<T>());
send_count[lr * MAX_CONNS + lw]++;
}
}
post_send_all<MAX_DIR>(sz, buff, n_wires);
buff++;
in_flight += 2 * MAX_DIR * n_wires;
}
// Main loop
//
// Keep going until we have no longer data in flight.
while (in_flight > 0) {
ibv_wc wc[WC_NUM];
int n = poll(left_, right_, WC_NUM, wc);
for (int i = 0; i < n; i++) {
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 / MAX_CONNS;
int lw = wire % MAX_CONNS;
in_flight--;
if (work_type == SEND_WR && send_count[wire] < n_steps) {
int64_t offset = lw * N + send_count[wire] * n_wires * N +
lr * n_wires * size_per_wire;
std::copy(
out_ptr + send_offset[lr] + offset,
out_ptr + send_offset[lr] +
std::max(offset, std::min(offset + N, send_limits[lr])),
send_buffer(sz, buff, lr, lw).begin<T>());
send_to(sz, buff, lr, lw);
in_flight++;
send_count[wire]++;
}
else if (work_type == RECV_WR) {
int64_t offset = lw * N + recv_count[wire] * n_wires * N +
lr * n_wires * size_per_wire;
reduce_op(
recv_buffer(sz, buff, lr, lw).begin<T>(),
out_ptr + recv_offset[lr] + offset,
std::max<int64_t>(0, std::min(N, recv_limits[lr] - offset)));
recv_count[wire]++;
if (recv_count[wire] + (PIPELINE - 1) < n_steps) {
recv_from(sz, buff, lr, lw);
in_flight++;
}
}
}
}
send_offset[0] = (send_offset[0] + chunk_multiple_size - chunk_size) %
chunk_multiple_size;
recv_offset[0] = (recv_offset[0] + chunk_multiple_size - chunk_size) %
chunk_multiple_size;
if constexpr (MAX_DIR == 2) {
send_offset[1] = (send_offset[1] + chunk_size) % chunk_multiple_size;
recv_offset[1] = (recv_offset[1] + chunk_size) % chunk_multiple_size;
send_limits[0] = std::min(
n_wires * size_per_wire, std::max<int64_t>(0, size - send_offset[0]));
send_limits[1] =
std::min(chunk_size, std::max<int64_t>(0, size - send_offset[1]));
recv_limits[0] = std::min(
n_wires * size_per_wire, std::max<int64_t>(0, size - recv_offset[0]));
recv_limits[1] =
std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[1]));
} else {
send_limits[0] =
std::min(chunk_size, std::max<int64_t>(0, size - send_offset[0]));
recv_limits[0] =
std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[0]));
}
for (int i = 0; i < MAX_DIR * MAX_CONNS; i++) {
send_count[i] = recv_count[i] = 0;
}
}
// Secondly all gather
//
// The offsets are correct from the scatter reduce
for (int k = 0; k < size_ - 1; k++) {
// Prefill the pipeline
int buff = 0;
while (buff < n_steps && buff < PIPELINE) {
post_recv_all<MAX_DIR>(sz, buff, n_wires);
for (int lr = 0; lr < MAX_DIR; lr++) {
for (int lw = 0; lw < n_wires; lw++) {
int64_t offset = lw * N +
send_count[lr * MAX_CONNS + lw] * n_wires * N +
lr * n_wires * size_per_wire;
std::copy(
out_ptr + send_offset[lr] + offset,
out_ptr + send_offset[lr] +
std::max(offset, std::min(offset + N, send_limits[lr])),
send_buffer(sz, buff, lr, lw).begin<T>());
send_count[lr * MAX_CONNS + lw]++;
}
}
post_send_all<MAX_DIR>(sz, buff, n_wires);
buff++;
in_flight += 2 * MAX_DIR * n_wires;
}
// Main loop
//
// Keep going until we have no longer data in flight.
while (in_flight > 0) {
ibv_wc wc[WC_NUM];
int n = poll(left_, right_, WC_NUM, wc);
for (int i = 0; i < n; i++) {
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 / MAX_CONNS;
int lw = wire % MAX_CONNS;
in_flight--;
if (work_type == SEND_WR && send_count[wire] < n_steps) {
int64_t offset = lw * N + send_count[wire] * n_wires * N +
lr * n_wires * size_per_wire;
std::copy(
out_ptr + send_offset[lr] + offset,
out_ptr + send_offset[lr] +
std::max(offset, std::min(offset + N, send_limits[lr])),
send_buffer(sz, buff, lr, lw).begin<T>());
send_to(sz, buff, lr, lw);
in_flight++;
send_count[wire]++;
}
else if (work_type == RECV_WR) {
int64_t offset = lw * N + recv_count[wire] * n_wires * N +
lr * n_wires * size_per_wire;
std::copy(
recv_buffer(sz, buff, lr, lw).begin<T>(),
recv_buffer(sz, buff, lr, lw).begin<T>() +
std::max<int64_t>(0, std::min(N, recv_limits[lr] - offset)),
out_ptr + recv_offset[lr] + offset);
recv_count[wire]++;
if (recv_count[wire] + (PIPELINE - 1) < n_steps) {
recv_from(sz, buff, lr, lw);
in_flight++;
}
}
}
}
send_offset[0] = (send_offset[0] + chunk_multiple_size - chunk_size) %
chunk_multiple_size;
recv_offset[0] = (recv_offset[0] + chunk_multiple_size - chunk_size) %
chunk_multiple_size;
if constexpr (MAX_DIR == 2) {
send_offset[1] = (send_offset[1] + chunk_size) % chunk_multiple_size;
recv_offset[1] = (recv_offset[1] + chunk_size) % chunk_multiple_size;
send_limits[0] = std::min(
n_wires * size_per_wire, std::max<int64_t>(0, size - send_offset[0]));
send_limits[1] =
std::min(chunk_size, std::max<int64_t>(0, size - send_offset[1]));
recv_limits[0] = std::min(
n_wires * size_per_wire, std::max<int64_t>(0, size - recv_offset[0]));
recv_limits[1] =
std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[1]));
} else {
send_limits[0] =
std::min(chunk_size, std::max<int64_t>(0, size - send_offset[0]));
recv_limits[0] =
std::min(chunk_size, std::max<int64_t>(0, size - recv_offset[0]));
}
for (int i = 0; i < MAX_DIR * MAX_CONNS; i++) {
send_count[i] = recv_count[i] = 0;
}
}
}
} // namespace mlx::core::distributed::jaccl
mlx/distributed/jaccl/ring.h
+178
View File
@@ -0,0 +1,178 @@
// Copyright © 2026 Apple Inc.
#pragma once
#include "mlx/distributed/distributed_impl.h"
#include "mlx/distributed/jaccl/utils.h"
using GroupImpl = mlx::core::distributed::detail::GroupImpl;
constexpr int MAX_CONNS = 4;
namespace mlx::core::distributed::jaccl {
/**
* The JACCL communication group for a ring where each node is connected to its
* two neighboring nodes. It should be the highest bandwidth communication
* group for large messages when many connections per peer are used.
*
* Like all JACCL groups it uses a side channel to exchange the necessary
* information and then configure the connections to be ready for RDMA
* operations.
*/
class RingGroup : public GroupImpl {
public:
RingGroup(
int rank,
int size,
const std::vector<std::string>& left_devices,
const std::vector<std::string>& right_devices,
const char* coordinator_addr);
Stream communication_stream(StreamOrDevice s) override {
return to_stream(s, Device::cpu);
}
int rank() override {
return rank_;
}
int size() override {
return size_;
}
void all_sum(const array& input, array& output, Stream stream) override;
void all_max(const array& input, array& output, Stream stream) override;
void all_min(const array& input, array& output, Stream stream) override;
void all_gather(const array& input, array& output, Stream stream) override;
void send(const array& input, int dst, Stream stream) override;
void recv(array& out, int src, Stream stream) override;
void sum_scatter(const array& input, array& output, Stream stream) override {
throw std::runtime_error("[jaccl] sum_scatter not supported.");
}
std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
throw std::runtime_error("[jaccl] Group split not supported.");
}
private:
template <typename T, typename ReduceOp>
void all_reduce(
const array& input,
array& output,
Stream stream,
ReduceOp reduce_op);
template <int MAX_DIR, typename T, typename ReduceOp>
void all_reduce_impl(
const T* in_ptr,
T* out_ptr,
int64_t size,
int n_wires,
ReduceOp reduce_op);
/**
* Performs the connection initialization. Namely, after this call all
* Connection objects should have a queue pair in RTS state and all buffers
* should have been allocated.
*/
void initialize();
/**
* Allocate all the buffers that we will use in the communication group.
*/
void allocate_buffers();
void send_to(int sz, int buff, int left_right, int wire) {
if (left_right) {
left_[wire].post_send(
send_buffer_left(sz, buff, wire),
SEND_WR << 16 | buff << 8 | (MAX_CONNS + wire));
} else {
right_[wire].post_send(
send_buffer_right(sz, buff, wire), SEND_WR << 16 | buff << 8 | wire);
}
}
void recv_from(int sz, int buff, int left_right, int wire) {
if (left_right) {
right_[wire].post_recv(
recv_buffer_right(sz, buff, wire),
RECV_WR << 16 | buff << 8 | (MAX_CONNS + wire));
} else {
left_[wire].post_recv(
recv_buffer_left(sz, buff, wire), RECV_WR << 16 | buff << 8 | wire);
}
}
SharedBuffer& send_buffer_right(int sz, int buff, int wire) {
return send_buffers_
[sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + wire];
}
SharedBuffer& send_buffer_left(int sz, int buff, int wire) {
return send_buffers_
[sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + MAX_CONNS +
wire];
}
SharedBuffer& send_buffer(int sz, int buff, int left_right, int wire) {
return send_buffers_
[sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 +
left_right * MAX_CONNS + wire];
}
SharedBuffer& recv_buffer_left(int sz, int buff, int wire) {
return recv_buffers_
[sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + wire];
}
SharedBuffer& recv_buffer_right(int sz, int buff, int wire) {
return recv_buffers_
[sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + MAX_CONNS +
wire];
}
SharedBuffer& recv_buffer(int sz, int buff, int left_right, int wire) {
return recv_buffers_
[sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 +
left_right * MAX_CONNS + wire];
}
template <int MAX_DIR>
void post_recv_all(int sz, int buff, int n_wires) {
for (int lr = 0; lr < MAX_DIR; lr++) {
for (int lw = 0; lw < n_wires; lw++) {
recv_from(sz, buff, lr, lw);
}
}
}
void post_recv_all(int sz, int buff) {
post_recv_all<2>(sz, buff, left_.size());
}
template <int MAX_DIR>
void post_send_all(int sz, int buff, int n_wires) {
for (int lr = 0; lr < MAX_DIR; lr++) {
for (int lw = 0; lw < n_wires; lw++) {
send_to(sz, buff, lr, lw);
}
}
}
void post_send_all(int sz, int buff) {
post_send_all<2>(sz, buff, left_.size());
}
int rank_;
int size_;
SideChannel side_channel_;
std::vector<Connection> left_;
std::vector<Connection> right_;
std::vector<SharedBuffer> send_buffers_;
std::vector<SharedBuffer> recv_buffers_;
};
} // namespace mlx::core::distributed::jaccl
mlx/distributed/jaccl/utils.cpp
+329
View File
@@ -0,0 +1,329 @@
// Copyright © 2025 Apple Inc.
#include <dlfcn.h>
#include <unistd.h>
#include <iostream>
#include <sstream>
#include "mlx/distributed/jaccl/utils.h"
#define LOAD_SYMBOL(symbol, variable) \
{ \
variable = (decltype(variable))dlsym(librdma_handle_, #symbol); \
char* error = dlerror(); \
if (error != nullptr) { \
std::cerr << IBV_TAG << " " << error << std::endl; \
librdma_handle_ = nullptr; \
return; \
} \
}
namespace {
void* page_aligned_alloc(size_t num_bytes) {
static size_t page_size = sysconf(_SC_PAGESIZE);
void* buf;
if (posix_memalign(&buf, page_size, num_bytes)) {
return nullptr;
}
return buf;
}
} // namespace
namespace mlx::core::distributed::jaccl {
IBVWrapper::IBVWrapper() {
librdma_handle_ = dlopen("librdma.dylib", RTLD_NOW | RTLD_GLOBAL);
if (librdma_handle_ == nullptr) {
return;
}
LOAD_SYMBOL(ibv_get_device_list, get_device_list);
LOAD_SYMBOL(ibv_get_device_name, get_device_name);
LOAD_SYMBOL(ibv_open_device, open_device);
LOAD_SYMBOL(ibv_free_device_list, free_device_list);
LOAD_SYMBOL(ibv_close_device, close_device);
LOAD_SYMBOL(ibv_alloc_pd, alloc_pd);
LOAD_SYMBOL(ibv_create_qp, create_qp);
LOAD_SYMBOL(ibv_create_cq, create_cq);
LOAD_SYMBOL(ibv_destroy_cq, destroy_cq);
LOAD_SYMBOL(ibv_destroy_qp, destroy_qp);
LOAD_SYMBOL(ibv_dealloc_pd, dealloc_pd);
LOAD_SYMBOL(ibv_query_port, query_port);
LOAD_SYMBOL(ibv_query_gid, query_gid);
LOAD_SYMBOL(ibv_modify_qp, modify_qp);
LOAD_SYMBOL(ibv_reg_mr, reg_mr);
LOAD_SYMBOL(ibv_dereg_mr, dereg_mr);
// Not really symbols but leaving them here in case they become symbols in
// the future.
//
// LOAD_SYMBOL(ibv_post_send, post_send);
// LOAD_SYMBOL(ibv_post_recv, post_recv);
// LOAD_SYMBOL(ibv_poll_cq, poll_cq);
}
IBVWrapper& ibv() {
static IBVWrapper wrapper;
return wrapper;
}
SharedBuffer::SharedBuffer(size_t num_bytes)
: data_(page_aligned_alloc(num_bytes)), num_bytes_(num_bytes) {}
SharedBuffer::SharedBuffer(SharedBuffer&& b) : data_(nullptr), num_bytes_(0) {
std::swap(data_, b.data_);
std::swap(num_bytes_, b.num_bytes_);
std::swap(memory_regions_, b.memory_regions_);
}
SharedBuffer::~SharedBuffer() {
for (auto& [pd, mr] : memory_regions_) {
ibv().dereg_mr(mr);
}
if (data_ != nullptr) {
std::free(data_);
}
}
void SharedBuffer::register_to_protection_domain(ibv_pd* protection_domain) {
auto [it, inserted] = memory_regions_.insert({protection_domain, nullptr});
if (!inserted) {
throw std::runtime_error(
"[jaccl] Buffer can be registered once per protection domain");
}
it->second = ibv().reg_mr(
protection_domain,
data_,
num_bytes_,
IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ |
IBV_ACCESS_REMOTE_WRITE);
if (!it->second) {
throw std::runtime_error("[jaccl] Register memory region failed");
}
}
Connection::Connection(ibv_context* ctx_)
: ctx(ctx_),
protection_domain(nullptr),
completion_queue(nullptr),
queue_pair(nullptr) {
src.local_id = -1;
}
Connection::Connection(Connection&& c) : Connection(nullptr) {
std::swap(ctx, c.ctx);
std::swap(protection_domain, c.protection_domain);
std::swap(completion_queue, c.completion_queue);
std::swap(queue_pair, c.queue_pair);
std::swap(src, c.src);
}
Connection::~Connection() {
if (queue_pair != nullptr) {
ibv().destroy_qp(queue_pair);
}
if (completion_queue != nullptr) {
ibv().destroy_cq(completion_queue);
}
if (protection_domain != nullptr) {
ibv().dealloc_pd(protection_domain);
}
if (ctx != nullptr) {
ibv().close_device(ctx);
}
}
void Connection::allocate_protection_domain() {
protection_domain = ibv().alloc_pd(ctx);
if (protection_domain == nullptr) {
throw std::runtime_error("[jaccl] Couldn't allocate protection domain");
}
}
void Connection::create_completion_queue(int num_entries) {
completion_queue = ibv().create_cq(ctx, num_entries, nullptr, nullptr, 0);
if (completion_queue == nullptr) {
throw std::runtime_error("[jaccl] Couldn't create completion queue");
}
}
void Connection::create_queue_pair() {
ibv_qp_init_attr init_attr;
init_attr.qp_context = ctx;
init_attr.qp_context = ctx;
init_attr.send_cq = completion_queue;
init_attr.recv_cq = completion_queue;
init_attr.srq = nullptr;
init_attr.cap.max_send_wr = MAX_SEND_WR;
init_attr.cap.max_recv_wr = MAX_RECV_WR;
init_attr.cap.max_send_sge = 1;
init_attr.cap.max_recv_sge = 1;
init_attr.cap.max_inline_data = 0;
init_attr.qp_type = IBV_QPT_UC;
init_attr.sq_sig_all = 0;
queue_pair = ibv().create_qp(protection_domain, &init_attr);
if (queue_pair == nullptr) {
throw std::runtime_error("[jaccl] Couldn't create queue pair");
}
}
const Destination& Connection::info() {
if (queue_pair == nullptr || src.local_id >= 0) {
return src;
}
ibv_port_attr port_attr;
ibv().query_port(ctx, 1, &port_attr);
ibv_gid gid;
ibv().query_gid(ctx, 1, 1, &gid);
src.local_id = port_attr.lid;
src.queue_pair_number = queue_pair->qp_num;
src.packet_sequence_number = 7; // TODO: Change to sth random
src.global_identifier = gid;
return src;
}
void Connection::queue_pair_init() {
ibv_qp_attr attr = {};
attr.qp_state = IBV_QPS_INIT;
attr.port_num = 1;
attr.pkey_index = 0;
attr.qp_access_flags =
IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE;
int mask =
IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS;
if (int status = ibv().modify_qp(queue_pair, &attr, mask); status != 0) {
std::ostringstream msg;
msg << "[jaccl] Changing queue pair to INIT failed with errno " << status;
throw std::invalid_argument(msg.str());
}
}
void Connection::queue_pair_rtr(const Destination& dst) {
ibv_qp_attr attr = {};
memset(&attr, 0, sizeof(attr));
attr.qp_state = IBV_QPS_RTR;
attr.path_mtu = IBV_MTU_1024;
attr.rq_psn = dst.packet_sequence_number;
attr.dest_qp_num = dst.queue_pair_number;
attr.ah_attr.dlid = dst.local_id;
attr.ah_attr.sl = 0;
attr.ah_attr.src_path_bits = 0;
attr.ah_attr.port_num = 1;
attr.ah_attr.is_global = 0;
if (dst.global_identifier.global.interface_id) {
attr.ah_attr.is_global = 1;
attr.ah_attr.grh.hop_limit = 1;
attr.ah_attr.grh.dgid = dst.global_identifier;
attr.ah_attr.grh.sgid_index = 1;
}
int mask = IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN |
IBV_QP_RQ_PSN;
if (int status = ibv().modify_qp(queue_pair, &attr, mask); status != 0) {
std::ostringstream msg;
msg << "[jaccl] Changing queue pair to RTR failed with errno " << status;
throw std::invalid_argument(msg.str());
}
}
void Connection::queue_pair_rts() {
ibv_qp_attr attr = {};
attr.qp_state = IBV_QPS_RTS;
attr.sq_psn = src.packet_sequence_number;
int mask = IBV_QP_STATE | IBV_QP_SQ_PSN;
if (int status = ibv().modify_qp(queue_pair, &attr, mask); status != 0) {
std::ostringstream msg;
msg << "[jaccl] Changing queue pair to RTS failed with errno " << status;
throw std::invalid_argument(msg.str());
}
}
std::vector<Connection> create_connections(
const std::vector<std::string>& device_names) {
std::vector<Connection> connections;
int num_devices = 0;
ibv_device** devices = ibv().get_device_list(&num_devices);
for (auto& name : device_names) {
// Empty so add a nullptr context
if (name.empty()) {
connections.emplace_back(nullptr);
continue;
}
// Search for the name and try to open the device
for (int i = 0; i < num_devices; i++) {
if (name == ibv().get_device_name(devices[i])) {
auto ctx = ibv().open_device(devices[i]);
if (ctx == nullptr) {
std::ostringstream msg;
msg << "[jaccl] Could not open device " << name;
throw std::runtime_error(msg.str());
}
connections.emplace_back(ctx);
break;
}
}
}
ibv().free_device_list(devices);
return connections;
}
SideChannel::SideChannel(int rank, int size, const char* addr)
: rank_(rank), size_(size) {
auto address = detail::parse_address(addr);
if (rank_ == 0) {
detail::TCPSocket server(IBV_TAG);
server.listen(IBV_TAG, address);
for (int i = 0; i < size - 1; i++) {
sockets_.push_back(server.accept(IBV_TAG));
}
std::vector<int> ranks(size - 1);
for (int i = 0; i < size - 1; i++) {
sockets_[i].recv(
IBV_TAG, reinterpret_cast<char*>(&ranks[i]), sizeof(int));
ranks[i]--;
}
for (int i = 0; i < size - 1; i++) {
while (i != ranks[i]) {
std::swap(sockets_[i], sockets_[ranks[i]]);
std::swap(ranks[i], ranks[ranks[i]]);
}
}
} else {
sockets_.push_back(
detail::TCPSocket::connect(
IBV_TAG, address, 4, 1000, [](int attempt, int wait) {
std::cerr << IBV_TAG << " Connection attempt " << attempt
<< " waiting " << wait << " ms" << std::endl;
}));
sockets_[0].send(IBV_TAG, reinterpret_cast<char*>(&rank_), sizeof(int));
}
}
SideChannel::SideChannel(SideChannel&& sc)
: rank_(sc.rank_), size_(sc.size_), sockets_(std::move(sc.sockets_)) {
sc.rank_ = -1;
sc.size_ = -1;
}
} // namespace mlx::core::distributed::jaccl
mlx/distributed/jaccl/utils.h
+342
View File
@@ -0,0 +1,342 @@
// Copyright © 2025 Apple Inc.
#pragma once
#include <infiniband/verbs.h>
#include <unordered_map>
#include <vector>
#include "mlx/distributed/utils.h"
constexpr const char* IBV_TAG = "[jaccl]";
constexpr int SEND_WR = 1;
constexpr int RECV_WR = 2;
constexpr int MAX_SEND_WR = 32;
constexpr int MAX_RECV_WR = 32;
constexpr int BUFFER_SIZES = 8;
constexpr int NUM_BUFFERS = 2;
constexpr int FRAME_SIZE = 4096;
namespace detail = mlx::core::distributed::detail;
namespace {
template <typename T, typename = void>
struct is_container : std::false_type {};
template <typename T>
struct is_container<
T,
std::void_t<typename T::value_type, typename T::iterator>>
: std::true_type {};
inline std::pair<int, int64_t> buffer_size_from_message(int64_t msg) {
if (__builtin_available(macOS 26.3, iOS 26.3, tvOS 26.3, visionOS 26.3, *)) {
for (int k = BUFFER_SIZES - 1; k > 0; k--) {
if (msg >= FRAME_SIZE * (1 << k)) {
return {k, FRAME_SIZE * (1 << k)};
}
}
}
return {0, FRAME_SIZE};
}
} // namespace
namespace mlx::core::distributed::jaccl {
/**
* Wrapper for the ibverbs API.
*/
struct IBVWrapper {
IBVWrapper();
bool is_available() {
return librdma_handle_ != nullptr;
}
// API
ibv_device** (*get_device_list)(int*);
const char* (*get_device_name)(ibv_device*);
ibv_context* (*open_device)(ibv_device*);
void (*free_device_list)(ibv_device**);
int (*close_device)(ibv_context*);
ibv_pd* (*alloc_pd)(ibv_context*);
ibv_qp* (*create_qp)(ibv_pd*, ibv_qp_init_attr*);
ibv_cq* (*create_cq)(ibv_context*, int, void*, ibv_comp_channel*, int);
int (*destroy_cq)(ibv_cq*);
int (*destroy_qp)(ibv_qp*);
int (*dealloc_pd)(ibv_pd*);
int (*query_port)(ibv_context*, uint8_t, ibv_port_attr*);
int (*query_gid)(ibv_context*, uint8_t, int, ibv_gid*);
int (*modify_qp)(ibv_qp*, ibv_qp_attr*, int);
ibv_mr* (*reg_mr)(ibv_pd*, void*, size_t, int);
int (*dereg_mr)(ibv_mr*);
private:
void* librdma_handle_;
};
IBVWrapper& ibv();
/**
* Contains the information that defines a destination to a remote device.
* Basically we can compute our own destination and share it with remote hosts
* over the side channel.
*/
struct Destination {
int local_id;
int queue_pair_number;
int packet_sequence_number;
ibv_gid global_identifier;
};
/**
* A buffer that can be registered to a number of protection domains.
*/
class SharedBuffer {
public:
SharedBuffer(size_t num_bytes);
SharedBuffer(SharedBuffer&& b);
~SharedBuffer();
SharedBuffer(const SharedBuffer&) = delete;
SharedBuffer& operator=(const SharedBuffer&) = delete;
void register_to_protection_domain(ibv_pd* protection_domain);
size_t size() const {
return num_bytes_;
}
uint32_t local_key(ibv_pd* protection_domain) const {
return memory_regions_.at(protection_domain)->lkey;
}
ibv_sge to_scatter_gather_entry(ibv_pd* protection_domain) const {
ibv_sge entry;
entry.addr = reinterpret_cast<uintptr_t>(data_);
entry.length = size();
entry.lkey = local_key(protection_domain);
return entry;
}
template <typename T>
T* data() {
return static_cast<T*>(data_);
}
template <typename T>
T* begin() {
return static_cast<T*>(data_);
}
template <typename T>
T* end() {
return static_cast<T*>(data_) + size() / sizeof(T);
}
private:
void* data_;
size_t num_bytes_;
std::unordered_map<ibv_pd*, ibv_mr*> memory_regions_;
};
/**
* Manipulates an RDMA connection. Enables (among other things)
*
* - Creating a queue pair
* - Sending and receiving
* - Checking completion
*/
struct Connection {
ibv_context* ctx;
ibv_pd* protection_domain;
ibv_cq* completion_queue;
ibv_qp* queue_pair;
Destination src; // holds the local information
Connection(ibv_context* ctx_);
Connection(Connection&& c);
Connection(const Connection&) = delete;
Connection& operator=(Connection&) = delete;
~Connection();
void allocate_protection_domain();
void create_completion_queue(int num_entries);
void create_queue_pair();
const Destination& info();
void queue_pair_init();
void queue_pair_rtr(const Destination& dst);
void queue_pair_rts();
void post_send(const SharedBuffer& buff, uint64_t work_request_id) {
ibv_send_wr work_request, *bad_work_request;
auto entry = buff.to_scatter_gather_entry(protection_domain);
work_request.wr_id = work_request_id;
work_request.sg_list = &entry;
work_request.num_sge = 1;
work_request.opcode = IBV_WR_SEND;
work_request.send_flags = IBV_SEND_SIGNALED;
work_request.next = nullptr;
if (int status =
ibv_post_send(queue_pair, &work_request, &bad_work_request);
status != 0) {
std::ostringstream msg;
msg << "[jaccl] Send failed with error code " << status;
throw std::invalid_argument(msg.str());
}
}
void post_recv(const SharedBuffer& buff, uint64_t work_request_id) {
ibv_recv_wr work_request, *bad_work_request;
auto entry = buff.to_scatter_gather_entry(protection_domain);
work_request.wr_id = work_request_id;
work_request.sg_list = &entry;
work_request.num_sge = 1;
work_request.next = nullptr;
if (int status =
ibv_post_recv(queue_pair, &work_request, &bad_work_request);
status != 0) {
std::ostringstream msg;
msg << "[jaccl] Recv failed with error code " << status;
throw std::invalid_argument(msg.str());
}
}
int poll(int num_completions, ibv_wc* work_completions) {
return ibv_poll_cq(completion_queue, num_completions, work_completions);
}
};
std::vector<Connection> create_connections(
const std::vector<std::string>& device_names);
inline int poll(
const std::vector<Connection>& connections,
int num_completions,
ibv_wc* work_completions) {
int completions = 0;
for (auto& c : connections) {
if (c.ctx == nullptr) {
continue;
}
if (completions >= num_completions) {
return completions;
}
int n = ibv_poll_cq(
c.completion_queue,
num_completions - completions,
work_completions + completions);
completions += n;
}
return completions;
}
inline int poll(
const std::vector<Connection>& connections_1,
const std::vector<Connection>& connections_2,
int num_completions,
ibv_wc* work_completions) {
int completions = 0;
completions += poll(connections_1, num_completions, work_completions);
completions += poll(
connections_2,
num_completions - completions,
work_completions + completions);
return completions;
}
/**
* Implement a TCP side channel to exchange information about the RDMA
* connections.
*
* Implements a simple all gather where every node sends to rank 0 and rank 0
* broadcasts to every node.
*/
class SideChannel {
public:
SideChannel(int rank, int size, const char* addr);
SideChannel(SideChannel&& sc);
SideChannel(const SideChannel&) = delete;
SideChannel& operator=(const SideChannel&) = delete;
template <typename T>
std::vector<T> all_gather(const T& v) {
std::vector<T> result(size_);
// T is a container of stuff like std::vector or std::string
if constexpr (is_container<T>::value) {
using U = typename T::value_type;
// Share the lengths first and set the communication size to be the
// maximum length of the containers.
auto lengths = all_gather<int>(v.size());
auto max_len = *std::max_element(lengths.begin(), lengths.end());
for (auto& s : result) {
s.resize(max_len);
}
// All gather of length max_len
if (rank_ == 0) {
std::copy(v.begin(), v.end(), result[rank_].begin());
for (int i = 1; i < size_; i++) {
sockets_[i - 1].recv(IBV_TAG, result[i].data(), sizeof(U) * max_len);
}
for (int i = 1; i < size_; i++) {
for (int j = 0; j < size_; j++) {
sockets_[i - 1].send(
IBV_TAG, result[j].data(), sizeof(U) * max_len);
}
}
} else {
std::copy(v.begin(), v.end(), result[rank_].begin());
sockets_[0].send(IBV_TAG, result[rank_].data(), sizeof(U) * max_len);
for (int i = 0; i < size_; i++) {
sockets_[0].recv(IBV_TAG, result[i].data(), sizeof(U) * max_len);
}
}
// Resize the outputs back to the original length
for (int i = 0; i < size_; i++) {
result[i].resize(lengths[i]);
}
}
// T is a scalar
else {
if (rank_ == 0) {
result[rank_] = v;
for (int i = 1; i < size_; i++) {
sockets_[i - 1].recv(IBV_TAG, &result[i], sizeof(T));
}
for (int i = 1; i < size_; i++) {
sockets_[i - 1].send(IBV_TAG, result.data(), size_ * sizeof(T));
}
} else {
sockets_[0].send(IBV_TAG, &v, sizeof(T));
sockets_[0].recv(IBV_TAG, result.data(), size_ * sizeof(T));
}
}
return result;
}
private:
int rank_;
int size_;
std::vector<detail::TCPSocket> sockets_;
};
} // namespace mlx::core::distributed::jaccl
python/mlx/_distributed_utils/common.py
+89 -49
View File
@@ -4,9 +4,9 @@ import argparse
import ipaddress
import json
import sys
from dataclasses import dataclass
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional
from typing import Optional, Union
@dataclass
@@ -14,7 +14,93 @@ class Host:
rank: int
ssh_hostname: str
ips: list[str]
rdma: list[Optional[str]]
rdma: list[Optional[Union[str, list[str]]]]
@dataclass
class Hostfile:
hosts: list[Host]
backend: str = ""
envs: list[str] = field(default_factory=list)
def to_json(self):
return {
"backend": self.backend,
"envs": self.envs,
"hosts": [
{"ssh": h.ssh_hostname, "ips": h.ips, "rdma": h.rdma}
for h in self.hosts
],
}
@classmethod
def from_file(cls, hostfile):
"""Parse the json hostfile that contains both the hostnames to ssh into and
the ips to communicate over when using the ring backend. It can also
contain the backend to be used and environment variables to set when
launching a distributed job.
Example:
{
"backend": "jaccl",
"envs": [
"MLX_METAL_FAST_SYNCH=1"
],
"hosts": [
{"ssh": "hostname1", "ips": ["123.123.123.1"], "rdma": [null, "rdma_en2", "rdma_en3"]},
{"ssh": "hostname2", "ips": ["123.123.123.2"], "rdma": ["rdma_en2", null, "rdma_en3"]},
...
{"ssh": "hostnameN", "ips": ["123.123.123.N"], "rdma": ["rdma_en2", "rdma_en3", null]},
]
}
Args:
hostfile (str): The path to the json file containing the host
information
"""
hostfile = Path(hostfile)
if not hostfile.exists():
raise ValueError(f"Hostfile {str(hostfile)} doesn't exist")
try:
data = json.load(open(hostfile))
backend = ""
envs = []
hosts = []
if isinstance(data, dict):
backend = data["backend"]
envs = data["envs"]
hosts = data["hosts"]
elif isinstance(data, list):
hosts = data
hosts = [
Host(i, h["ssh"], h.get("ips", []), h.get("rdma", []))
for i, h in enumerate(hosts)
]
return cls(hosts, backend, envs)
except Exception as e:
raise ValueError(
f"Failed to parse hostfile {str(hostfile)} ({str(e)})"
) from e
@classmethod
def from_list(cls, hostlist, repeats=1):
hosts = []
for i, h in enumerate(hostlist.split(",")):
if h == "":
raise ValueError("Hostname cannot be empty")
try:
ipaddress.ip_address(h)
ips = [h]
except ValueError:
ips = []
for i in range(repeats):
hosts.append(Host(i, h, ips, []))
return cls(hosts)
class OptionalBoolAction(argparse.Action):
@@ -47,49 +133,3 @@ def log_warning(*args, **kwargs):
def log_error(*args, **kwargs):
kwargs["file"] = sys.stderr
print("\033[31m[ERROR]", *args, "\033[0m", **kwargs)
def parse_hostlist(parser, hostlist, repeats):
hosts = []
for i, h in enumerate(hostlist.split(",")):
if h == "":
raise ValueError("Hostname cannot be empty")
try:
ipaddress.ip_address(h)
ips = [h]
except ValueError:
ips = []
for i in range(repeats):
hosts.append(Host(i, h, ips, []))
return hosts
def parse_hostfile(parser, hostfile):
"""Parse the json hostfile that contains both the hostnames to ssh into and
the ips to communicate over when using the ring backend.
Example:
[
{"ssh": "hostname1", "ips": ["123.123.123.1"], "rdma": [null, "rdma_en2", "rdma_en3"]},
{"ssh": "hostname2", "ips": ["123.123.123.2"], "rdma": ["rdma_en2", null, "rdma_en3"]},
...
{"ssh": "hostnameN", "ips": ["123.123.123.N"], "rdma": ["rdma_en2", "rdma_en3", null]},
]
Args:
hostfile (str): The path to the json file containing the host
information
"""
hostfile = Path(hostfile)
if not hostfile.exists():
parser.error(f"Hostfile {str(hostfile)} doesn't exist")
try:
hosts = []
with open(hostfile) as f:
for i, h in enumerate(json.load(f)):
hosts.append(Host(i, h["ssh"], h.get("ips", []), h.get("rdma", [])))
return hosts
except Exception as e:
parser.error(f"Failed to parse hostfile {str(hostfile)} ({str(e)})")
python/mlx/_distributed_utils/config.py
+113 -65
View File
@@ -14,12 +14,11 @@ import mlx.core as mx
from .common import (
Host,
Hostfile,
OptionalBoolAction,
log,
log_error,
log_warning,
parse_hostfile,
parse_hostlist,
)
@@ -70,9 +69,20 @@ def add_ips(hosts, verbose=False):
log_warning("Could not extract ip for", h.ssh_hostname)
def check_rdma(hosts, verbose=False):
def save_hostfile(args, hostfile):
if args.output_hostfile:
with open(args.output_hostfile, "w") as f:
json.dump(hostfile.to_json(), f, indent=4)
else:
print("Hostfile")
print("========")
print(json.dumps(hostfile.to_json(), indent=4))
def check_rdma(hosts, verbose=False, strict=True):
# Check whether the hosts are capable of RDMA over thunderbolt
warn = False
log_f = log_warning if not strict else log_error
failed = False
for h in hosts:
log(verbose, "Checking that", h.ssh_hostname, "supports RDMA")
rdma_devs = (
@@ -82,19 +92,20 @@ def check_rdma(hosts, verbose=False):
)
rdma_devs = [d for d in rdma_devs if d.startswith("rdma_")]
if not rdma_devs:
log_warning(h.ssh_hostname, "does not seem to have RDMA enabled")
warn = True
log_f(h.ssh_hostname, "does not seem to have RDMA enabled")
failed = True
if warn:
log_warning()
log_warning(
"Some of the hosts don't have RDMA enabled or they don't support RDMA."
)
log_warning()
log_warning(
"See https://ml-explore.github.io/mlx/build/html/usage/distributed.html"
)
log_warning("for instructions on how to enable RDMA.")
if failed:
log_f()
log_f("Some of the hosts don't have RDMA enabled or they don't support RDMA.")
log_f()
log_f("See https://ml-explore.github.io/mlx/build/html/usage/distributed.html")
log_f("for instructions on how to enable RDMA.")
if failed and strict:
sys.exit(1)
return not failed
def can_auto_setup(hosts, sshinfo, auto_setup=False):
@@ -340,6 +351,20 @@ def check_valid_mesh(hosts, connectivity, strict=True):
return True
def check_valid_ring(hosts, rings, strict=True):
has_ring = len(rings) > 0 and len(rings[0][0]) == len(hosts)
if strict and not has_ring:
log_error("Could not find a full ring.")
log_error()
log_error("Try passing --dot to visualize the connectivity")
if len(rings) > 0:
log_error("Rings found:")
for r in rings:
log_error(f" - {','.join(hosts[i].ssh_hostname for i in r)}")
sys.exit(1)
return has_ring
def check_ssh_connections(hosts):
results = [None] * len(hosts)
@@ -408,52 +433,38 @@ def prepare_ethernet_hostfile(args, hosts):
log(args.verbose, f"Preparing an ethernet hostfile")
add_ips(hosts, args.verbose)
hostfile = []
for h in hosts:
hostfile.append(dict(ssh=h.ssh_hostname, ips=h.ips))
hostfile = Hostfile(
[Host(i, h.ssh_hostname, h.ips, []) for i, h in enumerate(hosts)], "", args.env
)
if args.output_hostfile:
with open(args.output_hostfile, "w") as f:
json.dump(hostfile, f, indent=4)
else:
print("Hostfile")
print("========")
print(json.dumps(hostfile, indent=4))
save_hostfile(args, hostfile)
def configure_ring(args, hosts, ips, ring, sshinfo):
log(args.verbose, "Prepare a ring hostfile")
ring, count = ring
hostfile = []
ring_hosts = []
for i, node in enumerate(ring):
h = hosts[node]
peer = ring[i - 1]
hostfile.append(
{
"ssh": h.ssh_hostname,
"ips": [ips.ips[node, peer][c][1] for c in range(count)],
"rdma": [],
}
ring_hosts.append(
Host(
i, h.ssh_hostname, [ips.ips[node, peer][c][1] for c in range(count)], []
)
)
hostfile = Hostfile(ring_hosts, "ring", args.env)
has_sudo = can_auto_setup(hosts, sshinfo, args.auto_setup)
ips.setup(verbose=args.verbose, auto_setup=args.auto_setup and has_sudo)
if args.output_hostfile:
with open(args.output_hostfile, "w") as f:
json.dump(hostfile, f, indent=4)
else:
print("Hostfile")
print("========")
print(json.dumps(hostfile, indent=4))
save_hostfile(args, hostfile)
def configure_jaccl(args, hosts, ips, sshinfo):
log(args.verbose, "Prepare a jaccl hostfile")
check_rdma(hosts, args.verbose)
add_ips(hosts, args.verbose)
hostfile = []
jaccl_hosts = []
for i, h in enumerate(hosts):
rdma = []
for j in range(len(hosts)):
@@ -461,18 +472,42 @@ def configure_jaccl(args, hosts, ips, sshinfo):
rdma.append(None)
else:
rdma.append(f"rdma_{ips.ips[i, j][0][0]}")
hostfile.append({"ssh": h.ssh_hostname, "ips": h.ips, "rdma": rdma})
jaccl_hosts.append(Host(i, h.ssh_hostname, h.ips, rdma))
hostfile = Hostfile(jaccl_hosts, "jaccl", args.env)
has_sudo = can_auto_setup(hosts, sshinfo, args.auto_setup)
ips.setup(verbose=args.verbose, auto_setup=args.auto_setup and has_sudo)
if args.output_hostfile:
with open(args.output_hostfile, "w") as f:
json.dump(hostfile, f, indent=4)
else:
print("Hostfile")
print("========")
print(json.dumps(hostfile, indent=4))
save_hostfile(args, hostfile)
def configure_jaccl_ring(args, hosts, ips, ring, sshinfo):
log(args.verbose, "Prepare a jaccl-ring hostfile")
add_ips(hosts, args.verbose)
jaccl_hosts = []
num_nodes = len(hosts)
ring, count = ring
for i, node in enumerate(ring):
h = hosts[node]
peer_left = ring[i - 1]
peer_right = ring[(i + 1) % num_nodes]
rdmas = []
for j in range(len(hosts)):
if j not in (peer_left, peer_right):
rdmas.append(None)
else:
rdma = []
for c in range(count):
rdma.append(f"rdma_{ips.ips[i, j][c][0]}")
rdmas.append(rdma[0] if count == 1 else rdma)
jaccl_hosts.append(Host(i, h.ssh_hostname, h.ips, rdmas))
hostfile = Hostfile(jaccl_hosts, "jaccl-ring", args.env)
has_sudo = can_auto_setup(hosts, sshinfo, args.auto_setup)
ips.setup(verbose=args.verbose, auto_setup=args.auto_setup and has_sudo)
save_hostfile(args, hostfile)
def prepare_tb_hostfile(args, hosts, sshinfo):
@@ -489,37 +524,44 @@ def prepare_tb_hostfile(args, hosts, sshinfo):
if args.backend is None:
rings = extract_rings(connectivity)
has_mesh = check_valid_mesh(hosts, connectivity, False)
has_ring = len(rings) > 0 and len(rings[0][0]) == len(hosts)
has_ring = check_valid_ring(hosts, rings, False)
has_rdma = check_rdma(hosts, args.verbose, False)
if not has_ring and not has_mesh:
log_error("Neither thunderbolt mesh nor ring found.")
log_error("Perhaps run with --dot to generate a plot of the connectivity.")
sys.exit(1)
elif has_rdma and has_mesh:
configure_jaccl(args, hosts, ips, sshinfo)
elif has_rdma and has_ring:
configure_jaccl_ring(args, hosts, ips, rings[0], sshinfo)
elif has_ring:
configure_ring(args, hosts, ips, rings[0], sshinfo)
else:
configure_jaccl(args, hosts, ips, sshinfo)
log_error("RDMA is not available and ring is not found.")
log_error("Perhaps run with --dot to generate a plot of the connectivity.")
sys.exit(1)
elif args.backend == "ring":
rings = extract_rings(connectivity)
has_ring = len(rings) > 0 and len(rings[0][0]) == len(hosts)
if not has_ring:
log_error("Could not find a full ring.")
log_error()
log_error("Try passing --dot to visualize the connectivity")
if len(rings) > 0:
log_error("Rings found:")
for r in rings:
log_error(f" - {','.join(hosts[i].ssh_hostname for i in r)}")
sys.exit(1)
check_valid_ring(hosts, rings)
configure_ring(args, hosts, ips, rings[0], sshinfo)
elif args.backend == "jaccl":
check_valid_mesh(hosts, connectivity)
check_rdma(hosts, args.verbose)
configure_jaccl(args, hosts, ips, sshinfo)
elif args.backend == "jaccl-ring":
rings = extract_rings(connectivity)
check_valid_ring(hosts, rings)
check_rdma(hosts, args.verbose)
configure_jaccl_ring(args, hosts, ips, rings[0], sshinfo)
def main():
parser = argparse.ArgumentParser(
@@ -555,16 +597,22 @@ def main():
)
parser.add_argument(
"--backend",
choices=["ring", "jaccl"],
choices=["ring", "jaccl", "jaccl-ring"],
default=None,
help="Which distributed backend to configure",
)
parser.add_argument(
"--env",
action="append",
default=[],
help="Set environment variables for the jobs",
)
args = parser.parse_args()
if args.hostfile is not None:
hosts = parse_hostfile(parser, args.hostfile)
hosts = Hostfile.from_file(args.hostfile).hosts
else:
hosts = parse_hostlist(parser, args.hosts, 1)
hosts = Hostfile.from_list(args.hosts).hosts
# Check that we can ssh
log(
python/mlx/_distributed_utils/launch.py
+24 -12
View File
@@ -19,7 +19,7 @@ from subprocess import PIPE, Popen, run
import mlx.core as mx
from .common import log, log_warning, parse_hostfile, parse_hostlist, positive_number
from .common import Hostfile, log, log_warning, positive_number
class CommandProcess:
@@ -367,6 +367,7 @@ def launch_jaccl(parser, hosts, args, command):
if not hosts[0].ips:
raise ValueError("Rank 0 should have an IP reachable from all other ranks")
jaccl_ring = args.backend == "jaccl-ring"
have_rdmas = all(len(h.rdma) == len(hosts) for h in hosts)
have_nulls = all(h.rdma[i] is None for i, h in enumerate(hosts))
if not have_rdmas or not have_nulls:
@@ -376,6 +377,8 @@ def launch_jaccl(parser, hosts, args, command):
env = args.env
cwd = args.cwd
env.append(f"MLX_JACCL_COORDINATOR={coordinator}:{args.starting_port}")
if jaccl_ring:
env.append("MLX_JACCL_RING=1")
files = {"MLX_IBV_DEVICES": json.dumps([h.rdma for h in hosts])}
log(args.verbose, "Running", shlex.join(command))
@@ -474,8 +477,6 @@ def main():
parser.add_argument("--hostfile", help="The file containing the hosts")
parser.add_argument(
"--backend",
choices=["ring", "mpi", "nccl", "jaccl"],
default="nccl" if mx.cuda.is_available() else "ring",
help="Which distributed backend to launch",
)
parser.add_argument(
@@ -535,9 +536,16 @@ def main():
# Try to extract a list of hosts and corresponding ips
if args.hostfile is not None:
hosts = parse_hostfile(parser, args.hostfile)
hostfile = Hostfile.from_file(args.hostfile)
else:
hosts = parse_hostlist(parser, args.hosts, args.repeat_hosts)
hostfile = Hostfile.from_list(args.hosts, args.repeat_hosts)
# Extract extra arguments from the hostfile
if hostfile.backend != "" and args.backend is None:
args.backend = hostfile.backend
if args.backend is None:
args.backend = "nccl" if mx.cuda.is_available() else "ring"
args.env = hostfile.envs + args.env
# Check if the script is a file and convert it to a full path
if (script := Path(rest[0])).exists() and script.is_file():
@@ -549,10 +557,14 @@ def main():
# Launch
if args.backend == "ring":
launch_ring(parser, hosts, args, rest)
if args.backend == "mpi":
launch_mpi(parser, hosts, args, rest)
if args.backend == "nccl":
launch_nccl(parser, hosts, args, rest)
if args.backend == "jaccl":
launch_jaccl(parser, hosts, args, rest)
launch_ring(parser, hostfile.hosts, args, rest)
elif args.backend == "mpi":
launch_mpi(parser, hostfile.hosts, args, rest)
elif args.backend == "nccl":
launch_nccl(parser, hostfile.hosts, args, rest)
elif args.backend == "jaccl" or args.backend == "jaccl-ring":
launch_jaccl(parser, hostfile.hosts, args, rest)
else:
parser.error(
"The backend should be one of {'ring', 'mpi', 'nccl', 'jaccl', 'jaccl-ring'}"
)