Explore distributed inference and training with MLX
AI & Machine Learning 高级 22m

使用 MLX 探索分布式推理和训练

Explore distributed inference and training with MLX

2026年6月10日

在 Apple 官方观看视频

一句话判断

Mac 终于能名正言顺地当“穷人版” AI 算力集群了，这场 Session 最值得关注的是 Apple 用 Thunderbolt 5 上的 RDMA (Remote Direct Memory Access) 配合全新的 JACCL 通信库，把几台 Mac 物理缝合成了一个低延迟的巨型统一内存池。

这场 Session 讲了什么

大模型越来越吃内存和算力，单台 Mac 哪怕是 M3 Ultra 也扛不住万亿参数模型。Apple 给出的解法不是让你去买昂贵的云 GPU，而是让你把桌子上的几台 Mac 用 Thunderbolt 5 线连起来。通过底层的 RDMA 协议绕过 CPU 和系统开销，加上开源的 JACCL (集合通信库) 做节点协调，MLX 框架现在可以直接把推理和微调任务切分扔给这个“Mac 集群”。

对独立开发者和小团队来说，这绝对是省钱利器。以前跑个 27B 模型微调或者推理 1T 参数的 Kimi 2.6，要么忍受单机的龟速，要么去云上烧钱。现在只要写个 JSON 配置文件，用 mlx.launch 命令行工具一键分发，4 台 M3 Ultra 就能跑出接近 3 倍的 token 生成速度。API 层面也从 CLI 一路打通到了 Python、Swift 和 C++，让你能在自己的 App 里直接嵌入分布式推理逻辑。

值得深挖的点

RDMA over Thunderbolt 5 与 JACCL：把 Mac 变成真集群的底座

Apple 这次没搞虚的，直接下探到了硬件协议层。RDMA 允许一台 Mac 直接读写另一台 Mac 的内存，完全不需要 CPU 介入，这解决了分布式计算中最致命的通信延迟问题。配合 Apple 开源的 JACCL，开发者不需要自己去写底层的 Socket 通信或者处理复杂的网络拥塞控制。

这种设计的 trade-off 很明显：它极度依赖物理连接。Wi-Fi 或者普通的千兆以太网根本跑不满大模型张量并行 (Tensor Parallelism) 所需的带宽。你必须老老实实买 Thunderbolt 5 线缆，把机器物理连起来。JACCL 支持 Full Mesh (全网状) 和 Ring (环形) 两种拓扑。Mesh 延迟最低，但 4 台机器需要 6 根线，且 M3 Ultra 的接口数量可能不够用；Ring 只需要 4 根线连成一圈，适合大消息传输。Apple 提供 mlx.distributed_config 自动生成 Hostfile，算是把物理接线的痛苦降到了最低。

张量并行 vs 流水线并行：大模型切片的艺术

当单机内存装不下模型时，MLX LM 提供了两种切片策略。张量并行是把模型的每一层按宽度切开，分给不同的 Mac。优点是推理速度快，因为所有机器同时计算一层；缺点是每一层算完都要做 All-Reduce 通信，对 Thunderbolt 5 的带宽要求极高。流水线并行 (Pipeline Parallelism) 则是按深度切，比如前 10 层给 Mac A，后 10 层给 Mac B。通信只在层与层交接时发生一次，对带宽要求低，但会产生“气泡”(Bubble)，导致部分机器闲置。

在跑 1T 参数的 Kimi 2.6 时，Session 演示了加上 --pipeline 参数切换到流水线并行。如果你的集群只有 2-4 台机器且 TB5 连线完备，无脑用默认的张量并行；如果你搞了 8 台以上的 Mac 集群，或者网络拓扑受限，老老实实切流水线并行。

代码片段

1. 一键启动分布式推理

场景：用 4 台 M3 Ultra 集群跑 27B 的 Qwen 模型，体验 3 倍生成速度。

# 生成集群配置文件，自动配置 RDMA 和 JACCL 后端
mlx.distributed_config \
    --hosts m3-ultra-0,m3-ultra-1,m3-ultra-2,m3-ultra-3 \
    --output "m3-ultra-jaccl.json" \
    --env MLX_METAL_FAST_SYNCH=1 \
    --auto-setup \
    --backend jaccl

# 使用 mlx.launch 将任务分发到集群，默认使用张量并行
mlx.launch --hostfile "m3-ultra-jaccl.json" -- \
    /remote/path/to/mlx_lm.chat --model "Qwen/Qwen3.6-27B" --max-tokens 2048

坑：--auto-setup 会尝试自动配置 SSH 免密登录，如果 Mac 之间没配好 SSH Key，这里会直接报错卡死。

2. 在 Python 中手动切分模型层

场景：不想用 CLI，想在自己的 Python 训练脚本里精细控制某一层 Linear 的分布式切分。

import mlx.core as mx
import mlx.nn as nn

# 初始化 JACCL 分布式后端
group = mx.distributed.init(strict=True, backend="jaccl")

# 定义一个普通的 Linear 层，并按列切分给集群中的设备
layer = nn.Linear(1024, 1024)
sharded_layer = nn.layers.distributed.shard_linear(
    layer, strategy="all-to-sharded", group=group
)

# 生成随机数据并前向传播
data = mx.random.normal((1, 1, 1024))
output = sharded_layer(data)
# 触发 MLX 的惰性计算
mx.eval(output)

坑：strategy="all-to-sharded" 意味着输入数据会全量广播，输出会被切分。如果输入数据本身已经很大，这一步的内存开销会直接撑爆单台 Mac 的统一内存。

3. Swift 端的 All-Reduce 通信

场景：在 macOS 原生 App 中，让多个进程或设备对计算结果求和。

import MLX
import MLXDistributed

// 初始化 Ring 拓扑的分布式组
let group = try DistributedGroup(strict: .ring)

// 根据当前节点的 rank 生成不同的数据
let data = rank == 0
    ? MLXArray(converting: [1.0, 2.0, 3.0])
    : MLXArray(converting: [5.0, 6.0, 7.0])

// 跨设备求和，所有节点都会拿到总和
let result = try group.allSum(data)

坑：Swift API 目前对错误处理比较严格，如果集群中有一台 Mac 掉线或者 Thunderbolt 线缆松动，allSum 会直接抛出 fatal error，记得在外部做好 try-catch 和重试机制。

最佳实践

  • 已有项目的迁移：如果你之前用单台 Mac 跑 mlx_lm.lora 微调，现在只需要把 batch-size 乘以机器数量，然后用 mlx.launch 包装你的启动命令。代码一行不用改，MLX 会自动处理梯度的 All-Reduce 和平均。
  • 新项目的采用：强烈建议先用 CLI 工具跑通官方的 Qwen 或 Llama 模型，确认你的物理连线和 RDMA 网络配置没问题后，再切入 Python/Swift API 写自定义逻辑。物理层的坑比代码层的坑难排查得多。
  • 实战避坑：千万别用 Wi-Fi 或者普通的 USB-C 线去连集群。Thunderbolt 5 的 80Gbps 带宽是张量并行的生命线。另外，在 System Settings (系统设置) 里一定要检查 RDMA 是否被正确启用，macOS 默认有时会因为安全策略拦截直接的内存访问。

还有什么值得关注

  • JACCL 不仅限于 AI：这个开源通信库完全可以脱离 MLX 使用，如果你在做分布式的视频渲染或者物理模拟，JACCL 提供的全套原语同样适用。
  • 内置的分布式 Server：MLX LM 现在自带了 distributed server 模式，可以直接把 Mac 集群包装成一个兼容 OpenAI API 的本地服务端，对做 Agent 开发的同学非常友好。
  • 与 Agentic AI 的联动：Session 提到了配套的本地 Agent 视频，分布式推理提速后，本地跑复杂的多 Agent 协作工作流终于不会因为 token 生成太慢而卡脖子了。
MLX Distributed Inference RDMA JACCL Thunderbolt 5