AWS Trainium

AWS Trainium menjalankan PyTorch secara native melalui paket torch-neuronx, yang mendaftarkan perangkat neuron. Neuron mengompilasi model Anda menjadi satu atau lebih graf untuk setiap bentuk input yang dilihatnya, sehingga menjaga bentuk input tetap statis menghindari kompilasi ulang.

Instalasi

Instal torch-neuronx dengan mengikuti panduan setup AWS Neuron. Kemudian periksa bahwa PyTorch melihat perangkat Neuron.

import torch

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

Ia mencetak True dan jumlah perangkat Neuron, misalnya True 4.

Jalankan inferensi

device_map="auto" tidak mendeteksi Neuron. Penempatan otomatis bergantung pada Accelerate untuk mendaftar perangkat yang tersedia, dan Accelerate tidak mendaftar perangkat Neuron, sehingga model dimuat secara diam-diam di CPU. Sampaikan device_map="neuron" secara eksplisit ke from_pretrained(), atau device="neuron" ke Pipeline.

Lewatkan input generasi di CPU alih-alih memanggil .to(model.device). generate() hanya memindahkan tensor yang dikonsumsi oleh forward ke perangkat Neuron. Tensor yang bertambah setiap langkah (urutan output, attention_mask) tetap di CPU, sehingga graf yang dikompilasi tetap stabil. generate mengembalikan output di perangkat yang sama dengan input. Lihat Keep generation tensors on CPU untuk detail lebih lanjut.

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

Dengan cache statis, generate() mengompilasi forward pass secara otomatis menggunakan torch.compile di Neuron. Sampaikan disable_compile=True ke generate() untuk memilih keluar.

Latih dengan Trainer

TrainingArguments mendeteksi Neuron dan secara otomatis mengatur perangkat ke neuron. Semua generasi Trainium mendukung bf16.

from transformers import TrainingArguments

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

Compile model

Set torch_compile=True untuk membungkus model dengan torch.compile. Ini mengompilasi jalur maju (forward pass) model, sementara langkah optimizer berjalan secara eager. TrainingArguments mengatur torch_compile_backend menjadi "neuron" secara otomatis saat mendeteksi Neuron.

from transformers import TrainingArguments

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

Setiap bentuk input baru memicu kompilasi baru. Isi input hingga panjang tetap (padding="max_length" di tokenizer) dan atur dataloader_drop_last=True agar batch terakhir tidak memiliki ukuran yang berbeda.

Pelatihan terdistribusi

Neuron mendukung API terdistribusi Transformers. Buat DistributedConfig dengan ukuran paralelisme yang Anda inginkan dan teruskan ke from_pretrained(). Transformers memecah model saat dimuat, dan Trainer menggunakannya apa adanya.

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),
)

Jalankan skrip dengan torchrun.

torchrun --nproc_per_node=4 train.py

Lihat Paralelisme tensor, Paralelisme pakar, dan FSDP untuk opsi yang tersedia.

Keterbatasan

• device_map="auto" memuat model di CPU alih-alih Neuron. Sampaikan device_map="neuron" secara eksplisit.
• Bentuk dinamis memicu kompilasi ulang. Proses generasi dan pelatihan sama-sama melambat ketika bentuk input berubah dari satu langkah ke langkah berikutnya. Gunakan cache statis untuk generasi dan padding dengan panjang tetap untuk pelatihan.