На PyTorch Conference Amazon и Google объявили о выпуске открытых инструментов для запуска оптимизированного под Nvidia фреймворка PyTorch на своих собственных ускорителях ИИ. По словам представителей компаний, TorchNeuron от Amazon призван устранить ключевые барьеры совместимости PyTorch с чипами Trainium, а Google представила TorchTPU — аналогичный пакет для ускорителей TPU. В анонсах подчёркивалось, что речь идёт об открытом коде и прямой поддержке существующих проектов PyTorch, что может существенно упростить работу разработчиков, которые ранее сталкивались с необходимостью обходных решений для Trainium и TPU.

Согласно уточнениям в материалах, TorchNeuron позволяет запускать проекты PyTorch на ускорителях Trainium без крупных изменений в коде — достаточно указать, какой ускоритель будет использоваться. Отмечается, что в составе инструментария задействованы готовые операторы языка NKI, специально разработанного под чипы Trainium, а также имеется компилятор, который, по описаниям, оптимизирует модель на низком уровне, что может приводить к приросту скорости и уменьшению потребления ресурсов.

Параллельно Google с командой PyTorch в Meta✴ подготовила TorchTPU для запуска существующих рабочих нагрузок PyTorch на TPU с минимальными изменениями кода. Подчёркнута совместимость со стандартным API PyTorch, наличие инструментов диагностики и встроенных примитивов для распределённых вычислений, что, как заявляет Google Cloud, должно обеспечить высокую производительность на старте.

Особое внимание уделено поддержке проектов с открытыми весами через Hugging Face: теперь модели из этой экосистемы «из коробки» поддерживаются на TPU — достаточно указать соответствующий параметр. Также в релизе упоминаются совместимости с vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec и LiteRT, что призвано упрощать переход на оборудование Google и снижать порог входа для разработчиков при переходе на TPU и PyTorch как рабочую среду.

В контексте сообщения подчёркнуто, что новые инструменты не требуют перестройки всей программной системы и могут стать альтернативой для тех, кто хотел бы отказаться от Nvidia-инициатив в пользу собственных ускорителей. Это заявление сопровождается указанием на то, что TorchNeuron и TorchTPU — часть открытой экосистемы, ориентированной на совместимость и расширяемость в рамках PyTorch.

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