AWS Trainium2（エーダブリューエストレイニウムツー）
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学習基盤の選択肢を多様化する象徴的な存在となった。
主な特徴・仕組み
- Annapurna Labs設計: AWS子会社Annapurna Labs(2015年買収)が設計、TSMC N5プロセスで製造
- 演算性能: 1.3 PFLOPS BF16(BFloat16)、650 TFLOPS FP32、Trainium比4倍向上
- HBM3メモリ: 96GB HBM3、メモリ帯域幅2.9 TB/s、Trainium比3倍容量
- NeuronLink: 専用インターコネクト、640 GB/sチップ間帯域、UltraServer内Hot Path形成
- Trn2.48xlarge Instance: 16x Trainium2搭載EC2インスタンス、20.8PFLOPS BF16、$21.50/hr
- Trn2 UltraServer: 64x Trainium2(4 servers)を高速接続、83PFLOPS BF16、$172/hr
- EFA v3 Networking: Elastic Fabric Adapter v3、3.2 Tbps Inter-instance帯域、SRD2プロトコル
- Neuron SDK 2.20: PyTorch/TensorFlow/JAX対応、TorchTitan統合、Hugging Face transformers動作
- 省電力設計: 1チップ約350W、H100(700W)比で半分、データセンター熱密度に有利
- Project Rainier: 40万Trainium2チップで構成される世界最大級AI cluster、Anthropic Claude次世代モデル学習用
- Trainium3予告: 2025年末提供予定、TSMC N3、Trainium2比2倍性能・40%省電力
スペック比較表
| 項目 | 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 |
具体例・対応製品
- Trn2.48xlarge (16x Trainium2): 価格$21.50/hr、20.8PFLOPS BF16、Anthropic Claude学習基盤として正式採用
- Trn2 UltraServer (64x Trainium2): 価格$172/hr、83PFLOPS BF16、世界最大級トレーニング規模
- Project Rainier (400,000x Trainium2): Anthropicとの共同プロジェクト、全Trainium2 Cluster世界記録
- Inf2.48xlarge (Inferentia2 12chip): 推論専用、Trainium2の姉妹機、$13.81/hr
- 競合機種: NVIDIA H100/H200(p5.48xlarge $98.32/hr)、Google TPU v5p(v5p-8 $11.04/hr)、Intel Gaudi 3(DL1 $13.11/hr)
自作PCでの選び方・注意点
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インスタンスとして提供される閉鎖的環境。
よくある質問（FAQ）
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年完成予定。総投資額は数十億ドル規模と推定される。
まとめ
- AWS Trainium2 = 2024年12月公開Amazon自社設計AI学習チップ第2世代
- 1.3PFLOPS BF16+96GB HBM3+TSMC N5、Trainium比4倍性能
- Trn2.48xlarge $21.50/hr、H100比78%コスト削減
- Anthropic Claude学習基盤として正式採用、Project Rainier 40万チップ
- AWS EC2専用、Neuron SDK経由でPyTorch/JAX対応