TorchNeuron: запуск моделей HF на AWS Trainium став простішим, ніж будь-коли

Оновити на GitHub

TorchNeuron надає новий нативний бекенд для PyTorch на AWS Trainium. Transformers, Diffusers і Kernels нативно його підтримують, завдяки чому запуск моделей на Trainium та їх поширення на HF Hub стали простішими, ніж будь-коли.

TorchNeuron

TorchNeuron — це плагін PyTorch, який додає новий бекенд для Trainium із підтримкою eager mode, 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, а модель розподіляється на частини під час завантаження ваг. Trainer бачить, що модель уже розподілена, і не обгортає її повторно. Вам не потрібні прапорець --fsdp або конфігурація Accelerate FSDP.

У скрипті немає нічого, що було б прив'язане до Trainium: той самий код працює на будь-якому апаратному забезпеченні. Під час запуску на 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.

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, оскільки на Hub з'являтиметься дедалі більше ядер Neuron для прискорення ваших моделей. Ядра Neuron можна писати за допомогою Neuron Kernel Interface (NKI), який надає нижчорівневий контроль для оптимізації операцій. Таким чином користувачі можуть запускати власні операції та ділитися ними на Hub.

Наприклад, ви можете використати це ядро, щоб прискорити роботу ваших vision transformer у 2,4 раза на Trainium!

Подяки

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

Більше статей у нашому блозі