PyTorch 2.15 заметно меняет работу с MPS — бэкендом для GPU Apple Silicon. Команда PyTorch в Hugging Face вместе с мейнтейнерами и сообществом довела до релиза то, чего на Mac не хватало: FP8, быстрый проход декодирования для языковых моделей и замену части операций MPSGraph на собственные Metal-ядра. Если вы гоняете модели локально через transformers, diffusers или ComfyUI, обновление стоит попробовать — как и оркестрацию GPU в Kubernetes, локальный инференс упирается в те же ограничения по памяти и пропускной способности.
Что требуется для PyTorch 2.15
Нужен Mac на Apple silicon с macOS 15 или новее и Python 3.11 или новее. Если система слишком старая или Metal GPU недоступен, проверка доступности MPS теперь объясняет причину отказа. Попробовать релиз-кандидат в отдельном окружении можно так:
uv pip install torch==2.15.0 --index-url https://download.pytorch.org/whl/test/cpu
FP8 приходит на MPS
В 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: MoE-модели быстрее в 5–6,5 раза
У языковой модели есть фаза prefill и фаза декодирования. На декодировании каждый новый токен идёт через модель одной строкой, и крупные матричные умножения превращаются в умножения матрицы на вектор. В 2.15 ускорены оба типа слоёв, которые определяют этот шаг: плотные линейные слои и слои mixture-of-experts.
В MPS ядра GEMV для torch.mm были ещё с 2.14, но F.linear по-прежнему шёл через MPSGraph. Теперь они связаны: типичный вход декодирования [1, 1, hidden_size] использует их напрямую, а подходящие прибавления bias встраиваются в ядро. На 96 формах BF16-слоёв из реальных моделей это даёт геометрическое среднее ускорение 1,23×. Самый большой выигрыш — на небольших проекциях вроде гейтов в MoE-моделях, где старый путь не использовал большую часть пропускной способности памяти. Крупные слои и LM-головы уже упирались в память, поэтому у них около 1×.
Слой mixture-of-experts отправляет разное число токенов каждому эксперту. Grouped matmul обрабатывает такие неравные группы одной операцией, и в 2.15 у MPS появилась нативная реализация. В 2.14 grouped matmul копировал смещения групп на CPU и шёл по группам циклом. Новые Metal-ядра читают смещения прямо на GPU, используя Metal Performance Primitives там, где они доступны, и отдельный Metal-путь в остальных случаях. Для малых групп, обычного случая при декодировании, есть своё ядро.
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-модели декодируются в 5–6,5 раза быстрее, чем в 2.14: Qwen3.6-35B-A3B — с 9,4 до 61,3 токена в секунду, LFM2.5-8B-A1B — с 23,6 до 123,6 токена в секунду. Плотные модели тоже выигрывают за счёт более быстрых линейных слоёв и softmax для сэмплирования, который стал примерно в 8 раз быстрее. Похожий сдвиг в сторону компактных и быстрых моделей виден и в релизе MoE-модели Meta на 816B параметров, но там выигрыш идёт от архитектуры, а здесь — от бэкенда.
Линейная алгебра и обучение на GPU
Для комплексных тензоров complex64 LU-факторизация, LU-решение и решение треугольных систем теперь выполняются на MPS — вместе с операциями и градиентами, которые на них опираются, например комплексными определителями.
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, что делает надёжнее результаты псевдообращения и метода наименьших квадратов.
Для аудио, видео и циклов обучения в релизе тоже есть изменения: conv1d в прямом проходе получил специализированные ядра вместо MPSGraph и путь im2col-plus-matmul для коротких выходов; постоянный padding переехал на Metal-ядра; reflection и replication padding получили новые реализации на Metal, которые быстрее и исправляют корректность обратного прохода; out-of-place dropout в nn.Dropout и F.dropout теперь использует существующее слитое ядро; обратный проход upsampling заменил путь resize-gradient из MPSGraph на Metal-ядро; обратный проход conv2d в channels-last остаётся на эффективном пути channels-last вместо графа NCHW с дорогим доступом со страйдом.
Softmax и память при декодировании видео
softmax и log_softmax теперь работают на отдельных Metal-ядрах, и больше всего выигрывают большие словари: FP32 softmax по строке логитов [1, 128256] сократился с 57 µs до 7 µs, примерно в 8 раз быстрее версии на MPSGraph.
Самое заметное следствие — изменение постоянного padding. VAE-декодирование Wan-AI/Wan2.2-TI2V-5B-Diffusers делает 8 764 вызова padding на полное декодирование из 28 тайлов. С padding на Metal пиковая память драйвера падает примерно с 270 ГиБ до 8 ГиБ. Пик текущих аллокаций тензоров снижается с 4,422 до 4,095 ГиБ, пик зарезервированного аллокатором — не меняется (8,070 ГиБ), а общий пик памяти драйвера — с 270,135 до 8,108 ГиБ.
Больше контроля из Python
Теперь можно выбирать, в каком потоке 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, пополнился: torch.frexp, torch.logspace, torch.take и Tensor.put_, torch.heaviside, torch.special.laguerre_polynomial_l и сравнения беззнаковых целых. У logspace пока нет поддержки целочисленных и булевых выходов с более чем одним элементом. Исправлена и корректность attention: математическая реализация scaled dot-product attention на MPS теперь правильно обрабатывает широковещание по batch и головам.
Что готовят в PyTorch 2.16
Работа над 2.16 уже идёт. Часть изменений попала в main и доступна в nightly-сборках, часть ещё на ревью, так что детали могут поменяться. Когда в пути декодирования модели не остаётся операций MPSGraph, GPU начинает работу над токеном только после того, как CPU закончил кодировать все его ядра. В 2.16 командный буфер фиксируется каждые 16 ядер, и GPU включается, пока CPU кодирует остальные. Для скомпилированного BF16-декодирования с контекстом 2048 токенов LFM2.5-8B-A1B ускоряется с 108,8 до 126,0 токена в секунду (+16%), Qwen3.6-35B-A3B — с 52,9 до 62,8 (+19%), Qwen3.8-27B — с 12,1 до 12,6 (+4%).
Scaled dot-product attention полностью уходит с MPSGraph: любая форма будет считаться Metal-ядрами, с отдельными векторными ядрами для декодирования и Metal Performance Primitives для prefill в половинной точности на новых версиях macOS. Больше всего выиграют формы, которые раньше откатывались на MPSGraph. На M5 Pro декодирование с grouped-query attention (32 головы запросов, 8 голов KV, размер головы 80, контекст 4K) сокращается с 465 µs до 65 µs, а prefill 4K × 4K с размером головы 40 — с 42 мс до 6 мс. Продолжается и перенос операций с 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 МиБ.
Частые вопросы
Кого касается обновление?
Всех, кто запускает или обучает модели PyTorch на Mac с Apple silicon и использует MPS как основной ускоритель. Особенно заметна разница для тех, кто работает с MoE-моделями, генерацией видео и аудио.
Как проверить, подходит ли мой Mac?
Нужны Apple silicon, macOS 15 или новее и Python 3.11 или новее. Если система слишком старая или Metal GPU недоступен, проверка доступности MPS теперь сообщает причину.
Что делать, если macOS обновить нельзя?
Остаться на прежней ветке PyTorch: в 2.15 требование macOS 15 или новее указано явно. Обходного пути для старых систем в статье не описано.
Насколько быстрее станут MoE-модели?
В 5–6,5 раза по сравнению с 2.14. На M3 Ultra с torch.compile Qwen3.6-35B-A3B выдаёт 61,3 токена в секунду вместо 9,4, LFM2.5-8B-A1B — 123,6 вместо 23,6.
А обычные, плотные модели ускорятся?
Да, за счёт специализированных матрично-векторных ядер в F.linear и примерно восьмикратного ускорения softmax по словарю. Геометрическое среднее ускорение по 96 BF16-формам слоёв — 1,23×.
Хватит ли памяти для генерации видео?
Декодирование Wan2.2 с тайловым VAE теперь пикует около 8 ГиБ памяти драйвера вместо примерно 270 ГиБ. Это следствие перевода постоянного padding на Metal-ядра.
Что изменилось в FP8?
Тензоры torch.float8_e4m3fn на MPS занимают вдвое меньше памяти, чем BF16, и умножаются через F.scaled_mm с масштабированием на тензор или построчно.
Можно ли управлять потоками MPS из Python?
Да, через torch.mps.Stream и контекстный менеджер torch.mps.stream, а Tensor.record_stream помогает аллокатору отслеживать тензор в другом потоке.
Когда ждать attention полностью на Metal?
В PyTorch 2.16. Часть изменений уже в nightly-сборках, часть ещё на ревью, поэтому детали могут измениться до релиза.
Что делать, если операция на MPS не работает или тормозит?
Команда просит открыть issue на GitHub и описать проблему — это один из источников задач для следующих релизов.
Источник: huggingface.co