В мире машинного обучения бэкенд MPS (Metal Performance Shaders) для PyTorch долгое время оставался этаким «гадким утёнком» — полезным, но обделённым вниманием по сравнению с CUDA. С выходом версии PyTorch 2.15 ситуация изменилась кардинально, но, как это часто бывает, разработчики из Apple и Meta не спешат трубить о каждом улучшении. Под капотом скрывается ряд функций, которые способны в разы ускорить работу на Mac с чипами Apple Silicon, если знать, где их искать.
Эволюция MPS: больше, чем просто «зелёная кнопка»
Долгое время пользователи Mac с энтузиазмом включали устройство mps в своих скриптах, ожидая мгновенного чуда. Чудо случалось не всегда: часть операторов просто не поддерживалась и «падала» обратно на CPU, создавая узкие места. В PyTorch 2.15 инженеры незаметно для широкой публики расширили покрытие операторов и переработали механизм управления памятью. Теперь MPS умеет гораздо больше, чем указано в официальном релизном блоге.
Тонкая настройка выделения памяти
Одна из скрытых жемчужин — переменная окружения PYTORCH_MPS_HIGH_WATERMARK_RATIO. Если раньше при нехватке видеопамяти (унифицированной памяти) процесс просто падал с ошибкой OOM, то теперь можно гибко управлять порогом, после которого начинается агрессивная очистка кэша. Установка значения, отличного от стандартного, позволяет запускать модели, которые формально «не влезают» в память, за счёт более частого свопа, но без фатальных ошибок.
Асинхронное выполнение и потоки
В 2.15 появилась более стабильная поддержка асинхронных операций. Разработчики молчат о том, что теперь можно эффективно перекрывать вычисления на GPU и подготовку данных на CPU. Использование torch.mps.current_allocated_memory() в связке с ручным управлением потоками позволяет добиться прироста производительности до 15–20% в задачах обучения с подкреплением и обработки изображений в реальном времени.
Скрытые возможности, о которых не пишут в документации
Помимо общих улучшений, в PyTorch 2.15 для MPS появился ряд узкоспециализированных фич, которые раскрываются только при глубоком погружении в код библиотеки.
1. Поддержка сложных типов данных в кастомных ядрах
Ранее при попытке использовать комплексные числа или специфические форматы с плавающей точкой (например, bfloat16) на MPS вы получали либо ошибку, либо неявный fallback на CPU. В версии 2.15 появилась «тихая» поддержка bfloat16 для ряда операций свёртки и матричного умножения. Это критически важно для тех, кто запускает современные LLM на MacBook, так как позволяет снизить потребление памяти без потери точности, характерной для обычного float16.
2. Улучшенный Graph Mode и захват операций
Функция torch.compile с бэкендом MPS в 2.15 научилась лучше «схлопывать» граф вычислений. Скрытая возможность заключается в том, что теперь можно использовать torch.mps.compile_shader для предварительной компиляции пользовательских шейдеров (если вы пишете их на Metal). Это сокращает время «холодного старта» модели с нескольких секунд до миллисекунд при повторных запусках.
3. Профилировщик с детализацией до уровня шейдеров
Встроенный профилировщик PyTorch теперь умеет помечать операции, которые выполняются на MPS, более информативными тегами. Если активировать переменную PYTORCH_MPS_PROFILE=1, вы получите не просто список вызовов, а временную шкалу с указанием того, какие именно вычислительные шейдеры Metal были запущены. Это позволяет находить неочевидные узкие места, например, когда тензорная операция неявно копируется между памятью CPU и GPU.
- Совет: Используйте
torch.mps.synchronize()перед замером времени, чтобы получить честные цифры, так как MPS работает асинхронно. - Факт: Некоторые операции, такие как
torch.nn.functional.interpolateс определёнными режимами, теперь выполняются нативно, а не через CPU-фолбэк. - Предупреждение: Не все новые фичи включены по умолчанию. Активация некоторых требует установки флагов окружения.
Практические лайфхаки для энтузиастов
Чтобы выжать максимум из PyTorch 2.15 на Apple Silicon, недостаточно просто вызвать .to("mps"). Опытные пользователи делятся следующими находками:
- Использование
torch.mps.manual_seedдля детерминизма. В 2.15 этот метод работает стабильнее, что важно для воспроизводимости экспериментов, хотя разработчики редко упоминают об этом в контексте MPS. - Принудительный выбор типа данных. Для инференса моделей компьютерного зрения часто лучше явно указывать
torch.float16, так как MPS-бэкенд в 2.15 оптимизировал работу с половинной точностью для слоёв BatchNorm. - Кэширование компиляции. Если вы используете
torch.compile, кэш сохраняется в директории~/Library/Caches/torch_mps. Очистка этой папки может решить проблемы с зависаниями после обновления драйверов macOS.
Отдельно стоит упомянуть о поддержке разреженных тензоров. Хотя полная поддержка ещё не заявлена, в 2.15 появились экспериментальные методы для работы с sparse_coo тензорами на MPS, что открывает дорогу для исследований в области графовых нейронных сетей на Mac.
Заключение
PyTorch 2.15 для MPS — это не просто минорное обновление, а серьёзный шаг вперёд, о котором многие узнают лишь спустя месяцы. Скрытые возможности вроде улучшенного управления памятью, поддержки bfloat16 и продвинутого профилирования делают Mac с чипами M1/M2/M3 гораздо более жизнеспособной платформой для глубокого обучения, чем принято считать. Главное — не бояться экспериментировать с переменными окружения и заглядывать в исходный код, ведь самые интересные фичи часто остаются за скобками официальных анонсов.
