PyTorch 2.9: стабильный ABI, симметричная память и мультиплатформенность

Введение

Команда PyTorch объявила о выпуске PyTorch 2.9. В новой версии представлены обновления стабильного ABI libtorch для сторонних расширений C++/CUDA, поддержка симметричной памяти для многопроцессорных систем, расширенные возможности компиляции, а также улучшения производительности на платформах AMD, Intel, NVIDIA и различных чипах на архитектуре Arm.

Подробнее о PyTorch 2.9

Одним из ключевых нововведений стала поддержка стабильного ABI libtorch, позволяющего создавать C++/CUDA-расширения, совместимые между версиями PyTorch. Добавлены новые утилиты для работы с устройствами, включая Device Guard и Stream, а также API torch::stable::Tensor с поддержкой функций is_cpu, scalar_type и get_device_index. Дополнительно реализованы стабильные операции ATen — amax, narrow, pad, new_empty и new_zeros.

Вторым крупным улучшением стала система симметричной памяти (Symmetric Memory), упрощающая программирование ядер для нескольких GPU, соединенных через NVLink и RDMA. Symmetric Memory обеспечивает низкоуровневое взаимодействие между ядрами, односторонний удаленный доступ с низкой задержкой и возможность создания настраиваемых коммуникационных паттернов. Новая архитектура уже поддерживает симметричные тензоры для CUDA и NVSHMEM, а также ускоренные коллективные операции, включая one_shot_all_reduce и multimem_all_gather_out.

Еще одно важное обновление — возможность гибкого управления ошибками в torch.compile. Новый контекстный менеджер torch._dynamo.error_on_graph_break() позволяет пользователям задавать поведение при разрыве графа: выдавать ошибку или возобновлять выполнение. Эта функция делает процесс компиляции более предсказуемым и удобным при отладке.

PyTorch 2.9 также расширяет поддержку wheel-пакетов с добавлением платформ AMD ROCm, Intel XPU и NVIDIA CUDA 13, а также бинарных сборок Linux aarch64 для всех поддерживаемых версий CUDA. Это улучшает переносимость и упрощает установку PyTorch на различных аппаратных конфигурациях.

Дополнительно в релиз включена поддержка FlexAttention на графических процессорах Intel и оптимизация Flash-декодирования на процессорах x86. Новая реализация повышает производительность при генерации длинных последовательностей в LLM и улучшает загрузку CPU при больших контекстах.

Для архитектуры Arm реализованы оптимизации компилятора и операторов: ускорена работа сверток, активаций и квантования, а также добавлена поддержка CI-инстансов AWS Graviton 4 на базе Arm Neoverse V2.

Выводы

PyTorch 2.9 представляет собой значимое обновление экосистемы ИИ-фреймворка, направленное на повышение стабильности, совместимости и производительности. Стабильный ABI, симметричная память и улучшенная многоплатформенная поддержка делают фреймворк более гибким и удобным для разработки масштабируемых ИИ-систем.