PyTorch 2.15 представляет собой значительное обновление для пользователей Mac, использующих Metal Performance Shaders (MPS) в своих вычислительных задачах. Одним из наиболее ожидаемых нововведений является поддержка FP8, что позволяет значительно экономить память. Также улучшены пути декодирования языковых моделей, что делает их более быстрыми, и добавлены новые Metal-ядра, которые оптимизируют ресурсоемкие операции в области аудио, видео и обучения.

Поддержка FP8 на MPS

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

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

Каждая языковая модель проходит через два этапа: предварительное заполнение и декодирование. На этапе декодирования каждый новый токен обрабатывается моделью как отдельная строка, что превращает большие матричные умножения в матрично-векторные произведения. В PyTorch 2.15 ускорены как плотные линейные слои, так и слои смешанных экспертов. Например, на M3 Ultra с использованием `torch.compile`, модели смешанных экспертов теперь декодируют токены в 5-6.5 раз быстрее, чем в версии 2.14.

Увеличение производительности линейной алгебры

Обновление также включает улучшения в линейной алгебре. Теперь на MPS доступны операции с комплексными тензорами, такие как LU-разложение и решение систем уравнений. Малые матрицы, используемые в геометрии камеры и трансформациях изображений, теперь обрабатываются быстрее благодаря новым ядрам Metal. Кроме того, QR-разложение и треугольные решения также получили улучшения производительности.

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

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

Уменьшение объема памяти для декодирования видео

Одним из самых заметных результатов обновления является изменение в обработке паддинга. При декодировании видео с использованием VAE в Wan2.2 пиковая память, необходимая для обработки, снизилась с 270 ГиБ до 8 ГиБ, что делает работу с видео более эффективной.

Более точный контроль с помощью потоков MPS

Теперь пользователи могут выбирать, какой поток MPS будет выполнять их задачи, непосредственно из Python. Это позволяет более эффективно управлять вычислительными ресурсами и синхронизировать выполнение задач.

Новые операции на MPS

В версии 2.15 добавлено множество новых операций, которые теперь могут выполняться на MPS, включая `torch.frexp`, `torch.logspace`, и другие. Это расширяет возможности работы с данными и упрощает реализацию различных алгоритмов.

Взгляд на PyTorch 2.16

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

Заключение

PyTorch 2.15 предлагает множество улучшений для пользователей Mac, работающих с MPS. Поддержка FP8, ускорение декодирования языковых моделей и оптимизация операций линейной алгебры делают это обновление особенно важным для разработчиков и исследователей, работающих в области машинного обучения и обработки данных.