MPS у PyTorch 2.15

Оновити на GitHub

Якщо ви запускаєте моделі на Mac або використовуєте MPS як основний інструмент прискорення обчислень, у PyTorch 2.15 на вас чекає багато цікавого! На MPS з’являється FP8, декодування мовних моделей отримує швидші шляхи через лінійні шари та 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-розклад.
• Аудіо, відео та навчання також отримують увагу: швидші шляхи для згорток, доповнення, dropout, зворотного проходу інтерполяції та редукцій. Декодування відео VAE у Wan2.2 тепер використовує максимум 8 ГіБ пам’яті замість 270 ГіБ.
• Більше контролю з 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 токена/с. Щільні моделі також отримують переваги завдяки швидшим лінійним шарам і softmax, який у 8 разів швидше працює під час семплювання. (Докладніше)
• Генерація відео, що вміщується в пам’яті. Декодування відео за допомогою тайлового VAE Wan2.2 раніше збільшувало використання пам’яті драйвера до 270 ГіБ; тепер максимум становить близько 8 ГіБ. (Докладніше)
• Коректне редагування зображень за допомогою Qwen-Image-2.1. У версії 2.14 кодер VAE Qwen-Image-2.1 непомітно пошкоджував зображення на MPS, порушуючи робочі процеси редагування (про це повідомляли в ComfyUI). Та сама зміна доповнення це виправляє: тепер цикл кодування–декодування відповідає CPU зі значенням 29,95 дБ PSNR замість 10,06 дБ.
• Ваги FP8 на вашому Mac. Тензори FP8 займають удвічі менше пам’яті, ніж BF16, а F.scaled_mm безпосередньо перемножує їх на GPU. (Докладніше)

Зміст

• Початок роботи
• FP8 з’являється на MPS
• Швидше декодування мовних моделей
• Прискорення лінійної алгебри
• Згортки, доповнення та навчання
• Softmax і редукції
• Значно менше використання пам’яті під час декодування відео
• Більше контролю за допомогою потоків MPS
• Більше операцій працюють на MPS
• Попередній погляд на PyTorch 2.16
• Підсумки

Початок роботи

Вам знадобиться Mac з Apple silicon, macOS 15 або новішою та Python 3.11 або новішим (див. матрицю сумісності 2.15). Якщо ваша ОС застаріла або графічний процесор Metal недоступний, перевірка доступності MPS тепер повідомить причину.

Щоб випробувати кандидат на випуск в окремому середовищі:

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

FP8 з’являється на MPS

FP8 дає ще один спосіб заощадити пам’ять на Mac. PyTorch 2.15 додає на MPS тензори torch.float8_e4m3fn: один байт на значення, тобто вдвічі менше місця, ніж для 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] використовує їх безпосередньо, із об’єднанням придатних додавань зміщення. Для 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 отримують спеціальний швидкий шлях для пакетів трикутних систем FP32 розміром 16 × 16, 32 × 32 і 64 × 64 щонайменше з вісьмома правими частинами.

SVD також став надійнішим у крайніх числових випадках. Тепер він коректно обробляє дуже малі сингулярні значення, дуже великі або дуже малі значення вхідних даних та доповнення базису з full_matrices=True, що своєю чергою робить результати псевдообернення та методу найменших квадратів надійнішими.

Згортки, доповнення та навчання

Аудіомоделі, декодери відео та цикли навчання мають власні витратні операції. Цей випуск надає кільком із них нові шляхи через Metal:

Softmax і редукції

Семплювання токенів також пришвидшилося у версії 2.15. softmax і log_softmax тепер виконуються на спеціалізованих Metal-ядрах, а найбільшу користь отримують великі словники: FP32 softmax над рядком логітів [1, 128256] скорочується з 57 мкс до 7 мкс, тобто працює приблизно у 8 разів швидше, ніж версія MPSGraph. Прискорення зберігається й для інших поширених розмірів словника:

Значно менше використання пам’яті під час декодування відео

Один із найяскравіших результатів цього випуску пов’язаний зі зміною константного доповнення. Декодування VAE у Wan-AI/Wan2.2-TI2V-5B-Diffusers виконує 8764 виклики доповнення протягом повного декодування 28 тайлів. Завдяки доповненню на Metal пікове використання пам’яті драйвера зменшується приблизно з 270 ГіБ до 8 ГіБ:

Більше контролю за допомогою потоків 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 тепер коректно обробляє broadcasting пакетів і голів.

Попередній погляд на PyTorch 2.16

Робота над 2.16 уже активно триває. Деякі наведені нижче зміни вже потрапили до main і доступні у nightly-версіях; інші ще проходять перевірку, тому деталі можуть змінитися до виходу релізу.

Швидше декодування LLM

Щойно шлях декодування моделі більше не містить операцій MPSGraph, GPU починає обробляти токен лише після того, як CPU завершить кодування всіх її ядер. PyTorch 2.16 фіксує буфер команд кожні 16 ядер, тому GPU починає роботу, поки CPU ще кодує решту. Для скомпільованого декодування BF16 за контексту 2048 токенів:

Attention повністю працює на Metal

Масштабована скалярна увага повністю відмовляється від MPSGraph: кожна форма працюватиме на Metal-ядрах, зі спеціалізованими векторними ядрами для декодування та Metal Performance Primitives для попереднього заповнення зниженої точності в новіших версіях macOS. Найбільшу користь отримують форми, які раніше поверталися до MPSGraph. На M5 Pro декодування з grouped-query attention (32 голови запитів, 8 KV-голів, розмір голови 80, контекст 4K) скорочується з 465 мкс до 65 мкс, а попереднє заповнення 4K × 4K із розміром голови 40 — з 42 мс до 6 мс.

Більше міграцій з MPSGraph на Metal

Перехід від MPSGraph триває. Прямий і зворотний проходи batch norm перебувають у розробці, тоді як втрати MSE (до 12 разів швидші на малих входах), std_mean і var_mean (до 4 разів швидші) та bincount уже додані. bincount є найяскравішим прикладом: коли більшість значень потрапляє в один бін, підрахунок 4 мільйонів цілих чисел раніше займав 6,8 секунди, а тепер — пів мілісекунди.

Менший розмір інсталяції

Вбудовані шейдери за замовчуванням збиратимуться для Metal 3.1, а версії Metal 4 компілюватимуться лише для ядер, яким вони потрібні. Це зменшує розмір libtorch_cpu.dylib приблизно на 25% — із 384 МіБ до 287 МіБ — і створює основу для шейдерів Metal 4.1 зі швидшими матричними множеннями та нативним матричним множенням FP8.

Підсумки

PyTorch 2.15 додає FP8 на Mac, пришвидшує декодування LLM для щільних моделей і моделей MoE, переносить більше операцій лінійної алгебри на GPU та містить довгий список операцій, переміщених із MPSGraph до спеціалізованих Metal-ядер. Спробуйте кандидата на випуск!

Що буде далі

• Більше MPS: у версії 2.16 attention повністю працюватиме на Metal, а ми продовжимо переносити операції з MPSGraph.
• torch.compile: ми працюємо над усуненням розривів графа, щоб більше моделей компілювалися наскрізно.
• Ваші ідеї: над чим ще нам варто попрацювати в PyTorch? Якщо щось працює повільно, відсутнє або зламане — на MPS чи будь-де — створіть issue на GitHub і повідомте нам.

Щиро дякуємо Нікіті Шульзі, Курту Молеру та Аарону Гокаслану за перевірку значної частини цієї роботи, а також усім, хто долучився до розвитку MPS у цьому випуску.

Більше статей у нашому блозі