跳到正文
Apple Machine Learning·· 8 天前精选AI 评分61

如何用 MLX 与 RDMA 在 Mac Studio 上运行本地及分布式大语言模型推理

Local and Distributed Large LLM Inference with MLX and RDMA on Mac Studio

AI 导读

Apple 演示使用 MLX 与 Thunderbolt RDMA 在 Mac Studio 上运行本地及分布式推理,四台 M5 Ultra 以张量并行运行 Muse Spark Open,在 8,192 tokens 上下文下达到 36.73 tokens/s。

推荐理由

通过同一组模型的单机与多机测试,对比流水线并行和张量并行在首 token 延迟、生成速度及通信开销上的差异。

正文 · AI 翻译

内容类型 精选亮点

The MLX logo

搭载 Apple 芯片的 Mac 上的 MLX 让 AI 研究人员和开发者能够在私密环境中,本地试验最新的模型和方法。这个开源数组框架利用 Apple 芯片的统一内存架构,让任何人都能在自己的硬件上运行新的推理和微调方法,并探索其他 AI 技术。在此前的一篇文章中,我们重点介绍了 M5 架构及其 GPU 神经网络加速器的优势。

在本文中,我们将展示如何使用今天发布的最新版 MLX,在你的桌面上直接运行拥有近万亿参数的最新前沿级模型 Muse Spark Open。在由四台 Mac Studio 组成的集群上,Muse Spark Open 每秒可生成超过 36 个 token(见图 1),远高于人类的平均阅读速度。

从单台 Mac Studio(最高配备 512GB 统一内存)起步,再通过基于 Thunderbolt 的 RDMA 驱动程序扩展至多台 Mac,下面我们将概述如今可在 Mac Studio 上运行的各类模型。

图 1:在搭载 M5 Ultra 的 Mac Studio 上,上下文大小为 8,192 个 token 时,Muse Spark Open 的解码速度,单位为 token/秒。数值越高越好。我们报告了三种并行方案的性能:在两台 Mac Studio 电脑上进行流水线并行(2 M5 Pipeline),以及在两台(2 M5 Tensor)和四台 Mac Studio 电脑(4 M5 Tensor)上进行张量并行。

我们按照此前文章中介绍的方法,使用 MLX LM 软件包,在 Mac Studio 上对几款流行的开放权重大语言模型进行基准测试。我们比较了两代 Mac Studio:最近发布的 M5 Ultra 和上一代 M3 Ultra。M5 这一代配备 GPU 神经网络加速器,支持专用矩阵乘法运算,这对许多机器学习工作负载至关重要。

我们评估了三款前沿级模型(表 1),其规模跨度很大(参数量以十亿为单位,记作 B):新发布的 Muse Spark Open(816 B)、Mistral Medium 3.5(128 B)和 Qwen 3.8(27 B)。Muse Spark Open 是一款混合专家(MoE)模型,仅激活其中一小部分专家及其对应的参数,从而减少推理时的计算量,我们以 NVFP4 精度对其进行评估。Mistral Medium 3.5 和 Qwen 3.8 采用稠密架构,推理时所有参数均处于激活状态,我们以各自的原生浮点精度(分别为 FP8 和 BF16)对其进行评估。

模型总参数量激活参数量原生精度评估精度
Muse Spark Open816B42BBF16NVFP4
Mistral Medium 3.5128B128BFP8FP8
Qwen 3.827B27BBF16BF16
表 1:本文使用的语言模型。模型以多种精度格式发布;我们以原生精度评估 Qwen 3.8 和 Mistral Medium 3.5,以 4 位(NVFP4)精度评估 Muse Spark Open。

我们的评估在一系列上下文大小下衡量三项指标:首个 token 生成时间(TTFT,单位为秒)、生成吞吐量(单位为 token/秒)和峰值功耗(单位为瓦)。本文中的所有性能数据均使用 mlx_lm.benchmark 测得。例如,以 8,192 个 token 的上下文运行 Qwen 3.8 并生成 128 个 token,可通过以下命令行完成:

mlx_lm.benchmark --model Qwen/Qwen3.8-27B -p 8192 -g 128 -b 1

在 Mac Studio 上使用本地大语言模型的一种常见用途,是在私密环境中运行智能体体验,包括通常需要较大上下文的智能体编程应用。因此,我们报告了上下文大小分别为 8,192 和 65,536 时的性能。

原始 TTFT 和解码吞吐量分别见表 2和表 3。

Qwen 3.8 27B 和 Mistral Medium 3.5 128B 均可轻松装入一台配备 512GB 内存的 Mac Studio,在大上下文(65,536 个 token)配置下,内存占用分别最高为 68GB 和 158GB。在图 2中,我们比较了这两个模型在分别搭载 M3 Ultra 和 M5 Ultra、配备 512GB 内存的单台 Mac Studio 上的推理性能。

首 token 延迟(TTFT)主要受计算能力限制,而这正是 M5 Ultra GPU 中神经网络加速器大显身手的地方:GPU 神经网络加速器针对快速矩阵乘法提供的专用硬件支持,直接转化为 TTFT 的显著缩短,与 M3 Ultra 相比,速度最高可达其 4.6×(见图 2)。这使得搭载 M5 Ultra 的 Mac Studio 成为处理超大上下文任务时的首选台式机。

生成速度(以每秒生成的 token 数衡量)受内存带宽限制。M5 Ultra 芯片提供 1.2 TB/s 的统一内存带宽,高于 M3 Ultra 的 819 GB/s。同样,这也直接转化为生成速度的提升(见图 3),搭载 M5 Ultra 的 Mac Studio 每秒生成的 token 数达到前者的 1.6×。

图 2:在上下文长度为 8,192 和 65,536 个 token 时,Qwen 3.8 和 Mistral Medium 3.5 在 M5 Ultra 上相对于 M3 Ultra 基准的首 token 延迟(TTFT)加速比。M5 Ultra GPU 神经网络加速器在受计算能力限制的工作负载中表现出色。
图 3:在上下文长度为 8,192 和 65,536 个 token 时,Qwen 3.8 和 Mistral Medium 3.5 在 M5 Ultra 上相对于 M3 Ultra 基准的解码加速比。解码受内存带宽限制,M5 Ultra 的 1.2 TB/s 统一内存带宽(M3 Ultra 为 819 GB/s)是一项关键优势。

Mac Studio 在 LLM 推理方面也很节能(见图 4):搭载 M3 Ultra 的 Mac Studio 峰值功耗约为 200W,搭载 M5 Ultra 的 Mac Studio 则为 300W,仅为普通烤面包机功耗的一小部分。更重要的是,这种较低的功耗意味着 Mac Studio 可以部署在任何家庭或办公环境中,用于 AI 工作。

图 4:上下文长度为 65,536 个 token 时,Qwen 3.8 和 Mistral Medium 3.5 在搭载 M3 Ultra 和 M5 Ultra 的 Mac Studio 上进行推理期间的峰值功耗(单位:瓦)。

虽然 Qwen 和 Mistral 模型都能装入一台配备 512GB 内存的 Mac Studio,但 Muse Spark Open 在原生精度(BF16)下,仅模型参数就需要超过 1.6TB 的内存。采用 4 位量化(NVFP4)格式后,模型的内存占用降至约 493GB。再加上推理期间中间计算所需的内存,该模型便无法装入单台机器,必须在多台 Mac Studio 之间进行并行化。

像 Muse Spark Open 这样超大的前沿级模型无法装入单台机器。不过,可以将模型的一半存储在一台 Mac Studio 上,另一半存储在另一台上。然后,我们可以在两台计算机之间以流水线方式进行模型推理:第一台 Mac Studio 执行模型前一半的层,然后将输出发送给第二台 Mac Studio,由后者执行后一半的层。这种并行方案称为流水线并行(见图 5),实现起来极其简单,只需对代码进行少量修改。

图 5:MLX 支持两种将大模型拆分到多台计算机上的方法。流水线并行沿深度方向拆分模型,而张量并行则沿宽度方向拆分模型。

MLX 支持不同的并行方案,详见 MLX 分布式文档。这些方案依赖快速的机器间通信:macOS 26.2 引入了基于 Thunderbolt 5 的低延迟 RDMA 驱动,MLX(v0.30.1 及更高版本)基于该驱动,为不同的并行方案提供了新的集合通信操作(jaccl,读作“jackal”)。在我们的配置中,两台 Mac Studio 通过两根 Thunderbolt 线缆连接(以使链路带宽翻倍),提供 jaccl 所使用的基于 Thunderbolt 的 RDMA 链路。主机配置完成后,MLX 提供了用于跨多个节点启动作业的辅助命令 mlx.launch:

mlx.launch --hostfile hostfile.json --backend jaccl --env MLX_METAL_FAST_SYNCH=1 -- \
mlx_lm.benchmark --model mistralai/Mistral-Medium-3.5-128B -p 8192 -g 128 -b 1 --prefill-step-size 8192 --pipeline

在此命令中,--backend jaccl 指定 RDMA Thunderbolt 驱动后端。--env MLX_METAL_FAST_SYNCH=1 标志启用了一项实验性 MLX 功能,可实现快速 CPU-GPU 同步,降低 CPU(发起 Thunderbolt 通信)与 GPU(执行计算)之间的通信开销。

总体而言,流水线并行下的解码推理速度(见图 7)与单机非常接近,因为各层仍然按顺序计算。通信开销也很低,因为只需交换两组层之间的中间输出。

不过,值得注意的是,对于非常长的上下文(我们的基准测试中为 65,536 个 token),流水线并行的预填充速度最终超过了单节点,最高可达约 1.8×(见图 6)。这是因为 MLX 会将长上下文切分为较小的块(此处由 --prefill-step-size 8192 设置为 8,192 个 token),并在多台机器之间以流水线方式执行预填充:当第二台机器计算第 k 块的后面各层时,第一台机器会计算第 k+1 块的前面各层,因此相邻块的计算在时间上有所重叠。为了进一步提高计算并行度,接下来我们转向张量并行。

图 6:在上下文长度为 8,192 和 65,536 个 token 时,针对 Qwen 3.8 和 Mistral Medium 3.5,使用流水线并行的两颗 M5 Ultra 相对于单颗 M5 Ultra 基线的首 token 时间(TTFT)加速比。拆分上下文(此处拆分为每块 8,192 个 token)使 MLX 能够让两台机器上的计算在时间上重叠,因此在 65,536 个 token 时有加速,而在 8,192 个 token 时没有。
图 7:在上下文长度为 8,192 和 65,536 个 token 时,针对 Qwen 3.8 和 Mistral Medium 3.5,使用流水线并行的两颗 M5 Ultra 相对于单颗 M5 Ultra 基线的解码加速比。流水线并行使解码速度基本保持不变。

张量并行(见图 5)将 LLM 的宽度分布到多台计算机上,而流水线并行则沿模型的深度进行分布。此时,每台机器都持有模型各层的一部分,从而使每层的计算并行化。如果实现高效,张量并行可以提高推理速度。不过,张量并行比流水线并行更难实现:每一层都必须并行化,这增加了代码复杂度。该方案还涉及大量小规模通信操作:对于每个输入 token 和每一层,都必须广播输入并汇总输出。因此,要让这种方法带来任何加速,MLX jaccl 支持的 Thunderbolt RDMA 这类超低延迟通信技术至关重要。

借助 MLX 的高级辅助工具 mlx.launch,可以轻松切换到张量并行。Mac Studio 计算机通过 Thunderbolt 网状网络连接,每对机器之间均通过两根 Thunderbolt 线缆进行点对点连接。假设主机已设置完毕,我们只需使用以下命令在多台机器上启动基准测试:

mlx.launch --hostfile hostfile.json --backend jaccl --env MLX_METAL_FAST_SYNCH=1 -- \
mlx_lm.benchmark --model meta-models/Muse-Spark-Open-816B-A42B-NVFP4 -p 8192 -g 128 -b 1

图 8 和图 9 展示了两台或四台搭载 M5 Ultra 的 Mac Studio 计算机相对于单机基线(对于 Muse Spark Open,则为流水线并行基线)获得的加速比。值得注意的是,稠密架构在多台计算机上的扩展性更好。

在机器数量较少时,上下文非常长的情况下,张量并行的通信开销可能会抵消部分计算收益,不过这仅适用于预填充阶段。Muse Spark Open 就出现了这种情况:对各专家的输出分别进行归一化增加了一些预填充通信,足以使双机张量并行在上下文非常长(65,536 tokens)时落后于双机流水线并行。不过,随着机器数量增加,这种方法就能带来收益,使用四台 M5 Ultra 的结果便体现了这一点。解码不受影响:即使上下文很长,Muse Spark Open 在张量并行下仍能保持显著的解码加速,与 Qwen 3.8 和 Mistral Medium 3.5 等稠密模型相当。

双线缆网状网络在预填充阶段也能带来收益,额外的链路带宽在这一阶段发挥了作用。对于运行在四台 M5 Ultra 上、上下文长度为 65,536-token 的 Muse Spark Open,每条链路使用两根线缆可将 TTFT 从约 50.5s(单根线缆)缩短至 40.6s——表 2 报告的是双线缆结果,单线缆运行则是一项额外实验。

图 8:在上下文长度为 8,192 和 65,536 tokens 时,Qwen 3.8、Mistral Medium 3.5 和 Muse Spark Open 使用两台或四台 M5 Ultra 进行张量并行,相对于单台 M5 Ultra 基线的首 token 时间(TTFT)加速比。Muse Spark Open 无法装入单台机器,因此其基线是两台 M5 Ultra 的流水线配置。
图 9:在上下文长度为 8,192 和 65,536 tokens 时,Qwen 3.8、Mistral Medium 3.5 和 Muse Spark Open 使用两台或四台 M5 Ultra 进行张量并行,相对于单台 M5 Ultra 基线的解码加速比。Muse Spark Open 无法装入单台机器,因此其基线是两台 M5 Ultra 的流水线配置。

首 token 时间(秒)

越低越好

模型M3 UltraM5 Ultra2 M5 流水线并行2 M5 张量并行4 M5 张量并行
Muse Spark Open——8.48.24.8
Mistral Medium 3.595.924.624.514.47.5
Qwen 3.818.24.14.13.11.7
表 2:Qwen 3.8、Mistral Medium 3.5 和 Muse Spark Open 在上下文长度为 8,192 和 65,536 个 token 时的首 token 时间(TTFT),单位为秒。越低越好。我们报告了单台 M3 Ultra、单台 M5 Ultra、两台 M5 Ultra(流水线并行和张量并行)以及四台 M5 Ultra(张量并行)的结果。

解码速度(token/秒)

越高越好

模型M3 UltraM5 Ultra2 台 M5 流水线并行2 台 M5 张量并行4 台 M5 张量并行
Muse Spark Open——15.7424.9936.73
Mistral Medium 3.54.957.707.7013.6923.37
Qwen 3.812.4619.6919.4433.2750.89
表 3:Qwen 3.8、Mistral Medium 3.5 和 Muse Spark Open 在上下文长度为 8,192 和 65,536 个 token 时的解码速度,单位为 token/秒。越高越好。我们报告了单台 M3 Ultra、单台 M5 Ultra、两台 M5 Ultra(流水线并行和张量并行)以及四台 M5 Ultra(张量并行)的结果。解码速度通过测量上下文之后生成的 128 个 token 得出。

最近,Mac 上的智能体 AI 获得了广泛关注,智能体编程工具也已在开发者中得到广泛采用。借助 MLX,在本地为这些工具提供支持,只需使用 mlx_lm.server 工具启动一个服务器。例如,使用以下命令即可启动一个在端口 8080 上运行 Muse Spark Open 的服务器:

mlx.launch --hostfile hostfile.json --backend jaccl --env MLX_METAL_FAST_SYNCH=1 -- \
mlx_lm.server --model meta-models/Muse-Spark-Open-816B-A42B-NVFP4 --port 8080

与 mlx_lm.benchmark 一样,可以在 mlx_lm.server 辅助工具前加上 mlx.launch,以在多台 Mac Studio 上并行计算。在此基础上,即可通过 OpenCode 等工具启用智能体 AI,只需简单配置,就能让这些工具使用该本地 LLM 服务器:

{
  "$schema": "https://opencode.ai/config.json",
  "model": "meta-models/Muse-Spark-Open-816B-A42B-NVFP4",
  "provider": {
    "mlx": {
      "npm": "@ai-sdk/openai-compatible",
      "name": "Local MLX",
      "options": {
        "baseURL": "http://127.0.0.1:8080/v1"
      },
      "models": {
        "meta-models/Muse-Spark-Open-816B-A42B-NVFP4": {
          "name": "Muse Spark Open 816B (MLX)"
        }
      }
    }
  }
}

通过类似方法,还可以实现许多其他本地智能体 AI 体验。

过去一年,我们在 MLX 中实现了多项关键创新,让用户能够更轻松、更高效地在近期推出的 Apple 芯片上本地运行最先进的 LLM。这些创新包括:为 M5 系列芯片 GPU 中的神经网络加速器提供软件支持,通过雷雳上的 RDMA 实现多台 Mac 之间的低延迟通信,以及高效实现混合专家(MoE)等近期热门架构。每一项创新本身都能带来帮助,但正是它们与最新 Apple 芯片及相关技术的结合——配备 M5 Ultra 的 Mac Studio、M5 Ultra GPU 中的神经网络加速器、512GB 统一内存,以及雷雳 5 上的 RDMA——使我们如今能够在本地运行 Muse Spark Open 这类参数量超过 800B 的模型,启动速度(首 token 时间)最高可达上一代 Mac Studio 的 4.6×。

相关阅读与最新动态。

Apple 将在一年一度的国际学习表征会议(ICLR)上展示新的研究成果。本届会议将于 April 23 至 27 在巴西里约热内卢线下举行。我们很荣幸再次赞助这一汇聚深度学习领域学术界与产业界研究群体的会议。以下是 Apple 参与 ICLR 2026 的概览。

阅读更多

搭载 Apple 芯片的 Mac 越来越受到 AI 开发者和研究人员的欢迎,他们希望使用自己的 Mac 试验最新的模型和技术。借助 MLX,用户可以在 Mac 上高效探索和运行 LLM。它让研究人员能够在自己的硬件上试验新的推理或微调技术,或在私密环境中研究 AI 技术。MLX 适用于所有搭载 Apple 芯片的系统,而借助最新的……

阅读更多

来源:Apple Machine Learning · machinelearning.apple.com