Amazon Web Services は火曜日のre:Inventで、Trainium3 と呼ばれる次世代 AI アクセラレータを初公開しました。同社によれば、このアクセラレータは来年末にリリースされる時点で、前世代機よりも 4 倍高いパフォーマンスを発揮するという。
詳細はまだほとんど明らかになっていないが、re:Invent に先立ってThe Registerのインタビューに応じた AWS の Annapurna Labs チームの製品および顧客エンジニアリング担当ディレクターの Gadi Hutt 氏は、Trainium3 が 3nm プロセスノード上に構築される初の専用機械学習アクセラレータとなり、論文発表から 1 年後に一般提供が開始される Trainium2 と比較して 40 パーセントの効率向上を達成すると予想している。これについては後ほど詳しく説明する。
パフォーマンスに関しては、Amazonは実際のパフォーマンス数値について曖昧な情報を提供しています。Trainium3の4倍のパフォーマンス向上は、完全な「UltraServer」構成に基づいていますが、これはまだ開発中とのことです。
わかっているのは、合計64個のアクセラレータを搭載したTrainium2 UltraServerが、83.2ペタフロップスの高密度FP8性能を発揮するということです。したがって、理論上、Trainium3 UltraServerは332.8ペタフロップスの演算性能を発揮するはずですが、その精度は不明です。
AWS に再度問い合わせて詳細を確認しましたが、推測するに、6 ビットまたは 4 ビットの浮動小数点演算を検討していると思われます。これは、Nvidia が Blackwell で市場に投入し、AMD が来年中に MI355X で導入する予定のものです。
スパース性を考慮すると、Trainium3 もそのプロセッサと同じ 4 倍の乗数をサポートしていると仮定すると、Amazon の次世代 UltraServer は 1.3 エクサフロップスを超える AI コンピューティングを実現できる可能性があります。
また、これらのパフォーマンスに関する主張は、ピーク時の演算性能（FLOPS）を指しており、漠然としたAIベンチマークではないことも保証されています。これは重要な詳細です。AIワークロードによっては、パフォーマンスはFLOPSだけでなく、様々な要因に左右されるからです。例えば、メモリ帯域幅の増加は、大規模言語モデル（LLM）推論性能の大幅な向上につながる可能性があります。これは、NVIDIAの帯域幅強化型H200チップで既に確認されています。
アマゾンはパフォーマンスと効率の指標を公表することには前向きだが、チップのメモリ搭載量についてはまだ詳細を明らかにしていない。
推測するに、Amazon が次世代 AI ASIC を発表する準備が整った頃に、この部分に関するより詳しい情報が得られるだろう。
Trainium2が戦闘準備完了
Trainium3 の詳細を待つ間、Amazon は第 2 世代の Trainium コンピューティング サービスを一般市場に投入する予定です。
昨年の re:Invent で紹介された Trainium2 は、その名前にもかかわらず実際にはトレーニングと推論の両方を行うチップであり、1.3 ペタフロップスの高密度 FP8 コンピューティングと、それぞれ 2.9 TBps の帯域幅を提供できる 96 ギガバイトの高帯域幅メモリを備えています。
参考までに、単一の Nvidia H100 は、約 2 petaFLOPS の高密度 FP8 パフォーマンス、80 GB の HBM、および 3.35 TBps の帯域幅を誇ります。
チップ自体は、TSMC のチップオンウェーハオンサブストレート (CoWoS) パッケージング技術を使用して統合された 5nm コンピューティングダイのペアと 4 つの 24GB HBM スタックで構成されています。
アンナプルナ研究所のTrainium2加速器を詳しく見てみましょう - クリックして拡大
GoogleのTensor Processing Unit（TPU）と同様に、これらのアクセラレータはラックスケールのクラスタにまとめられています。64個のTrainium2コンポーネントが、相互接続された2つのラックに分散配置されています。
前述したように、この Trn2 UltraServer 構成は、高密度 FP8 パフォーマンスで 83.2 ペタFLOPS、または 4 倍スパース モードを有効にすると 332.8 ペタFLOPS を出力できます。
2つのラックに64個のTranium2チップを搭載したAWSの新しいTrn2 UltraServerを詳しく見てみましょう。 - クリックして拡大
これ以上のコンピューティング能力が必要であれば、Amazon では 16 個のアクセラレータと約 20.8 ペタフロップスの高密度コンピューティングを備えた Trn2 インスタンスも提供しています。
Amazon によれば、これらのインスタンスは、EC2 で利用可能な現在の世代の GPU ベースのインスタンス (具体的には Nvidia H200 ベースの P5e および P5en ベースのインスタンス) に比べて 30 ～ 40 パーセント優れた価格性能を提供します。
モデルの学習にチップを使用する場合、Trainium2は10万個以上のチップを搭載したさらに大規模なクラスターに拡張できます。これはまさにAWSとモデルビルダーのAnthropicがProject Rainierで計画していることで、AI学習用のクラスターに「数十万個」のTrainium2チップを投入する予定です。このスーパーコンピューティングシステムは2025年にフル稼働が予定されており、「[Anthropicの]最新世代AIモデルの学習に使用されるエクサフロップスの5倍」の演算能力を発揮できると言われています。
Trn2インスタンスは現在AWSの米国東部（オハイオ）で利用可能で、近日中に他のリージョンでも利用可能になる予定です。また、より大規模なTrn2 UltraServer構成は現在プレビュー版としてご利用いただけます。
- GenAIは、かつて自動化から「安全」と考えられていた仕事に参入する
- バイデン政権、AIアクセラレーターに不可欠なHBMチップの中国による購入を禁止
- 「オープン」なAIの主張はしばしば明らかな嘘であると研究は主張している
- AIへの野望が銅を限界まで押し上げている
賭けをヘッジする
AWSのAnnapurna Labsチームはカスタムシリコンの開発を推進していますが、すべてを一つのバスケットに詰め込んでいるわけではありません。このクラウド大手は既にH200、L40S、L4アクセラレータなど、幅広いインスタンスをサポートしており、Project Ceibaの下でBlackwellの大規模クラスターを展開中です。
Nvidia の Grace-Blackwell スーパーチップ (GB200) をベースにしたこの大規模な AI スーパーコンピューターは、約 20,736 個の Blackwell GPU を誇り、各 GPU は 800 Gbps (スーパーチップあたり 1.6 Tbps) の Elastic Fabric Adapter 帯域幅で接続されます。
このマシンは、合計で約414エクサフロップスの超低精度スパースFP4演算能力を発揮すると予想されています。ただし、FP4精度はほぼ推論にのみ使用され、学習にはより高精度なFP8とFP/BFが使用されます。学習においては、Ceibaは依然として驚異的な51エクサフロップスの高密度BF16演算能力を発揮すると予想されており、FP8まで落とせばその2倍の性能を発揮するでしょう。
いずれにせよ、AWS は Trainium シリコンを推進しているかもしれませんが、Nvidia との取り組みはまだ終わっていません。®
