资讯动态

MLX 分布式程序启动实战指南:mlx.launch 与 mlx.distributed_config 完全解析

发布时间:2026/9/10 23:06:08 来源:尧图企业网站定制
MLX 分布式程序启动实战指南mlx.launch 与 mlx.distributed_config 完全解析【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlxMLXApple silicon 上的数组框架支持将训练与推理的计算负载分摊到多台 Mac 乃至多 GPU CUDA 机器上。本文围绕 launching_distributed.rst 文档系统讲解 MLX Python 包提供的两大分布式启动工具负责自动化配置网络接口与生成 hostfile 的mlx.distributed_config以及负责跨节点拉起、监控与转发进程输出的mlx.launch。读完本文你将掌握 thunderbolt/以太网两种物理链路下 ring、JACCL、NCCL、MPI 四种后端的完整配置流程、hostfile 的 JSON 结构与各后端专属启动参数并能结合 launch.py 与 config.py 的源码理解每一步的实际行为从而独立完成多节点分布式环境的搭建与排障。概览两个工具各司其职MLX 的 Python 包提供两个配套工具来支撑分布式计算mlx.launch负责把脚本分发到多个节点或多进程单节点上运行并维护进程生命周期、转发标准输入输出。它对应mlx.core.distributed模块的运行时启动入口。mlx.distributed_config负责自动化配置 Mac 的网络接口尤其是 thunderbolt 通信并生成供mlx.launch使用的 hostfile。两者的命令行入口在 setup.py 中通过console_scripts注册entry_points { console_scripts: [ mlx.launch mlx._distributed_utils.launch:main, mlx.distributed_config mlx._distributed_utils.config:main, ], }也就是说安装 MLX Python 包后这两个命令即可直接使用。完整的后端体系ring、JACCL、NCCL、MPI入门可参考 分布式通信概览文档。mlx.distributed_config自动化配置 Mac 集群除非你只是在本地做开发测试或运行多 GPU 的 CUDA 环境否则通常都需要先配置多台 Mac 的分布式通信环境。mlx.distributed_config的目的就是自动化完成两件事配置网络接口尤其是 thunderbolt 通信创建供mlx.launch使用的 hostfile。以下分析该工具的三种典型使用场景JACCL over thunderbolt、ring over thunderbolt、ring over ethernet。其主流程位于 config.py 的main()函数先校验 ssh 连通性再根据--over thunderbolt或--over ethernet分支处理。场景一RDMA over thunderboltJACCL在按照启用 RDMA 的步骤macOS 26.2 起 thunderbolt 5 支持 RDMA需在 macOS 恢复模式中执行rdma_ctl enable并重启完成 RDMA 使能后运行如下命令即可配置节点并生成 hostfilemlx.distributed_config --verbose --backend jaccl \ --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 --over thunderbolt \ --auto-setup --output m3-ultra-jaccl.json脚本自动完成的七个步骤与 config.py 实现一一对应ssh 到所有节点验证可达性check_ssh_connections并发执行ssh -o BatchModeyes -o ConnectTimeout5探活同时探测各节点是否具备免密 sudoconfig.py提取 thunderbolt 连接关系在每个节点上运行system_profiler SPThunderboltDataType -json与networksetup -listallhardwareports解析出每个端口连接的对方 UUID构建拓扑extract_connectivityconfig.py验证存在合法的全连接 meshcheck_valid_mesh检查任意两节点之间是否有直连不满足则报错并提示用--dot可视化config.py。如上图所示JACCL 要求任意两台 Mac 之间都有 thunderbolt 线缆直连检查 RDMA 已启用check_rdma在各节点运行ibv_devices确认存在rdma_*设备config.py提取 en0 接口的以太网 IPadd_ips依次尝试ipconfig getifaddr en0、ipconfig getifaddr en1取得各节点 IPconfig.py禁用 thunderbolt 桥接并建立每条线缆的点对点网络IPConfigurator.setup生成sudo ifconfig bridge0 down、为每条连接分配192.168.x.y子网 IP 并配置路由config.py写出 hostfilesave_hostfile将结果写入--output指定的 JSON 文件未指定则打印到 stdoutconfig.py。理解上述步骤后你既可以手工复现配置也可以在出问题时定位故障点。例如把以太网 IP 换成别的接口直接改配置也是可行的——只要该 IP 从所有节点可达即可。--auto-setup参数要求各节点具备免密 sudosudo ls无需密码可执行。如果不可用脚本不会执行配置命令而是把需要在每个节点上手动执行的命令打印出来IPConfigurator.setup中auto_setupFalse的分支会逐条输出并等待回车确认。场景二ring over thunderbolt在 thunderbolt 上搭 ring 后端只需把--backend从jaccl改成ringmlx.distributed_config --verbose --hosts host1,host2,host3,host4 --backend ring执行步骤与 JACCL 高度相似主要区别在于不再校验全连接 mesh而是通过extract_rings用深度优先搜索在连接矩阵中寻找环状拓扑或若干个环再用check_valid_ring确认存在覆盖全部节点的完整环config.py。默认不指定--backend时脚本会同时探测 mesh 与 ring并按RDMA 且 mesh → jacclRDMA 且 ring → jaccl-ring仅 ring → ring的优先级自动选择配置方案config.py。场景三ring over ethernet如果走以太网则不需要配置网络接口脚本只做两件事提取各节点en0的 IP然后写出 hostfilemlx.distributed_config --verbose --hosts host1,host2,host3,host4 --over ethernet对应的prepare_ethernet_hostfile只调用了add_ips并直接保存 hostfileconfig.py。由于--over默认值是thunderbolt走以太网时必须显式传入--over ethernet。调试线缆连接--dot 导出拓扑图mlx.distributed_config可以把节点间的 thunderbolt 连接关系导出为 GraphViz 图帮助快速定位哪根线缆没有插对mlx.distributed_config --verbose \ --hosts host1,host2,host3,host4 \ --over thunderbolt --dot--dot分支调用tb_connectivity_to_dot输出 DOT 格式描述每个节点一个矩形框每条边标注两端接口名如en2/en2config.py。配合dot命令即可渲染出可视化图形例如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这也是分布式文档中先接线、再可视化、确认无误后自动配置的标准工作流。mlx.launch跨节点启动分布式程序最小用法与核心行为最简用法是在多节点上启动脚本mlx.launch --hosts ip1,ip2 my_script.py或在本机多进程测试mlx.launch -n 2 my_script.pymlx.launch的行为要点实现于 launch.py 的_launch_with_io与RemoteProcess连接给定主机在每个主机上启动输入脚本监控每个已启动进程一旦某个进程异常退出或mlx.launch本身被终止如 CtrlC其余进程也会被终止_launch_with_io中通过stop标志和RemoteProcess.terminate实现pid 通过 mktemp 文件记录以便远程 killlaunch.py把每个远程进程的stdout / stderr 分别转发到本地对应流广播 stdin 到所有进程这使得交互式程序与pdb调试器在分布式模式下也能正常工作_launch_with_io用select轮询 stdin 并写入每个进程的 stdin 队列launch.py。此外--env可在远端设置环境变量如--env MLX_METAL_FAST_SYNCH1--cwd可指定各节点的工作目录--python可指定远端使用的解释器路径。提供 Hosts命令行参数与 JSON hostfileHosts 既可以像上面那样直接作为命令行参数给出也可以通过 JSON hostfile 完整定义。hostfile 的 schema 很简单一个对象列表每个对象用ssh指定用于 ssh 的主机名、用ips给出该节点用于通信的 IP 列表[ {ssh: hostname1, ips: [123.123.1.1, 123.123.2.1]}, {ssh: hostname2, ips: [123.123.1.2, 123.123.2.2]} ]这里ssh 的主机名与通信 IP是解耦的你可以 ssh 到一个主机名但让它监听另一个 IP。解析逻辑在 common.pyHostfile.from_file支持两种顶层结构——纯列表如上或带backend/envs/hosts字段的对象Hostfile.from_list则把命令行--hosts拆分为 Host并自动识别是 IP 则作为 ips、是主机名则仅作 ssh 用。需要强调Hostfile.from_file还能从 hostfile 中读取backend与envs字段mlx.launch会在命令行未显式指定后端时采用 hostfile 内声明的后端并把其中声明的环境变量一并注入launch.py。用mlx.distributed_config --over ethernet即可生成 IP 对应en0接口的 hostfile。设置远程主机三件套检查清单要在每台主机上启动脚本必须满足ssh hostname无需密码或主机确认即可连通Python 解释器在所有主机上位于相同路径——用mlx.launch --print-python查看当前使用的路径该参数直接打印sys.executable后退出launch.py要运行的脚本在所有主机上位于相同路径。不满足以上任一条件分布式启动就会失败这通常也是排查 rank 启动失败 的第一站。ring 后端专属参数ring 是默认后端在无 CUDA 环境下由main()自动选择nccl if mx.cuda.is_available() else ringlaunch.py也可用--backend ring显式指定。它与其他后端不同的要求与参数--hosts只接受IP而不接受主机名若 ssh 用的主机名与要绑定的 IP 不一致则必须提供 hostfile--starting-port别名-p默认32323定义远端主机绑定的起始端口。rank 0 的第一个 IP 使用该端口之后每个 IP 或 rank 依次加 1launch.py 的launch_ring按此逐 IP、逐连接累加端口并生成MLX_HOSTFILE环境变量--connections-per-ip增加相邻节点间的连接数以提升带宽等价于mpirun的--mca btl_tcp_links 2。从源码看ring 后端在 C 侧ring.cpp通过 TCP socket 实现环形拓扑的 all reduce 与 all gather节点只与环上的邻居通信MLX_RING_VERBOSE1可开启[ring]前缀的调试日志。JACCL 后端专属参数--backend jaccl选择 JACCL。启动该后端必须有 hostfile因为 hostfile 需要包含连接各节点的 RDMA 设备列表每行rdma字段null表示自身。launch_jaccl会严格校验每个 host 的rdma数量必须与节点数一致、自身对应的设备必须是null、且任意两节点之间不能缺失 RDMA 设备否则直接报错拒绝启动launch.py。启动时mlx.launch会为各进程注入MLX_JACCL_COORDINATORrank0_ip:port与MLX_IBV_DEVICES由 hostfile 的 rdma 矩阵序列化而来。JACCL 的 hostfile 形如对应上图的四节点全连接 mesh[ { ssh: m3-ultra-1, ips: [123.123.123.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] } ]注意即使通信走 Thunderbolt RDMA不经过 TCP/IP禁用 thunderbolt 桥接和为每条 thunderbolt 连接建立隔离的本地子网这两步依然必需。这些都可以由mlx.distributed_config自动完成。关于 JACCL 初始化时在 rank 之间通过 side channel 交换 RDMA 元数据的机制all_gather_factory自定义 side-channel all-gather以及低延迟场景下可选的MLX_METAL_FAST_SYNCH环境变量默认关闭存在死锁风险建议保持不设置详见分布式通信概览文档。NCCL 后端专属参数NCCL 是 CUDA 环境的默认后端。当从 Mac 启动到带 CUDA 的 Linux 机器时需用--backend nccl显式指定。多节点 多 GPU 任务要用--repeat-hosts别名-nmlx.launch --backend nccl --hosts linux-1,linux-2 -n 8 -- ./my-job.sh上面的命令会在每个节点上启动 8 个进程共 16 个进程运行my-job.sh。launch_nccl的实现细节rank 0 的第一个 IP 作为NCCL_HOST_IP端口由--nccl-port默认12345指定MLX_WORLD_SIZE设为进程总数每个 rank 的CUDA_VISIBLE_DEVICES取rank % repeat_hosts即各节点上按序分配本地 GPU 编号launch.py。加--verbose时自动注入NCCL_DEBUGINFO。MPI 后端专属参数--backend mpi让mlx.launch成为mpirun的薄封装。此时hostfile 中的IP 被忽略只使用其中的 ssh 主机名并按重复次数生成slots写入临时 hostfilelaunch.py对 ssh 连通性的要求更高每个节点都要能连到其他所有节点mpirun必须在每个节点上以相同路径可用。mlx.launch还代为处理了 macOS 上的库路径问题通过otool -L探测libmpi动态库名自动注入DYLD_LIBRARY_PATH与MLX_MPI_LIBNAMEget_mpi_libnamelaunch.py兼容 Homebrew / pip 安装的 MPI。向mpirun透传参数用--mpi-arg。例如为 MPI 的 byte-transfer-layer 指定网卡mlx.launch --backend mpi --mpi-arg --mca btl_tcp_if_include en0 --hostfile hosts.json my_script.py也可以使用--mpi-arg --mca btl_tcp_links N增加每对主机间的 TCP 连接数以提升带宽。需要更高带宽的 all reduce 时分布式文档建议优先考虑 ring 后端Thunderbolt 或以太网均可。实战建议与排查要点综合文档与源码以下是落地分布式任务时的实用建议先本地小规模验证用mlx.launch -n 2 -- my_script.py在单机先跑通脚本内建议直接调用mx.distributed.all_sum等集合通信无需写if world.size() 1的分支判断——MLX 在单进程 group 下所有分布式操作都是 noop善用可视化排障mlx.distributed_config --hosts h1,h2,h3 --over thunderbolt --dot检查线缆连接是否正确再决定是否--auto-setup善用交互调试mlx.launch广播 stdin、聚合 stdout配合pdb调试远端进程非常方便善用--verbosemlx.launch --verbose与mlx.distributed_config --verbose会打印每个步骤的 INFO 日志绿色[INFO]前缀是理解脚本行为和定位问题的第一手线索核对三件套ssh 免密、Python 路径一致--print-python、脚本路径一致是分布式启动失败时最常被忽略的原因。如果项目由集群调度器如 SLURM而非mlx.launch来拉起进程则可以直接设置各后端所需的环境变量如 ring 的MLX_RANKMLX_HOSTFILE、JACCL 的MLX_RANKMLX_JACCL_COORDINATORMLX_IBV_DEVICES、NCCL 的MLX_RANKMLX_WORLD_SIZENCCL_HOST_IPNCCL_PORTCUDA_VISIBLE_DEVICES具体清单与示例见分布式通信概览文档。至此从线缆接线、节点配置到跨机启动的完整链路已经打通你可以基于 examples/python/distributed_data_parallel.py 与 examples/python/distributed_tensor_parallel.py 中的示例脚本把训练或推理负载正式扩展到整个集群。【免费下载链接】mlxMLX: An array framework for Apple silicon项目地址: https://gitcode.com/GitHub_Trending/ml/mlx创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价