Перейти к основному содержимому
← Все записи

AWS выпускает TorchNeuron — нативный бэкенд PyTorch для Trainium с поддержкой Hugging Face

TorchNeuron делает Trainium нативным устройством PyTorch. Библиотеки Hugging Face уже работают на нём, а Amazon заявляет, что также работает над ослаблением зависимости PyTorch от Nvidia.

Amazon хочет убедить вас, что самая сложная часть отказа от Nvidia — изменить одну строку конфигурации. Возможно, это даже правда — и именно поэтому анонс интересен. AWS выпустила TorchNeuron — плагин для PyTorch, который превращает её чип Trainium в обычное устройство PyTorch, а Hugging Face заявляет, что её основные библиотеки теперь работают на нём напрямую. Обещание — простая миграция. Подтекст — настоящим барьером для конкурентов Nvidia всегда были не чипы, а привычки в работе с ПО.

Что такое TorchNeuron на самом деле

TorchNeuron — плагин для PyTorch, добавляющий новый бэкенд для Trainium. Он поддерживает режим eager, torch.compile и встроенные API распределённых вычислений PyTorch, включая FSDP и тензорный параллелизм. Neuron отображается как устройство — так же, как GPU.

Инструмент выйдет в публичную бету 20 октября 2026 года. Исследователи и разработчики смогут установить его командой pip install torch-neuronx.

По словам Amazon, которые приводят СМИ, приложения на PyTorch можно запускать на Trainium, изменив одну строку конфигурации и не переписывая код. Примеры Hugging Face честнее раскрывают детали, и до них мы ещё доберёмся. Утверждение об одной строке отлично подходит для заголовка. А вот настоящая работа скрывается в сносках.

Почему мишенью стала Nvidia

PyTorch — бесплатный инструментарий с открытым исходным кодом для создания, обучения и запуска моделей ИИ. Его создала Facebook, а с 2022 года он находится под управлением Linux Foundation. По данным Amazon, на нём написано около 85% приложений ИИ.

Nvidia годами добивалась того, чтобы её программная платформа CUDA хорошо работала с PyTorch. В результате оборудование Nvidia стало выбором по умолчанию — прекрасная позиция для производителя оборудования. Никто не выбирает вариант по умолчанию. Его просто никогда не меняют.

Amazon не первая пытается потеснить Nvidia. Пользователи PyTorch уже могут переносить приложения на GPU AMD или чипы Apple, изменив конфигурацию. Так что новизна здесь не в самой идее, а в том, кто её реализует и что к ней прилагается.

Что предлагают: не просто чип

Amazon утверждает, что перенос на Trainium даёт ИИ-стартапам доступ ко всему облаку AWS, а не только к более дешёвому ускорителю. Заодно компания критикует неооблака, многие из которых привязаны к оборудованию Nvidia. В качестве примеров были названы Coreweave, Nebius и Firmus.

Колин Брейс, вице-президент по технологическому ПО в AWS Annapurna Labs — подразделении, которое разрабатывает собственные чипы Amazon, — сформулировал это так: «Вы строите на базе AWS. Вы не делаете ставку на узкоспециализированного провайдера вычислительных ресурсов и не пытаетесь в спешке выстроить вокруг него всё остальное». Он добавил, что инфраструктура для перехода от прототипа к промышленной эксплуатации уже есть.

Экономический аргумент состоит из двух частей:

  • Один чип, две задачи. Trainium используется и для обучения (когда модель учат), и для инференса (когда готовую модель пускают в дело — например, она отвечает на вопрос пользователя). AWS заявляет, что единая платформа делает выполнение всей рабочей нагрузки экономически выгоднее.
  • Соотношение цены и производительности, а не цена по прайсу. AWS признаёт, что её цена за токен может быть не самой низкой. Компания утверждает, что интеграция с остальной облачной инфраструктурой всё равно снижает итоговый счёт.

В качестве доказательства приводится Deepgram — поставщик голосовых ИИ-технологий и конкурент ElevenLabs. По словам Брейса, Deepgram перешла на Trainium с использованием нативного PyTorch и теперь обрабатывает на 70% больше аудио на каждый доллар, сохраняя точность, достаточную для промышленной эксплуатации. Это пример одного клиента, приведённый представителем самого поставщика, поэтому стоит считать его отдельным наблюдением, а не бенчмарком.

Ставка на апстрим

Более амбициозная часть вообще не связана с Trainium. Amazon заявляет, что вносит изменения в основной проект PyTorch, вместе с сообществом устраняя код, который привязывает фреймворк к конкретным чипам. По словам Amazon, цель — сделать PyTorch более независимым от оборудования для всего сообщества разработчиков, а не только для пользователей Trainium.

Это хорошая новость для всех, кто продаёт оборудование или облачные сервисы, конкурирующие с Nvidia. Правда, с самим назначением PyTorch это сочетается не без противоречий. Значительная часть ценности фреймворка — в том, что он скрывает неприятные подробности запуска ИИ-кода на оборудовании. Даст ли ослабление зависимости от Nvidia больше выбора и более дешёвые варианты, как обещает Amazon, или лишь добавит головной боли с конфигурацией — ещё предстоит выяснить.

Сам Брейс резюмировал это так: «Главное здесь — устранить привязку к поставщику, из-за которой у разработчиков вообще не было настоящего выбора». Если эти изменения попадут в основной проект, для Nvidia это станет более серьёзной проблемой, чем любой отдельный чип конкурента.

Что меняется для пользователей Hugging Face

До сих пор Hugging Face работала на Trainium и Inferentia через библиотеку optimum-neuron. Она опиралась на пакет torch-neuronx, построенный на PyTorch XLA. В своё время это был разумный выбор, но он заставлял мыслить как компилятор. Вы получали ленивые тензоры вместо немедленного выполнения, только статические формы и повторную компиляцию при каждом разрыве графа. Optimum Neuron скрывала эту сложность, а для этого поддерживала собственный специализированный код, включая классы моделей и тренеры.

TorchNeuron устраняет необходимость в этой прослойке. Neuron становится нативным устройством с немедленным выполнением, DTensor, FSDP2 и torch.compile. Hugging Face заявляет, что Transformers, Diffusers, Accelerate, PEFT и TRL теперь работают на нём с теми же классами и скриптами, которые вы использовали бы на любом другом ускорителе. Пользователям Neuron компания рекомендует работать с этими библиотеками напрямую, а не через Optimum Neuron.

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

Инференс

В Transformers появилась нативная поддержка генерации текста для широкого спектра моделей. Главная цифра: с TorchNeuron из коробки работают 89% из 100 самых используемых архитектур моделей для генерации текста, по сравнению с 46% в Optimum Neuron v0.4.5.

Чтобы перенести модель на устройства Neuron, нужно передать device_map=neuron. Пример Hugging Face идёт дальше и демонстрирует генерацию текста с тензорным параллелизмом, который, по словам компании, позволяет полностью использовать возможности инстанса Neuron. В рецепте есть несколько составляющих:

  • Группа процессов. Инициализируйте распределённую группу процессов с бэкендом neuron.
  • Тензорный параллелизм. Передайте DistributedConfig с размером TP, равным общему числу процессов, а также DeviceMesh для устройств Neuron.
  • Статический кеш. Используйте статический кеш для генерации: Hugging Face настоятельно рекомендует это на Neuron, чтобы избежать частых повторных компиляций.
  • Дополнение и разбиение на чанки. Дополните промпт слева до длины, кратной размеру чанка предзаполнения (512 в примере), и задайте максимальную длину кеша (1024).
  • Настройка компиляции. Используйте CompileConfig с бэкендом neuron, fullgraph=False и dynamic=False.

В примере указана модель Qwen/Qwen3.8-27B, загружаемая в формате bfloat16.

Обратите внимание, что подразумевает рекомендация о статическом кеше. Предпочтения компилятора никуда не исчезли — они просто переехали из специализированной библиотеки во флаги конфигурации. Это прогресс, но не магия.

Обучение

В части обучения Hugging Face обещает, что вам больше не нужны ни специальный тренер, ни отдельный форк кода модели для Neuron. Стандартный скрипт дообучения на TRL запускается на инстансе Trainium с несколькими настройками, специфичными для Neuron.

В примере используется скрипт дообучения с учителем на базе SFTTrainer из TRL. За параллелизм отвечает сам Transformers. В model_init_kwargs передаётся DistributedConfig с размерами групп тензорного параллелизма и FSDP, SFTTrainer передаёт его в from_pretrained, и модель шардируется по мере загрузки её весов. Trainer видит, что модель уже распределена, и не оборачивает её повторно, поэтому не нужны ни флаг --fsdp, ни конфигурация FSDP для Accelerate.

Для запуска используется torchrun — по одному процессу на NeuronCore. В примере TP_SIZE=4 и FSDP_SIZE=8 дают 32 процесса, организованных в сетку 4 на 8. Тензорный параллелизм распределяет каждый слой по 4 ядрам группы TP, а FSDP шардирует параметры, градиенты и состояние оптимизатора между 8 группами TP.

В этом запуске Qwen/Qwen3-8B дообучается на наборе данных HuggingFaceH4/ultrachat_200k (часть train_sft) в формате bfloat16. Среди настроек — максимальная длина 2048, упаковка последовательностей в режиме wrapped, скорость обучения 2e-5, одна эпоха, размер батча 4 на устройство и 4 шага накопления градиентов. Hugging Face выполнила запуск на 8 чипах Trainium2 (32 NeuronCores), что соответствует половине инстанса trn2.48xlarge.

Запуск завершается, и trainer.save_model записывает стандартную контрольную точку Transformers. В ней ничто не привязано к Neuron, так что её можно как есть отправить на Hub и загрузить на другом оборудовании. Запасные выходы, оказывается, — тоже фича.

Сноски

Сам скрипт не привязан к Trainium, но Hugging Face перечисляет несколько особенностей, которые нужно учитывать при запуске на нём:

Тема Что делать
Статические формы Отключите автоматические динамические формы в torch.compile и используйте упаковку в режиме wrapped для последовательностей длиной до 2048 токенов. В TRL 1.16 и новее батчи на Neuron дополняются до длины, кратной 128 токенам, что ограничивает число форм входных данных. Скомпилированные графы повторно используются на последующих шагах с такими же формами.
Компиляция Передайте --torch_compile true. На инстансе Trainium torch.compile определяет устройство и автоматически выбирает бэкенд Neuron.
Загрузка весов Встроенная в Transformers асинхронная загрузка весов не поддерживается на Neuron, поэтому пока задайте HF_DEACTIVATE_ASYNC_LOAD=1.

В скрипте также есть комментарий о том, что бэкенд Neuron пока не умеет компилировать автоматические динамические формы Dynamo, и ссылка на открытый issue в TRL. Так что у «одной строки» есть компания.

Diffusers

Многие пайплайны Diffusers совместимы с TorchNeuron как в режиме eager, так и в режиме компиляции. FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 и H3 также работают с тензорным параллелизмом, что позволяет полнее использовать возможности инстанса Neuron.

В примере Qwen-Image-2.1 генерирует изображение размером 1024 на 1024 за 40 шагов инференса со степенью тензорного параллелизма 8. Несколько деталей:

  • Текстовый энкодер остаётся на CPU, а полученные эмбеддинги переносятся на устройство Neuron.
  • Повторяющиеся блоки трансформера компилируются с бэкендом neuron, с параметрами fullgraph=True и dynamic=False.
  • Начальное значение генератора случайных чисел задаётся на CPU, чтобы каждый ранг получал один и тот же шум.
  • Запуск выполняется командой NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py.

Hugging Face заявляет, что поддержка тензорного параллелизма появится и в других пайплайнах, и указывает на руководство по её добавлению для понравившегося вам пайплайна, если такой поддержки пока нет.

Вычислительные ядра

Hugging Face также обращает внимание на вычислительные ядра от команды Neuron: они и дальше будут появляться на Hub для ускорения моделей. Их можно писать с помощью Neuron Kernel Interface (NKI), который даёт более низкоуровневый контроль над операциями. Пользователи также могут писать собственные операции и делиться ими на Hub.

Одно из опубликованных ядер, как утверждается, ускоряет трансформеры для компьютерного зрения на Trainium до 2,4 раза. Hugging Face поблагодарила команду AWS за работу над TorchNeuron и за быструю реакцию на обратную связь в ходе интеграции.

Где источники расходятся

Стоит отметить два небольших расхождения.

  • Сроки. В одном источнике говорится, что AWS выпустила инструмент в понедельник, 19 октября. Hugging Face сообщает, что TorchNeuron станет доступен в публичной бете с 20 октября. Самое вероятное объяснение — в один день состоялся анонс, а на следующий вышла бета, но источники этого не уточняют.
  • Трудозатраты. Amazon описывает переход как изменение одной строки конфигурации без переписывания кода. В собственных примерах Hugging Face фигурируют группа процессов, сетка устройств, распределённые конфигурации, статические кеши, настройки форм и переменная окружения. Оба утверждения могут быть верны, если одна строка — это базовый вариант, а всё остальное — настройка для масштабирования, но разработчику стоит ожидать, что для любого распределённого сценария понадобится больше одной строки.

Вердикт — какой есть

В этом релизе три слоя. Есть плагин — он существует и доступен для установки. Есть интеграция с Hugging Face — конкретная, документированная, с измеримым охватом архитектур (89% против 46%). И есть работа над основным проектом, призванная сделать PyTorch менее заточенным под Nvidia, — пока это обещание.

Первые два упрощают знакомство с Trainium. Именно третье может придать всему этому значение. Конкуренты Amazon в сегменте неооблаков напрямую заинтересованы в противоположном исходе, а Nvidia давно получила фору в качестве выбора по умолчанию. Настройки по умолчанию живучи, потому что никто их не замечает, а теперь Amazon просит разработчиков обратить на них внимание.

Если всё сработает, лучший результат — аргумент device_map, о котором никто не задумывается. Если нет, второй по привлекательности исход — очень хорошо документированный список переменных окружения.

Источники

В сыре есть дырки. В вашей инфраструктуре их быть не должно.