TorchNeuron: стартирането на HF модели на AWS Trainium е по-лесно от всякога

Актуализиране в GitHub

TorchNeuron предоставя нов собствен бекенд за PyTorch върху AWS Trainium. Transformers, Diffusers и Kernels го поддържат нативно, което прави стартирането на модели върху Trainium и споделянето им в HF Hub по-лесно от всякога.

TorchNeuron

TorchNeuron е плъгин за PyTorch, който добавя нов бекенд за Trainium и позволява поддръжка на eager режим, torch.compile и собствени разпределени API като FSDP и TP. Прочетете статията на PyTorch, за да научите повече за TorchNeuron.

От 20 октомври 2026 г. TorchNeuron ще бъде достъпен като публична бета версия. Изследователи и разработчици ще могат да го инсталират с pip install torch-neuronx.

Нов бекенд, собствен интерфейс в библиотеките на HF

Досега изживяването с Hugging Face върху AWS Trainium и Inferentia — AI чипове, проектирани съответно за рентабилно обучение и инференция в голям мащаб — се предоставяше чрез 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 върху графични процесори, той би трябвало да ви изглежда познат.

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 или 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, тъй като в Hub ще стават достъпни все повече Neuron ядра за ускоряване на вашите модели. Neuron ядрата могат да бъдат написани чрез Neuron Kernel Interface (NKI), който предоставя по-ниско ниво на контрол за оптимизиране на операциите. По този начин потребителите могат да изпълняват собствени операции и да ги споделят в Hub.

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

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

Специални благодарности на екипа на AWS за работата му по TorchNeuron и за бързата реакция на обратната ни връзка, с която направи интеграцията с HF възможно най-добра.

Още статии от нашия блог