TorchNeuron: запуск моделей HF на AWS Trainium стал проще, чем когда-либо

Вернуться к статьям

Опубликовано 19 октября 2026 г.

TorchNeuron предоставляет новый нативный бэкенд для PyTorch на AWS Trainium. Transformers, Diffusers и Kernels нативно его поддерживают, что упрощает запуск моделей на Trainium и их публикацию на Hub Hugging Face.

TorchNeuron

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

Начиная с 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 во избежание частых перекомпиляций.

import os

import torch
import torch.distributed as dist
import torch_neuronx  
from torch.distributed import DeviceMesh
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
from transformers.distributed import DistributedConfig
from transformers.generation.configuration_utils import CompileConfig

MODEL_ID = "Qwen/Qwen3.8-27B"
MAX_SEQ_LEN = 1024
PREFILL_CHUNK_SIZE = 512

dist.init_process_group(backend="neuron")
world_size = dist.get_world_size()

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=world_size),
    device_mesh=DeviceMesh("neuron", list(range(world_size))),
).eval()

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"

inputs = tokenizer.apply_chat_template(
    [{"role": "user", "content": "The future of AI is"}],
    add_generation_prompt=True,
    padding=True,
    tokenizer_kwargs={"pad_to_multiple_of": PREFILL_CHUNK_SIZE},
    return_tensors="pt",
)

generation_config = GenerationConfig(
    max_new_tokens=128,
    do_sample=False,
    use_cache=True,
    pad_token_id=tokenizer.pad_token_id,
    eos_token_id=tokenizer.eos_token_id,
    cache_implementation="static",
    max_cache_len=MAX_SEQ_LEN,
    prefill_chunk_size=PREFILL_CHUNK_SIZE,
    compile_config=CompileConfig(backend="neuron", fullgraph=False, dynamic=False),
)

with torch.no_grad():
    output_ids = model.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        generation_config=generation_config,
    )

if dist.get_rank() == 0:
    prompt_len = inputs["input_ids"].shape[1]
    print(tokenizer.decode(output_ids[0, prompt_len:], skip_special_tokens=True))

Обучение

Теперь, когда Neuron является нативным устройством PyTorch, обучение на Trainium использует те же библиотеки Hugging Face, что и любой другой ускоритель: Transformers, Accelerate, PEFT и TRL. Вам больше не нужны специализированный тренер или версия кода модели для Neuron. Стандартный скрипт тонкой настройки TRL запускается на экземпляре Trainium всего с несколькими настройками, специфичными для Neuron.

Ниже приведён скрипт обучения с учителем на основе SFTTrainer из TRL. Если вы раньше дообучали модель с помощью TRL на GPU, он покажется вам знакомым.

import os

import torch
from datasets import load_dataset
from transformers.distributed import DistributedConfig

from trl import ModelConfig, ScriptArguments, SFTConfig, SFTTrainer, TrlParser


def main(script_args, training_args, model_args):
    
    torch._dynamo.config.automatic_dynamic_shapes = False

    training_args.model_init_kwargs = {
        "revision": model_args.model_revision,
        "trust_remote_code": training_args.trust_remote_code,
        "attn_implementation": model_args.attn_implementation,
        "dtype": model_args.dtype,
        "device_map": None,
    }
    if int(os.environ.get("WORLD_SIZE", "1")) > 1:
        
        training_args.model_init_kwargs["distributed_config"] = DistributedConfig(
            tp_size=int(os.environ.get("TP_SIZE", "1")),
            fsdp_size=int(os.environ.get("FSDP_SIZE", "1")),
        )

    
    train_dataset = load_dataset(
        script_args.dataset_name,
        name=script_args.dataset_config,
        split=script_args.dataset_train_split,
    )
    
    train_dataset = train_dataset.select_columns(["messages"])

    
    trainer = SFTTrainer(
        model=model_args.model_name_or_path,
        args=training_args,
        train_dataset=train_dataset,
    )

    trainer.train()
    trainer.save_model(training_args.output_dir)
    trainer.end()


if __name__ == "__main__":
    parser = TrlParser((ScriptArguments, SFTConfig, ModelConfig))
    script_args, training_args, model_args = parser.parse_args_and_config()
    main(script_args, training_args, model_args)

Параллелизм нативно поддерживается Transformers. Вы добавляете DistributedConfig с размерами тензорного параллелизма (TP) и FSDP в model_init_kwargs, SFTTrainer передаёт его в from_pretrained, а модель распределяется по устройствам во время загрузки весов. Тренер видит, что модель уже распределена, и не оборачивает её повторно. Вам не нужны флаг --fsdp или конфигурация FSDP для Accelerate.

В скрипте нет ничего привязанного к Trainium: тот же код работает на любом оборудовании. При запуске на Trainium стоит учитывать несколько моментов:

Статические формы. Скрипт отключает автоматические динамические формы в torch.compile и использует упаковку с обёрткой для формирования последовательностей длиной до 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.

TP_SIZE=4 FSDP_SIZE=8 HF_DEACTIVATE_ASYNC_LOAD=1 uv run torchrun --nproc_per_node=32 sft_neuron_example.py \
    --model_name_or_path Qwen/Qwen3-8B \
    --dataset_name HuggingFaceH4/ultrachat_200k \
    --dataset_train_split train_sft \
    --output_dir Qwen3-8B-sft-tp4-fsdp8 \
    --dtype bfloat16 \
    --max_length 2048 \
    --packing true \
    --packing_strategy wrapped \
    --learning_rate 2e-5 \
    --warmup_steps 0.03 \
    --num_train_epochs 1 \
    --logging_steps 2 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 4 \
    --save_only_model true \
    --torch_compile true

Мы запустили его на 8 чипах Trainium2 (32 NeuronCore), то есть на половине экземпляра 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:

import torch, torch.distributed as dist, torch_neuronx 
from torch.distributed.device_mesh import DeviceMesh 
from diffusers import QwenImage21Pipeline, QwenImage21Transformer2DModel, TensorParallelConfig


dist.init_process_group(backend="neuron") 
dev = torch.neuron.current_device()
mesh = DeviceMesh("neuron", list(range(8))) 


model = "Qwen/Qwen-Image-2.1"
transformer = QwenImage21Transformer2DModel.from_pretrained(
    model,
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
    parallel_config=TensorParallelConfig(mesh=mesh), 
)
pipe = QwenImage21Pipeline.from_pretrained(
    model,
    transformer=transformer,
    torch_dtype=torch.bfloat16,
)
pipe.vae.to(dev)
pipe.transformer.compile_repeated_blocks(backend="neuron", fullgraph=True, dynamic=False)


emb, mask, _ = pipe.encode_prompt(prompt="A capybara wearing a wizard hat, oil painting", device="cpu")  
image = pipe(
    prompt_embeds=emb.to(dev),
    prompt_embeds_mask=mask.to(dev) if mask is not None else None,
    height=1024,
    width=1024,
    num_inference_steps=40,
    generator=torch.Generator("cpu").manual_seed(0), 
).images[0]
if dist.get_rank() == 0:
    image.save("out.png")

Этот фрагмент кода можно запустить с помощью команды: NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py.

Всё больше конвейеров будут совместимы с TP. Если нужный вам конвейер пока не поддерживается, вы также можете добавить поддержку самостоятельно, следуя этому руководству.

Kernels

Обязательно ознакомьтесь с ядрами от команды Neuron, поскольку всё больше ядер Neuron будет становиться доступно на Hub для ускорения ваших моделей. Ядра Neuron можно писать с помощью Neuron Kernel Interface (NKI), который предоставляет более низкоуровневый контроль для оптимизации операций. Благодаря этому пользователи могут запускать собственные операции и делиться ими на Hub.

Например, вы можете использовать это ядро, чтобы ускорить работу vision-трансформеров в 2,4 раза на Trainium!

Благодарности

Особая благодарность команде AWS за работу над TorchNeuron и оперативную реакцию на наши отзывы, благодаря чему интеграция с HF стала максимально качественной.

Другие статьи в нашем блоге