AWS открыла бету TorchNeuron для запуска PyTorch на Trainium

Подробно

Что такое TorchNeuron и чем он отличается от прежнего подхода AWS?

TorchNeuron - не отдельная версия и не форк PyTorch. Это нативный аппаратный backend: он добавляет в обычный upstream PyTorch устройство neuron через официальный механизм PrivateUse1. Пакет и репозиторий называются torch-neuronx. Раньше под этим названием выходила XLA-реализация, поэтому AWS отдельно уточняет, что версии 2.13.3.0. и 2.14.3.0. относятся уже к новому нативному backend.

Продукт представили 19 октября 2026 года AWS и авторы блога PyTorch. Публичную бету и открытие исходников приурочили к 20 октября - дню конференции PyTorch Conference North America в Сан-Хосе (Калифорния). К моменту проверки репозиторий aws-neuron/torch-neuronx уже открыт, а в документации Neuron указан выпуск публичной беты в составе Neuron SDK 2.33.0.

До этого PyTorch на Trainium запускали через стек Neuron и PyTorch XLA со специфическими интерфейсами. Инстансы Trn1 на первом поколении Trainium стали общедоступны 10 октября 2022 года. В декабре 2024 года Neuron 2.21 добавил поддержку Trainium2 и Trn2, а 2 декабря 2025 года AWS запустила Trn3 UltraServers и пообещала нативную интеграцию с PyTorch без изменения кода моделей. После этого TorchNeuron проходил закрытое бета-тестирование.

«Если вы знаете PyTorch, значит, вы знаете Trainium», - заявил Колин Брейс, вице-президент по технологическому ПО Annapurna Labs в AWS.

Как перенести PyTorch-модель на Trainium и правда ли хватит одной строки?

Для базового переноса модель отправляют на устройство вызовом .to("neuron"). По словам AWS, код, который уже использует аппаратно-независимый интерфейс torch.accelerator, может заработать и без этой правки. Для компилируемого режима backend выбирают командой torch.compiler.set_default_backend("neuron") или прямо в вызове torch.compile(model, backend="neuron").

Формула «без переписывания кода» и «одна строка» - это описание типового сценария от Amazon, а не гарантия для любого приложения. Она подходит для простого или изначально аппаратно-независимого кода. Если проект завязан на CUDA, полная совместимость не доказана: часть API отсутствует, а для высокой производительности могут понадобиться компиляция, профилирование и собственные низкоуровневые ядра.

По данным AWS, TorchNeuron поддерживает:

• eager execution и torch.compile;
• torch.distributed, DDP, FSDP2, DTensor и DeviceMesh;
• pipeline и context parallelism;
• профилирование через torch.profiler;
• собственные ядра на Neuron Kernel Interface (NKI).

Hugging Face и AWS добавили поддержку Trainium в Transformers, Diffusers, TRL и библиотеку kernels. Также TorchNeuron интегрирован с TorchTitan, vLLM Neuron и vLLM Omni Neuron.

«TorchNeuron делает Trainium нативным backend для Transformers и Diffusers и позволяет работать со знакомыми определениями моделей через полностью открытый стек», - отметил Джефф Будье, вице-президент по продукту Hugging Face.

Какие ограничения есть у публичной беты TorchNeuron?

Бета работает только с Trainium2 и Trainium3: первого поколения Trainium (Trn1) в матрице поддержки нет. Официально поддерживаются Linux x86-64, Ubuntu 24.04 или 26.04, PyTorch 2.13.0 или 2.14.0 и Python 3.12-3.14. Ubuntu 22.04 для новых образов больше не поддерживается.

Ключевые ограничения из release notes:

• torch.compile требует статических форм и параметра dynamic=False, а новый размер входных данных может вызвать повторную компиляцию;
• первая операция или граф выполняются с задержкой на компиляцию, если в кеше нет готового результата;
• не поддерживаются torch.sparse, torch.func, torch.nested и torch.distributions, FFT поддерживается частично;
• TorchAO и составной тип MXFP8 в публичный релиз не вошли;
• неподдерживаемые eager-операции могут выполняться на CPU;
• float64 фактически считается как float32 с предупреждением, и в задачах, которым нужна 64-битная точность, это может изменить численные результаты.

Стабильную версию (GA), SLA поддержки и отдельный стабильный API-контракт AWS пока не объявила.

Насколько TorchNeuron совместим с PyTorch по тестам AWS?

AWS заявляет, что постоянно прогоняет 100 000 моделей с Hugging Face, на которые приходится 99,9% скачиваний. Модели тестируются в уменьшенных репрезентативных конфигурациях: проверяются загрузка, прямой и обратный проход, компиляция и случаи перехода на CPU.

В августе 2026 года AWS провела внутренний прогон upstream-тестов PyTorch и получила около 154 000 результатов. Примерно 69 000 тестов были пропущены, в том числе около 39 000 - намеренно для Neuron. Около 22 000 из них касались операций, для которых ещё нет нативной реализации (lowering) под Neuron.

AWS также отправила в основной репозиторий PyTorch серию RFC и pull request, которые должны убрать жёсткую привязку к CUDA и конкретному железу в тестах. Не все изменения уже приняты: как минимум PR #193439 на момент проверки оставался открытым.

Amazon утверждает, что 85% исследователей и разработчиков ИИ пишут код на PyTorch. Ни методика, ни выборка, ни исходное исследование в публикации не приводятся. Эту цифру нельзя трактовать как долю всех ИИ-приложений.

Сможет ли TorchNeuron конкурировать с CUDA от Nvidia?

В AWS прямо называют цель релиза - уменьшить зависимость разработчиков от одной платформы.

«Главное здесь - устранить привязку, из-за которой у разработчиков вообще не было реального выбора», - сказал Колин Брейс.

«У разработчиков должен быть одинаковый опыт работы на разных аппаратных платформах», - заявила Яна ван Грёнен, директор по разработке PyTorch в Meta.

В закрытой бете TorchNeuron использовали более 30 компаний и внутренних команд Amazon. Среди названных пользователей и партнёров - Cartesia, Deepgram, Lieber Institute for Brain Development, Hugging Face, Anthropic, Databricks, Meta, Rivian и Salesforce.

Редактор TechCrunch Джули Борт называет TorchNeuron самой прямой атакой Amazon на программное преимущество Nvidia и CUDA. При этом, по её мнению, пока неясно, даст ли отвязка PyTorch от Nvidia реальный выбор и экономию или только добавит разработчикам проблем с настройкой. Сооснователь GeekWire Тодд Бишоп пишет, что именно CUDA и накопленный вокруг неё софт мешают компаниям переходить на альтернативные чипы. Он также отмечает, что крупнейшие пользователи Trainium пока в основном финансово связаны с Amazon - среди них Anthropic и OpenAI. Автор TechsCurrent считает бету подходящей для контролируемых тестов переносимости, но не универсальной заменой CUDA: то, что модель запустилась, ещё не доказывает её производительность или экономическую выгоду.

Какие результаты показал тест Deepgram на Trainium?

Deepgram и AWS провели совместный тест: модель распознавания речи Nova-3 на одном инстансе trn2.3xlarge обрабатывала на 11-70% больше аудио на доллар, чем вычислительные конфигурации, которые Deepgram использовала раньше. Это результат заинтересованных сторон, полученный на пакетной нагрузке, где важна пропускная способность. При переносе применялись готовые NKI-ядра, а инженеры AWS помогали с оптимизацией.

Независимых воспроизводимых сравнений TorchNeuron с CUDA на репрезентативном наборе production-нагрузок к моменту публикации не найдено.

Сколько стоит работа с TorchNeuron на AWS EC2 Trn2?

У самого TorchNeuron нет лицензионной платы: он распространяется под Apache 2.0. Платить нужно за инфраструктуру EC2. По текущему прайсу EC2 Capacity Blocks инстанс trn2.3xlarge с одним чипом Trainium2 стоит $2,235 за инстанс-час в регионах Hyderabad, Melbourne и São Paulo. Инстанс trn2.48xlarge с 16 чипами Trainium2 - $35,7608 за инстанс-час в Ohio и Hyderabad.

Это ставки Capacity Blocks, а не стандартная цена On-Demand. ОС и другие ресурсы могут оплачиваться отдельно. Полный список регионов и On-Demand-цены для всех вариантов Trn2 и Trn3 в материалах запуска не приведены.

Доступен ли TorchNeuron в России?

Исходный код TorchNeuron открыт и доступен для скачивания. Но чтобы запустить его на Trainium, нужна инфраструктура AWS. В России у AWS нет ни региона, ни дата-центров, ни офисов. С марта 2022 года компания не регистрирует новые аккаунты из России и Беларуси. При этом AWS упоминала, что у некоторых зарубежных клиентов есть команды разработчиков в России. Сможет ли конкретный существующий аккаунт пользоваться Trn2 или Trn3, зависит от его статуса и доступных зарубежных регионов.