Explore distributed inference and training with MLX

Highlight

Apple 在 macOS 26.2 中通过 Thunderbolt 5 支持 RDMA，配合开源通信库 JACCL 和 MLX，让几台 Mac 组成集群就能本地运行万亿参数大模型，推理和微调速度提升约 3 倍，无需云服务器。

核心内容

单台 Mac 跑不动大模型了

本地大模型发展很快。模型越来越大，能力越来越强，开发者用它处理更长的上下文、更复杂的任务。但内存、算力、带宽迟早会撞上天花板。一台 M3 Ultra 能跑的模型有上限，更大的模型根本装不下。

WWDC26 的 session 232 展示了如何用 MLX 在单台 Mac 上跑本地 AI Agent。但如果桌面上有多台 Mac，能不能把它们连起来，把负载分摊出去？

这场演讲回答了这个问题的完整路径。

从硬件到软件的全栈方案

Apple 没有只做框架层的高 level 抽象，而是从底层通信协议开始构建。

（02:29）macOS 26.2 开始支持通过 Thunderbolt 5 走 RDMA（Remote Direct Memory Access）。RDMA 让数据直接从一台机器的内存搬到另一台，绕过 CPU 和操作系统的大部分开销，实现高带宽、低延迟的机器间通信。

（03:17）光有 RDMA 还不够，分布式程序需要更高层的通信原语。Apple 开源了 JACCL（Just Another Collective Communication Library），基于 RDMA over Thunderbolt 提供集合通信操作，比如跨机器发送数据、汇总结果。JACCL 不限于机器学习，任何 Apple Silicon 上的分布式负载都能用。

（03:57）最上层是 MLX。MLX 利用 JACCL 做低延迟分布式通信，并提供工具来编排集群上的分布式任务。

把 4 台 M3 Ultra 变成集群

（04:22）演讲用 4 台 M3 Ultra 演示了完整流程。

第一步是选拓扑。JACCL 支持两种：

• Mesh（全连接）：每台机器直连其他所有机器，延迟最低，适合需要频繁通信的场景。
• Ring（环）：每台机器只连两个邻居，省端口、省线缆，可以通过多根线缆增加单链路带宽，适合带宽敏感的大消息传输。

（06:46）JACCL 会自动根据消息大小和操作类型选择最优拓扑：小消息走 mesh 降延迟，大消息走 ring 提带宽。

第二步启用 RDMA。在系统设置里搜索 “RDMA”，打开 “Enable RDMA over Thunderbolt”，重启即可。

（07:59）第三步用 mlx.launch 启动分布式任务。它通过 SSH 连到每台机器，启动可执行文件，之后所有机器直接通过 Thunderbolt 链路通信。

mlx.launch 需要一个 hostfile 描述集群。MLX 提供了 mlx.distributed_config 工具自动生成：

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

关键点：

• --hosts：集群中所有机器的主机名。
• --output：生成的 hostfile 路径。
• --env MLX_METAL_FAST_SYNCH=1：启用更快的 GPU-CPU 同步，分布式任务中计算在 GPU、通信在 CPU，同步速度很关键。
• --auto-setup：自动禁用 Thunderbolt Bridge 并配置 RDMA。
• --backend jaccl：mesh 拓扑；用 jaccl-ring 则是 ring 拓扑。

脚本会先检查 SSH 连通性，再探测 Thunderbolt 端口发现物理连接关系，最后输出 JSON hostfile。

分布式推理：单条命令跑集群

（10:39）集群就绪后，用 MLX LM 跑分布式推理几乎和单机一样简单。

单机推理：

mlx_lm.chat --model "Qwen/Qwen3.6-27B" --max-tokens 2048

分布式推理：

mlx.launch --hostfile "m3-ultra-jaccl.json" -- \
    /remote/path/to/mlx_lm.chat --model "Qwen/Qwen3.6-27B" --max-tokens 2048

（12:21）演讲现场对比了单机 vs 4 机集群跑 Qwen 3.6（270 亿参数）。集群的 token 生成速度接近单机的 3 倍。

（13:18）速度不是唯一的理由。有些模型大到单机根本装不下。Kimi 2.6 有 1 万亿参数，8-bit 量化后权重约 1TB，一台 M3 Ultra 放不下，但 4 台可以。

Pipeline vs Tensor 并行

MLX 支持两种模型切分策略：

（13:47）Pipeline 并行按深度切分。每台机器保存一组层，数据依次流过各机器。它不加速推理（每个 token 仍需串行经过所有层组），但通信简单，只在层组边界交换激活值。

（14:13）Tensor 并行按宽度切分。每台机器保存每层的一部分，所有机器同时处理同一个 token。它加速推理，但每层、每个 token 都要通信，对延迟要求高。这也是 mesh 拓扑重要的原因——任意两台机器一跳直达。

Tensor 并行是 MLX LM 的默认策略。切到 pipeline 并行只需加一个 --pipeline 标志：

mlx.launch --hostfile "m3-ultra-jaccl.json" -- \
    /remote/path/to/mlx_lm.chat --model "moonshotai/Kimi-K2.6" \
                                 --max-tokens 2048 \
                                 --pipeline

（15:03）演讲演示了用 tensor 并行在 4 台 M3 Ultra 上跑 Kimi 2.6——万亿参数模型，一条命令，本地运行。

分布式微调：数据并行

（16:11）微调时，单机把训练数据分成 batch，逐个计算梯度、更新权重。多机加速的思路是数据并行：每台机器复制一份模型，各自处理不同的 batch，然后平均梯度，让更新融合所有 batch 的信息。

N 台机器理论上可以把数据处理速度提升 N 倍。

分布式微调和单机命令几乎一样：

# 单机微调
mlx_lm.lora --model "Qwen/Qwen3.5-9B" \
             --data "mlx-community/wikisql" \
             --train --batch-size 4

# 分布式微调（batch-size 按设备数放大）
mlx.launch --hostfile "hostfile.json" -- \
    /remote/path/to/mlx_lm.lora --model "Qwen/Qwen3.5-9B" \
                                  --data "mlx-community/wikisql" \
                                  --train --batch-size 16

（17:55）现场对比：单机 M3 Ultra 约 180 tokens/秒，4 机集群约 600 tokens/秒，微调速度提升超过 3 倍。

详细内容

Python API 控制分布式推理

（19:01）CLI 之外，MLX 提供了细粒度的 Python API。

import mlx.core as mx
from mlx_lm import stream_generate
from mlx_lm.utils import sharded_load

# 初始化分布式后端
group = mx.distributed.init(strict=True, backend="jaccl")
# 定义并行策略
tensor_group, pipeline_group = group, None

# 加载并切分模型
model, tokenizer = sharded_load("moonshotai/Kimi-K2.6", pipeline_group, tensor_group)
for response in stream_generate(model, tokenizer, prompt, max_tokens=1024):
    if group.rank() == 0:
        print(response.text, end="", flush=True)

关键点：

• mx.distributed.init(backend="jaccl")：初始化 JACCL 分布式组。
• tensor_group = group, pipeline_group = None：启用 tensor 并行，不启用 pipeline 并行。
• sharded_load()：自动按指定的并行策略切分并加载模型。
• group.rank() == 0：只有 rank 0 的节点输出结果，避免重复打印。
• 加载后的模型用法和单机完全一致，分布式通信由 MLX 内部处理。

手动切分层

（19:31）如果需要更底层的控制，可以直接用 MLX 的 shard_linear：

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

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

# 定义层并按列切分
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)
mx.eval(output)

关键点：

• strategy="all-to-sharded"：将完整层按列切分到各设备。
• shard_linear 返回的层在调用时会自动处理跨设备通信。
• mx.eval(output)：触发实际计算和同步。

跨语言的原语：All-reduce

（19:47）MLX 在 Python、Swift、C++ 中暴露了相同的分布式原语。以 all-reduce（跨设备求和）为例：

Python：

import mlx.core as mx
world = mx.distributed.init(strict=True, backend="jaccl")
data = mx.full((4,), float(world.rank()), dtype=mx.float32)
result = mx.distributed.all_sum(data, group=world)
mx.eval(result)

Swift：

let group = try DistributedGroup(strict: .ring)
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)

C++：

namespace mx = mlx::core;
auto world = mx::distributed::init(/* strict */ true, "jaccl");
mx::array data = mx::full({4}, static_cast<float>(world.rank()), mx::float32);
mx::array result = mx::distributed::all_sum(data, world);
mx::eval(result);

关键点：

• 三种语言的 API 结构一致：初始化 group → 准备数据 → 执行 all_sum → eval 同步。
• world.rank() 获取当前节点编号，用于构造差异化测试数据。
• all_sum 将各设备上的数组按元素求和，结果广播回所有设备。

脱离 MLX，直接用 JACCL

（20:14）JACCL 可以独立使用，不依赖 MLX。它提供 C++ API：

#include <jaccl/jaccl.h>
#include <iostream>

int main() {
    // 初始化 JACCL 组
    auto group = jaccl::init();
    std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl;
    // 执行 all-reduce 求和
    float data[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(data, output, sizeof(data), jaccl::Float32);
    std::cout << "Result: " << output[0] << std::endl;
    return 0;
}

关键点：

• jaccl::init()：自动发现集群配置并初始化通信组。
• group->rank() / group->size()：获取当前节点编号和总节点数。
• all_sum 的参数：输入数组、输出数组、字节数、数据类型。
• 任何需要跨机器同步数据的程序都可以基于 JACCL 构建。

核心启发

1. 本地部署企业级知识库

• 做什么：用 2-4 台 Mac Studio 组成集群，本地部署 Qwen 3.6 或更大模型，作为公司内部知识库问答系统。
• 为什么值得做：数据不出公司网络，满足合规要求；分布式推理速度接近云端 A100 实例，成本是一次性硬件投入。
• 怎么开始：先在一台 Mac 上用 mlx_lm.chat 验证模型效果，确认后买第二台 Mac，用 mlx.distributed_config 生成配置，再用 mlx.launch 启动分布式服务。

2. 私有领域模型微调

• 做什么：收集公司内部文档、客服对话、产品手册，用 MLX 的 LoRA 微调一个基础模型，打造专属助手。
• 为什么值得做：session 展示了 4 机集群微调速度提升 3 倍以上，本地微调意味着敏感训练数据不会上传到任何第三方。
• 怎么开始：用 mlx_lm.lora 在单机上用小数据集验证微调流程，确认效果后把 --batch-size 按设备数放大，用 mlx.launch 扩展到集群。

3. 多 Agent 协作系统的推理后端

• 做什么：构建一个本地多 Agent 系统，不同 Agent 负责代码审查、文档生成、测试用例编写，共享同一个大模型后端。
• 为什么值得做：session 232 讲了本地 Agentic AI，session 233 解决了”模型不够大、推理不够快”的瓶颈。两者结合可以做出真正可用的本地 Agent 工作流。
• 怎么开始：先用 session 232 的方法搭好单 Agent 原型，再用本 session 的 sharded_load + stream_generate 把推理层换成分布式，支撑多 Agent 并发请求。

4. 用 JACCL 做非 ML 的分布式计算

• 做什么：利用 JACCL 的 C++ API 做图像批量处理、音频分析、科学计算等分布式任务。
• 为什么值得做：JACCL 不绑定 MLX，任何需要跨机器同步数据的程序都能用。Thunderbolt 5 的带宽和延迟比 WiFi/以太网好一个数量级。
• 怎么开始：写一个简单的 C++ 程序，包含 <jaccl/jaccl.h>，调用 jaccl::init() 和 all_sum 验证通信，再把业务逻辑接入。

5. 在 Swift App 中嵌入分布式推理

• 做什么：在 macOS App 中集成大模型推理能力，用户的多台 Mac 自动组成集群加速响应。
• 为什么值得做：MLX 的 Swift API 和 Python API 结构一致，可以把实验好的 Python 模型直接移植到 App 中。
• 怎么开始：先用 Python API 验证模型和切分策略，再用 MLX Swift 的 DistributedGroup 和对应 API 重写核心推理逻辑，嵌入 App。

关联 Session

• Session 232: Run local agentic AI on the Mac using MLX — 单机本地 AI Agent 的入门指南，和本 session 的分布式能力形成完整拼图
• Session 328: MLX Swift — 在 Swift 中使用 MLX 的详细指南，适合想把分布式推理嵌入 App 的开发者
• Session 324: Core AI — Apple 平台 AI 核心框架概览，了解 MLX 在整个生态中的位置
• Session 241: Foundation Models — 苹果基础模型的设计理念，帮助理解为什么本地运行大模型越来越重要