PyTorch 2.15 с поддержкой MPS: обзор обновлений

Выход релиза PyTorch 2.15 приносит множество важных изменений для разработчиков и исследователей, использующих аппаратное ускорение на компьютерах Mac. Ключевым направлением обновления стало внедрение поддержки FP8, оптимизация процесса генерации текста языковыми моделями и перевод ряда тяжелых операций с интерфейса MPSGraph на специализированные ядра Metal.

Для стабильной работы актуальной версии программного обеспечения требуется операционная система macOS 15 или новее, а также установленный интерпретатор Python версии 3.11 либо выше. Проверка доступности ускорителя теперь выдает детальные пояснения при выявлении несовместимости оборудования или устаревшей версии системы.

Что изменилось в практическом применении

Пользователи трансформеров, диффузионных моделей и графических интерфейсов вроде ComfyUI получают заметный прирост производительности и улучшение стабильности. В первую очередь это касается локального запуска языковых моделей, где модели типа mixture-of-experts демонстрируют существенно возросшую скорость генерации токенов.

Внедрение формата FP8 на платформе MPS

Новая версия добавляет поддержку тензоров torch.float8_e4m3fn, занимающих всего один байт на значение. Это вдвое сокращает объем памяти, требуемый для хранения данных по сравнению с форматами FP16 или BF16. Умножение таких тензоров выполняется напрямую на графическом процессоре с помощью функции F.scaled_mm, поддерживающей как постраничное, так и тензорное масштабирование.

Ускорение декодирования языковых моделей

Процесс генерации текста в языковых моделях состоит из фаз предварительного заполнении и последовательного декодирования. На этапе декодирования обработка каждого нового токена сводится к матрично-векторным операциям. Обновление затрагивает оба доминирующих типа слоев.

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

Интерфейс F.linear теперь напрямую задействует специализированные матрично-векторные ядра, что позволяет эффективно обрабатывать стандартные входные последовательности декодирования с автоматическим слиянием смещений. Для типичных конфигураций это обеспечивает существенное ускорение вычислений, особенно на небольших проекциях в моделях смешанных экспертов.

Слои смешанных экспертов

Архитектуры mixture-of-experts распределяют токены по разным экспертам неравномерно. В версии 2.15 реализована нативная поддержка группированного умножения матриц. Если ранее смещения групп копировались на центральный процессор с организацией циклов, то новые ядра Metal считывают их непосредственно на графическом чипе, существенно повышая итоговую пропускную способность токенов.

Развитие линейной алгебры на графическом процессоре

Работа с комплексными тензорами стала значительно проще благодаря переносу на ускоритель факторизации LU, а также процедур решения систем линейных уравнений для комплексных чисел. Подобные изменения открывают возможность корректного вычисления комплексных детерминантов и градиентов.

Малые матрицы и сложные вычисления

Для задач компьютерной графики и геометрии камер добавлена быстрая инверсия пакетов матриц размером до 4×4 в рамках единого ядра Metal. Алгоритмы QR-разложения получили блочные обновления, а сингулярный анализ стал устойчивее к экстремальным числовым значениям, обеспечивая высокую точность псевдообратных вычислений.

Оптимизация сверток, дополнений и обучения

Ресурсоемкие операции в аудиомоделях, видеодекодерах и циклах обучения получили новые оптимизированные пути выполнения через Metal. Так, специализированные ядра заменили старые граф-процессы для одномерных сверток, а постоянное дополнение избавило систему от критической утечки памяти в процессе декодирования видеопотоков.

Снижение потребления памяти при декодировании видео

Одним из наиболее заметных результатов оптимизации постоянного дополнения стало кардинальное сокращение потребления ресурсов видеопамяти. В процессе декодирования видеофайлов с использованием VAE пиковая нагрузка на драйвер снижается в десятки раз, что позволяет комфортно запускать ресурсоемкие генеративные цепочки.

Управление потоками MPS из Python

Пользователи получили программный интерфейс для прямого выбора потока обработки задач через класс torch.mps.Stream и соответствующий контекстный менеджер. Специальный метод позволяет распределителю памяти отслеживать использование тензоров в параллельных потоках.

Дополнительные улучшения операций

Список поддерживаемых на графическом процессоре операций пополнился функциями логарифмического распределения, поиска индексов, специальными полиномами и операциями сравнения беззнаковых целых чисел. Кроме того, скорректировано масштабирование внимания для корректной обработки пакетов и голов.

Предварительный взгляд на PyTorch 2.16

Разработка следующего мажорного выпуска уже идет полным ходом. В нем планируется перенести вычисления механизма внимания полностью на Metal, продолжить миграцию операций с устаревшего фреймворка и уменьшить размер базового дистрибутива библиотеки за счет разделения версий шейдеров.