AWS Trainium

AWS Trainium örnekleri, neuron cihazını kaydeden torch-neuronx paketi aracılığıyla PyTorch'u yerel olarak çalıştırır. Neuron, gördüğü her giriş şekli için modelinizi bir veya daha fazla grafiğe derler; bu nedenle giriş şekillerini statik tutmak yeniden derlemeyi önler.

Kurulum

AWS Neuron kurulum kılavuzunu izleyerek torch-neuronx paketini yükleyin. Ardından PyTorch'un Neuron cihazlarını gördüğünü kontrol edin.

import torch

print(torch.neuron.is_available(), torch.neuron.device_count())

True ve Neuron cihazlarının sayısını, örneğin True 4 olarak yazdırır.

Çıkarım çalıştır

device_map="auto" Neuron’u algılamaz. Otomatik yerleştirme, kullanılabilir cihazları listelemek için Accelerate’e dayanır ve Accelerate Neuron cihazlarını listelemediğinden model sessizce CPU’ya yüklenir. Bunu önlemek için from_pretrained() çağrısına açıkça device_map="neuron" veya Pipeline çağrısına device="neuron" parametresini iletin.

Girişleri .to(model.device) çağrısı yapmak yerine CPU üzerinde iletin. generate() yalnızca forward tarafından tüketilen tensörleri Neuron cihazına taşır. Her adımda büyüyen tensörler (çıktı dizisi, attention_mask) CPU üzerinde kalır, böylece derlenmiş grafik kararlılığını korur. generate, çıktıyı girdilerle aynı cihazda döndürür. Daha fazla ayrıntı için Üretim tensörlerini CPU üzerinde tutma bölümüne bakın.

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", dtype="bfloat16", device_map="neuron"
)

# leave the inputs on CPU
inputs = tokenizer("The French Bread Law states", return_tensors="pt")
output = model.generate(
    **inputs, do_sample=False, max_new_tokens=50, cache_implementation="static"
)
print(tokenizer.decode(output[0], skip_special_tokens=True))

The French Bread Law states that the French government must provide a certain amount of bread for the citizens of the country. The law is passed by the French parliament. The French parliament is a body of the French government. The French government is a body of the French state. The

Statik bir önbellekle, generate() Neuron üzerinde ileri geçişi otomatik olarak torch.compile ile derler. Bu davranıştan vazgeçmek için generate() fonksiyonuna disable_compile=True parametresini iletin.

Trainer ile Eğitim

TrainingArguments, Neuron'u algılar ve cihazı otomatik olarak neuron olarak ayarlar. Tüm Trainium nesilleri bf16'yi destekler.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
)

Modeli derle

torch_compile=True değerini ayarlayarak modeli torch.compile ile sarmalayın. Bu işlem, modelin ileri geçişini derlerken optimizasyon adımının aceleci (eager) olarak çalışmasını sağlar. TrainingArguments, Neuron algıladığında torch_compile_backend değerini otomatik olarak "neuron" olarak ayarlar.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
    torch_compile=True,
)

Her yeni giriş şekli yeni bir derlemeyi tetikler. Girdileri sabit bir uzunluğa doldurun (tokenizörde padding="max_length") ve son partinin farklı bir boyuta sahip olmaması için dataloader_drop_last=True ayarlayın.

Dağıtık eğitim

Neuron, Transformers dağıtık API'sini destekler. İstediğiniz paralellik boyutlarıyla bir DistributedConfig oluşturun ve bunu from_pretrained() fonksiyonuna iletin. Transformers, modeli yükleme sırasında bölümlere ayırır ve Trainer bunu olduğu gibi kullanır.

import torch
from transformers import AutoModelForCausalLM, DistributedConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-1.7B",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=2, fsdp_size=2),
)

Betik dosyasını torchrun ile başlatın.

torchrun --nproc_per_node=4 train.py

Mevcut seçenekler için Tensor parallelism, Expert parallelism ve FSDP konularına bakın.

Kısıtlamalar

• device_map="auto" modeli Neuron yerine CPU üzerinde yükler. device_map="neuron" parametresini açıkça belirtin.
• Dinamik şekiller yeniden derlemeyi tetikler. Girdi şekilleri adım adım değiştiğinde hem üretim hem de eğitim yavaşlar. Üretim için statik bir önbellek, eğitim için sabit uzunlukta dolgu kullanın.