AWS Trainium2（エーダブリューエストレイニアムツー）
Amazon Web Services が開発した第2世代 AI 学習専用チップ Trainium2。前世代比で性能 4 倍、エネルギー効率 2 倍を実現し、EC2 UltraCluster で最大 10 万チップの超大規模学習クラスタを構成できる。
AWS Trainium2 とは
AWS Trainium2 は、Amazon Web Services が AI モデルの学習ワークロード専用に設計した第2世代カスタムチップです。2024年に発表され、前世代 Trainium1 と比較して学習性能 4 倍、エネルギー効率 2 倍を達成しています。EC2 UltraCluster 構成では最大 10 万チップを束ねた超大規模分散学習が可能で、兆パラメータ規模のモデル学習に対応します。
アーキテクチャと主な特徴
- NeuronCore-v3: 第3世代ニューロンコア、1 チップあたり 2 コア搭載で並列演算能力を強化
- 演算精度: FP32 / BF16 / FP8 / cFP8（Configurable FP8）をサポート
- メモリ: 96 GB HBM3e 搭載、帯域幅 3.2 TB/s で大規模パラメータの高速アクセス
- NeuronLink: チップ間高速インターコネクト、Trn2 サーバー内 16 チップを直接接続
- EFA v3（Elastic Fabric Adapter）: ノード間 3.2 Tbps のネットワーク帯域
- UltraCluster: 最大 100,000 チップを Petabit 級ファブリックで接続する大規模構成
- ソフトウェア: AWS Neuron SDK（PyTorch / JAX 対応）、SageMaker Training 統合
スペック比較
| チップ | 世代 | HBM 容量 | ピーク BF16 | チップ間帯域 | 最大クラスタ |
|---|---|---|---|---|---|
| Trainium1 | 第1世代 | 32 GB HBM2e | 210 TFLOPS | 800 Gbps | 30,000 チップ |
| Trainium2 | 第2世代 | 96 GB HBM3e | 840 TFLOPS | 1.6 Tbps | 100,000 チップ |
| NVIDIA H100 | Hopper | 80 GB HBM3 | 989 TFLOPS | 900 GB/s NVLink | 32,768 GPU |
| Google TPU v6e | Trillium | 32 GB HBM3 | 918 TFLOPS | ICI | 256 チップ |
EC2 インスタンスタイプ
- trn2.48xlarge: Trainium2 × 16 チップ、768 GB アクセラレータメモリ、192 vCPU
- trn2n.48xlarge: trn2.48xlarge + EFA v3 強化ネットワーキング（UltraCluster 構成用）
主な採用事例
- Anthropic: Claude 3.5 / Claude 4 シリーズの事前学習に UltraCluster を使用
- Meta: Llama 3 の学習ワークロードの一部を Trainium2 で実行
- Amazon 社内: Alexa LLM、Amazon Q、Rufus（ショッピング AI）の学習
- スタートアップ: AI21 Labs、Databricks が SageMaker 経由で利用
NVIDIA GPU との比較
Trainium2 は NVIDIA H100 と同等クラスの BF16 演算性能を持ちながら、AWS エコシステム内での TCO を最大 50% 削減できると AWS は主張しています。ただし、CUDA エコシステムの豊富なライブラリ群（cuDNN、NCCL、DeepSpeed 等）が使えないため、Neuron SDK への移植コストが発生します。
よくある質問（FAQ）
Q1: Trainium2 と Inferentia2 の違いは何ですか？ A: Trainium2 はモデル学習に最適化された ASIC で、Inferentia2 は推論に特化しています。Trainium2 は HBM 容量が大きく勾配計算に必要なメモリ帯域を確保、Inferentia2 はレイテンシとスループットのバランスを重視した設計です。
Q2: 既存の PyTorch コードをそのまま Trainium2 で動かせますか？ A: AWS Neuron SDK が PyTorch の torch_xla バックエンドとして動作するため、多くのケースでコード変更は最小限です。ただしカスタム CUDA カーネルや一部の動的制御フローは対応が必要です。
Q3: Trainium2 のコストパフォーマンスは NVIDIA H100 より良いですか？ A: AWS の公式ベンチマークでは、Llama 3.1 70B の学習で H100 比 30-40% のコスト削減を報告しています。ただし、ワークロードの種類や最適化度合いによって差は変動します。
まとめ
- AWS Trainium2 は第2世代 AI 学習専用 ASIC、前世代比 4 倍の性能
- 96 GB HBM3e + NeuronLink で大規模モデル学習に対応
- UltraCluster で最大 10 万チップの超大規模分散学習が可能
- Anthropic Claude や Meta Llama の学習に採用実績あり