ИИтог

Hugging Face добавила поддержку TorchNeuron для AWS Trainium в Transformers и Diffusers

инфраструктураРедакция ИИтогHugging Face🔥62Рейтинг ИИтог62Влияние55Новизна60Релевантность95Неожиданность30Участники55Достоверность85
Hugging Face добавила поддержку TorchNeuron для AWS Trainium в Transformers и Diffusers

Hugging Face сообщила о поддержке TorchNeuron, нового нативного бэкенда PyTorch для чипов AWS Trainium. TorchNeuron, это плагин для PyTorch, который добавляет бэкенд для Trainium с поддержкой eager-режима (немедленного выполнения), torch.compile и нативных распределённых API, таких как FSDP и TP (тензорный параллелизм). Начиная с 20 октября 2026 года TorchNeuron будет доступен как публичная бета; установить его можно командой pip install torch-neuronx.

Раньше работа с Hugging Face на Trainium и Inferentia шла через библиотеку optimum-neuron, которая опиралась на пакет torch-neuronx на базе PyTorch XLA. По словам авторов, это означало «мыслить как компилятор»: ленивые тензоры вместо немедленного выполнения, только статические формы, перекомпиляция при каждом разрыве графа. Optimum Neuron скрывала эту сложность и держала собственный код (классы моделей, тренеры и т. д.).

TorchNeuron встраивается прямо в PyTorch и делает neuron нативным устройством: с eager-выполнением, штатными распределёнными примитивами PyTorch (DTensor, FSDP2) и torch.compile. Для библиотек Hugging Face Trainium становится обычным устройством: Transformers, Diffusers, Accelerate, PEFT и TRL работают на нём с теми же классами и скриптами, что и на любом другом ускорителе. Hugging Face рекомендует пользователям Neuron отныне работать с этими библиотеками напрямую, а не через Optimum Neuron.

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

Transformers, обучение. Отдельный тренер или Neuron-специфичный форк кода модели больше не нужны: стандартный скрипт дообучения на TRL (SFTTrainer) запускается на Trainium с несколькими специальными настройками. Параллелизм обрабатывает сам Transformers: в model_init_kwargs передаётся DistributedConfig с размерами TP и FSDP, модель шардируется при загрузке весов, флаг --fsdp и конфигурация Accelerate FSDP не нужны. Нюансы запуска: нужны статические формы (автоматические динамические формы в torch.compile отключены; упаковка wrapped даёт последовательности до 2048 токенов, а при TRL ≥ 1.16 батчи на Neuron дополняются до кратности 128 токенов); достаточно --torch_compile true, бэкенд Neuron выбирается автоматически; встроенная асинхронная загрузка весов в Transformers на Neuron не поддерживается, поэтому пока её отключают через HF_DEACTIVATE_ASYNC_LOAD=1. В примере дообучается Qwen3-8B на датасете HuggingFaceH4/ultrachat_200k через torchrun с 32 процессами (TP_SIZE=4 × FSDP_SIZE=8), на 8 чипах Trainium2 (32 NeuronCores), то есть на половине инстанса trn2.48xlarge. Обучение дошло до конца, а trainer.save_model записал обычный чекпойнт Transformers, который можно выложить на Hub и загрузить на любом другом железе.

Diffusers. Многие пайплайны совместимы с TorchNeuron в eager- и скомпилированном режимах. FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 и H3 дополнительно поддерживают тензорный параллелизм. В примере с Qwen-Image-2.1 используется TP=8, текстовый энкодер остаётся на CPU, а запуск идёт через torchrun с NEURON_LOGICAL_NC_CONFIG=2.

Kernels. На Hub будет появляться всё больше ядер от команды Neuron для ускорения моделей. Их можно писать на Neuron Kernel Interface (NKI), который даёт более низкоуровневый контроль над операциями; пользователи могут запускать собственные операции и делиться ими на Hub. Одно из таких ядер ускоряет vision-трансформеры на Trainium до x2,4 (база сравнения не указана). В благодарностях авторы отмечают команду AWS за работу над TorchNeuron и оперативную реакцию на замечания.

Ключевые факты

  • TorchNeuron, плагин PyTorch с нативным бэкендом для AWS Trainium: eager-режим, torch.compile, FSDP и тензорный параллелизм; публичная бета с 20 октября 2026 года, установка через pip install torch-neuronx.
  • Transformers, Diffusers, Accelerate, PEFT, TRL и Kernels работают с Trainium напрямую; Hugging Face рекомендует использовать их вместо Optimum Neuron.
  • С TorchNeuron «из коробки» поддержано 89% из 100 самых используемых архитектур генерации текста против 46% с Optimum Neuron v0.4.5.
  • Дообучение Qwen3-8B через TRL SFTTrainer запущено на 8 чипах Trainium2 (32 NeuronCores) с раскладкой TP=4 × FSDP=8 и завершилось сохранением стандартного чекпойнта.
  • Для Diffusers FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 и H3 поддерживают тензорный параллелизм; ядро Neuron ускоряет vision-трансформеры до x2,4.

Почему это важно

Раньше запуск моделей Hugging Face на Trainium и Inferentia требовал отдельной библиотеки optimum-neuron с собственными классами моделей и тренерами, поверх PyTorch XLA со статическими формами и перекомпиляциями. Теперь Trainium становится обычным устройством PyTorch, а стандартные библиотеки экосистемы работают на нём без Neuron-специфичных форков. Охват архитектур для генерации текста вырос с 46% до 89% из топ-100.

Кому это важно

Разработчикам и исследователям, которые обучают или запускают модели на инстансах AWS Trainium и Inferentia; командам, уже использующим Transformers, TRL, PEFT, Diffusers; авторам кастомных операций, которые хотят писать ядра на NKI и делиться ими на Hub.

Как это применить

После выхода публичной беты 20 октября 2026 года TorchNeuron устанавливается командой pip install torch-neuronx. Для инференса в Transformers достаточно передать device_map=neuron; для генерации рекомендован статический кэш. Для обучения стандартный скрипт TRL запускают через torchrun, задавая TP_SIZE и FSDP_SIZE (их произведение равно числу процессов), включая --torch_compile true и отключая асинхронную загрузку весов через HF_DEACTIVATE_ASYNC_LOAD=1. Для Diffusers пример запускается как NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py. Цены и доступность инстансов в тексте не указаны.

Можно ли доверять

Источник, собственный блог Hugging Face, то есть заявление самого разработчика интеграции. Конкретные цифры (89% против 46%, ускорение до x2,4) приведены без методики; для x2,4 база сравнения не названа. Замеров пропускной способности, задержки, стоимости или скорости обучения в сравнении с GPU или Optimum Neuron в тексте нет. Для обучения сказано лишь, что запуск Qwen3-8B дошёл до конца, метрик качества нет.

Риски и подводные камни

TorchNeuron выходит как публичная бета. Есть известные ограничения: бэкенд пока не умеет компилировать автоматические динамические формы Dynamo, а встроенная асинхронная загрузка весов в Transformers не поддерживается и отключается. Нужны статические формы: батчи дополняются до кратности 128 токенов, а для генерации нужен статический кэш, иначе возможны частые перекомпиляции. Какие именно архитектуры не вошли в 89%, не перечислено. Не все пайплайны Diffusers поддерживают тензорный параллелизм.