TorchNeuron предоставляет новый встроенный бэкенд для PyTorch на AWS Trainium. Библиотеки Transformers, Diffusers и Kernels поддерживают его «из коробки», что делает запуск моделей на Trainium и их публикацию в хабе HF проще, чем когда-либо.
TorchNeuron
TorchNeuron — это плагин для PyTorch, который добавляет новый бэкенд для Trainium, обеспечивая поддержку режима eager, torch.compile и встроенных распределенных API, таких как FSDP и TP. Прочитайте blog по PyTorch, чтобы узнать больше о TorchNeuron.
Начиная с 20 октября 2026 года, TorchNeuron будет доступен в режиме публичной бета-версии. Исследователи и разработчики смогут установить его с помощью команды pip install torch-neuronx.
Новый бэкенд, нативный интерфейс в библиотеках HF
До сих пор взаимодействие с Hugging Face на AWS Trainium и Inferentia — чипах ИИ, созданных соответственно для экономичного обучения и инференса в масштабе — осуществлялось через optimum-neuron. Эта библиотека опирается на пакет torch-neuronx, который основан на PyTorch XLA. В то время это был естественный выбор, но он требовал мышления компилятора: ленивые тензоры вместо eager-выполнения, только статические формы, рекомпиляция при каждом разрыве графа. Optimum Neuron существовал для того, чтобы скрывать эту сложность, и поддерживал для этого собственный код (классы моделей, тренеры и т. д.).
TorchNeuron меняет это. Он интегрируется напрямую в PyTorch и представляет neuron как нативное устройство с eager-выполнением, собственными примитивами распределения PyTorch (DTensor, FSDP2) и torch.compile.
Для библиотек Hugging Face Trainium становится устройством, как и любое другое: Transformers, Diffusers, Accelerate, PEFT и TRL запускаются на нем напрямую, используя те же классы и скрипты, которые вы применяли бы на любом другом ускорителе. Отныне мы рекомендуем пользователям Neuron работать с этими библиотеками напрямую, а не через Optimum Neuron.
Transformers
Инференс
Генерация текста теперь нативно поддерживается в Transformers для широкого спектра моделей. С TorchNeuron поддерживается 89% из 100 самых популярных архитектур моделей генерации текста «из коробки» по сравнению с 46% в Optimum Neuron v0.4.5. Вам нужно лишь передать аргумент device_map=neuron в своем коде, и ваша модель будет перенесена на устройства Neuron.
Ниже приведен пример генерации текста с тензорным параллелизмом, который позволит в полной мере задействовать возможности вашего экземпляра Neuron. Обратите внимание, что здесь для генерации используется статический кэш, так как это настоятельно рекомендуется для Neuron во избежание частых рекомпиляций.
Обучение
Теперь, когда Neuron является нативным устройством PyTorch, обучение на Trainium использует те же библиотеки Hugging Face, что и любой другой ускоритель: Transformers, Accelerate, PEFT и TRL. Вам больше не нужен специализированный тренер или форк кода модели, специфичный для Neuron. Стандартный скрипт дообучения TRL запускается на экземпляре Trainium всего с несколькими настройками, специфичными для Neuron.
Вот скрипт контролируемого дообучения (supervised fine-tuning), созданный на базе SFTTrainer из TRL. Если вы раньше дообучали модели с помощью TRL на графических процессорах, он покажется вам знакомым.
Параллелизм обрабатывается нативно силами Transformers. Вы указываете DistributedConfig с размерами тензорного параллелизма (TP) и FSDP в model_init_kwargs, SFTTrainer передает его в from_pretrained, и модель разбивается на части (шардируется) по мере загрузки ее весов. Тренер видит, что модель уже распределена, и не оборачивает ее повторно. Вам не нужны флаг --fsdp или конфигурация Accelerate FSDP.
Ничто в скрипте не привязано к Neuron: тот же код выполняется на любом оборудовании. При запуске на Trainium полезно знать о нескольких вещах:
- Статические формы. Скрипт отключает автоматические динамические формы в torch.compile и использует обернутую упаковку (wrapped packing) для создания последовательностей длиной до 2048 токенов. Начиная с TRL ≥ 1.16, батчи на Neuron дополняются до кратного 128 токенам значения, что ограничивает количество входных форм. Скомпилированные графы повторно используются для последующих шагов с теми же формами.
- Включенный torch.compile. Достаточно передать --torch_compile true: на экземпляре Trainium torch.compile обнаруживает устройство и автоматически выбирает бэкенд Neuron без необходимости дополнительных настроек.
- Отключение асинхронной загрузки весов. Встроенная асинхронная загрузка весов в Transformers не поддерживается на Neuron, поэтому пока мы отключаем ее с помощью HF_DEACTIVATE_ASYNC_LOAD=1.
Команда ниже запускает обучение с помощью torchrun, который инициирует один процесс на один NeuronCore. TP_SIZE и FSDP_SIZE задают схему параллелизма, а их произведение равно количеству процессов: здесь 4 x 8 = 32, что соответствует параметру --nproc_per_node. 32 процесса образуют сетку 4 x 8. Тензорный параллелизм распределяет каждый слой между 4 ядрами группы TP, а FSDP шардирует параметры, градиенты и состояние оптимизатора между 8 группами TP.
Мы запустили его на 8 чипах Trainium2 (32 NeuronCores), что составляет половину экземпляра trn2.48xlarge.
Обучение доходит до конца, и trainer.save_model записывает стандартную контрольную точку (чекпоинт) Transformers. В ней нет ничего привязанного к Neuron: вы можете отправить ее в Hugging Face Hub как есть и загрузить на любом другом оборудовании.
Diffusers
Инференс
Многие пайплайны из библиотеки Diffusers совместимы с TorchNeuron в eager-режиме и режиме компиляции. FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 и H3 также совместимы с тензорным параллелизмом, чтобы вы могли максимально эффективно использовать свои устройства Neuron. Ниже приведен пример, показывающий, как можно использовать TP:
Вы можете запустить этот фрагмент кода с помощью команды: NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py.
Все больше пайплайнов будут получать поддержку TP. Если нужный вам пайплайн еще не поддерживается, вы также можете добавить эту поддержку, следуя данному руководству.
Ядра (Kernels)
Обязательно ознакомьтесь с ядрами от команды Neuron, поскольку на Hub будет появляться все больше ядер Neuron для ускорения ваших моделей. Ядра Neuron могут быть написаны с использованием Neuron Kernel Interface (NKI), который предоставляет более низкоуровневый контроль для оптимизации операций. Таким образом пользователи могут запускать собственные кастомные операции и делиться ими в хабе.
Вы можете использовать, например, это ядро для ускорения ваших Vision Transformers до 2.4 раз на Trainium!
Благодарности
Особая благодарность команде AWS за их работу над TorchNeuron и за оперативный отклик на наши отзывы, позволивший сделать интеграцию с HF максимально качественной.
