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 благодаря нативному grouped matmul со специализированным путем декодирования, меньшему числу синхронизаций 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). Если ваша ОС слишком старая или GPU 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 есть этап предварительного заполнения (prefill) и этап декодирования. Во время декодирования каждый новый токен проходит через модель как одна строка, поэтому большие матричные умножения превращаются в матрично-векторные операции. PyTorch 2.15 ускоряет оба типа слоев, определяющих производительность этого этапа: плотные линейные слои и слои mixture-of-experts.

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

В MPS уже с версии 2.14 были GEMV-ядра для torch.mm, но F.linear по-прежнему использовал MPSGraph. PyTorch 2.15 подключает их к F.linear, поэтому распространенный вход для декодирования размерности [1, 1, hidden_size] использует их напрямую, с объединением подходящих операций добавления смещения. Для 96 форм BF16-слоев из реальных моделей это дает геометрическое среднее ускорения в 1,23 раза. Наибольший выигрыш наблюдается у небольших проекций, например у вентилей в моделях MoE, где старый путь оставлял большую часть пропускной способности памяти неиспользованной. Большие слои и головы LM уже были близки к ограничению пропускной способностью, поэтому их результат остается около 1×.

Слои mixture-of-experts

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

В версии 2.14 grouped matmul копировал смещения групп на 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 раза больше токенов в секунду:

Ускорение линейной алгебры

Вычисления с комплексными числами

Работать с комплексными тензорами стало проще. LU-факторизация, решение LU-систем и решение треугольных систем для complex64 теперь выполняются на 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-ядрах, причем наибольшую пользу получают большие словари: softmax FP32 для строки логитов [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 также исправлена корректность: математическая реализация scaled dot-product attention в MPS теперь корректно обрабатывает broadcasting по пакетам и головам.

Предварительный обзор 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 мкс до 65 мкс, а предварительное заполнение 4K × 4K с размерностью головы 40 — с 42 мс до 6 мс.

Новые миграции с MPSGraph на Metal

Переход от MPSGraph продолжается. Прямой и обратный проходы batch norm находятся в разработке, а MSE loss (до 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 в этом выпуске.

Другие статьи в нашем блоге