# AWS Trainium Amazon Web Services(AWS)が提供する AI トレーニング向けアクセラレータ。Trainium 1(EC2 インスタンス trn1.32xlarge)と Trainium 2(trn2.48xlarge)の2世代がある。単一コアあたりの理論ピーク性能は、Trainium 1 が PeakBW 440.2 GB/s・PeakMM 23.75 TFLOPS・PeakVec 286.8 GFLOPS、Trainium 2 が PeakBW 640.0 GB/s・PeakMM 19.75 TFLOPS・PeakVec 550.0 GFLOPS([[@2026__MLSys2026__AccelOpt - A Self-Improving LLM Agentic System for AI Accelerator Kernel Optimization]] Table 4)。1コアはテンソルエンジン・ベクタエンジン・スカラエンジン・GPSIMD エンジンが並行動作し、オンチップメモリ経由で HBM と通信する構成を持つ。 Trainium は [[Neuron Kernel Interface (NKI)]] という Python 埋め込みカーネル言語でプログラムされる。ハードウェアとプログラミングモデルの双方が比較的新しいため、GPU のような成熟した最適化レシピや性能ヒューリスティクスが確立されておらず、開発者は限られた性能直感しか持たない([[@2026__MLSys2026__AccelOpt - A Self-Improving LLM Agentic System for AI Accelerator Kernel Optimization]] §1)。 ## 「TPUのファストフォロワー」という位置づけ [[AI Chip Architectures]]は、Trainiumの中核設計を「[[Google TPU]]の思想(128×128 weight-stationaryシストリックアレイ・ソフトウェア管理スクラッチパッド・全プログラムコンパイル)をほぼそのまま再実装したもの」と位置づけ、コンパイラもGoogleのOpenXLAを共用する(neuronx-ccがXLA HLOグラフをNEFFバイナリへ低レベル変換)点を明示する。独自性はTensor Engineの周囲に集まる: Vector Engine(層正規化・softmax等の縮約)・Scalar Engine(要素ごとの演算)・GPSIMD Engine(8基のプログラマブルベクトルプロセッサ、C言語で記述)という3つの補助エンジンと、専用の集団通信ハードウェア**CC-Cores**(Trn2で20個/チップ、NeuronLinkポートに直結してall-reduce/all-to-all等をTensor/Vector Engineと並行実行)。GPU上ではNCCLカーネルが行列演算と同じSM上で通信と計算を奪い合うのに対し、Trainiumはこれを専用シリコンへ切り出している。Trn1/Trn2はTPUと同じトーラストポロジ(NeuronLink、4×4 2Dトーラス→4×4×4 3Dトーラス)を採用したが、Trn3ではMoEのall-to-allトラフィックがトーラスの最悪ケースになることから、スイッチ型all-to-allファブリックNeuronSwitch-v1へ転換した(TPUがBoardflyへ転換したのと同じ動機)。スケールアウトはAWS既存のEFA/SRD(Nitroオフロード、最大64パスへのフロー分散でヘッドオブラインブロッキングを回避)。Anthropicは Project Rainier(2025年後半に稼働開始、約50万個のTrainium2、2026年初頭には100万個超)でClaudeの訓練・推論の双方にTrainiumを用いており、独自のNKIカーネルをNeuronスタックへ直接アップストリームしている点が、GPUの成熟したCUTLASSエコシステムに対するTrainiumのソフトウェア面での若さを補う実例とされる。(Source: [[AI Chip Architectures]]) ## 関連 - [[Amazon Web Services]] — 提供元 - [[Neuron Kernel Interface (NKI)]] — Trainium 向けカーネル言語 - [[NKIBench]] — Trainium カーネルのベンチマークスイート - [[AIアクセラレータ]] — AI アクセラレータ全般の横断概念 - [[Google TPU]] — 設計思想の元になったアーキテクチャ - [[@2026__MLSys2026__AccelOpt - A Self-Improving LLM Agentic System for AI Accelerator Kernel Optimization]] — 初出ソース - [[AI Chip Architectures]] — Trn2/Trn3世代、CC-Cores、NeuronSwitch、Project Rainier ## 出典 - [[@2026__MLSys2026__AccelOpt - A Self-Improving LLM Agentic System for AI Accelerator Kernel Optimization]] - [[AI Chip Architectures]]