跳到正文
Apple Machine Learning Research·· 11 天前精选AI 评分63

Apple 使用 MLX 和 RDMA 在 Mac Studio 上实现本地与分布式大模型推理

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

AI 导读

Apple 机器学习研究团队介绍了如何利用最新版 MLX 框架和 Thunderbolt 5 上的 RDMA 驱动程序,在 Mac Studio(搭载 M5 Ultra 芯片)上运行近万亿参数的 Muse Spark Open 等前沿大模型。

推荐理由

文章详细展示了利用最新 MLX 和 RDMA 技术在 Mac Studio 集群上运行近万亿参数大模型的方法与性能数据,为端侧大模型部署提供了具体的工程参考。

正文 · AI 翻译

content type 精选亮点

The MLX logo

MLX on Mac with Apple silicon lets AI researchers and developers experiment locally with the latest models and methods in a private environment. The open-source array framework leverages the unified memory architecture of Apple silicon, enabling anyone to run new inference and fine-tuning methods on their own hardware and explore other AI techniques. In a previous post, we highlighted the advantages of the M5 architecture and its GPU Neural Accelerators.

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

从单台 Mac Studio(最高 512GB 统一内存)开始,然后通过 RDMA over Thunderbolt 驱动程序扩展到多台 Mac,我们在下文中分享了你现在可以在 Mac Studio 上运行的各种模型的概览。

Figure 1: Muse Spark Open decode speed in tokens per second, on Mac Studio with M5 Ultra, at a context size of 8,192 tokens. Higher is better. We report performance for three parallelization schemes: pipeline parallelization over two Mac Studio computers (2 M5 Pipeline), and tensor parallelization over two (2 M5 Tensor) and four Mac Studio (4 M5 Tensor) computers.

We benchmark several popular open-weight LLMs on Mac Studio, using the MLX LM package, as described in our prior post. We compare two generations of Mac Studio: the recently released M5 Ultra and the previous generation, M3 Ultra. The M5 generation features the GPU Neural Accelerator, which supports dedicated matrix multiplication operations, essential for many machine learning workloads.

We evaluate three frontier-class models (Table 1) spanning a wide range of sizes (in billions of parameters, B): the newly released Muse Spark Open (816 B), Mistral Medium 3.5 (128 B), and Qwen 3.8 (27 B). Muse Spark Open is a mixture of experts (MoE) model, which keeps only a fraction of its experts — and thus parameters — active, reducing compute at inference, and it is evaluated in NVFP4 precision. Mistral Medium 3.5 and Qwen 3.8 are dense architectures, where all parameters are active at inference, and they are evaluated in their native floating-point precision (FP8 and BF16, respectively).

ModelTotal ParametersActive ParametersNative PrecisionEvaluation Precision
Muse Spark Open816B42BBF16NVFP4
Mistral Medium 3.5128B128BFP8FP8
Qwen 3.827B27BBF16BF16
Table 1: Language models used in this post. Models ship in various precision formats; we evaluate Qwen 3.8 and Mistral Medium 3.5 in their native precision, and Muse Spark Open in 4-bit (NVFP4) precision.

Our evaluation measures three metrics across a range of context sizes: time to first token (TTFT, in seconds), generation throughput (in tokens/second), and peak power usage (in watts). All the performance numbers in this post were obtained with mlx_lm.benchmark. For example, running Qwen 3.8 with 8,192 tokens of context and generating 128 tokens is performed with the command line:

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

A popular use of local LLMs on Mac Studio is running agentic experiences in a private environment, including agentic coding applications that often need large contexts. We therefore report performance at both 8,192 and 65,536 context sizes.

原始首字延迟(TTFT)和解码吞吐量分别见 Table 2 和 Table 3。

Qwen 3.8 27B 和 Mistral Medium 3.5 128B 可以轻松容纳在配备 512GB 内存的单个 Mac Studio 上,在大上下文(65,536 个 token)设置下分别占用高达 68GB 和 158GB 的 RAM。在 Figure 2 中,我们在配备 512GB RAM 的单台 Mac Studio 上(分别在 M3 Ultra 和 M5 Ultra 上)对比了它们的推理性能。

首字延迟(TTFT)本质上受限于计算能力,这也是 M5 Ultra GPU 中的神经加速器大显身手的地方:GPU 神经加速器对快速矩阵乘法的专用硬件支持,直接转化为 TTFT 的显著提速,与 M3 Ultra 相比最高可快 4.6 倍(见 Figure 2)。这使得搭载 M5 Ultra 的 Mac Studio 成为处理超大上下文任务时的首选桌面设备。

生成速度(以每秒 token 数计算)受限于内存带宽。M5 Ultra 芯片提供 1.2 TB/s 的统一内存带宽,高于 M3 Ultra 的 819 GB/s。同样,这直接转化为生成速度(见 Figure 3),搭载 M5 Ultra 的 Mac Studio 的 token/s 速度提升了 1.6 倍。

Figure 2: 在 8,192 和 65,536 token 的上下文长度下,Qwen 3.8 和 Mistral Medium 3.5 在 M5 Ultra 上相比 M3 Ultra 基准的首字延迟(TTFT)加速比。M5 Ultra GPU 神经加速器在计算密集型工作负载中表现亮眼。
Figure 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)推理也具有极高的能效(见 Figure 4):搭载 M3 Ultra 的 Mac Studio 峰值功耗约为 200W,搭载 M5 Ultra 的 Mac Studio 约为 300W——仅为普通烤面包机耗电量的一小部分。更重要的是,这种适度的功耗意味着 Mac Studio 可以部署在任何家庭或办公室环境中的 AI 工作场景下。

Figure 4: 在搭载 M3 Ultra 和 M5 Ultra 的 Mac Studio 上,Qwen 3.8 和 Mistral Medium 3.5 在上下文大小为 65,536 token 时的推理峰值功耗(以瓦特为单位)。

虽然 Qwen 和 Mistral 模型都可以放入配备 512GB RAM 的单个 Mac Studio 中,但 Muse Spark Open 仅模型参数在其原生精度(BF16)下就需要超过 1.6TB 的 RAM。在 4 位量化(NVFP4)格式下,模型内存占用降至约 493GB。再加上推理过程中中间计算所需的内存,该模型无法容纳在单台机器上。它需要跨多台 Mac Studio 计算机进行并行化。

Muse Spark Open 等超大型前沿级模型无法放入单台机器中。不过,可以将一半的模型存储在一台 Mac Studio 上,另一半存储在另一台。然后,我们可以在两台计算机之间流水线化模型的推理:第一台 Mac Studio 执行模型的前半部分层,然后将输出发送到第二台 Mac Studio,由其执行后半部分。这种并行化方案称为流水线并行(见 Figure 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:对于 Qwen 3.8 和 Mistral Medium 3.5,在上下文大小为 8,192 和 65,536 个 token 时,两台 M5 Ultra 采用流水线并行相比单台 M5 Ultra 基准的首字延迟(TTFT)加速比。拆分上下文(此处拆分为 8,192 token 的块)使 MLX 能够在两台机器之间重叠计算——因此在 65,536 时有加速,而在 8,192 时则没有。
图 7:对于 Qwen 3.8 和 Mistral Medium 3.5,在上下文大小为 8,192 和 65,536 个 token 时,两台 M5 Ultra 采用流水线并行相比单台 M5 Ultra 基准的解码加速比。流水线并行使解码速度基本保持不变。

张量并行(见 Figure 5)将大语言模型的宽度分布在多台计算机上,这与沿模型深度分布的流水线并行形成对比。现在每台机器持有模型各层的一部分,从而使逐层计算得以并行。当高效实现时,张量并行可以提高推理速度。然而,张量并行的实现比流水线并行更复杂:必须对每一层进行并行化,这增加了代码的复杂性。该方案还涉及许多小型通信操作:对于每个输入 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

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

在机器数量较少时,张量并行的通信开销可能会在非常大的上下文下抵消其部分计算优势,尽管这仅适用于预填充(prefill)。这在 Muse Spark Open 中表现得很明显,专家输出的单独归一化增加了一些预填充通信,以至于在非常大的上下文(65,536 个 Token)下,两机张量并行的表现落后于两机流水线并行。然而,随着机器数量的增加,这种方法会带来回报,正如四台 M5 Ultra 所示。解码(Decode)不受影响:即使在大型上下文中,Muse Spark Open 在张量并行下也能保持显著的解码加速,这与 Qwen 3.8 和 Mistral Medium 3.5 等稠密模型相当。

双电缆网状网络在预填充阶段也大有裨益,额外的链路带宽在此发挥了作用。对于在四台 M5 Ultra 上运行、上下文为 65,536 个 Token 的 Muse Spark Open 而言,每条链路两根电缆将 TTFT 从约 50.5 秒(单根电缆)缩短至 40.6 秒——这是 Table 2 中报告的双电缆数值,单电缆运行是一项额外的实验。

Figure 8: 在上下文大小为 8,192 和 65,536 个 Token 时,采用张量并行的两台或四台 M5 Ultra 相对于单台 M5 Ultra 基准的首字延迟(TTFT)加速比,针对 Qwen 3.8、Mistral Medium 3.5 和 Muse Spark Open。Muse Spark Open 无法在单机上运行,因此其基准是两台 M5 Ultra 流水线配置。
Figure 9: 在上下文大小为 8,192 和 65,536 个 Token 时,采用张量并行的两台或四台 M5 Ultra 相对于单台 M5 Ultra 基准的解码加速比,针对 Qwen 3.8、Mistral Medium 3.5 和 Muse Spark Open。Muse Spark Open 无法在单机上运行,因此其基准是两台 M5 Ultra 流水线配置。

首字延迟(秒)

越低越好

模型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 时的首字延迟(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 中神经加速器(Neural Accelerators)的软件支持、利用 Thunderbolt 上的 RDMA 实现多台 Mac 之间低延迟通信,以及对混合专家模型(MoE)等近期热门架构的高效实现。这些改进每一项都大有裨益,但正是它们与最新 Apple 芯片的结合——搭载 M5 Ultra 的 Mac Studio、M5 Ultra GPU 中的神经加速器、512GB 统一内存以及 Thunderbolt 5 上的 RDMA——让我们现在能够在本地运行参数量超过 800B 的 Muse Spark Open 等模型,其启动速度(首字延迟)最高可达上一代 Mac Studio 的 4.6 倍。

相关阅读与更新。

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

阅读全文

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

阅读全文

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