# AWS Trainium

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

## インストール

[AWS Neuron セットアップガイド](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/setup/pytorch/index.html)に従って`torch-neuronx`をインストールします。その後、PyTorchがNeuronデバイスを確認できることを確認してください。

```py
import torch

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

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

## 推論を実行する

`device_map="auto"` では Neuron が検出されません。自動配置は [Accelerate](https://huggingface.co/docs/accelerate/index) が利用可能なデバイスを列挙する機能に依存していますが、Accelerate は Neuron デバイスを列挙しないため、モデルは CPU にサイレントにロードされます。[from_pretrained()](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel.from_pretrained) に `device_map="neuron"` を、または [Pipeline](/docs/transformers/main/ja/main_classes/pipelines#transformers.Pipeline) に `device="neuron"` を明示的に指定してください。

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

```py
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))
```

```txt
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()](/docs/transformers/main/ja/main_classes/text_generation#transformers.GenerationMixin.generate) は Neuron 上で `torch.compile` によるフォワードパスのコンパイルを自動的に実行します。[generate()](/docs/transformers/main/ja/main_classes/text_generation#transformers.GenerationMixin.generate) に `disable_compile=True` を渡すことで、この動作を無効にできます。

## Trainerで学習

[TrainingArguments](/docs/transformers/main/ja/main_classes/trainer#transformers.TrainingArguments)はNeuronを検出し、デバイスを自動的に`neuron`に設定します。すべてのTrainium世代がbf16をサポートしています。

```py
from transformers import TrainingArguments

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

### モデルをコンパイルする

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

```py
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](/docs/transformers/main/ja/expert_parallelism#transformers.DistributedConfig)を作成し、希望する並列サイズを指定して[from_pretrained()](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel.from_pretrained)に渡します。Transformersはモデルの読み込み時にシャード化を行い、[Trainer](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer)はそれをそのまま使用します。

```py
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`で起動します。

```bash
torchrun --nproc_per_node=4 train.py
```

利用可能なオプションについては、[Tensor parallelism](./tensor_parallelism)、[Expert parallelism](./expert_parallelism)、および[FSDP](./fsdp)を参照してください。

## 制限事項

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