Amazon открыла исходный код TorchNeuron — нативного бэкенда PyTorch для облачных ускорителей AWS Trainium. Проект уже доступен на GitHub под лицензией Apache 2.0 и позволяет переносить модели на устройство neuron через стандартные механизмы PyTorch.
Бэкенд работает в режиме eager execution, поддерживает torch.compile, распределённое обучение, DTensor, FSDP и пользовательские ядра NKI. Amazon заявляет, что существующие проекты можно адаптировать без переписывания модели: в базовом сценарии меняется устройство, на котором размещаются модель и входные данные. Поддерживаются ускорители Trainium2 и Trainium3.
Публичная версия пока имеет статус beta и привязана к определённому окружению. В репозитории указаны Python 3.12–3.14, PyTorch 2.14 и Neuron SDK 2.33. Для запуска нужны облачные инстансы Trn2 или Trn3 либо соответствующая среда AWS Neuron.
Есть и технические ограничения. Компиляция через torch.compile работает только со статическими формами при dynamic=false; неконтинуальные тензоры копируются перед выполнением; вычисления float64 переводятся в float32, что способно повлиять на точность. Первый запуск операции или графа включает время компиляции и может быть заметно медленнее последующих.
TorchNeuron снижает барьер для команд, которые уже используют PyTorch, но рассматривают альтернативу привычным GPU. При этом совместимость конкретной модели всё равно нужно проверять по таблице поддерживаемых операций и ограничениям beta-версии.
Пока никто не ответил. Начните обсуждение.