2024年AWS公開のCustom AI Training Accelerator第2世代。1.3PFLOPS BF16+96GB HBM3+UltraServer 64chip+Trn2 EC2 Instance+EFA v3+Neuron SDK 2.20を搭載。
AWS Trainium2は2024年12月のre:Invent 2024でAWSが発表したAmazon自社設計のAI学習向けカスタムアクセラレータ第2世代。前世代Trainium(2020年12月)から4倍の演算性能と3倍のメモリ容量を実現し、大規模言語モデル(LLM)の学習・ファインチューニングに最適化された設計。Anthropic Claude 3.5 Sonnet・Claude 3.7 Sonnet等の大規模モデル学習で実際に採用され、Project Rainierとして40万個のTrainium2チップを使った専用クラスターをAnthropic向けに構築中。NVIDIA H100/H200の代替としてコスト効率30-40%改善を謳い、AI学習基盤の選択肢を多様化する象徴的な存在となった。
| 項目 | Trainium2 | Trainium (前世代) | NVIDIA H100 SXM | Google TPU v5p |
|---|
| プロセス | TSMC N5 | TSMC N7 | TSMC N4 | TSMC N5 |
| BF16性能 | 1.3 PFLOPS | 0.21 PFLOPS | 0.99 PFLOPS | 0.46 PFLOPS |
| メモリ | 96GB HBM3 | 32GB HBM2e | 80GB HBM3 | 95GB HBM2e |
| 帯域幅 | 2.9 TB/s | 0.82 TB/s | 3.35 TB/s | 2.76 TB/s |
| TDP | ~350W | ~400W | 700W | ~270W |
| ネットワーク | NeuronLink 640GB/s | NeuronLink 192GB/s | NVLink 900GB/s | ICI 3D-Torus |
| 提供形態 | AWS EC2のみ | AWS EC2のみ | OEM/Cloud | Google Cloudのみ |
| 価格(参考) | $21.50/hr Trn2 | $13.10/hr Trn1 | $98.32/hr p5.48xl | $11.04/hr v5p-8 |
Trainium2はAWS EC2でのみ利用可能なクラウド専用ハードウェアで、自作PCには搭載できない。代わりに、AWS Neuron SDKを使ってTrn2インスタンス上でPyTorch/TensorFlow/JAXコードを実行する形となる。NVIDIA CUDAとの最大の違いは、PyTorchのtorch_neuronxプラグインを介した自動コンパイル方式で、CUDAコードの書き換えなしでTrainium2移行が可能(ただしカスタムCUDAカーネルは要書き換え)。コスト面ではNVIDIA H100(p5.48xlarge $98.32/hr)比でTrn2.48xlarge($21.50/hr)は78%安価で、同等規模のLLM学習で大幅なコスト削減が可能。一方、PyTorch/Hugging Face transformersは公式対応するが、JAX独自のFlash Attention等の最新研究実装はNVIDIAより数か月遅れる傾向あり。学習基盤の選択肢として、Anthropic等の大手LLM企業がTrainium2を主力採用しているのは大規模学習でのコスト優位性が決定打となっている。
AWS Inferentia2との違い: Inferentia2は推論専用、Trainium2は学習専用という分業設計。Inferentia2は190 TFLOPS BF16+32GB HBM、Trainium2は1.3PFLOPS+96GB HBMで性能差は大きい。同じNeuron SDK上で動作し、コード移植は容易。 NVIDIA H100との違い: H100は汎用CUDA Ecosystem+柔軟なソフトウェアスタック、Trainium2はAWS Neuron SDK限定+大規模クラスタ最適化。NVIDIAはPCIe/SXM等で多様な接続形態、Trainium2はEC2インスタンスとして提供される閉鎖的環境。
Q1: Trainium2とNVIDIA H100でどちらを選ぶべき？ A: 大規模LLM学習(70B+)+コスト重視ならTrainium2。研究・実験+最新CUDA研究実装ならH100。Anthropic等の大手はコスト優位性でTrainium2に流れているが、PyTorch標準コード以外を多用する場合はH100が無難。
Q2: Neuron SDKの学習曲線はどれくらい？
A: PyTorch標準コードならtorch_neuronxプラグイン+xla_device()の数行追加で動作する。ただし、カスタムCUDAカーネル(Triton等)を多用する場合はTrainium2向けに書き換えが必要で、ここで2-4週間の開発工数が発生する。
Q3: Project Rainierクラスターの規模はどれほどか？ A: 40万個のTrainium2チップで構成、520 EFLOPS BF16級、世界最大規模のAI学習基盤。Anthropic Claude次世代モデル(Claude 3.7/4.0系)の学習専用に構築中で、2025年完成予定。総投資額は数十億ドル規模と推定される。