AWS Trainium

Trn3 UltraServers は、スケールアップ UltraServer テクノロジーにおける最新のイノベーションを体現しており、NeuronSwitch-V1 を使用して
最大 144 枚の Trainium3 チップにわたって全対全のコレクティブを高速化しています。まとめると、1 台の Trn3 UltraServer で
最大 20.7 TB の HBM3e、706 TB/秒のメモリ帯域幅、362 FP8 PFLOP を実現し、最大 4.4 倍のパフォーマンスを発揮し、
Trn2 UltraServer の 4 倍以上のエネルギー効率を実現しています。Trn3 はトレーニングと推論のコストを最小限に抑えながら最も高いパフォーマンスを発揮し、
最新の 1T 以上のパラメータ (MoE) と推論タイプモデルを備え、Trainium2 ベースのインスタンスと比較して GPT-OSS を大規模に提供する場合のスループットが
大幅に向上しています。

Trn2 UltraServers は、最大 1T のパラメータのモデルの生成 AI トレーニングと推論のための高性能で費用対効果の高いオプションで
あり続けています。Trn2 インスタンスには最大 16 個の Trainium2 チップが搭載され、Trn2 UltraServer には独自のチップ間相互接続である
NeuronLink に接続された最大 64 個の Trainium2 チップが搭載されています。

Trn1 インスタンスには最大 16 個の Trainium チップが搭載され、最大 3 つの FP8 PFLOP、9.8 TB/秒のメモリ帯域幅を備えた 512 GB の HBM、
および最大 1.6 Tbps の EFA ネットワークを実現します。

研究と実験用に構築

AWS Neuron SDK は、Trn3 インスタンス、Trn2 インスタンスと Trn1 インスタンスから最大限のパフォーマンスを引き出すのに役立つため、モデルの構築とデプロイに集中し、市場投入までの時間を短縮できます。AWS Neuron は、PyTorch、Jax、および Hugging Face、vLLM、PyTorch Lightning などの重要なライブラリとネイティブに統合されています。プロファイリングとデバッグのための深いインサイトを提供しながら、すぐに分散型トレーニングと推論用にモデルを最適化します。AWS Neuron は、Amazon SageMaker、Amazon SageMaker Hyerpod、Amazon Elastic Kubernetes Service (Amazon EKS)、Amazon Elastic Container Service (Amazon ECS)、AWS ParallelCluster、AWS Batch などのサービスのほか、Ray (Anyscale)、Domino Data Lab、Datadog などのサードパーティーサービスと統合されています。

精度の目標を達成しながら高いパフォーマンスを実現するために、AWS Trainium は
BF16、FP16、FP8、MXFP8、MXFP4 などのさまざまな混合精度データタイプをサポートしています。生成 AI による急速なイノベーションを支援するために、
Trainium2 と Trainium3 は、4 倍のスパース性 (16:4)、マイクロスケーリング、確率的丸め、専用の集合エンジン向けの
ハードウェア最適化を実現しています。

Neuron を使用すると、デベロッパーはカーネル開発に Neuron Kernel Interface (NKI) を使用してワークロードを最適化できます。NKI は完全な Trainium ISA を公開しているため、命令レベルのプログラミング、メモリ割り当て、実行スケジューリングを完全に制御できます。デベロッパーは、独自のカーネルを構築するほかに、最適化されたカーネルをすぐにデプロイできるオープンソースの Neuron Kernel Library を利用できます。そして最後に、Neuron Explore はスタック全体を可視化し、デベロッパーのコードからハードウェアのエンジンに至るまで接続します。