ITPub博客

首页 > 人工智能 > 人工智能 > 在 Mac Studio 上使用 MLX 和 RDMA 进行本地与分布式大语言模型推理

在 Mac Studio 上使用 MLX 和 RDMA 进行本地与分布式大语言模型推理

人工智能 作者:ITPUB社群 时间:1970-01-01 08:33:46 0 删除 编辑

搭载 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 上运行的模型范围。

Muse Spark Open 解码速度

上下文大小为 8,192 个 token

使用 MLX 运行大语言模型

我们使用 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 Open | 816B | 42B | BF16 | NVFP4 |
| Mistral Medium 3.5 | 128B | 128B | FP8 | FP8 |
| Qwen 3.8 | 27B | 27B | BF16 | BF16 |

我们的评估在一系列上下文大小上测量三个指标:首 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 中。

单台 Mac Studio 上的推理性能

Qwen 3.8 27B 和 Mistral Medium 3.5 128B 可轻松运行在配备 512GB 内存的单台 Mac Studio 上,在大上下文尺寸(65,536 tokens)配置下分别占用最多 68GB 和 158GB 内存。在图 2 中,我们比较了二者在配备 512GB 内存的单台 Mac Studio 上的推理性能,分别基于 M3 Ultra 和 M5 Ultra。

首 token 延迟(TTFT)本质上受计算能力限制,而这正是 M5 Ultra GPU 中神经加速器(Neural Accelerators)的优势所在:GPU 神经加速器对快速矩阵乘法的专用硬件支持,直接转化为 TTFT 的显著加速,相比 M3 Ultra 最高可提升 4.6 倍(见图 2)。这使得配备 M5 Ultra 的 Mac Studio 成为处理超大规模上下文任务的首选桌面设备。

生成速度(以 tokens/秒计)受内存带宽限制。M5 Ultra 芯片提供 1.2 TB/s 的统一内存带宽,高于 M3 Ultra 的 819 GB/s。这同样直接转化为生成速度的提升(见图 3),配备 M5 Ultra 的 Mac Studio 报告的 tokens/s 提升 1.6 倍。

TTFT 加速比

单台 Mac Studio,M5 Ultra 相对于 M3 Ultra 基线(1.0×)

解码加速比

单台 Mac Studio,M5 Ultra 相对于 M3 Ultra 基线(1.0×)

Mac Studio 在大语言模型(LLM)推理方面也十分节能(见图 4):配备 M3 Ultra 的 Mac Studio 峰值功耗约为 200W,配备 M5 Ultra 的 Mac Studio 约为 300W——仅为普通烤面包机功耗的一小部分。更重要的是,如此适中的功耗意味着 Mac Studio 可以部署在任何家庭或办公环境中用于 AI 工作。

峰值功耗

上下文尺寸为 65,536 tokens

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

跨两台 Mac Studio 计算机的流水线推理

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

将大型模型拆分到多台计算机

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 k 块的后续层时,第一台机器计算第 k + 1 k+1 k + 1 块的前序层,因此连续的块在时间上重叠。为了实现更好的计算并行化,我们接下来转向张量并行化。

TTFT 加速比

采用流水线并行化的两台 Mac Studio,相对于单台 M5 Ultra 基线(1.0×)

解码加速比

采用流水线并行化的两台 Mac Studio,相对于单台 M5 Ultra 基线(1.0×)

用于更快推理的张量并行化

张量并行化(见图 5)将 LLM 的宽度分布到多台计算机上,而流水线并行化则将模型按其深度进行分布。现在每台机器持有模型每一层的一部分,从而使逐层计算得以并行化。在高效实现的情况下,张量并行化可提升推理速度。然而,张量并行化的实现比流水线并行化更复杂:每一层都必须并行化,这增加了代码复杂度。该方案还涉及许多小型通信操作:对于每个输入 token 和每一层,都必须广播输入并聚合输出。因此,像 Thunderbolt RDMA 这样由 MLX jaccl 支持的极低延迟通信技术,对于该方法能否带来任何加速至关重要。

借助 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 中报告的双线缆数值,单线缆运行是额外的实验。

TTFT 加速比

使用张量并行化的两台或四台 Mac Studio,相对于单台 M5 Ultra 基线(1.0×;Muse Spark Open:双机流水线)

Decode 加速比

使用张量并行化的两台或四台 Mac Studio,相对于单台 M5 Ultra 基线(1.0×;Muse Spark Open:双机流水线)

首 token 时间(秒)

越低越好

| 模型 | M3 Ultra | M5 Ultra | 2 M5 流水线 | 2 M5 张量 | 4 M5 张量 |
| --- | --- | --- | --- | --- | --- |
| Muse Spark Open | — | — | 8.4 | 8.2 | 4.8 |
| Mistral Medium 3.5 | 95.9 | 24.6 | 24.5 | 14.4 | 7.5 |
| Qwen 3.8 | 18.2 | 4.1 | 4.1 | 3.1 | 1.7 |

解码速度(每秒 token 数)

越高越好

| 模型 | M3 Ultra | M5 Ultra | 2 M5 Pipeline | 2 M5 Tensor | 4 M5 Tensor |
| --- | --- | --- | --- | --- | --- |
| Muse Spark Open | — | — | 15.74 | 24.99 | 36.73 |
| Mistral Medium 3.5 | 4.95 | 7.70 | 7.70 | 13.69 | 23.37 |
| Qwen 3.8 | 12.46 | 19.69 | 19.44 | 33.27 | 50.89 |

Mac 上的本地智能体 AI

近期,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——使我们如今能够在本地运行参数量超过 800B 的 Muse Spark Open 等模型,启动(首 token 时间)速度最高可达上一代 Mac Studio 的 4.6 倍。


转载声明:本文为转载整理,原文版权归原作者及原站点所有。

原文来源:Apple ML Research

原创地址:https://machinelearning.apple.com/research/mlx-distributed-inference

原文作者:Apple ML Research;如有侵权,请联系管理员删除。

请登录后发表评论 登录
全部评论

注册时间:2022-11-21

  • 博文量
    1199
  • 访问量
    549965