跳到正文
Apple Machine Learning Research·· 18 小时前精选AI 评分60

Apple 演示用 MLX 与 RDMA 在 Mac Studio 集群本地运行 816B 大模型

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

AI 导读

Apple 机器学习研究团队演示如何在 Mac Studio 上本地运行近万亿参数的 Muse Spark Open 等前沿模型,使用今天发布的新版 MLX(v0.30.1 起)与 macOS 26.2 的 RDMA over Thunderbolt 驱动。

推荐理由

Apple 给出了在 Mac Studio 集群上跑 816B MoE 模型的并行方案与基准数据,可迁移到本地大模型部署实践。

正文 · 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 统一内存)开始,再通过 RDMA over Thunderbolt 驱动扩展到多台 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 上对几个流行的开放权重 LLM 进行了基准测试,方法如我们之前的文章所述。我们比较了两代 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 上本地 LLM 的一个流行用途是在私密环境中运行智能体体验,包括通常需要大上下文的智能体编程应用。因此,我们报告了 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 中,我们比较了它们在一台配备 512GB 内存的 Mac Studio 上的推理性能,分别使用 M3 Ultra 和 M5 Ultra。

首 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 所使用的 RDMA-over-Thunderbolt 链路。主机设置完成后,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 将大上下文切分为更小的块(此处为 8,192 个 token,由 --prefill-step-size 8192 设置),并在各机器之间流水线化预填充:当第二台机器计算块 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 时获得加速,而在 8,192 时没有。
图 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,则是流水线并行化基线)所获得的加速。值得注意的是,稠密架构在多台计算机上的扩展性更好。

在机器数量较少时,张量并行化的通信开销在非常大的上下文下可能会抵消其部分计算优势,不过这只适用于 prefill。这在 Muse Spark Open 上有所体现:对专家输出进行单独归一化会增加一些 prefill 通信,足以导致在非常大的上下文(65,536 个 token)下,双机张量并行化落后于双机流水线并行化。然而,随着机器数量的增加,这种方法会带来回报,正如四台 M5 Ultra 所显示的那样。Decode 不受影响:即使在大型上下文下,Muse Spark Open 在张量并行化下也能保持显著的 decode 加速,与 Qwen 3.8 和 Mistral Medium 3.5 等稠密模型相当。

双线缆网状网络在 prefill 期间也能带来回报,此时额外的链路带宽得以体现。对于在四台 M5 Ultra 上以 65,536 token 上下文运行的 Muse Spark Open,每条链路两根线缆将 TTFT 从约 50.5 秒(一根线缆)缩短至 40.6 秒——即 表 2 中报告的双线缆数值,单线缆运行是额外的实验。

图 8:在上下文大小为 8,192 和 65,536 个 token 时,对于 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 个 token 时,对于 Qwen 3.8、Mistral Medium 3.5 和 Muse Spark Open,使用张量并行化的两台或四台 M5 Ultra 相对于单台 M5 Ultra 基线的 decode 加速。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 中神经加速器的软件支持、使用基于 Thunderbolt 的 RDMA 在多台 Mac 之间进行低延迟通信,以及对混合专家(MoE)等近期流行架构的高效实现。每一项创新本身都有帮助,但正是它们与最新 Apple 芯片的结合——搭载 M5 Ultra 的 Mac Studio、M5 Ultra GPU 中的神经加速器、512GB 统一内存以及基于 Thunderbolt 5 的 RDMA——让我们现在能够在本地运行像 Muse Spark Open 这样参数超过 800B 的模型,启动(首 token 时间)速度最高可达上一代 Mac Studio 的 4.6 倍。

相关阅读和更新。

Apple 将在年度 国际学习表征会议(ICLR)上展示新的研究成果,该会议将于 4 月 23 日至 27 日在巴西里约热内卢线下举行。我们很自豪再次赞助该会议,它汇聚了专注于深度学习的科学和工业研究社区。以下是 Apple 参与 ICLR 2026 的概览。

阅读更多

搭载 Apple 芯片的 Mac 在有兴趣使用 Mac 试验最新模型和技术的 AI 开发者和研究人员中越来越受欢迎。借助 MLX,用户可以在 Mac 上高效地探索和运行 LLM。它允许研究人员试验新的推理或微调技术,或在自己的硬件上、在私密环境中研究 AI 技术。MLX 适用于所有 Apple 芯片系统,并且与最新的…

阅读更多

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