diff --git a/mlx/distributed/jaccl/CMakeLists.txt b/mlx/distributed/jaccl/CMakeLists.txt index dbd60c8e..567ed672 100644 --- a/mlx/distributed/jaccl/CMakeLists.txt +++ b/mlx/distributed/jaccl/CMakeLists.txt @@ -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() diff --git a/mlx/distributed/jaccl/jaccl.cpp b/mlx/distributed/jaccl/jaccl.cpp index a9d28f4d..5935a98f 100644 --- a/mlx/distributed/jaccl/jaccl.cpp +++ b/mlx/distributed/jaccl/jaccl.cpp @@ -1,1160 +1,128 @@ // Copyright © 2025 Apple Inc. -#include -#include -#include #include -#include +#include #include -#include "mlx/backend/cpu/encoder.h" #include "mlx/distributed/distributed_impl.h" -#include "mlx/distributed/reduction_ops.h" -#include "mlx/distributed/utils.h" -#include "mlx/dtype_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; \ - } \ - } - -constexpr const char* IBV_TAG = "[jaccl]"; -constexpr int NUM_BUFFERS = 2; -constexpr int BUFFER_SIZE = 4096; -constexpr int MAX_SEND_WR = 32; -constexpr int MAX_RECV_WR = 32; -constexpr int SEND_WR = 1; -constexpr int RECV_WR = 2; -constexpr int MAX_PEERS = 8; +#include "mlx/distributed/jaccl/mesh.h" +#include "mlx/distributed/jaccl/ring.h" +#include "mlx/distributed/jaccl/utils.h" using GroupImpl = mlx::core::distributed::detail::GroupImpl; using json = nlohmann::json; -namespace detail = mlx::core::distributed::detail; -namespace allocator = mlx::core::allocator; -struct IBVWrapper { - IBVWrapper() { - librdma_handle_ = dlopen("librdma.dylib", RTLD_NOW | RTLD_GLOBAL); - if (librdma_handle_ == nullptr) { - return; - } +namespace { - 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); - } - - bool is_available() { - return librdma_handle_ != nullptr; - } - - void* librdma_handle_; - - // 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*); -}; - -IBVWrapper& ibv() { - static IBVWrapper wrapper; - return wrapper; -} - -template -struct is_container : std::false_type {}; - -template -struct is_container< - T, - std::void_t> - : std::true_type {}; - -std::ostream& operator<<(std::ostream& os, const ibv_gid& gid) { - os << std::hex << std::setfill('0'); - for (int i = 0; i < 16; i += 2) { - uint16_t part = (gid.raw[i] << 8) | gid.raw[i + 1]; - os << std::setw(4) << part; - if (i < 14) - os << ":"; - } - os << std::dec; - return os; -} - -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; -} - -/** - * 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; -}; - -std::ostream& operator<<(std::ostream& os, const Destination& dst) { - os << dst.local_id << " " << dst.queue_pair_number << " " - << dst.packet_sequence_number << " " << dst.global_identifier; - return os; -} - -/** - * A buffer that can be registered to a number of protection domains. - */ -class SharedBuffer { - public: - SharedBuffer(size_t num_bytes) - : data_(page_aligned_alloc(num_bytes)), num_bytes_(num_bytes) {} - ~SharedBuffer() { - for (auto& [pd, mr] : memory_regions_) { - ibv().dereg_mr(mr); - } - if (data_ != nullptr) { - std::free(data_); - } - } - - SharedBuffer(const SharedBuffer&) = delete; - SharedBuffer& operator=(const SharedBuffer&) = delete; - 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_); - } - - void register_to_protection_domain(ibv_pd* protection_domain) { - auto [it, inserted] = memory_regions_.insert({protection_domain, nullptr}); - if (!inserted) { +struct DeviceFile { + DeviceFile(const char* dev_file) { + std::ifstream f(dev_file); + json devices = json::parse(f); + if (!devices.is_array()) { throw std::runtime_error( - "[jaccl] Buffer can be registered once per protection domain"); + "[jaccl] The device file should start with an array"); } - 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"); - } - } - - 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(data_); - entry.length = size(); - entry.lkey = local_key(protection_domain); - return entry; - } - - template - T* data() { - return static_cast(data_); - } - - template - T* begin() { - return static_cast(data_); - } - - template - T* end() { - return static_cast(data_) + size() / sizeof(T); - } - - private: - void* data_; - size_t num_bytes_; - std::unordered_map 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_) - : ctx(ctx_), - protection_domain(nullptr), - completion_queue(nullptr), - queue_pair(nullptr) { - src.local_id = -1; - } - - 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(const Connection&) = delete; - Connection& operator=(Connection&) = delete; - - ~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 allocate_protection_domain() { - protection_domain = ibv().alloc_pd(ctx); - if (protection_domain == nullptr) { - throw std::runtime_error("[jaccl] Couldn't allocate protection domain"); - } - } - - void 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 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& 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 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 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 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()); - } - } - - 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()); - } - } -}; - -/** - * 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) : 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)); + devices_.resize(devices.size()); + for (int rank = 0; rank < devices.size(); rank++) { + auto conn = devices[rank]; + if (!conn.is_array()) { + throw std::runtime_error( + "[jaccl] The device file should have an array of arrays"); + } + if (conn.size() != devices_.size()) { + std::ostringstream msg; + msg << "[jaccl] The device file should contain the connectivity of each rank to " + << "all other ranks but rank " << rank << " contains only " + << conn.size() << " entries."; + throw std::runtime_error(msg.str()); } - std::vector ranks(size - 1); - for (int i = 0; i < size - 1; i++) { - sockets_[i].recv( - IBV_TAG, reinterpret_cast(&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]]); + devices_[rank].resize(conn.size()); + for (int dst = 0; dst < conn.size(); dst++) { + auto names = conn[dst]; + if (names.is_string()) { + devices_[rank][dst].push_back(names); + } else if (names.is_array()) { + for (auto name_it = names.begin(); name_it != names.end(); + name_it++) { + devices_[rank][dst].push_back(*name_it); + } + } else if (!names.is_null()) { + throw std::runtime_error( + "[jaccl] Device names should be null, a string or array of strings."); } } - } 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(&rank_), sizeof(int)); } } - SideChannel(const SideChannel&) = delete; - SideChannel& operator=(const SideChannel&) = delete; - - SideChannel(SideChannel&& sc) - : rank_(sc.rank_), size_(sc.size_), sockets_(std::move(sc.sockets_)) { - sc.rank_ = -1; - sc.size_ = -1; + int size() { + return devices_.size(); } - template - std::vector all_gather(const T& v) { - std::vector result(size_); - - // T is a container of stuff like std::vector or std::string - if constexpr (is_container::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(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); + bool is_valid_mesh() { + for (int src = 0; src < size(); src++) { + for (int dst = 0; dst < size(); dst++) { + if (devices_[src][dst].size() != static_cast(src != dst)) { + return false; } - 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); + } + } + + return true; + } + + bool is_valid_ring() { + int num_connections = devices_[0][1].size(); + if (num_connections == 0) { + return false; + } + + for (int src = 0; src < size(); src++) { + int left = (src + size() - 1) % size(); + int right = (src + 1) % size(); + for (int dst = 0; dst < size(); dst++) { + if (dst != left && dst != right) { + if (devices_[src][dst].size() != 0) { + return false; + } + } else { + if (devices_[src][dst].size() != num_connections) { + return false; } } - } 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; + return true; } - private: - int rank_; - int size_; - std::vector sockets_; + std::vector extract_mesh_connectivity(int rank) { + std::vector devices(size()); + for (int dst = 0; dst < size(); dst++) { + if (dst != rank) { + devices[dst] = devices_[rank][dst][0]; + } + } + return devices; + } + + std::pair, std::vector> + extract_ring_connectivity(int rank) { + int left = (rank + size() - 1) % size(); + int right = (rank + 1) % size(); + + return std::make_pair(devices_[rank][left], devices_[rank][right]); + } + + std::vector>> devices_; }; -/** - * Manages a set of connections. Among other things it uses a side channel to - * exchange the necessary information and then configure the connections to be - * ready for RDMA operations. - */ -class ConnectionManager { - public: - ConnectionManager( - int rank, - const std::vector& device_names, - const char* coordinator_addr) - : rank_(rank), - size_(device_names.size()), - side_channel_(rank_, size_, coordinator_addr) { - create_contexts(device_names); - if (connections_[rank_].ctx != nullptr) { - throw std::runtime_error("[jaccl] Malformed device file"); - } - } - - int rank() const { - return rank_; - } - - int size() const { - return size_; - } - - /** - * Performs the connection initialization. Namely, after this call all - * Connection objects should have a queue pair in RTS state. - */ - void initialize(int num_buffers, size_t num_bytes) { - // 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(num_buffers, num_bytes); - - // 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 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 allocate_buffers(int num_buffers, size_t num_bytes) { - // Deregister any buffers and free the memory - buffers_.clear(); - - // Allocate the memory - for (int i = 0; i < num_buffers; i++) { - for (int j = 0; j < size_; j++) { - buffers_.emplace_back(num_bytes); - } - } - - 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_[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_[i * size_ + j].register_to_protection_domain( - connections_[j].protection_domain); - } - } - } - } - - void send_to(int rank, int buff) { - connections_[rank].post_send( - buffers_[buff * size_ + rank_], SEND_WR << 16 | buff << 8 | rank); - } - - void recv_from(int rank, int buff) { - connections_[rank].post_recv( - buffers_[buff * size_ + rank], RECV_WR << 16 | buff << 8 | rank); - } - - /** - * Poll all connections and save the work completions and return the - * corresponding length. - */ - int poll(int num_completions, ibv_wc* work_completions) { - int completions = 0; - for (int r = 0; r < size_; r++) { - if (r == rank_) { - continue; - } - if (completions >= num_completions) { - return completions; - } - - int c = ibv_poll_cq( - connections_[r].completion_queue, - num_completions - completions, - work_completions + completions); - - completions += c; - } - return completions; - } - - /** - * - */ - int poll(int rank, int num_completions, ibv_wc* work_completions) { - return ibv_poll_cq( - connections_[rank].completion_queue, num_completions, work_completions); - } - - SharedBuffer& send_buffer(int buff) { - return buffers_[buff * size_ + rank_]; - } - - SharedBuffer& buffer(int rank, int buff) { - return buffers_[buff * size_ + rank]; - } - - void barrier() { - side_channel_.all_gather(0); - } - - private: - void create_contexts(const std::vector& device_names) { - 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); - } - - int rank_; - int size_; - SideChannel side_channel_; - std::vector connections_; - std::vector buffers_; -}; - -std::vector load_device_names(int rank, const char* dev_file) { - std::vector device_names; - std::ifstream f(dev_file); - - json devices = json::parse(f); - devices = devices[rank]; - for (auto it = devices.begin(); it != devices.end(); it++) { - std::string n; - if (!it->is_null()) { - n = *it; - } - device_names.emplace_back(std::move(n)); - } - - return device_names; -} +} // namespace namespace mlx::core::distributed::jaccl { -class IBVGroup : public GroupImpl { - public: - IBVGroup(ConnectionManager cm) - : cm_(std::move(cm)), rank_(cm.rank()), size_(cm.size()) {} - - Stream communication_stream(StreamOrDevice s) override { - return to_stream(s, Device::cpu); - } - - int rank() override { - return cm_.rank(); - } - - int size() override { - return cm_.size(); - } - - void all_sum(const array& input, array& output, Stream stream) override { - dispatch_all_types(output.dtype(), [&](auto type_tag) { - using T = MLX_GET_TYPE(type_tag); - all_reduce(input, output, stream, detail::SumOp{}); - }); - } - - void all_max(const array& input, array& output, Stream stream) override { - dispatch_all_types(output.dtype(), [&](auto type_tag) { - using T = MLX_GET_TYPE(type_tag); - all_reduce(input, output, stream, detail::MaxOp{}); - }); - } - - void all_min(const array& input, array& output, Stream stream) override { - dispatch_all_types(output.dtype(), [&](auto type_tag) { - using T = MLX_GET_TYPE(type_tag); - all_reduce(input, output, stream, detail::MinOp{}); - }); - } - - void all_gather(const array& input, array& output, Stream stream) override { - auto in_ptr = input.data(); - auto out_ptr = output.data(); - 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; - constexpr int64_t N = BUFFER_SIZE; - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * MAX_PEERS * 2; - int64_t total = static_cast(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(buff); - std::copy( - our_data + read_offset, - our_data + std::min(read_offset + N, total), - cm_.send_buffer(buff).begin()); - post_send_all(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 = cm_.poll(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), - cm_.send_buffer(buff).begin()); - post_send_all(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( - cm_.buffer(rank, buff).begin(), - cm_.buffer(rank, buff).begin() + - 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) { - cm_.recv_from(rank, buff); - in_flight++; - } - } - } - } - }); - } - - void send(const array& input, int dst, Stream stream) override { - auto data = input.data(); - 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; - constexpr int N = BUFFER_SIZE; - - 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), - cm_.send_buffer(buff).begin()); - cm_.send_to(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 = cm_.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), - cm_.send_buffer(buff).begin()); - cm_.send_to(dst, buff); - - read_offset += N; - in_flight++; - } - } - } - }); - } - - void recv(array& out, int src, Stream stream) override { - auto data = out.data(); - 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; - constexpr int N = BUFFER_SIZE; - - int in_flight = 0; - int64_t write_offset = 0; - - // Prefill the pipeline - int buff = 0; - while (N * buff < n_bytes && buff < PIPELINE) { - cm_.recv_from(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 = cm_.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( - cm_.buffer(src, buff).begin(), - cm_.buffer(src, buff).begin() + - std::min(n_bytes - write_offset, static_cast(N)), - data + write_offset); - write_offset += N; - - if (write_offset + (PIPELINE - 1) * N < n_bytes) { - cm_.recv_from(src, buff); - - in_flight++; - } - } - } - }); - } - - void sum_scatter(const array& input, array& output, Stream stream) override { - throw std::runtime_error("[jaccl] sum_scatter not supported."); - } - - std::shared_ptr split(int color, int key = -1) override { - throw std::runtime_error("[jaccl] Group split not supported."); - } - - private: - void post_recv_all(int buffer) { - for (int i = 0; i < size_; i++) { - if (i == rank_) { - continue; - } - cm_.recv_from(i, buffer); - } - } - - void post_send_all(int buffer) { - for (int i = 0; i < size_; i++) { - if (i == rank_) { - continue; - } - cm_.send_to(i, buffer); - } - } - - template - void all_reduce( - const array& input, - array& output, - Stream stream, - ReduceOp reduce_op) { - auto in_ptr = input.data(); - auto out_ptr = output.data(); - 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; - constexpr int64_t N = BUFFER_SIZE / sizeof(T); - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * MAX_PEERS * 2; - int64_t total = static_cast(size); - 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 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(buff); - std::copy( - data + read_offset, - data + std::min(read_offset + N, total), - cm_.send_buffer(buff).begin()); - post_send_all(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 = cm_.poll(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), - cm_.send_buffer(buff).begin()); - post_send_all(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( - cm_.buffer(r, buff).begin(), - data + w, - std::min(N, total - w)); - w += N; - s++; - if (w + (PIPELINE - 1) * N < total) { - cm_.recv_from(r, buff); - in_flight++; - } - } - completed_recv_begin[r] = s; - } - } - }); - } - - ConnectionManager cm_; - int rank_; - int size_; -}; - bool is_available() { return ibv().is_available(); } @@ -1163,6 +131,7 @@ std::shared_ptr init(bool strict /* = false */) { const char* dev_file = std::getenv("MLX_IBV_DEVICES"); const char* coordinator = std::getenv("MLX_JACCL_COORDINATOR"); const char* rank_str = std::getenv("MLX_RANK"); + const char* ring = std::getenv("MLX_JACCL_RING"); if (!is_available() || !dev_file || !coordinator || !rank_str) { if (strict) { @@ -1179,20 +148,31 @@ std::shared_ptr init(bool strict /* = false */) { } auto rank = std::atoi(rank_str); - auto device_names = load_device_names(rank, dev_file); + bool prefer_ring = ring != nullptr; + DeviceFile devices(dev_file); - auto cm = ConnectionManager(rank, device_names, coordinator); - if (cm.size() > MAX_PEERS) { + if (rank >= devices.size() || rank < 0) { std::ostringstream msg; - msg << "[jaccl] The maximum number of supported peers is " << MAX_PEERS - << " but " << cm.size() << " was provided"; + msg << "[jaccl] Invalid rank " << rank << ". It should be between 0 and " + << devices.size(); throw std::runtime_error(msg.str()); } - cm.initialize(NUM_BUFFERS, BUFFER_SIZE); - cm.barrier(); - - return std::make_shared(std::move(cm)); + if (prefer_ring && devices.is_valid_ring()) { + auto [left, right] = devices.extract_ring_connectivity(rank); + return std::make_shared( + rank, devices.size(), left, right, coordinator); + } else if (devices.is_valid_mesh()) { + auto device_names = devices.extract_mesh_connectivity(rank); + return std::make_shared(rank, device_names, coordinator); + } else if (devices.is_valid_ring()) { + auto [left, right] = devices.extract_ring_connectivity(rank); + return std::make_shared( + rank, devices.size(), left, right, coordinator); + } else { + throw std::runtime_error( + "[jaccl] The device file should define a valid mesh or a valid ring."); + } } } // namespace mlx::core::distributed::jaccl diff --git a/mlx/distributed/jaccl/mesh.cpp b/mlx/distributed/jaccl/mesh.cpp new file mode 100644 index 00000000..c8df4e67 --- /dev/null +++ b/mlx/distributed/jaccl/mesh.cpp @@ -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& 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(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 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(input, output, stream, detail::SumOp{}); + }); +} + +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(input, output, stream, detail::MaxOp{}); + }); +} + +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(input, output, stream, detail::MinOp{}); + }); +} + +void MeshGroup::all_gather(const array& input, array& output, Stream stream) { + auto in_ptr = input.data(); + auto out_ptr = output.data(); + 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(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()); + 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()); + 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(), + recv_buffer(sz, buff, rank).begin() + + 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(); + 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()); + 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()); + send_to(sz, dst, buff); + + read_offset += N; + in_flight++; + } + } + } + }); +} + +void MeshGroup::recv(array& out, int src, Stream stream) { + auto data = out.data(); + 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(), + recv_buffer(sz, buff, src).begin() + + std::min(n_bytes - write_offset, static_cast(N)), + data + write_offset); + write_offset += N; + + if (write_offset + (PIPELINE - 1) * N < n_bytes) { + recv_from(sz, src, buff); + + in_flight++; + } + } + } + }); +} + +template +void MeshGroup::all_reduce( + const array& input, + array& output, + Stream stream, + ReduceOp reduce_op) { + auto in_ptr = input.data(); + auto out_ptr = output.data(); + 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(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()); + 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()); + 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(), + 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 diff --git a/mlx/distributed/jaccl/mesh.h b/mlx/distributed/jaccl/mesh.h new file mode 100644 index 00000000..6f779e9c --- /dev/null +++ b/mlx/distributed/jaccl/mesh.h @@ -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& 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 split(int color, int key = -1) override { + throw std::runtime_error("[jaccl] Group split not supported."); + } + + private: + template + 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 connections_; + std::vector buffers_; +}; + +} // namespace mlx::core::distributed::jaccl diff --git a/mlx/distributed/jaccl/ring.cpp b/mlx/distributed/jaccl/ring.cpp new file mode 100644 index 00000000..fe27a2e9 --- /dev/null +++ b/mlx/distributed/jaccl/ring.cpp @@ -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& left_devices, + const std::vector& 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(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 left_info; + for (auto& conn : left_) { + left_info.emplace_back(conn.info()); + } + std::vector 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(input, output, stream, detail::SumOp{}); + }); +} + +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(input, output, stream, detail::MaxOp{}); + }); +} + +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(input, output, stream, detail::MinOp{}); + }); +} + +void RingGroup::all_gather(const array& input, array& output, Stream stream) { + auto in_ptr = input.data(); + auto out_ptr = output.data(); + 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()); + 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()); + 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(), + recv_buffer(sz, buff, lr, lw).begin() + + std::max(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(); + 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()); + 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()); + 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(); + 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(), + recv_buffer(sz, buff, dir, lw).begin() + + std::max( + 0, std::min(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 +void RingGroup::all_reduce( + const array& input, + array& output, + Stream stream, + ReduceOp reduce_op) { + auto in_ptr = input.data(); + auto out_ptr = output.data(); + 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 +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(0, size - send_offset[0])); + send_limits[1] = + std::min(chunk_size, std::max(0, size - send_offset[1])); + recv_limits[0] = std::min( + n_wires * size_per_wire, std::max(0, size - recv_offset[0])); + recv_limits[1] = + std::min(chunk_size, std::max(0, size - recv_offset[1])); + } else { + send_limits[0] = + std::min(chunk_size, std::max(0, size - send_offset[0])); + recv_limits[0] = + std::min(chunk_size, std::max(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(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()); + send_count[lr * MAX_CONNS + lw]++; + } + } + post_send_all(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()); + 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(), + out_ptr + recv_offset[lr] + offset, + std::max(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(0, size - send_offset[0])); + send_limits[1] = + std::min(chunk_size, std::max(0, size - send_offset[1])); + recv_limits[0] = std::min( + n_wires * size_per_wire, std::max(0, size - recv_offset[0])); + recv_limits[1] = + std::min(chunk_size, std::max(0, size - recv_offset[1])); + } else { + send_limits[0] = + std::min(chunk_size, std::max(0, size - send_offset[0])); + recv_limits[0] = + std::min(chunk_size, std::max(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(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()); + send_count[lr * MAX_CONNS + lw]++; + } + } + post_send_all(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()); + 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(), + recv_buffer(sz, buff, lr, lw).begin() + + std::max(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(0, size - send_offset[0])); + send_limits[1] = + std::min(chunk_size, std::max(0, size - send_offset[1])); + recv_limits[0] = std::min( + n_wires * size_per_wire, std::max(0, size - recv_offset[0])); + recv_limits[1] = + std::min(chunk_size, std::max(0, size - recv_offset[1])); + } else { + send_limits[0] = + std::min(chunk_size, std::max(0, size - send_offset[0])); + recv_limits[0] = + std::min(chunk_size, std::max(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 diff --git a/mlx/distributed/jaccl/ring.h b/mlx/distributed/jaccl/ring.h new file mode 100644 index 00000000..a59ceb3d --- /dev/null +++ b/mlx/distributed/jaccl/ring.h @@ -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& left_devices, + const std::vector& 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 split(int color, int key = -1) override { + throw std::runtime_error("[jaccl] Group split not supported."); + } + + private: + template + void all_reduce( + const array& input, + array& output, + Stream stream, + ReduceOp reduce_op); + + template + 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 + 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 + 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 left_; + std::vector right_; + std::vector send_buffers_; + std::vector recv_buffers_; +}; + +} // namespace mlx::core::distributed::jaccl diff --git a/mlx/distributed/jaccl/utils.cpp b/mlx/distributed/jaccl/utils.cpp new file mode 100644 index 00000000..dc87d428 --- /dev/null +++ b/mlx/distributed/jaccl/utils.cpp @@ -0,0 +1,329 @@ +// Copyright © 2025 Apple Inc. + +#include +#include +#include +#include + +#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 create_connections( + const std::vector& device_names) { + std::vector 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 ranks(size - 1); + for (int i = 0; i < size - 1; i++) { + sockets_[i].recv( + IBV_TAG, reinterpret_cast(&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(&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 diff --git a/mlx/distributed/jaccl/utils.h b/mlx/distributed/jaccl/utils.h new file mode 100644 index 00000000..5b700f49 --- /dev/null +++ b/mlx/distributed/jaccl/utils.h @@ -0,0 +1,342 @@ +// Copyright © 2025 Apple Inc. + +#pragma once + +#include + +#include +#include + +#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 +struct is_container : std::false_type {}; + +template +struct is_container< + T, + std::void_t> + : std::true_type {}; + +inline std::pair 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(data_); + entry.length = size(); + entry.lkey = local_key(protection_domain); + return entry; + } + + template + T* data() { + return static_cast(data_); + } + + template + T* begin() { + return static_cast(data_); + } + + template + T* end() { + return static_cast(data_) + size() / sizeof(T); + } + + private: + void* data_; + size_t num_bytes_; + std::unordered_map 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 create_connections( + const std::vector& device_names); + +inline int poll( + const std::vector& 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& connections_1, + const std::vector& 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 + std::vector all_gather(const T& v) { + std::vector result(size_); + + // T is a container of stuff like std::vector or std::string + if constexpr (is_container::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(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 sockets_; +}; + +} // namespace mlx::core::distributed::jaccl diff --git a/python/mlx/_distributed_utils/common.py b/python/mlx/_distributed_utils/common.py index 16bf3f2b..d57747a3 100644 --- a/python/mlx/_distributed_utils/common.py +++ b/python/mlx/_distributed_utils/common.py @@ -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)})") diff --git a/python/mlx/_distributed_utils/config.py b/python/mlx/_distributed_utils/config.py index e0dc7f08..1c0bbfe6 100644 --- a/python/mlx/_distributed_utils/config.py +++ b/python/mlx/_distributed_utils/config.py @@ -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( diff --git a/python/mlx/_distributed_utils/launch.py b/python/mlx/_distributed_utils/launch.py index 7f144799..d866eaca 100644 --- a/python/mlx/_distributed_utils/launch.py +++ b/python/mlx/_distributed_utils/launch.py @@ -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'}" + )