Hugging Face来源发布时间

PyTorch 2.15 中的 MPS

如果你在 Mac 上运行模型,或使用 MPS 作为主要的计算加速工具,那么 PyTorch 2.15 有很多值得期待的内容!FP8 登陆 MPS,语言模型解码在 linear 层和 softmax 中获得了更快的路径,新的 Metal 内核则解决了音频、视频和训练工作负载中一些开销出乎意料地高的部分。

如果你在 Mac 上运行模型,或使用 MPS 作为主要的计算加速工具,那么 PyTorch 2.15 有很多值得期待的内容!FP8 登陆 MPS,语言模型解码在 linear 层和 softmax 中获得了更快的路径,新的 Metal 内核则解决了音频、视频和训练工作负载中一些开销出乎意料地高的部分。

简要介绍一下我们是谁:我们是 Hugging Face 的 PyTorch 团队。我们全面参与 PyTorch 的贡献,其中一个重点领域是 MPS 后端,目标是让您的 Mac 成为一等公民的场所,用于运行和训练 PyTorch 模型。在2026年初,在 MPS 上运行哪怕很小的模型,也常常意味着要绕过不支持的操作、性能断崖和意外的内存增长。因此,我们将来自 Hugging Face Hub 的语言、图像和视频领域的真实模型实际运行起来,找出 MPS 的不足之处,并在 PyTorch 上游进行修复——越来越多的情况下,是用专用的 Metal 内核替换 MPSGraph。本文的大部分内容就是这项工作的成果,它是与 PyTorch 维护者和社区贡献者共同完成的。

简而言之

  • FP8 来到 Mac: 商店 float8_e4m3fn 以 BF16 一半的空间表示数值,并使用缩放矩阵乘法以及 CUDA/CPU 上支持的所有其他操作!
  • MoE 模型的解码速度比 2.14 快 5–6.5 倍: Qwen3.6-35B-A3B 在 M3 Ultra 上的速度从 9.4 提升到 61.3 tokens/s,这得益于原生的分组矩阵乘法(grouped matmul)配有一条专用的解码路径、更少的 GPU-CPU 同步以及更快的 top-k 路由。
  • 每个 LLM 每个 token 所需的时间更短: 单行线性层使用专门的矩阵-向量内核,而词汇表上的softmax运算速度约快8倍。
  • 在 GPU 上进行更多线性代数运算: 复杂的LU分解和三角求解、更快速的小矩阵求逆,以及更快速的QR分解。
  • 音频、视频和培训也受到了关注: 针对卷积、填充、dropout、插值反向传播和归约的更快路径。Wan2.2 的视频 VAE 解码现在峰值内存为 8 GiB,而不是 270 GiB。
  • 从 Python 获得更多控制: 公开的MPS流以及更可靠的事件计时。
  • 更新后的要求: MPS 现在要求 macOS 15 或更高版本;PyTorch 2.15 需要 Python 3.11 或更高版本.
  • 2.16版抢先看: 全力聚焦于 Metal,更多从 MPSGraph 到 Metal 的迁移以及更快的 LLM 解码。

这对您意味着什么

如果你在 Mac 上使用 transformers、diffusers 或 ComfyUI 等工具运行模型,以下是 2.15 在实际使用中的变化:

  • 本地LLM速度大幅提升。 专家混合模型的解码速度比 2.14 快 5–6.5 倍:在配备 M3 Ultra 的 torch.compile, Qwen3.6-35B-A3B 从 9.4 tokens/s 提升到 61.3 tokens/s,LFM2.5-8B-A1B 从 23.6 tokens/s 提升到 123.6 tokens/s。稠密模型也从中受益,得益于更快的线性层和采样时快 8 倍的 softmax。(更多)
  • 能装进显存的视频生成。 使用 Wan2.2 的分块 VAE 解码视频过去会把驱动内存推高到 270 GiB;现在峰值约为 8 GiB。(更多)
  • 使用 Qwen-Image-2.1 的正确图像编辑。 在 2.14 版本中,Qwen-Image-2.1 的 VAE 编码器会在 MPS 上悄无声息地损坏图像,破坏图像编辑工作流(已在 ComfyUI 中报告)。 same padding change 这一修复解决了该问题:现在一次编码-解码往返在 29.95 dB PSNR 下与 CPU 结果一致,而此前仅为 10.06 dB。
  • 在你的 Mac 上使用 FP8 权重。 FP8 张量占用的内存是 BF16 的一半,并且 F.scaled_mm 可以直接在 GPU 上对它们进行乘法运算。(更多)

目录

入门指南

你需要一台搭载Apple芯片的Mac,并且装有 macOS 15 或更高版本 和 Python 3.11 或更高版本 (参见 2.15 兼容性矩阵)。如果你的操作系统过旧或没有可用的 Metal GPU,MPS 可用性检查现在会告诉你原因。

要在单独的环境中试用该候选版本:

uv pip install torch==2.15.0 --index-url https://download.pytorch.org/whl/test/cpu

FP8 登陆 MPS

FP8 为你提供了另一种在 Mac 上节省内存的方式。PyTorch 2.15 添加了 torch.float8_e4m3fn MPS 上的张量:每个值一个字节,存储量是 FP16 或 BF16 的一半。

你可以将这些张量与 F.scaled_mm相乘,它同时支持 tensorwise 和 rowwise 缩放。下面是一个每个张量一个缩放因子的小示例:

import torch
import torch.nn.functional as F

def quantize_fp8(x):
    limit = torch.finfo(torch.float8_e4m3fn).max
    scale = x.abs().amax().clamp_min(1e-12) / limit
    return (x / scale).to(torch.float8_e4m3fn), scale

x = torch.randn(1, 1024, device="mps", dtype=torch.float32)
weight = torch.randn(1024, 1024, device="mps", dtype=torch.float32)
x_fp8, x_scale = quantize_fp8(x)
w_fp8, w_scale = quantize_fp8(weight)

y = F.scaled_mm(
    x_fp8,
    w_fp8.t(),
    scale_a=x_scale,
    scale_recipe_a=F.ScalingType.TensorWise,
    scale_b=w_scale,
    scale_recipe_b=F.ScalingType.TensorWise,
    output_dtype=torch.bfloat16,
)
print(y.shape, y.dtype)

更快的语言模型解码

每个 LLM 都有一个预填充阶段和一个解码阶段。在解码期间,每个新 token 以单行的形式通过模型,因此大矩阵乘法变成了矩阵-向量乘积。PyTorch 2.15 加速了主导这一步骤的两类层:稠密线性层和混合专家层。

线性层

MPS 自 2.14 起就已为 torch.mm 自 2.14 起,但 F.linear 仍然要经过 MPSGraph。PyTorch 2.15 将它们直接连接到 F.linear,因此常见的 [1, 1, hidden_size] 解码输入会直接使用它们,并将符合条件的偏置加法融合进来。在来自真实模型的 96 个 BF16 层形状上,这带来了 1.23× 的几何平均加速。最大的收益落在小型投影上,例如 MoE 模型中的门控,旧路径在这些地方浪费了大部分内存带宽。大型层和 LM 头本已接近带宽瓶颈,因此它们保持在接近 1× 的水平。

混合专家层

混合专家层会向每个专家发送不同数量的 token。分组 matmul 可以在单个操作中处理这些不均匀的分组,2.15 为 MPS 提供了 原生实现.

在 2.14 中,分组 matmul 会将分组偏移量复制到 CPU 并循环遍历各分组。新的 Metal 内核直接在 GPU 上读取偏移量,在可用时使用 Metal Performance Primitives,在其他地方则使用单独的 Metal 路径。小型分组(解码期间的常见情况)获得了 自己的内核 也一样。

例如,64 个 token 可以拆分为 16、32 和 16 的组:

import torch
import torch.nn.functional as F

tokens = torch.randn(64, 64, device="mps", dtype=torch.bfloat16)
weights = torch.randn(3, 128, 64, device="mps", dtype=torch.bfloat16)
ends = torch.tensor([16, 48, 64], device="mps", dtype=torch.int32)

out = F.grouped_mm(tokens, weights.transpose(-2, -1), offs=ends)
print(out.shape)  # torch.Size([64, 128])

解码内核直接体现在 token 吞吐量上。在 M3 Ultra 和 torch.compile上,MoE 模型每秒生成的 token 数约为 1.3 倍:

线性代数获得提升

复数计算

处理复数张量变得更加容易。 complex64 LU 分解, LU 求解 和 三角求解 现在可以在 MPS 上运行,这也解锁了构建于其上的运算和梯度,例如复数行列式。以下是一个直接在 GPU 上求解的小型系统:

import torch

a = torch.randn(8, 8, device="mps", dtype=torch.complex64)
a = a + 8 * torch.eye(8, device="mps", dtype=torch.complex64)
b = torch.randn(8, 2, device="mps", dtype=torch.complex64)

lu, pivots = torch.linalg.lu_factor(a)
x = torch.linalg.lu_solve(lu, pivots, b)
print((a @ x - b).abs().max().item())

小型矩阵与更大规模的求解

小型矩阵也拥有快速内核。对于相机几何和图像变换,最多可达 32 个的 FP32 矩阵批次 4 × 4 的 FP32 矩阵批次可以 在单个 Metal 内核中完成求逆 ,该内核将分解和求解结合在一起。更大规模的问题也得到了关注: QR 使用分块更新,而 三角求解 则将其大部分工作转化为高效的矩阵乘法。

Gated DeltaNet 工作负载获得了一条 专用快速路径 适用于以下规模的批次: 16 × 16、32 × 32 和 64 × 64 的 FP32 三角系统 且至少有八个右侧项。

SVD 在数值极端情况下也更加稳健。它现在能正确处理 非常小的奇异值, 非常大或非常小的输入幅值 以及 基补全与 full_matrices=True,这反过来使伪逆和最小二乘结果更加可靠。

卷积、填充与训练

音频模型、视频解码器和训练循环都有各自开销高昂的操作。本次发布为其中多项操作提供了经由 Metal 的新路径:

操作 2.15 中的变化
conv1d forward 专用内核取代了 MPSGraph,并针对短输出增加了 im2col 加 matmul 的路径。适用于音频和流式模型的形状。
常数填充 迁移到 Metal 内核,修复了分块视频 VAE 解码中驱动程序内存失控增长的问题(下文详述)。
反射填充和复制填充 全新的 Metal 实现,速度更快并修复了反向传播的正确性问题。
非就地 dropout nn.Dropout 和 F.dropout 现在使用现有的融合内核。
上采样反向传播 Metal gather 内核取代了 MPSGraph 的缩放梯度路径。
Channels-last conv2d 反向传播 保留在高效的 channels-last 路径上,而不是使用代价高昂的跨步访问的 NCHW 图。

Softmax 与归约

在 2.15 中,token 采样也变得更快。 softmax 和 log_softmax 现在运行在专用的 Metal 内核上,大型词表受益最大:对一个 FP32 softmax 而言, [1, 128256] 一行 logits 从 57 µs 降至 7 µs,大约 快 8 倍 于 MPSGraph 版本。这一加速在其他常见词表大小下同样成立:

视频解码的内存占用大幅缩小

本次发布中最引人注目的结果之一来自常数填充(constant-padding)的改动。 Wan-AI/Wan2.2-TI2V-5B-Diffusers 中的 VAE 解码在完整的 28 瓦片解码过程中会发起 8,764 次填充调用 超过一次完整的 28 瓦片解码。在 Metal 上启用填充后,驱动程序的峰值内存从大约 270 GiB 降至 8 GiB:

内存指标 MPSGraph 填充 Metal 填充
当前张量分配峰值 4.422 GiB 4.095 GiB
分配器保留峰值 8.070 GiB 8.070 GiB
驱动程序总内存峰值 270.135 GiB 8.108 GiB

通过 MPS 流获得更多控制权

你现在可以直接从 Python 中选择由哪个 MPS 流来运行你的工作。 torch.mps.Stream 以及 torch.mps.stream 上下文管理器 提供了接口,而 Tensor.record_stream 让分配器能够跟踪一个张量在另一个流上的使用。

下面是一个小示例,它将工作发送到一个单独的流,并在读取结果之前进行同步:

import torch

x = torch.randn(1024, device="mps")
torch.mps.synchronize()
worker = torch.mps.Stream()

with torch.mps.stream(worker):
    y = x.square()
    x.record_stream(worker)

del x
worker.synchronize()
print(y.cpu().shape)

更多操作支持 MPS

现在有更多操作可以在 MPS 上运行: torch.frexp, torch.logspace, torch.take 以及 Tensor.put_, torch.heaviside, torch.special.laguerre_polynomial_l 以及 无符号整数比较. logspace 尚不支持包含多个元素的整数或布尔输出。

注意力机制也得到了一个正确性修复:scaled dot-product attention 的 MPS math 实现现在 能够正确处理 batch 和 head 的广播.

PyTorch 2.16 抢先看

2.16 的工作已经在顺利进行中。下面的一些更改已经合入 main 并可在 nightly 版本中使用;其他更改仍在审查中,因此在发布前细节可能会有所变化。

更快的 LLM 解码

一旦模型的解码路径不再包含任何 MPSGraph 操作,GPU 只有在 CPU 完成对一个 token 的所有内核编码之后才会开始处理该 token。PyTorch 2.16 每 16 个内核提交一次命令缓冲区,因此 GPU 可以在 CPU 仍在编码其余部分时开始工作。对于 2048-token 上下文下的编译 BF16 解码:

模型 之前 之后 变化
LFM2.5-8B-A1B 108.8 tokens/s 126.0 tokens/s +16%
Qwen3.6-35B-A3B 52.9 tokens/s 62.8 tokens/s +19%
Qwen3.8-27B 12.1 tokens/s 12.6 tokens/s +4%

注意力机制完全在 Metal 上运行。

缩放点积注意力(Scaled dot-product attention)是 完全放弃 MPSGraph: 每种形状都将运行在 Metal 内核上,在较新的 macOS 版本上配有专用于解码的向量内核以及用于半精度预填充的 Metal Performance Primitives。以前回退到 MPSGraph 的形状受益最大。在 M5 Pro 上,使用分组查询注意力(32 个查询头、8 个 KV 头、头维度 80、4K 上下文)的解码从 465 µs 降至 65 µs,头维度 40 的 4K × 4K 预填充则从 42 ms 降至 6 ms。

更多 MPSGraph 到 Metal 的迁移

告别MPSGraph的迁移仍在继续。批量归一化 向前 和 向后 正在进行中,而 MSE 损失 (在小输入上最高快 12 倍), std_mean 和 var_mean (最高快 4 倍)以及 bincount 已经落地。 bincount 是最引人注目的:当大多数数值落入单个区间时,对 4 百万个整数进行计数过去需要 6.8 秒,现在只需半毫秒。

更小的安装体积

捆绑的着色器将 默认针对 Metal 3.1 构建,Metal 4 版本仅针对需要它们的内核编译。这将 libtorch_cpu.dylib 缩减了约 25%,从 384 MiB 降至 287 MiB,并为 Metal 4.1 着色器奠定了基础,后者将带来更快的矩阵乘法和原生 FP8 矩阵乘法。

总结

PyTorch 2.15 为 Mac 带来了 FP8,为稠密模型和 MoE 模型带来更快的 LLM 解码,增加了更多 GPU 上的线性代数功能,并将一长串操作从 MPSGraph 迁移到专用的 Metal 内核。快来试试这个候选版本吧!

我们的下一步计划

  • 更多 MPS: 2.16 将把注意力机制完全放到 Metal 上运行,我们会继续将操作从 MPSGraph 迁出。
  • torch.compile: 我们正在努力消除 graph break,以便更多模型能够端到端编译。
  • 你的想法: 在 PyTorch 中我们还应该做些什么?如果有什么地方在 MPS 或其他任何地方运行缓慢、缺失或出错, 请在 GitHub 上提交一个 issue 并告诉我们。

特别感谢 Nikita Shulga, Kurt Mohler 以及 Aaron Gokaslan 审阅了这项工作的大部分内容,也感谢本版本中为 MPS 做出贡献的所有人。

原始出处

Hugging Face

内容说明

原始发布及相关权利归来源方。

机器翻译 · 请以原文为准