
MLX 中的 JACCL基于 Thunderbolt 5 RDMA 的低延迟分布式通信库实战指南【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx导读JACCLJack and Angelos Collective Communication Library读作 Jackal是 MLX 框架内置的、面向 macOS 平台的低延迟分布式通信库专为搭载 Thunderbolt 5 接口的 Mac 集群设计。它利用 Apple 在 macOS 26.2 中引入的 RDMA over Thunderbolt 技术实现比传统 TCP 方案低一个数量级的通信延迟可支撑多机张量并行推理、高性能分布式训练以及 Mac 之间的低延迟集合通信。本文将以仓库中的 JACCL 文档 为主体结合 JACCL 源码 与 MLX 集成层 的实现细节完整讲解 RDMA 环境准备、独立库构建、C API 使用、MLX Python 集成与集群自动配置帮助你在多台 Mac 上快速搭建可用的分布式通信环境。JACCL 概述与适用场景JACCL 的定位是 macOS 上的类 NCCL通信库它基于 verbs APIrdma.h 中封装了ibv_*系列函数在 Thunderbolt 链路上直接操作 RDMA绕开 TCP/IP 协议栈的拷贝与调度开销从而获得远低于 TCP 的端到端延迟。按照官方文档的定位JACCL 主要面向三类场景大型模型推理的张量并行Tensor Parallelism多机切分模型权重在每一层计算后通过all_sum/all_reduce同步中间结果低延迟直接决定推理吞吐高性能分布式训练梯度同步all-reduce是分布式训练的热点路径Mac 之间的低延迟集合操作all-sum、all-max、all-min、all-gather 等原生集合原语。需要强调的是JACCL 依赖 Apple 的 RDMA over Thunderbolt 技术该技术随 macOS 26.2 提供因此在硬件与系统版本上都有硬性门槛详见下文 Requirements 一节。核心特性根据文档与源码JACCL 提供以下能力Mesh 拓扑全连接拓扑任意两个节点可直接通信适合小消息、低延迟场景Ring 拓扑环形拓扑每个节点只与左右邻居相连通过流水线化 reduce-scatter all-gather 实现大消息的高带宽 all-reducering.h 中RingGroup的注释明确指出当每个 peer 使用多条连接时它是大消息下带宽最高的通信组集合操作all_sum、all_max、all_min、all_gather点对点操作send、recv同步原语barrier阻塞直到组内所有节点到达数据类型Bool、Int8-64、UInt8-64、Float16、BFloat16、Float32、Float64、Complex64。以上类型映射在 group.h 的Dtype枚举中完整定义并由 types.h 中的dispatch_all_types在编译期分发到具体类型的归约实现。值得注意的源码细节是JACCL 自带与 MLX 兼容的float16_t、bfloat16_t与complex64_t实现保证库可以独立于 MLX 单独编译同时通过has_native_bf16_support()在运行时检测 CPU 是否支持FEAT_BF16从而一次编译、按机器能力启用原生 bf16。环境要求RequirementsJACCL 对运行环境有明确且严格的要求macOS SDK 26.2这一限制同时体现在文档与构建脚本中。独立库的 CMakeLists.txt 会先通过xcrun --sdk macosx --show-sdk-version探测 SDK 版本若低于 26.2 直接跳过构建MLX 侧的 jaccl/CMakeLists.txt 则要求MACOS_SDK_VERSION与CMAKE_OSX_DEPLOYMENT_TARGET均大于等于 26.2 才编译 JACCL 后端否则回退到no_jaccl.cpp节点间 Thunderbolt 5 互连启用 RDMA over Thunderbolt需在 macOS 恢复模式中配置见下节。此外构建脚本要求 CMake 3.24、C20CMAKE_CXX_STANDARD 20并会通过 FetchContent 拉取 nlohmann/jsonv3.11.3用于解析设备配置文件——这是设备文件格式依赖 JSON 的原因。启用 RDMA over ThunderboltRDMA over Thunderbolt 默认未开启需要在 macOS 恢复模式Recovery Mode中执行一次配置步骤如下将 Mac 启动进入恢复模式在实用工具 - 终端中打开终端执行rdma_ctl enable重启系统。启用后可通过ibv_devices验证 RDMA 设备是否可见正常输出形如device node GUID ------ ---------------- rdma_en2 8096a9d9edbaac05 rdma_en3 8196a9d9edbaac05 rdma_en5 8396a9d9edbaac05设备名如rdma_en5会作为后续设备配置文件中的连接标识使用请在各节点上记录自己实际可见的设备名。从源码角度看rdma.h 通过运行时动态加载 librdma 句柄来判断 RDMA 是否可用is_available()这也是 MLX 侧jaccl::is_available()的底层依据。构建 JACCL独立库构建cd mlx/distributed/jaccl/lib mkdir build cd build cmake .. make构建脚本默认使用 Release 编译类型CMakeLists.txt 中专门注释说明不设构建类型时 CMake 会用空类型即-O0会严重影响归约和 memcpy 热点路径性能。构建产物会安装到lib、include头文件安装到include/jaccl并导出jaccl::jaccl的 CMake 目标。在自己的 CMake 工程中引入FetchContent_Declare( jaccl GIT_REPOSITORY https://github.com/ml-explore/mlx.git GIT_TAG main SOURCE_SUBDIR mlx/distributed/jaccl/lib ) FetchContent_MakeAvailable(jaccl)仓库内的 examples/CMakeLists.txt 正是模拟了这种用法把../即 lib 目录当作 FetchContent 依赖引入然后链接jaccl库来构建minimal_env、minimal_cfg、minimal_barrier与allreduce_bench四个示例程序。使用方式环境变量初始化推荐JACCL 支持通过环境变量完成初始化这也是jaccl::init()无参调用的配置来源对应 jaccl.cpp 中的Config::from_env()。每个变量都提供JACCL_*与MLX_*两种写法优先读取前者环境变量别名含义JACCL_RANKMLX_RANK本进程的 rank从 0 开始的整数JACCL_IBV_DEVICESMLX_IBV_DEVICES描述设备连接关系的 JSON 文件路径JACCL_COORDINATORMLX_JACCL_COORDINATOR协调者rank 0 监听端的 IP:portJACCL_RINGMLX_JACCL_RING可选设置了即优先使用 ring 拓扑而非 mesh源码实现上from_env()会按上述优先级依次读取JACCL_RING存在时调用prefer_ring(true)最终通过Config::is_valid()校验配置完整性——若 rank、设备文件或 coordinator 缺失且以stricttrue初始化会抛出带完整提示信息的运行时错误。设备文件格式Device File设备文件是一个 JSON 数组每个条目描述某个 rank 到其余所有 rank 所使用的 RDMA 设备名[ [null, rdma_en5, rdma_en4, rdma_en3], [rdma_en5, null, rdma_en3, rdma_en4], [rdma_en4, rdma_en3, null, rdma_en5], [rdma_en3, rdma_en4, rdma_en5, null] ]格式约定对mesh拓扑devices[i][j]应保存连接 rank i 到 rank j 的设备名i j时为null对ring拓扑只有相邻节点间应填设备名其余位置为null。解析逻辑在 jaccl.cpp 的parse_devices_json中它要求顶层必须是数组且每个 rank 的连接条目数量必须等于节点总数否则抛出包含具体 rank 与缺失数量的错误每个元素可以是null、单个设备名字符串或字符串数组一条链路上允许多个设备。校验方面is_valid_mesh()要求每个非对角位置恰好有一个设备、对角位置为空is_valid_ring()则要求每个节点到左右邻居的设备数量一致。基础示例环境变量模式以下代码直接对应仓库 examples/minimal_env.cpp#include iostream #include jaccl/jaccl.h int main() { // Initialize JACCL group auto group jaccl::init(); if (!group) { std::cerr Failed to initialize JACCL std::endl; return 1; } std::cout Rank group-rank() of group-size() std::endl; // Perform all-reduce sum float input[10] {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f}; float output[10]; group-all_sum(input, output, sizeof(input), jaccl::Float32); std::cout Result: output[0] std::endl; return 0; }运行前需要为每个进程设置好JACCL_RANK、JACCL_IBV_DEVICES、JACCL_COORDINATOR环境变量。手动配置模式你也可以不依赖环境变量直接用Config对象显式配置对应 examples/minimal_cfg.cpp#include iostream #include jaccl/jaccl.h int main() { auto cfg jaccl::Config() .set_rank(0) // 每个节点应设置为不同的值 .set_coordinator(192.168.1.1:32132) // rank 0 将在此地址监听 .set_devices({ {{}, {rdma_en5}, {rdma_en4}, {rdma_en3}}, {{rdma_en5}, {}, {rdma_en3}, {rdma_en4}}, {{rdma_en4}, {rdma_en3}, {}, {rdma_en5}}, {{rdma_en3}, {rdma_en4}, {rdma_en5}, {}} }); auto group jaccl::init(cfg); if (!group) { std::cerr Failed to initialize JACCL std::endl; return 1; } std::cout Rank group-rank() of group-size() std::endl; // Perform all-reduce sum float input[10] {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f}; float output[10]; group-all_sum(input, output, sizeof(input), jaccl::Float32); std::cout Result: output[0] std::endl; return 0; }与 MLX 配合使用JACCL 作为 MLX 的分布式后端之一可以直接在 Python 侧使用import mlx.core as mx # Initialize with JACCL backend world mx.distributed.init(backendjaccl) # Perform distributed operations x mx.ones((10,)) result mx.distributed.all_sum(x, groupworld)集成层位于 mlx/distributed/jaccl/jaccl.cpp其中的JACCLGroup适配器把独立库的jaccl::Group包装成 MLX 的GroupImpl通过dtype_to_jaccl_dtype完成 MLX 数据类型到 JACCLDtype的映射并把集合操作通过 CPU command encoder 调度到 MLX 的流stream上执行。该文件还实现了all_max、all_min、all_gather、sum_scatter、send/recv等接口需要留意split组内再分片目前不支持会抛出 Group split not supported 错误。启动脚本可使用mlx.launchmlx.launch --backend jaccl --hostfile hosts.json my_script.pyHostfile 示例供mlx.launch使用的 hostfile 是一个 JSON 文件每个 host 节点通过ssh别名、ips与rdma设备行描述{ backend: jaccl, hosts: [ { ssh: m3-ultra-1, ips: [192.168.1.1], rdma: [null, rdma_en5, rdma_en4, rdma_en3] }, { ssh: m3-ultra-2, ips: [], rdma: [rdma_en5, null, rdma_en3, rdma_en4] }, { ssh: m3-ultra-3, ips: [], rdma: [rdma_en4, rdma_en3, null, rdma_en5] }, { ssh: m3-ultra-4, ips: [], rdma: [rdma_en3, rdma_en4, rdma_en5, null] } ] }其中每个 host 的rdma数组即对应设备文件矩阵中的一行null表示自己到自己的连接。自动配置mlx.distributed_configMLX 提供mlx.distributed_config工具可自动探测并配置各节点的 Thunderbolt 连接关系省去手工编写设备矩阵的麻烦# 可视化连接拓扑生成 DOT 图并用 Preview 打开 mlx.distributed_config --verbose \ --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \ --over thunderbolt --dot | dot -Tpng | open -f -a Preview # 自动配置并生成 hostfile mlx.distributed_config --verbose \ --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \ --over thunderbolt --backend jaccl \ --auto-setup --output m3-ultra-jaccl.json第一条命令以 DOT 图形形式展示节点间的实际物理连接第二条命令在--auto-setup模式下自动完成 RDMA 相关配置并把探测到的连接矩阵写入--output指定的 hostfile可直接交给mlx.launch使用。通信组 APIJACCL 的核心 API 是通信组Group。重要约定JACCL 自身不做任何内存分配所有输出指针必须指向已分配且足以容纳结果的内存区域。class Group { public: virtual ~Group() {} // 查询本进程在组中的身份 virtual int rank() 0; virtual int size() 0; // All-reduce 实现。输入输出大小相同 // 归约按 dtype 语义在组内进行。 virtual void all_sum(const void* input, void* output, size_t n_bytes, int dtype) 0; virtual void all_max(const void* input, void* output, size_t n_bytes, int dtype) 0; virtual void all_min(const void* input, void* output, size_t n_bytes, int dtype) 0; // All-gather 实现。输出大小为 group-size() * n_bytes。 virtual void all_gather(const void* input, void* output, size_t n_bytes) 0; // 简单的 send/recv 原语。 virtual void send(const void* input, size_t n_bytes, int dst) 0; virtual void recv(void* output, size_t n_bytes, int src) 0; // 阻塞直到组内每个 rank 都到达此点。 virtual void barrier() 0; };接口定义与 group.h 保持一致仓库当前实现中还额外提供了sum_scatterreduce-scatter with sum输入为size()个连续的n_bytes分块归约后 rank r 的输出为各 rank 第 r 块之和。创建通信组只需调用initstd::shared_ptrGroup init(bool strict false); std::shared_ptrGroup init(const Config cfg, bool strict false);无参init从环境变量构建配置并创建组stricttrue时初始化失败会抛出异常否则返回nullptr带Config的版本允许通过代码而非环境变量配置另有支持自定义 all-gather 工厂的init(bool strict, std::functionAllGatherFn(int, int) factory)重载见 jaccl.h用于替换默认的 TCP 侧信道来交换 RDMA 连接元数据。init的拓扑选择逻辑jaccl.cpp为若设置了prefer_ring且配置是合法 ring则创建RingGroup否则优先创建MeshGroup两者都非法时按strict决定返回空指针或抛错。Config 配置类class Config { public: Config(); Config set_rank(int rank); Config set_coordinator(std::string coordinator); Config set_devices(std::vectorstd::vectorstd::vectorstd::string devices); Config prefer_ring(bool prefer true); bool is_valid_mesh() const; bool is_valid_ring() const; }源码中还提供了更多便捷方法set_rank(const char*)、set_coordinator(const char*)、set_devices_from_file(const char* dev_file)直接读取 JSON 设备文件、set_all_gather(...)/set_all_gather_factory(...)定制侧信道以及static Config from_env()与is_valid()。set_devices会推导组大小size_ devices_.size()并校验矩阵必须是方阵否则抛出std::invalid_argument。示例与基准程序examples 目录 提供了四个可直接编译运行的参考程序minimal_env.cpp最小环境变量模式示例等价于文档中的 Basic Exampleminimal_cfg.cpp最小手动配置示例minimal_barrier.cpp演示barrier()的栅栏语义——各 rank 按100ms * rank错峰到达 barrier退出后用一次all_sum(Int32)验证组仍然健康并校验结果为size * (size 1) / 2allreduce_bench.cppNCCL 风格的 all-reduce 基准对 1K256M 字节的消息扫描带宽与延迟支持-w预热次数默认 5、-n计时迭代默认 20、-b/-e最小/最大消息字节数、-f倍增步长默认 2、-dfloat32/float16/bfloat16与-c正确性检查等参数输出算法带宽与总线带宽bus BW 按 ring 的2*(n-1)/n因子折算也可用mlx.launch --hostfile hosts.json ./jaccl_allreduce_bench启动。版本与平台注意事项本文描述的能力以当前仓库main 分支为准JACCL 依赖 macOS SDK 26.2 与部署目标 26.2非 Darwin 平台会在构建阶段被跳过MLX 主库在满足上述平台条件时编译 JACCL 后端否则编译no_jaccl.cpp占位实现见 jaccl/CMakeLists.txt设备文件、hostfile 与Config::set_devices三种表达连接矩阵的方式语义一致均可用于 mesh 或 ring 拓扑选择 ring 时非相邻节点位置必须留空null/空数组。许可证与致谢JACCL 是 MLX 的一部分采用与 MLX 相同的许可证发布见仓库根目录 LICENSE。名称 JACCL发音 Jackal代表 Jack and Angelos Collective Communication Library既是对 NVIDIA NCCL 的戏仿式致敬也纪念主导 Apple RDMA over Thunderbolt 技术开发的 Jack Beasley。【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考