AWS Trainium

AWS Trainium インスタンスは、torch-neuronx パッケージを介して PyTorch をネイティブで実行し、neuron デバイスを登録します。Neuron は、検出する各入力形状に対してモデルを1つ以上のグラフにコンパイルするため、入力形状を静的に保つことで再コンパイルを回避できます。

インストール

AWS Neuron セットアップガイドに従ってtorch-neuronxをインストールします。その後、PyTorchがNeuronデバイスを確認できることを確認してください。

import torch

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

TrueとNeuronデバイスの数が出力されます。例えば、True 4のように表示されます。

推論を実行する

device_map="auto" は Neuron を検出しません。自動配置は Accelerate が利用可能なデバイスを列挙する機能に依存していますが、Accelerate は Neuron デバイスを列挙しないため、モデルは CPU にサイレントにロードされます。from_pretrained() に device_map="neuron" を明示的に渡すか、Pipeline に device="neuron" を渡してください。

.to(model.device) を呼び出す代わりに、生成の入力を CPU に渡します。generate() は、forward が消費するテンソルのみを Neuron デバイスに移動します。各ステップで増加するテンソル（出力シーケンス、attention_mask）は CPU に保持されるため、コンパイルされたグラフは安定した状態を維持します。generate は、入力と同じデバイスで出力を返します。詳細については、生成テンソルを CPU に保持するを参照してください。

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

静的キャッシュを使用すると、generate() は Neuron 上で torch.compile によるフォワードパスのコンパイルを自動的に実行します。generate() に disable_compile=True を渡すことで、この動作を無効にできます。

Trainerで学習

TrainingArgumentsはNeuronを検出し、デバイスを自動的にneuronに設定します。すべてのTrainium世代がbf16をサポートしています。

from transformers import TrainingArguments

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

モデルをコンパイルする

torch_compile=True を設定すると、モデルが torch.compile でラップされます。これによりモデルのフォワードパスがコンパイルされ、オプティマイザーステップはイージリーに実行されます。Neuron が検出された場合、TrainingArguments は torch_compile_backend を自動的に "neuron" に設定します。

from transformers import TrainingArguments

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

新しい入力形状のたびに新しいコンパイルがトリガーされます。入力を固定長にパディングし（トークナイザーでpadding="max_length"）、最後のバッチが異なるサイズにならないようにdataloader_drop_last=Trueを設定してください。

分散学習

NeuronはTransformersの分散APIをサポートしています。DistributedConfigを作成し、希望する並列サイズを指定してfrom_pretrained()に渡します。Transformersはモデルの読み込み時にシャード化を行い、Trainerはそれをそのまま使用します。

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

スクリプトをtorchrunで起動します。

torchrun --nproc_per_node=4 train.py

利用可能なオプションについては、Tensor parallelism、Expert parallelism、およびFSDPを参照してください。

制限事項

device_map="auto" を指定すると、モデルは Neuron ではなく CPU にロードされます。device_map="neuron" を明示的に指定してください。
動的な形状は再コンパイルをトリガーします。ステップごとに入力の形状が変化すると、生成と学習の両方が遅くなります。生成には静的なキャッシュを、学習には固定長のパディングを使用してください。

Update on GitHub