Обзор PyTorch 2.11: ключевые изменения и новые возможности версии

PyTorch – один из главных инструментов, на которых сегодня строится обучение нейросетей. Его используют как академические исследователи, так и команды в крупных компаниях. 23 марта вышла версия 2.11, собравшая 2723 коммита от 432 участников с момента предыдущего релиза. Это не революция, но вполне весомое накопительное обновление – особенно если работать с несколькими GPU одновременно или обучать модели на Mac.

Дифференцируемые коллективные операции для распределенного обучения

Когда тысячи GPU должны работать как один

Одно из центральных изменений релиза касается распределённого обучения – то есть ситуаций, когда модель обучается не на одном ускорителе, а на десятках или сотнях одновременно. В таких системах устройства постоянно обмениваются данными между собой: синхронизируют градиенты, перераспределяют тензоры, согласовывают состояния. Эти операции называются коллективными.

До сих пор одна из практических проблем заключалась в том, что такие операции были своего рода «чёрным ящиком» с точки зрения обратного распространения ошибки. Проще говоря: когда модель обучается, она не просто делает предсказание – она «понимает», как изменить свои веса, чтобы прогноз стал точнее. Для этого нужно уметь «смотреть назад» сквозь все выполненные операции. Коллективные операции это прохождение блокировали.

В PyTorch 2.11 эта блокировка снята: коллективные операции стали дифференцируемыми. Теперь обратный проход работает сквозь них без дополнительных ухищрений. Для большинства пользователей это изменение останется незаметным в повседневной работе, зато исследователи, строящие нестандартные архитектуры обучения на кластерах, получают заметно больше свободы.

Ускорение механизма внимания и интеграция FlashAttention-4

Механизм внимания стал быстрее – и это ощутимо

Большинство современных языковых и мультимодальных моделей используют механизм внимания – это ключевая операция, которая позволяет модели «понимать», какие части входных данных важны в контексте друг друга. От того, насколько быстро она работает, напрямую зависит скорость обучения и работы (инференса) модели.

В PyTorch есть инструмент FlexAttention, позволяющий гибко настраивать этот механизм под конкретные задачи. В версии 2.11 он получил новый бэкенд на основе FlashAttention-4 – более эффективную реализацию, оптимизированную под современные GPU архитектур Hopper и Blackwell (это флагманские линейки NVIDIA последних поколений).

Прирост производительности на нагруженных задачах составляет от 1.2× до 3.2× по сравнению с предыдущей реализацией. Если коротко: та же модель, то же железо, но обучение идёт заметно быстрее. Пока это экспериментальная функция, и её поведение может меняться по мере стабилизации, но протестировать её можно уже сейчас.

Улучшение поддержки Apple Silicon и бэкенда MPS в PyTorch 2.11

Mac становится более серьёзной платформой для ML

Разработчики, работающие на MacBook с чипами Apple Silicon (M-серия), получают в этом релизе ощутимое расширение возможностей. MPS – это бэкенд, через который PyTorch использует графический процессор внутри чипов Apple для ускорения вычислений. Раньше поддержка была неполной: часть операций просто не работала на этом бэкенде, и PyTorch либо переходил в запасной режим, либо выдавал ошибки.

В 2.11 добавлены новые математические операции, расширена поддержка типов данных (целых и комплексных чисел), а также появилась асинхронная отчётность об ошибках. Последнее – важная деталь: теперь PyTorch может сообщать об ошибках, которые происходят непосредственно на GPU, например о выходе за границы массива при индексировании. Раньше такие ошибки было сложно отловить, и зачастую они приводили к неявным сбоям.

Это не значит, что Mac превращается в полноценную замену серверному кластеру. Но для локальной разработки, экспериментов и запуска моделей на машинах без внешней видеокарты ситуация заметно улучшается.

Экспорт рекуррентных нейросетей LSTM и GRU на GPU

Экспорт моделей – теперь и для рекуррентных сетей

Рекуррентные сети (LSTM, GRU и им подобные) – это архитектуры, которые хорошо работают с последовательными данными: текстом, временными рядами, аудио. Они менее популярны, чем трансформеры, но всё ещё активно используются во многих задачах, особенно там, где важна вычислительная экономность.

До сих пор у этих архитектур была проблема: их было неудобно экспортировать для последующего развёртывания. Экспорт – это процесс, при котором обученная модель «упаковывается» в формат, пригодный для запуска в рабочей среде (продакшене) без лишних зависимостей. В 2.11 модели LSTM и GRU наконец получили полноценную поддержку экспорта на GPU, включая динамические размеры входных данных. Это делает их полноправными участниками стандартного цикла разработки – от обучения до деплоя.

Оптимизация работы PyTorch на графических процессорах Intel и AMD

Intel и AMD тоже не забыты

Для GPU от Intel добавлена поддержка XPUGraph – механизма, позволяющего «записать» последовательность операций и затем воспроизводить её многократно, исключая накладные расходы на повторный запуск. Проще говоря: Python запускает граф один раз, фиксирует его, а далее выполнение идёт напрямую на железе без промежуточных шагов. Это снижает нагрузку на центральный процессор и ускоряет повторяющиеся задачи инференса.

Для GPU от AMD появилась поддержка отладочных утверждений (assertions) непосредственно на устройстве, а также был улучшен оператор TopK – операция поиска наибольших значений в тензоре, которая часто встречается в задачах ранжирования и сэмплинга.

Поддержка вычислений в формате FP16 на центральных процессорах

FP16 на CPU: полезно для периферийных устройств

Ещё одно небольшое, но практически значимое изменение: PyTorch теперь умеет выполнять матричные умножения в формате FP16 (арифметика половинной точности) на обычных процессорах. Раньше FP16 был почти исключительно прерогативой GPU. На CPU это открывает возможность более быстрого инференса на устройствах без видеокарты – серверах с CPU-инференсом, edge-устройствах или встроенных системах.

Переход от TorchScript к современному механизму torch.export

Что уходит: TorchScript официально устарел

TorchScript – это прежний механизм сериализации и запуска PyTorch-моделей. Он был удобен в своё время, но накопил ряд ограничений. Устаревшим его объявили ещё в версии 2.10, а теперь это подтверждено официально: разработчикам рекомендуется переходить на более современный подход через torch.export.

Для тех, кто только начинает: это означает, что новые проекты лучше сразу строить на актуальном инструментарии. Существующие проекты на TorchScript пока продолжат работать, но инвестировать время в этот метод уже не стоит.

Изменение графика релизов и переход на двухмесячный цикл

Релизы теперь будут чаще

Отдельно стоит отметить изменение в расписании выпусков: начиная с этого года PyTorch переходит с квартального цикла на двухмесячный. Это значит, что улучшения будут доходить до пользователей быстрее. С одной стороны, это хорошо: меньше ждать нужных исправлений и новых возможностей. С другой – у разработчиков, тщательно тестирующих ПО перед обновлением, прибавится работы по отслеживанию изменений.

В целом PyTorch 2.11 – это аккуратный, насыщенный релиз без громких анонсов, но с реальными улучшениями по всем направлениям: от ускорения механизмов внимания до расширения поддержки платформ. Именно такие обновления делают инструмент по-настоящему зрелым.