MPS в PyTorch 2.15

Актуализиране в GitHub

Ако изпълнявате модели на своя Mac или използвате MPS като основен инструмент за ускоряване на изчисленията, в PyTorch 2.15 има много интересни новини! FP8 идва в MPS, декодирането на езикови модели получава по-бързи пътища през линейните слоеве и softmax, а новите Metal ядра се справят с някои изненадващо скъпи части от работните натоварвания за аудио, видео и обучение.

Накратко за това кои сме ние: ние сме екипът на PyTorch в Hugging Face. Допринасяме за PyTorch в различни области, а една от основните ни насоки е бекендът MPS, чиято цел е да превърне вашия Mac в пълноценна платформа за изпълнение и обучение на модели на PyTorch. В началото на 2026 г. изпълнението дори на малки модели в MPS често изискваше заобикаляне на неподдържани операции, спадове в производителността и неочаквано нарастване на използваната памет. Затова изпълняваме реални модели от Hugging Face Hub в областите на езика, изображенията и видеото, откриваме къде MPS не достига и отстраняваме проблемите директно в PyTorch, все по-често чрез замяна на MPSGraph със специализирани Metal ядра. Голяма част от тази публикация е резултат от тази работа, извършена съвместно с поддържащите PyTorch и участниците от общността.

Накратко

• FP8 идва в Mac: съхранявайте стойности от тип float8_e4m3fn с наполовина по-малко пространство от BF16 и използвайте мащабирано матрично умножение и всички други операции, поддържани в CUDA/CPU!
• MoE моделите се декодират 5–6,5 пъти по-бързо от 2.14: Qwen3.6-35B-A3B преминава от 9,4 на 61,3 токена/сек на M3 Ultra благодарение на нативно групирано матрично умножение със специален път за декодиране, по-малко синхронизации между GPU и CPU и по-бързо маршрутизиране top-k.
• По-малко време за всеки токен при всеки LLM: линейните слоеве с един ред използват специализирани матрично-векторни ядра, а softmax върху речника е приблизително 8 пъти по-бърз.
• Повече линейна алгебра върху GPU: комплексни LU и триъгълни решения, по-бързи инверсии на малки матрици и по-бързо QR разлагане.
• Аудиото, видеото и обучението също получават внимание: по-бързи пътища за свертки, padding, обратно разпространение при интерполация и редукции. Декодирането на видео VAE на Wan2.2 вече достига пик от 8 GiB памет вместо 270 GiB.
• Повече контрол от Python: публични MPS потоци и по-надеждно измерване на времето на събитията.
• Актуализирани изисквания: MPS вече изисква macOS 15 или по-нова версия; PyTorch 2.15 изисква Python 3.11 или по-нова версия.
• Поглед към 2.16: attention изцяло върху Metal, още миграции от MPSGraph към Metal и по-бързо декодиране на LLM.

Какво означава това за вас

Ако изпълнявате модели на Mac с transformers, diffusers или инструменти като ComfyUI, ето какво променя 2.15 на практика:

• Значително по-бързи локални LLM. Моделите от типа mixture-of-experts се декодират 5–6,5 пъти по-бързо от 2.14: на M3 Ultra с torch.compile Qwen3.6-35B-A3B преминава от 9,4 на 61,3 токена/сек, а LFM2.5-8B-A1B — от 23,6 на 123,6 токена/сек. Плътните модели също се възползват от по-бързите линейни слоеве и 8 пъти по-бързия softmax за семплиране. (Още)
• Генериране на видео, което се побира в паметта. Декодирането на видео с тайлинга VAE на Wan2.2 преди можеше да увеличи паметта на драйвера до 270 GiB; сега пикът е около 8 GiB. (Още)
• Коректно редактиране на изображения с Qwen-Image-2.1. В 2.14 VAE енкодерът на Qwen-Image-2.1 тихомълком повреждаше изображенията в MPS, което нарушаваше работните процеси за редактиране на изображения (съобщено в ComfyUI). Същата промяна в padding отстранява проблема: двупосочното кодиране и декодиране вече съвпада с CPU при 29,95 dB PSNR спрямо 10,06 dB.
• FP8 тегла на вашия Mac. FP8 тензорите използват наполовина по-малко памет от BF16, а F.scaled_mm ги умножава директно върху GPU. (Още)

Съдържание

• Първи стъпки
• FP8 идва в MPS
• По-бързо декодиране на езикови модели
• Подобрения в линейната алгебра
• Свертки, padding и обучение
• Softmax и редукции
• Много по-малък отпечатък в паметта при декодиране на видео
• Повече контрол с MPS потоци
• Повече операции работят в MPS
• Поглед към PyTorch 2.16
• В заключение

Първи стъпки

Ще ви е необходим Mac с Apple silicon, 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, който поддържа както мащабиране по тензори, така и мащабиране по редове. Ето малък пример с един мащаб за всеки тензор:

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 има фаза на предварително запълване и фаза на декодиране. По време на декодирането всеки нов токен преминава през модела като един ред, така че големите матрични умножения се превръщат в матрично-векторни произведения. PyTorch 2.15 ускорява и двата вида слоеве, които доминират в тази стъпка: плътните линейни слоеве и слоевете mixture-of-experts.

Линейни слоеве

MPS разполага с GEMV ядра за torch.mm още от 2.14, но F.linear все още преминаваше през MPSGraph. PyTorch 2.15 свързва тези ядра с F.linear, така че обичайният вход за декодиране [1, 1, hidden_size] ги използва директно, като подходящите операции за добавяне на bias са слети. При 96 BF16 форми на слоеве от реални модели това дава геометрично средно ускорение от 1,23 пъти. Най-големите подобрения са при малките проекции, например гейтовете в MoE моделите, където старият път оставяше по-голямата част от пропускателната способност на паметта неизползвана. Големите слоеве и главите на езиковите модели вече бяха близо до ограничението от пропускателната способност, така че остават около 1 път.

Слоеве mixture-of-experts

Слой mixture-of-experts изпраща различен брой токени към всеки експерт. Групираното матрично умножение обработва тези неравномерни групи в една операция, а 2.15 предоставя на MPS нативна реализация.

В 2.14 групираното матрично умножение копираше отместванията на групите към CPU и обхождаше групите в цикъл. Новите Metal ядра четат отместванията директно от GPU, като използват Metal Performance Primitives, когато са налични, и отделен Metal път в останалите случаи. Малките групи, които са обичайни при декодиране, също получават собствено ядро.

Например 64 токена могат да бъдат разделени на групи от 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)  

Ядрото за декодиране се проявява директно в пропускателната способност на токените. На M3 Ultra с torch.compile MoE моделите генерират приблизително 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())

Малки матрици и по-големи решения

Малките матрици също получават бързи ядра. За геометрия на камерата и трансформации на изображенията пакети от FP32 матрици с размер до 4 × 4 се обръщат в едно Metal ядро, което комбинира факторизацията и решаването. По-големите задачи също получават внимание: QR използва блокови обновявания, а триъгълните решения превръщат голяма част от работата си в ефективно матрично умножение.

Работните натоварвания на Gated DeltaNet получават специален бърз път за пакети от 16 × 16, 32 × 32 и 64 × 64 FP32 триъгълни системи с поне осем десни страни.

SVD също е по-устойчив при числени крайности. Вече коректно обработва много малки сингулярни стойности, много големи или много малки входни величини и допълване на базата с full_matrices=True, което от своя страна прави резултатите от псевдообратните матрици и метода на най-малките квадрати по-надеждни.

Свертки, padding и обучение

Аудио моделите, видео декодерите и циклите за обучение имат свои скъпи операции. Това издание предоставя на няколко от тях нови пътища през Metal:

Softmax и редукции

Семплирането на токени също става по-бързо в 2.15. softmax и log_softmax вече се изпълняват върху специализирани Metal ядра, а големите речници се възползват най-много: FP32 softmax върху ред [1, 128256] от логити намалява от 57 µs на 7 µs, тоест е приблизително 8 пъти по-бърз от версията с MPSGraph. Ускорението се запазва и при други често срещани размери на речника:

Много по-малък отпечатък в паметта при декодиране на видео

Един от най-впечатляващите резултати в това издание идва от промяната при константния padding. VAE декодирането в Wan-AI/Wan2.2-TI2V-5B-Diffusers извършва 8764 извиквания на padding при пълно декодиране от 28 плочки. С padding върху Metal пиковата памет на драйвера спада от около 270 GiB на 8 GiB:

Повече контрол с MPS потоци

Вече можете да избирате кой MPS поток да изпълнява работата ви директно от Python. 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 все още не поддържа целочислени или булеви изходи с повече от един елемент.

Attention също получава корекция за правилност: MPS математическата реализация на scaled dot-product attention вече обработва правилно broadcasting между batch и head.

Поглед към PyTorch 2.16

Работата по 2.16 вече е в разгара си. Някои от промените по-долу са влезли в main и са налични в nightly версиите; други все още се преглеждат, така че подробностите може да се променят преди изданието.

По-бързо декодиране на LLM

Когато пътят за декодиране на даден модел вече не съдържа операции на MPSGraph, GPU започва работа върху токен едва след като CPU приключи с кодирането на всичките му ядра. PyTorch 2.16 изпраща командния буфер на всеки 16 ядра, така че GPU започва работа, докато CPU все още кодира останалите. При компилирано BF16 декодиране и контекст от 2048 токена:

Attention се изпълнява изцяло върху Metal

Scaled dot-product attention напълно премахва MPSGraph: всяка форма ще се изпълнява върху Metal ядра, със специализирани векторни ядра за декодиране и Metal Performance Primitives за предварително запълване с половинна точност в по-новите версии на macOS. Формите, които преди преминаваха към MPSGraph, ще имат най-голяма полза. На M5 Pro декодирането с grouped-query attention (32 query глави, 8 KV глави, размер на главата 80, контекст 4K) спада от 465 µs на 65 µs, а предварителното запълване 4K × 4K с размер на главата 40 преминава от 42 ms на 6 ms.

Още миграции от MPSGraph към Metal

Преминаването от MPSGraph продължава. Напред и назад при batch norm са в процес на разработка, докато MSE загубата (до 12 пъти по-бърза при малки входове), std_mean и var_mean (до 4 пъти по-бързи) и bincount вече са добавени. bincount е най-впечатляващият пример: когато повечето стойности попадат в един bin, преброяването на 4 милиона цели числа преди отнемаше 6,8 секунди, а сега отнема половин милисекунда.

По-малка инсталация

Включените шейдъри ще се компилират за Metal 3.1 по подразбиране, а версиите за Metal 4 ще се компилират само за ядрата, които се нуждаят от тях. Това намалява размера на libtorch_cpu.dylib с около 25% — от 384 MiB на 287 MiB — и подготвя основата за шейдъри за Metal 4.1 с по-бързи matmul операции и нативно FP8 матрично умножение.

В заключение

PyTorch 2.15 пренася FP8 на Mac, осигурява по-бързо декодиране на LLM за плътни и MoE модели, повече линейна алгебра върху GPU и дълъг списък от операции, преместени от MPSGraph към специализирани Metal ядра. Изпробвайте кандидата за издание!

Какво следва за нас

• Още MPS: 2.16 ще изпълнява attention изцяло върху Metal и ще продължим да преместваме операции извън MPSGraph.
• torch.compile: работим по премахването на прекъсванията на графа, за да могат повече модели да се компилират от край до край.
• Вашите идеи: върху какво друго трябва да работим в PyTorch? Ако нещо е бавно, липсва или е повредено в MPS или където и да е другаде, отворете issue в GitHub и ни уведомете.

Голямо благодаря на Nikita Shulga, Kurt Mohler и Aaron Gokaslan за преглеждането на голяма част от тази работа и на всички, допринесли за MPS в това издание.

Още статии от нашия блог