re:Invent 2025: AWS Trainium3 UltraServersの性能とAnthropicの大規模AI基盤構築事例
はじめに
海外の様々な講演を日本語記事に書き起こすことで、隠れた良質な情報をもっと身近なものに。そんなコンセプトで進める本企画で今回取り上げるプレゼンテーションはこちら！
re:Invent 2025 の書き起こし記事については、こちらの Spreadsheet に情報をまとめています。合わせてご確認ください
📖 re:Invent 2025: AWS re:Invent 2025 - AWS Trn3 UltraServers: Power next-generation enterprise AI performance(AIM3335)
この動画では、AWSのInferentiaとTrainiumチップのEC2プロダクトマネージャーJoe Senerschiaが、TrainiumチーフアーキテクトRon Diamantと、AnthropicのTrainium推論リードJonathan Grayを招き、AIインフラストラクチャの構築について解説しています。Trainium3チップの発表を中心に、360ペタフロップスのFP8コンピュート、20テラバイトのHBM容量、144チップまでスケールするUltraServersなど具体的なスペックが紹介されました。Ronはmicroscaling量子化やaccelerated SoftMax命令などのマイクロアーキテクチャ最適化により、持続パフォーマンスをピークに近づける設計思想を説明し、GPT-3.5ベンチマークでTrainium2比5倍のメガワットあたりトークン生成を実現したことを示しました。Jonathan GrayはAnthropicがProject Rainierで100万個のTrainium2チップ上でClaudeモデルを提供している実例を挙げ、Neuron Kernel InterfaceとNeuron Explorerを使った実際のカーネル最適化手法を詳述しました。
※ こちらは既存の講演の内容を最大限維持しつつ自動生成した記事になります。誤字脱字や誤った内容が記載される可能性がありますのでご留意下さい。
本編
AIがもたらす地殻変動的シフト：科学からソフトウェアエンジニアリングまで
皆さん、ようこそ。私の名前はJoe Senerschiaです。InferentiaとTrainiumチップのEC2プロダクトマネージャーを務めています。皆さんにお越しいただき、大変嬉しく思います。では、ちょっと手を挙げていただきたいのですが、InferentiaとTrainiumをご存知の方はどれくらいいらっしゃいますか?なるほど。では、Anthropic Claudeモデルはいかがでしょうか?なるほど、もう少し多いですね。さて、本日は大変嬉しいことに、この両方の分野のエキスパートをお招きしています。TrainiumのチーフアーキテクトであるRon Diamantと、AnthropicでTrainium推論のリードを務めるJonathan Grayです。彼はTrainium上でClaudeモデルを最適化することについて考えています。
では、簡単に本日の内容をご紹介します。まず私から、AWSがどのようにAWSのAIインフラストラクチャを構築し、考えているかについてお話しします。次にRonから、パフォーマンス、スケール、使いやすさを実現するためにTrainiumをどのように構築したかについてお話しします。そして、Jonathan Grayに登壇していただき、実際にTrainium上で効果的に動作するよう、さまざまなカーネルをどのように最適化しているかをご覧いただきます。それでは、始めましょう。
まず、なぜAIに関するニュースがこれほど多いのでしょうか?なぜこれほど盛り上がっているのでしょうか?それは、私たちが世界を構築し、デプロイし、そして相互作用する方法において、本当に地殻変動的なシフトだからだと思います。つまり、これは単なる段階的な変化ではありません。AIがそれを可能にしたことで、新しい機能が次々と登場しているのです。そして、Andy Jassyは最近、こう述べました。「私たちは人生で最大の技術変革の始まりにいる」と。
そこで、私が一歩下がって見てみたい分野の一つは、AIが本当に再構築している場所、それは科学分野です。これを見てみると、例えばタンパク質生物学のような分野では、 モデルが今や数分で新しいタンパク質を予測し設計できるようになっています。これは従来は何時間もかかっていたものです。あるいは数学の分野では、AlphaGeometryのようなモデルが オリンピアードレベルで競い合い、形式的証明も解いています。もう一つ、それ自体が独自の科学分野になっている領域が ソフトウェアエンジニアリングです。ここではAIが今や独自の画期的な力となっており、自らコードを解決し、開発し、デプロイし、コード内のバグを解決し、さらには大規模なコードベース全体にわたって推論することもできます。
そして、これはまだ始まりに過ぎません。しかし、これらのイノベーションを総合すると、もはや科学的発見を支援するだけではなくなっています。実際に科学的発見を推進するエンジンになっているのです。では、これが実際にどのように起こっているのか、先ほど述べたソフトウェアエンジニアリングを例に見てみましょう。
時間の経過とともに、私たちは従来のプログラミングを見てきました。そしてこの数年間で、コード補完、チャットベースのプログラミング、あるいはvibe codingとのコラボレーションといったものが普及し始めるのを目にしてきました。しかし、これらが普及し始めるのと同じ時期に、SWE-Bench Liteのようなベンチマークでも記録が更新されるのを見てきました。モデルが実際のGitHubのissueの最大80%を完了したり、SWE-Bench Verifiedのようなより難しいベンチマークでも、一部のモデルが完全な正確性を持って最大50%を完了できるようになっています。
そして、これが可能にするのが、ソフトウェアエンジニアリングにおけるAIの次のフェーズです。つまり、エージェントやエージェントフリート、エージェントクラスターをデプロイして、自律的に動作し、ソフトウェアの問題を解決できるようにすることです。この未来が正確にどのような形になるかはわかりませんが、ソフトウェア開発者がエージェントのフリート全体と密接に協力する世界を想像することができます。そしてこれが実現することで、ソフトウェアの機能を提供するための新しいスピードとスケールが開かれるのです。
AWSの包括的AIインフラストラクチャスタックとTrainium2の急速な展開
では一歩下がって、なぜこれらすべてが重要なのでしょうか。重要なポイントは、これらすべてのAIを支える基盤となるインフラストラクチャなしには、これらのどれも実現しないということです。そしてAWSでは、10年以上かけて、最も包括的で深く統合されたAIスタックを構築してきました。一番上のコンピュートから始めると、幅広いポートフォリオのアクセラレーテッドGPUインスタンス、または最新のInferentiaとTrainiumを提供しており、これらはAIワークロードに対してコスト効率を提供します。
ネットワーク層では、数万、あるいは数十万のチップまでスケールできるUltraClustersをデプロイしており、すべて低レイテンシ、低ジッターのElastic Fabric Adapterで接続されています。そしてストレージでは、高スループットストレージオプションを増やしており、FSx for LustreやS3 Express One Zoneを使って、GPUにデータを供給し続けることができます。以前の10倍の速度でデータにアクセスできるようになりました。
そして重要なのがセキュリティです。これはAWSのすべてのインフラストラクチャにとって重要ですが、AIにとっても重要です。Nitroシステムを備えており、ワークロードを分離して顧客データを保護することができます。そして一番下には、CloudWatchのような管理サービスと可観測性ツールを提供しており、ノードを監視して、健全かつ効率的に動作していることを確認できます。
そして、これらすべてが統合されて、フロンティアモデルのトレーニングと推論のためのフルスタックプラットフォームとなります。そして、私たちがこれを実現できる理由は、 舞台裏で10年以上にわたってシリコンを開発してきたからです。ここの上部にあるNitroシステムでは、現在6世代以上が利用可能になっており、仮想化を専用ハードウェアにオフロードすることで、より高いパフォーマンスとワークロードのより強固な分離を実現しています。
私たちはGravitonを構築しました。これは現在、多数のワークロードと数万のお客様をサポートしています。そして、AIの成長を予測し、InferentiaとTrainiumチップの構築を早期に開始しました。2019年にInferentiaをリリースし、このフルスタック全体でイノベーションを続けています。これは、AIインフラストラクチャの次のフェーズを推進するために本当に重要なことです。
昨年Trainium2を発表した際、チップとそのスペックについて多くお話ししましたが、それがチップだけの話ではないことも示しました。サーバーレベルとネットワークレベルでもたらしているイノベーションについてもです。 チップレベルでは、1,300 FP8 dense teraflopsのようなコンピュートとflopsを推進するイノベーションがあります。 あるいは、サーバーレベルでは、1テラバイト毎秒の接続性を持つNeuronLinkを介して最大64チップまでスケールアップ可能な、初のTrainium2 UltraServerをリリースしました。または、ネットワークレベルでは、これらのチップを数万個、すべてElastic Fabric Adapterで接続して展開しました。
そして実際、これらすべてのエンジニアリングとエンドツーエンドの設計を見ると、それらが一体となることで、これまでやったことのないことを可能にしてくれます。その一部がここに示されているように、 メーカーからチップを受け取ってからお客様の手に届けるまでの時間を短縮することです。ここでは、Trainium2のライフサイクルを通じて、それを70%短縮したことがわかります。その結果、Trainium2を他のAWS AIインスタンスの4倍の速さでランプアップし、 他のどのインスタンスよりも33倍大きな容量のフットプリントにすることができました。そして、その容量はすべて完全に契約されています。
2025年のAIトレンドと次世代システム要件：Trainium3の登場
これこそが、世界最大の公式発表されたコンピュートクラスターであるProject Rainierのようなものを構築するために必要なエンドツーエンドのイノベーションなのです。しかし、より大きなスケールを構築する中で、次に何がスケールするのかに目を向け続けたいと思っています。そこで、お客様と継続的にトレンドを見て、業界がどこに向かっているのかを見ています。 ここでは、2025年に見られたいくつかのトレンドを見ていきます。ポストトレーニングへの重点の高まりです。モデル開発者がモデルを実際の環境に投入し、仮想的に生成されたものであれ、ロボティクスのような実際の実環境であれ、フィードバックを得ようとする中で、強化学習がより重要になってきています。
次に推論モデルがありますが、これらのモデルは少し時間をかけています。レイテンシーは増えますが、複数のステップにわたって推論を行うことで、深い質問に対してより正確な回答を生成できるんです。そして最後に、先ほどソフトウェアのところでお話ししたagenticワークロードに戻りますが、複数のエージェントが自律的に協調し、ツール呼び出しを行って、幅広い問題に対して独立したソリューションを実現しているのが見られます。では、もう少し深く掘り下げて、その影響は何でしょうか？これがAWS AIインフラストラクチャで次に構築するものにどう影響するのでしょうか？
これは本当に、私たちが検討しているいくつかの新しいシステム要件に集約されると思います。チップだけでなくシステムと言っているのは、これがより大きな全体像に関わることだからです。そして、これらのシステムは今、推論モデルがより長いコンテキストで動作するのを見ていますが、コンテキスト長が100万トークンを超えるようになっており、その機能をサポートするシステムが必要です。mixture of expertモデルのサポートも必要で、これは通信負荷が高く、疎に活性化されたmixture of expertモデルがスケールアップドメイン全体で通信を行います。また、事前トレーニング、事後トレーニング、推論に使用できるインフラストラクチャのサポートも必要で、お客様がそれぞれを独立してスケールする際に、利用可能なコンピュートを本当に最適化できるようにします。そして最後に、多数の並行エージェントが自律的に独立して動作できる、非常に高いバッチサイズ、高スループットのシステムのサポートが必要です。
つまり、ここでの重要なテーマは、次世代のAIインフラストラクチャはコンピュートフロップスだけの話ではないということです。それ以上のものなんです。バランスの取れたコンピュートを持つことが重要で、つまりより多くのメモリ、より多くのメモリ帯域幅を持ち、また、モデルがスケールするにつれて幅広いエキスパート並列設計をサポートできる、より大きなスケールアップドメインを持つことです。そしてそれこそが、Trainium3をご紹介できることを嬉しく思う理由なんです。これは、次世代のagenticワークロード、推論ワークロード、そして次世代AIシステムのコンピュート需要を牽引するビデオ生成ワークロードのために構築されたチップです。
そして先ほど申し上げたように、これはチップだけの話ではなく、システムの話でもあります。ですので、Mattの基調講演をご覧になった方はご存知かと思いますが、最近、144チップまでスケールアップするTrainium3 UltraServersを発表しました。
これらの統計をすべて説明することはしませんので、それは私たちのチーフアーキテクトであるRonに任せますが、覚えておいていただきたい重要なことは、これらのそれぞれにイノベーションがあり、それが次世代AIシステムの機能を推進しているということです。それでは、Trainium3について説明してもらうため、Ronにバトンタッチします。
Trainium3のパフォーマンス設計：持続パフォーマンスを重視したアーキテクチャ
さて皆さん、Joe、どうもありがとう。そして皆さん、本日はご参加いただきありがとうございます。 このトークの次のパートでは、私たちがどのようにTrainium3を構築したかについて、もう少し深く掘り下げていきたいと思います。具体的には、高性能で、スケールに対応でき、使いやすいものにするためにどのように構築したかについてです。まずはパフォーマンスから始めましょう。Joeが少し触れたように、パフォーマンスというのは実は単一の指標ではありません。実際には複数の指標の組み合わせなんです。もちろん、コンピュート、つまり1秒あたりの浮動小数点演算がありますが、メモリ帯域幅やメモリ容量、そしてこれらのチップ間を接続するインターコネクトも重要で、これらすべてがバランスよく整っていないと最大のパフォーマンスを達成できないんです。
実は昨年のトークでこれについて詳しく触れていまして、右上にQRリンクがあります。ちなみに、このトーク全体を通して、オフラインで自己学習できる機会があるたびに、右上にQRリンクが表示されます。Trainium3 UltraServersは、これらのパフォーマンス次元のそれぞれにおいて大幅な飛躍を遂げました。 360ペタフロップスのマイクロスケールドFP8コンピュートを実現しました。これが正確に何を意味するのかは後ほど説明します。これはTrainium2 UltraServersと比べて4.4倍です。HBM容量は20テラバイトで3.4倍、そしてHBMメモリ帯域幅は毎秒700テラバイトで、Trainium2 UltraServersの3.9倍です。
また、インターコネクトは2倍高速になっています。そして、ラックの中央にあるこれらのスイッチに注目していただきたいんです。 これらは新しいコンポーネントで、NeuronSwitchesと呼んでいます。これらはTrainium2コンピュートとTrainium3コンピュートスレッドをフルメッシュトポロジーで接続します。各スレッドは単一ホップで他のすべてのスレッドに接続されています。そして、これらはトップレベルのスペックには現れないようなシステム最適化なんですが、実際のワークロードパフォーマンスには絶対的に影響を与えます。なぜなら、異なるトポロジーをデプロイする柔軟性が高まり、各Trainium3デバイスのペア間のレイテンシーを削減し、all-to-all通信において本当に高いパフォーマンスを提供してくれるからです。
all-to-all通信のパフォーマンスを重視する理由、少なくともその理由の一つは、mixture of expertモデル、略してMOEと呼ばれるものです。このようなモデル、MOEモデルでは、異なるエキスパートを異なるチップに配置し、トークンをリアルタイムで関連するエキスパートにルーティングして、その特定のチップだけでコンピュートを実行する傾向があります。そして、これには超高速のall-to-all通信が必要で、それこそがまさにNeuronSwitchesが提供するものなんです。
これで次のポイントに移りますが、それはピークパフォーマンスと持続パフォーマンスの違いです。実際のところ、1スライド前に私が引用した数字について考えてみると、あれらはスペックパフォーマンスの数字、つまりピークパフォーマンスの数字でした。しかし実際には、パフォーマンスの話はそこから始まるだけで、そこで終わるわけではありません。これについて私が思いつく良いアナロジーの一つは、マラソンレースで勝つのはどちらだと思いますか、短距離走者かマラソンランナーか？明らかにマラソンランナーですよね？でも考えてみてください、短距離走者の方がスペック速度、つまりピーク速度は高いんです。ただ、マラソンレース全体を通してそれを維持できないだけなんです。つまり、少なくともいくつかの状況では、短いピークパフォーマンスよりも持続パフォーマンスの方が実際には重要だということがわかります。そしてAIチップにおいても、実際に同じことが言えます。私たちは、特定のスペック数値よりも、達成可能で持続可能なパフォーマンスを非常に重視しているんです。
さて、私たちがTrainium3チップの開発を始めたとき、ソフトウェアチームから課題を投げかけられました。 持続的なパフォーマンスをピークパフォーマンスに可能な限り近づけるチップを作るには何が必要か？お金を払った分の浮動小数点演算を全て活用できるようにする、ということです。そしてそれが、マイクロアーキテクチャの改善リストにつながりました。 これらの改善は、パフォーマンスの最後の1パーセントまで引き出すことを目的としています。そして、これがどのようなものか、そして私たちがこのワークロードをいかにエンドツーエンドで最適化しているかを理解していただくために、いくつかの例をご紹介したいと思います。
microscaling技術：低精度トレーニングにおけるモデル精度の維持
まずはmicroscalingから始めましょう。 低精度トレーニングと推論の動機は非常に明確ですよね？純粋に物理の問題です。
より小さなデータ型、つまり低精度のデータ型を使用すれば、より小さな回路で計算を実行でき、チップ上でより小さなデータを移動できます。これにより、より高いパフォーマンスとより優れたエネルギー効率が得られます。しかし、人生の多くの良いことと同様に、少し注意書きが付いてきます。
例えば、高精度のデータ型、 例えばBFloat16から、低精度のデータ型、例えばFP8に単純にキャストするだけだと、実はモデルが完全に壊れてしまうことがわかります。その理由は、BFloat16はFP8と比較して、はるかに高いダイナミックレンジ、つまり表現できる数値の範囲を持っているためです。つまり、大きな数値は無限大にオーバーフローし、小さな数値はゼロに押しつぶされてしまう傾向があります。
これは量子化と呼ばれる技術で修正できます。 ここでは、ABSmaxと呼ばれる量子化技術をお見せしています。テンソル内の最大絶対値を計算し、FP8のダイナミックレンジ全体を正確に捉えられるようにテンソル全体をスケーリングします。これは実際、テンソル内に分布の外れ値という興味深いケースに到達するまでは、かなりうまく機能します。テンソル内の要素の1つが他の値よりも100倍大きいケースを想像してみてください。それが そこにある緑色の要素です。
そのため、テンソルをスケーリングして、緑の要素がFP8で表現可能な最大値にマッピングされるようにするわけですが、そうすると他のすべての要素がゼロまたはゼロに近い値に押しつぶされてしまい、FP8へのキャストや量子化の後、表現能力が完全に失われてしまいます。これもmicroscalingと呼ばれる技術で解決できます。microscalingでは、ABSmax量子化をもう一度行いますが、今回は要素の小さなグループごとに行います。
ここでは、緑と黄色の要素を持つ1つのグループと、オレンジ、ピンク、青の要素を持つもう1つのグループがあります。緑が外れ値であることがわかります。緑の要素は他のどの要素よりもはるかに大きいです。これが何を引き起こすかというと、最初のmicroscalingグループでは、量子化後、緑は表現可能な最大値になり、黄色はゼロに押しつぶされます。しかし、2番目のmicroscalingグループでは、新しい分布でゼロから量子化を行うため、青、ピンク、オレンジの要素は、緑の要素、つまり外れ値の影響を全く受けずに、非常にうまく量子化されていることがわかります。
これがまさにmicroscalingの仕組みであり、低精度のトレーニングと推論においてモデルの精度を保つのに非常に効率的であることが示されています。しかし、microscalingは実行が難しいのです。なぜなら、テンソルを取り出してグループに分割し、各グループでスケーリングを計算し、スケールを算出し、スケールを適用し、そして逆量子化する際にこれらすべてを逆の順序で行う必要があるからです。そこで、Trainium3で私たちが行ったのは、microscaling量子化と逆量子化を完全にオフロードするためにハードウェア回路を構築したことです。
accelerated SoftMax命令とGPT-3.5ベンチマークでの5倍の効率向上
基本的に、コンピュートエンジンに一切のオーバーヘッドをかけることなく、microscaled量子化の精度上のメリットをすべて得ることができます。これは、私がお見せしたピーク値には現れませんが、エンドツーエンドのワークロードパフォーマンスを確実に向上させます。では、別の例を見てみましょう。この場合は、高速化されたsoftmax命令です。背景を説明すると、すべての最新の、あるいはほとんどの最新のAIモデルの中核には、self-attentionと呼ばれる演算子があります。
これは、Claudeなどのモデルが今日のように優れた動作をする要因となっているtransformerアーキテクチャにおけるブレークスルーの1つでした。self-attention計算の中核では、ここでQとKという2つの行列を掛け合わせ、その結果に対してsoftmaxを計算し、最後にそれを別の行列Vと掛け合わせます。タイムラインでお見せすると、行列乗算を行い、その後softmax演算を行い、その後もう1つの行列乗算を行います。
これを複数の計算タイルにわたってパイプライン化すると、非常にクリーンなパイプラインが得られることがわかります。テンソルエンジン、つまり行列乗算を行うエンジン、システムの中で最も貴重なリソースが、常にビジー状態で100%稼働しています。これは素晴らしいことです。では次に、先ほどお話しした最適化、microscaled FP8を適用してみましょう。すべての行列乗算が今や遥かに高速に実行されますが、全体的なself-attentionの計算はそれほど加速されませんでした。これはsoftmaxがFP8を活用していないためです。実際にはより高い精度で実行されます。モデルの精度を保つためには、これを行う必要があります。
これはML分野ではよく知られた秘訣です。もし私たちがこれに注意を払っていなかったら、ここでいくつかの問題が発生していたかもしれませんよね？まず第一に、先ほどお見せした素晴らしい最適化にもかかわらず、エンドツーエンドの高速化は期待したほどではなく、そしてテンソルエンジン、システムの中で最も貴重なリソースが、今や十分に活用されていない状態になっています。
幸いなことに、私たちのチームはこれを早い段階で見抜いていました。そしてmicro-scale FP8の最適化に取り組む際に、テンソルエンジンを常に稼働させ続けるための別の最適化リストも導入しました。このケースでは、同じ精度で精度の損失ゼロで、SoftMaxを4倍高速に実行できる、accelerated SoftMax命令でした。accelerated SoftMax命令を使うとこのようになります。これで期待していたエンドツーエンドの高速化が得られ、テンソルエンジンは再び常に100%の稼働率で動作するようになりました。達成されたパフォーマンスは、ピークパフォーマンスに可能な限り近づいています。
さて、私たちはこれらの最適化の膨大なリストを持っています。実際にこれらを文書化していて、オンラインで自己学習できます。繰り返しになりますが、右上にリンクがあります。そしてこれらすべての最適化は互いに積み重なって、Trainium 3デバイスが提供するすべての浮動小数点演算を確実に使用できるようにしています。
すべてをまとめてみましょう。ここでは1200億パラメータのGPT-3.5というモデルをベンチマークしています。これはOpenAIによるオープンソース、オープンウェイトのモデルです。x軸では、私たちがインタラクティビティと呼んでいるものを測定しています。これはユーザーごとの体験、つまりどれだけ速く出力トークンを生成できるかです。そしてy軸では、全体的なスループットを測定しています。サーバーが同時に複数のリクエストを処理している場合、1秒あたりに生成できるトークンの総数はどれくらいか、ということです。
そして本当に公平な比較、つまりリンゴとリンゴを比べるために、Y軸をメガワットで正規化しました。これでTrainium 2とTrainium 3を同じ土俵で比較できるわけです。どちらがより効率的か？結果は本当に素晴らしいものだと思います。Trainium 3では、Trainium 2と比較して、メガワットあたり5倍のトークンを生成できます。同時にインタラクティビティも向上させています。私たちはこの結果を本当に誇りに思っています。皆さんに真の価値を提供できると考えています。
初日からスケールに対応：モジュラー設計とProject Rainierの実現
次はスケールについてお話しします。昨年、このグラフをお見せしましたが、これはML分野における採用曲線が他のテクノロジーで見慣れた採用曲線とは大きく異なることを示しています。これが典型的な採用曲線です。アーリーアダプターのフェーズがあり、その後立ち上がりが始まり、最終的に大量導入に至ります。そしてMLでは、今日Trainium 3で行っているように新しいテクノロジーを導入すると、すぐにこの新しいテクノロジー、この新世代で巨大なクラスターを構築したいという顧客の需要が発生するのです。
そしてこれにより、私たちはTrainium 3を初日からスケールに対応できるように構築する必要がありました。考えてみてください、それこそがまさにAnnapurnaとAWSが出会い、互いに補完し合う場所なのです。私たちAnnapurnaでは、初日からスケールに対応できるように構築されたTrainium 3を作りました。その理由と方法を正確にお見せします。そしてそれを、世界中の誰よりも速く大規模なコンピュートクラスターを展開する数十年の専門知識を持つAWSと組み合わせることで、Project Rainierのようなプロジェクトや、Trainium 3で構築しようとしているものを実現するのです。
ここにお見せしているのはTrainium 3のコンピュートスレッドです。非常にモジュラーな設計になっていますが、これは単なるエレガントなデザインの選択ではなく、重要なことなのです。つまり、すべてのコンポーネントを個別にテストしてからシステムに組み込むことができるということです。そしてすべてのコンポーネントが上部からアクセス可能で交換可能です。これは非常に重要です。なぜなら、生産ラインを自動化し、組み立てを完全にロボット化できるからです。つまり、手動や複雑な組み立てと比較して、はるかに速くスケールできるということです。また、本番環境でこれらのカードをサービスする必要がある場合、非常に迅速かつ効率的に行うことができ、インフラストラクチャを稼働し続けることができます。これは私たち全員が望んでいることです。
これを分解してみましょう。このコンピュートスレッドの背面には、4つのTrainium 3デバイスがあります。そして前面には、EFAによるスケールアウトネットワーキング用の2つのNitroデバイスがあります。そして中央には、入出力と全体の管理を担当するGraviton CPUがあります。これらのチップはすべて社内で構築されており、深い専門知識があります。最適化の方法も、デバッグの方法も、サービスの方法も熟知しています。
そしてそれらの間での深いco-optimizationが行われています。これは再び、最大のパフォーマンスを提供するために非常に重要です。達成可能なパフォーマンスはピークパフォーマンスにできるだけ近づける必要があり、それを実現するためにはスタック全体にわたって最適化する必要があります。
Joeがこれらのグラフをお見せしました。Trainium 2では、 AWSの他のどのAIチップよりも4倍速く、3倍大きな容量でデプロイしました。例えば、彼はProject Rainierについて言及しました。では、Project Rainierについてお話ししましょう。実は昨年、私たちがこのステージに立っていた時に、Project Rainierを発表しました。私たちはAnthropicのために巨大なAIクラスターを構築すると言いました。そして今、12ヶ月後、100万個のチップが稼働し、最先端のClaudeモデルをプロダクションでトレーニングし、サービング提供しています。これは将来の発表の話ではありません。今日、実際に稼働しているのです。そしてこれが12ヶ月で実現しました。
3つの顧客ペルソナへの対応とPyTorchネイティブサポート
今お見せしたTrainium 3では、Trainium 2でスケールした時よりも速く、 実際にはるかに速く、そしてはるかに大規模にスケールすることを期待しています。最後に、使いやすさについてお話ししましょう。 私たちはここで非常に洗練されたインフラストラクチャを構築しており、お客様が簡単に使用でき、そこから最大の価値を得られるようにする必要があります。そして、使いやすさを最適化したいのであれば、 お客様を深く理解する必要があることがわかっていたので、彼らと多くの対話を重ねました。そして明らかになったのは、実際には3つの顧客ペルソナがあり、それぞれ異なるニーズを持っているということです。
ここの上部には、既存のモデルに基づいてAIアプリケーションを構築しているML developersがいます。彼らが最も重視するのは、非常に強力で堅牢なサードパーティライブラリの統合と、すぐに使える事前最適化されたモデルです。次にresearchersがいます。researchersは新しいモデルや新しいオペレーターを発明しています。彼らは非常に迅速に反復したいと考えており、実際にはパフォーマンスよりも、堅牢でフリクションレスな体験を重視しています。彼らが気にするのは開発サイクルです。実験は非常に迅速に行われる必要があります。
そして最後にperformance engineersがいます。これはJonathan Grayのような、ハードウェア最適化を呼吸し生きている人たちです。Jonathan Grayの話はすぐに聞けます。ちなみに彼はこの分野で最高の一人なので、彼が現時点でTrainium 2に対してどのように最適化しているか、Trainium 3も来ますが、特に彼らが最も重視しているのは、ハードウェアを完全にコントロールできるツールです。これについても話します。
それでは、ML開発者から順番に見ていきましょう。NeuronはPyTorch LightningやvLLM、Hugging Faceといったサードパーティライブラリと深く統合されており、これらのライブラリからモデルを取得して、シームレスに、摩擦なくTrainium上で実行することができます。また、大学のコースやハッカソンを通じてコミュニティとも関わっています。そこにある例を見ていただくと、参加者がHugging Faceのモデルを取得して、特定のタスクを実行するためにファインチューニングしているのがわかります。チェスをプレイするためのモデルのファインチューニングを行うハッカソンのQRリンクがそこに表示されていますが、最終的にはTrainium上でモデルを提供することもでき、これまでに得られているフィードバックは圧倒的にポジティブなものです。
研究者向けには、PyTorchとJAXとの深い統合があり、そして私が本当に共有したいと思っているのは、TrainiumがPyTorchネイティブにもなりつつあるということです。それについてお話ししましょう。PyTorchチームがここで行った最近の進歩により、Private Use Oneと呼ばれるものが導入され、カスタムAIバックエンドをPyTorchに統合できるようになりました。私たちはTrainiumをPyTorchでネイティブにサポートされるようにしました。つまり、CPUやGPU上で実行できるPyTorchで書いたコードが、まったく同じコードでTrainium上でシームレスに実行できるということです。これについては後ほどお見せします。
これは、皆さんがPyTorchで知っていて愛用しているeager executionの体験がTrainiumデバイス上で得られるということを意味します。また、PyTorchがtorch.compileを通じて導入した自動コード最適化も、Trainium上でシームレスに実行されるということも意味します。ここでの良い副次効果、ポジティブな副次効果は、皆さんが知っていて愛用しているPyTorch上で動作するすべてのツールやライブラリも一緒についてくるということです。ワークロードを分散するためにFSDPやDTensorを使用している場合、それもTrainium上でシームレスに実行されます。そして、大規模トレーニングを行うためにTorch Titanのようなライブラリを使用している場合も、それも同様にTrainium上でシームレスに実行されます。
コードでどのように見えるかをお見せします。左側にはGPU上で実行されるPyTorchコードがあり、右側にはTrainium上で実行される対応するPyTorchコードがあります。違いを見つけるのは難しいはずです。なぜなら、違いがあまりないからです。文字通り一つの単語です。
「to CUDA」と言う代わりに、「to Neuron」と書くだけで、残りは私たちが処理します。ただそれだけで動作します。繰り返しになりますが、ここでPyTorchチームに多大な感謝を伝えたいと思います。彼らがPyTorchフレームワークを拡張した方法により、私たちがここでお見せしていることが可能になりました。私たちはすでに選ばれた顧客グループとこの機能をパイロット運用しています。非常に良いフィードバックを得ており、Q1に一般提供する予定です。
Neuron Kernel InterfaceとNeuron Explorer：パフォーマンスエンジニアのための完全なコントロール
最後になりますが、パフォーマンスエンジニアについてお話ししましょう。このカテゴリーのお客様、つまりカスタマーペルソナに対して、私たちは2つの新しい機能を導入しました。1つ目はNeuron Kernel Interface、略してNKIと呼んでいます。私はNKIと呼ぶのに慣れています。これはTrainiumデバイスを直接プログラムするための低レベルのプログラミングインターフェースです。これは昨年から存在していましたが、かなり進化させました。これについては後ほど詳しくお話しします。そして2つ目はNeuron Explorerです。これはTrainiumデバイス上でパフォーマンス最適化を行うためのツールキットで、Neuron Profilerの上に構築されており、Trainium上で実行されるワークロードに対して非常に深いインサイトと可観測性を提供します。そしてこの2つを組み合わせることで、Trainium上でワークロードを最適化するための完全なコントロールが得られます。
それでは1つずつ見ていきましょう。NKIはPythonに組み込まれたDSLですが、非常にユニークな特徴があります。単一のプログラミング環境の中で2つのレベルの抽象化を組み合わせています。タイルベースの言語でコードを実装でき、サブマトリクス間の計算を行うだけで、これは特にNumPyやTritonから来られた方にとっては非常に習得しやすいでしょう。しかし、本当に最適化したい領域を特定した場合は、タイルベースのセマンティクスと非常に似たセマンティクスで、アセンブリレベルまで深く掘り下げることができます。そしてこの組み合わせにより、非常に迅速に習得し、非常に深く最適化する能力が得られます。
今年、私たちはNKIにいくつかの新機能を導入しています。スケジューリングとアロケーションAPIを含んでおり、マシン上で実行されているさまざまな命令のスケジューリングや、さまざまなテンソルをどこに割り当てるかについて、きめ細かい制御を行うことができます。そしてこれにより、先ほどself-attentionの例でお見せした非常に構造化されたパイプラインを構築することができます。これは実際に一部のお客様からの機能リクエストでした。私たちは耳を傾けました。これはすでに利用可能です。今すぐ使い始めることができます。
さらに、エラーメッセージが大幅に改善された新しいフロントエンドも導入しました。これにより、基本的にセルフサービスで、NKIプログラミング体験をより迅速に反復でき、Trainium上で最適化されたカーネルを得るまでの時間を改善できます。そして最後になりますが、実は私はこれにかなり興奮しています。私たちはNKIコンパイラをオープンソース化することを決定しました。これは今後数ヶ月以内に提供されます。そしてこれを行うことにした理由は、NKIがコントロールと可観測性を提供することがすべてだからです。つまり、今やコードがどのようにTrainiumハードウェアにコンパイルされるかについて完全な透明性を提供し、Trainiumスタック全体にわたる業界とコミュニティの貢献も歓迎します。
ここに良い例があります。これはDescartesという会社です。彼らはリアルタイムの生成AIビデオ生成を行うクールなアプリケーションを持っています。ビデオを取り込み、編集し、ビデオを生成して返すことができます。ここに例が見えます。彼らはNKIに基づいて全体のモデルを構築することを決定し、驚異的な利用率の数値を達成しました。実際、チームが3〜4ヶ月でできると私が期待していた以上のものでした。
次に、Neuron Explorerについてお話しします。高度に最適化されたコードを書いたことがある方ならご存知だと思いますが、最良の友となるのは、ハードウェア上で何が実行されているのか、そしてボトルネックがどこにあるのかを教えてくれる強力なプロファイラーやトレーサーです。Trainiumには、業界をリードするNeuron Profilerがあり、実際に実行されているワークロードに対してパフォーマンスへの影響をゼロにしながら、ハードウェア上で実行されている内容の命令レベルのトレースを取得できます。これはワークロードを遅くすることなく、ナノ秒レベルの可観測性を実現します。
そこで私たちはNeuron Profilerを大幅に拡張し、その上にNeuron Explorerと呼ぶツール群を構築しました。まず、インタラクティブ性が4倍向上しており、これによってデバッグをはるかに高速に行え、全体的なデバッグ体験が向上します。
さらにその上で、開発者間で簡単に共有できるようWebアプリケーション経由で利用可能にし、VS CodeのようなIDEとも深く統合しました。これは実際かなり重要なことです。画面に表示されているのは、私がNeuron Kernelコードの1行をハイライトすると、Neuron Explorerが自動的にプロファイラー内の関連する命令をハイライトしている様子です。これにより、書いているコードと実際にハードウェア上で実行されている内容との間に、はるかに緊密な接続が生まれ、何を最適化する価値があるのかという感覚が得られます。
また、システムレベルのプロファイリングも導入しています。これはまだ準備中で、1ヶ月後に提供されます。システムレベルのプロファイリングを導入することで、複数のデバイスでの完全な実行を確認でき、それらが緊密に同期しているか、あるいは遅いマシンが1台あるかを確認できます。大規模なトレーニング実行のような高度に分散されたコードをデバッグする際に、本当に役立ちます。
他にもいくつか機能を追加しました。階層ビューを導入したので、Neuron Explorerを起動すると、self-attentionや全結合層といったフレームワークレベルのオペレーターが表示され、そこからクリックして命令レベルまでドリルダウンできます。これによってデバッグ体験がはるかに段階的になります。高レベルから始めてボトルネックがどこにあるかを理解しようとし、本当に何かにズームインしたいときには、そこまでドリルダウンできます。私の見解では、これによってデバッグ体験がはるかに快適になります。
また、異なるエンジンがどのように活用されているかを示すサマリーページも提供します。ここでは左側にtensor engineが表示されています。tensor engineは非常によく活用されていて、60から70%のMFUとなっており、他のエンジンは軽く利用されている程度です。これによってワークロードがどのように実行されているかがわかります。右上では、メモリ帯域幅をどのように活用しているか、どの部分が読み取りに使われているか、どの部分が書き込みに使われているか、そしてメモリが実際にアイドル状態で待機している時間の割合を確認できます。このトップレベルのビューを見ると、ワークロードがハードウェア上でどれだけうまく実行されているかを本当に把握することができます。
また、統計情報とビジュアライゼーションも提供しています。右下のものは私が特に気に入っているものです。ここでは、このケースでは推論実行全体を通じてcollective communicationを示しており、それらの散布図を表示しています。ここで見えているのは良い状態です。ほとんどの通信がほぼ正確に同じ期間で発生しており、これはパフォーマンスが非常に一貫性があり予測可能であることを意味します。しかし、もしここで大きなばらつきが見られた場合は、外れ値があることがわかり、デバッグして何が起こったのかを理解する必要があります。
そして最後に、これはクールなものです。私たちはPerformance Insightsと呼ぶものを導入しました。サマリーページには、パフォーマンスのボトルネックがどこにあると私たちが考えているか、そして実際にそれを解決するために何ができるかを示す一連のボックスが表示されます。これはAIベースの技術と人間ベースの技術の組み合わせによって実現しています。もし私たちが何度かデバッグした場合、ここにルールを導入して、これがパフォーマンス向上に役立つかもしれないというヒントを提供するようにしています。
さて、これをAnthropicの方々にお見せしました。そこにTristanという素晴らしいパフォーマンスエンジニアがいて、彼がこれを見たとき、これはすべてのパフォーマンスエンジニアの夢だと言ってくれました。これは近年私が最も気に入っている言葉の一つです、特にTristanのような人からの言葉ですからね。
さて、まとめに入りますが、私たちは異なる顧客ペルソナに対してこの使いやすさを提供しており、今日お見せしたもののほとんどはオープンソース化される予定です。Neuron Kernel compilerについてお話ししました。それに加えて、Torch nativeのトレーニングバックエンドもオープンソース化される予定で、さらにNeuron Kernel libraryもオープンソース化します。これは私たちのユースケースのために構築した事前最適化されたカーネルのスイートで、世界中で利用できるようにしたいと考えています。
Jay Grayに渡す前に、ご想像の通り、私たちは今Trainium3の実装に深く取り組んでいます。詳細な事実を共有するにはまだ少し早いのですが、私たちが目指しているものは時間とともに加速しており、おそらくFP8で6倍のパフォーマンス向上、4倍のメモリ帯域幅の向上、そして2倍のメモリ容量の向上を超えるでしょう。エネルギー効率の向上は非常に大きなものになりますが、それについてはまだ共有する準備ができていません。Jay Gray、実際にこれらのチップをどのように使っているか話してもらえますか。ありがとう、みんな。
AnthropicのClaude推論最適化：Flash Attentionカーネルの実践的改善
素晴らしい、ありがとうRon、ありがとうJoe。あなた方とステージを共にできて光栄です。ここに来られて超興奮しています。皆さんこんにちは、私の名前はJay Grayで、AnthropicでTraining Inference Leadを務めています。Anthropicはその規模において史上最も急成長しているビジネスです。私たちのClaude 4とClaude 4.5モデルは、世界中の企業から最も信頼されているAIであり、特に先週リリースしたばかりのClaude Opus 4.5により、Claudeは世界最高のコーディングモデルであり、agenticワークフローにおいて最高のモデルとなっています。
そして、これらすべての鍵となるのは、私たちのすべてのプロダクトサーフェス全体で、ファーストパーティAPIとAWS Bedrockを通じて、Claude codeのすべての使用、私たちのウェブアプリ、モバイルアプリにおいて、今日のトラフィックの大部分がTrainium 2で提供されているということです。では、このようなモデル推論のスケールを可能にしているものは何でしょうか。私のチームの仕事は、このような前例のない速度でスケールすることを可能にする、Trainium上のコア推論技術を提供することです。そして今日は、このスケールを可能にする私たちが行っているパフォーマンスエンジニアリングの仕事について深く掘り下げていきます。
では、私たちは実際に何をしているのでしょうか。私たちの仕事は、基本的に、指数関数的に成長する顧客群に対して、できるだけ効率的にサービスを提供しながら、モデルをできるだけ高速に実行することです。シンプルな仕事ですね。モデルのpre-fill時間を10%削減するたびに、新しいプロダクトのユースケースが開かれます。より長いコンテキストのより人間工学的な使用により、コードベース全体をコンテキストに入れることができ、より速い応答時間により、より人間工学的なインタラクティブなユースケースが可能になります。そして、トークン生成速度を上げるたびに、Claudeがもう少し長く考えることができるようになり、コードがもう少し速く書けるようになり、あるいはバックエンドでサンプリングバッチサイズを増やして、あなたのトラフィックをもう少し効率的に静かに処理できるようになります。
Anthropicでは、モデル推論のすべてのオペレーションとすべてのカーネルが、Trainiumチップから最高のパフォーマンスを引き出すように設計されています。そこで今日は、私たちが日常的に行っているパフォーマンス作業について、皆さんを深く掘り下げた旅にお連れするのが面白いと思いました。では、まず始めに、Anthropicのカスタムモデルアーキテクチャとカスタムカーネルの概要を見てみましょう。さて、これは少し冗談です。私は、まだいくつかの企業秘密を持っていますし、文字通り私たちのモデルアーキテクチャを説明するつもりはありませんが、現実的な大規模LLM推論カーネルで私たちが行ったいくつかの実際の最適化についてお話しします。
それではこれから5分か10分ほど、これを使って説明していきます。これは実際のfused flash attention kernelを3つのパートに分けたものです。まず最初は大規模な行列乗算で、これがqueries、keys、valuesを生成します。これらはRonが先ほど説明したself-attention operationへの入力となります。そして実際のself-attention operationがあり、最後にもう一つの大きな行列乗算があって、attentionの出力をresidual stream spaceに射影して戻します。
本題に入る前に、Neuronアーキテクチャについて非常に簡単に概要を説明します。もしすでにTrainiumでプログラミングをされている方には復習になりますね。他のアーキテクチャでのプログラミングに慣れている方には、NeuronCoreアーキテクチャの簡単で興味深い概要になればと思います。すべてのTrainiumチップの中核には、NeuronCoresのセットがあり、各コアには異なる線形代数演算に特化した複数の異なるエンジンがあります。
この中心にあるのがTensor Engineで、これは小さなタイルの行列乗算を行います。もしMLパフォーマンスやカーネル最適化の話から一つだけ持ち帰るとしたら、これです：カーネルの目標、そしてカーネルエンジニアの目標は、Tensor Engineが常に行列乗算を行っているようにすることです。それ以外のすべては本質的に補助的なデータ移動と追加の操作であり、Tensor Engineが一つの行列乗算を終えたときに、次の行列乗算に必要なデータがすぐに入力できる状態になっていて、行列乗算を密にパックできるようにするためのものです。
Vector Engineは、ベクトルの総和のような、データのストリームに対する削減や処理を専門とするエンジンで、Scalar Engineは活性化関数やsoftmaxの指数部分のような要素ごとの演算を専門としています。最後のエンジンは、Trainiumアーキテクチャの面白いイノベーションで、GPSIMDまたはGeneral Purpose SIMD Engineと呼ばれるもので、基本的に任意のCコードを書いてデータを操作でき、他のエンジンに当てはまらないような変わった操作をカスタムアーキテクチャに組み込むことができます。
これらすべてのエンジンは、エンジンの近くにある高速SRAMメモリバンクのセット、SBUFとPSUMと呼ばれるものに読み書きします。ここではこの2つの違いには触れません。そしてDMAエンジンのセットがあり、これがエンジンに近い高速SRAMメモリと、チップ上のより大きなHBMとの間でデータを往復させます。
それでは、Flash Attentionカーネルに戻りましょう。ここでご覧いただいているのは、実際のカーネルの実際のプロファイラービューです。そして、ここの各行は、先ほど説明したエンジンの1つに対応しており、各線はそれらのエンジンの1つで実際に発生している操作を表しています。そして、ここで数字に深く入り込まなくても分かることは、私たちはここでかなり良い結果を出しているということです。視覚的に見ても、Tensor Engineがこれらの青い行列乗算操作で密にパックされていることが分かります。これはかなり良い感じですが、どうやってここまで到達したのでしょうか？
それでは、最初の最適化として、大きな行列乗算の最初のものであるQKVに飛び込んでいきましょう。そして、Tensor Matrixで発生している単一の操作を見ることから始めましょう。ここで少し立ち止まって、この点を本当に強調したいと思います。
なぜなら、これはTrainiumでプログラミングする上で、私が最も気に入っていることかもしれないからです。ここでご覧いただいているのは、128×128の行列乗算操作1つの実際のISA読み出しです。これは、カーネル内で、完全なフォワードパス内で発生する多くの操作の1つです。ここでご覧いただいているのは、ナノ秒単位、読み書きされている個々のメモリ空間のバイト単位まで、完全な読み出しです。これはまさに実際に起こっていることであり、つまり、他のアクセラレータアーキテクチャでプログラミングをしたことがある方なら、これがどれほど素晴らしいことかすぐに理解できるでしょう。これは、カーネルのパフォーマンスに対する可視性のレベルであり、他ではなかなか得られないものです。
すべてのflop、すべてのナノ秒、すべてのカーネルのすべての操作におけるメモリのすべてのバイトが、このレベルの詳細まで追跡できます。そして、これこそが、Trainiumチップから最大のパフォーマンスを引き出すことを可能にしているのです。それでは、ここから始めましょう。標準的なBFloat16フォーマットで密にパックされた行列乗算ですが、最近の多くのLLM推論、特にデコードでは、Ronが言及したように、より小さく、より効率的なデータフォーマットを使用することが重要です。そして、Trainium2は、完全幅のBFloat16から得られる速度の2倍の速度を、より小さなFP8フォーマットから引き出すように設計されています。そして、これらの操作を遅いBFloat16から、より速い、この場合はFP8 E4M3フォーマットに移行することで、この行列乗算で即座に2倍の速度向上、つまり2倍のスピードアップを得ることができます。
次に、実際のself attention操作に飛び込んでみましょう。これは、もうお分かりのように、少し複雑なカーネルです。そして、attentionの最適化は、最も興味深い問題の1つだと思います。現代のLLM推論において。これは、単一の行列乗算を扱うよりもはるかに複雑な最適化です。なぜなら、その中にはより多くの最適化、つまりより多くの操作があり、カーネルが行列乗算にすべての時間を費やすことを妨げるボトルネックの機会がたくさんあるからです。ここで視覚的に分かることは、先ほど見ていた行列乗算とは異なり、緑色のtensor matrix、上から3行目を見ると、私たちが望むように、常に行列乗算を密にパックして実行しているわけではないということです。
私たちが目にするのは、これらの行列乗算のバーストですが、その間にギャップが散在していて、実際には多数の小さなベクトルコア演算を行っているように見えます。そして、先ほどお見せしたプロファイラービューを使ってここを深く掘り下げ、ISAビューを読むと、ここでのボトルネックは私たちが望むような行列乗算を行うことではないことがわかります。ここでのボトルネックは、実際には行列乗算の結果を1つのメモリバンクから別のメモリバンクへシャトルする際に、非効率的に大量の小さなベクトルコア演算を使用していることです。そしてそれに気づいたとき、私たちはタイリングを書き直して、より少ない数のより大きなベクトルコア演算を使用してメモリを1つのバンクから別のバンクへ移動させ、命令起動のオーバーヘッドを償却し、命令をより有効に活用することで、行列乗算により多くの時間を費やせるようにしました。そしてこれに手を加えるだけで、attentionで13%の高速化を実現しました。
通信最適化とTrainium3への移行：100万チップ規模での運用実績
大したことないように聞こえるかもしれませんが、私たちが運用している規模では、これは膨大な量のチップの節約であり、膨大な量の追加トラフィックを節約できます。それでは通信について話しましょう。現代のLLMが単一のチップに収まらないほど大きくなってから、もう何年も経っています。そして、パフォーマンスエンジニアとして私たちが持つ興味深い設計空間の多くは、完全なLLMのフォワードパスのデータと計算を複数のチップに分割してシャーディングし、その後collectiveを使用してそれらの間で通信を行い、正しい結果に到達する方法です。
Trainiumは、ほとんどのチップアーキテクチャと同様に、先ほど述べたように、より大きなHBMと通信する少量の高速SRAMメモリバンクで動作します。そしてデフォルトでここで起こることは、collective操作を行うために、操作の1つの結果を取得し、高速SRAMからHBMへシャトルし、異なるチップのHBMからHBMへcollectiveを実行し、その結果をSRAMに戻すということです。特にトークンをできるだけ速くストリーミングしようとしているtoken decodeでは、この3ステップのメモリ移動はレイテンシにとって非常に悪く、通信を他の計算とオーバーラップできない場合、このような通信に時間を費やすことは、低レイテンシカーネルにとってまさに致命的です。
そこでTrainiumが私たちに可能にしてくれるこのクールな最適化は、異なるチップ間でSRAMからSRAMへ直接collectiveを実行できるクールなハードウェア機能の1つを利用し、SRAMとHBM間のメモリの余分なホップを節約することです。そしてここで見ることができるのは、それほど明白ではありませんが、私が赤い円で注釈を付けたGPSIMD操作で、左側ではすべての時間をSRAMとHBM間のメモリ移動DMAのディスクリプタを書き込むことに費やしています。
これは右側では消えており、より高速なSRAMからSRAMへのcollectiveにより、通信に費やす時間が少なくなり、decodeのレイテンシが速くなります。
最後の最適化は、もちろん、このカーネルをTrainium3上で実行することです。今日ご説明したすべての演算が、Trainium3上ではより高速になります。最初の最適化で見た倍速のFP8行列乗算は、Trainium3では4倍高速になり、Ronが説明したmicroscalingアーキテクチャを活用して、より効率的なブロック単位の量子化と逆量子化を実現します。attentionのような複雑な実際のワークロードのボトルネックになりやすいベクトル演算やスカラー演算も、Trainium3ではより高速になっています。通信も高速化されています。ICIドメインあたりのHBM容量が増加しており、単一のICIドメイン上でより大きなモデルを提供できるようになっています。まだまだ挙げられますが、この場合、過去5分から10分間取り組んできたこのカーネルは、最適化後にTrainium2で約60%のテンソルエンジン使用率を達成していましたが、Trainium3では90%以上に到達します。
それでは、ここで締めくくりたいと思います。1年前、私たちはProject Rainierを発表し、AnthropicがTrainiumチップの初期使用を発表しました。1年後の今、私たちはほぼ100万個のTrainium2チップ上でAnthropicのモデルを提供しており、2026年以降、Trainium3でどこまで到達できるか、とても楽しみにしています。それでは、Joeにお返しします。
素晴らしいですね。改めてJay GrayとRonの両名に感謝します。彼らは、スタック全体にわたってTrainiumの最適化についてどのように考えているか、詳細に深く掘り下げて本当に素晴らしい仕事をしてくれました。終わる前にいくつかポイントをお伝えします。Trainium3は一般提供開始されています。このチップは昨日のMattの基調講演で発表されました。Trainium3については、チップだけでなく、システムとしても考えてください。Jay Grayが説明したように、最適化にとってシステムがいかに重要か、フロップスだけでなく、ICIバンド幅、つまり私たちがNeuronLinkと呼んでいるものも重要だということをご覧いただきました。私たちは144チップで真にスケールするシステム、Trainium3 UltraServerを構築しています。
そしてもう一つ、始めるのは簡単です。私たちは本当に簡単にしていますので、Neuron SDKについて多くの情報を提供しており、皆さんが外に出て学べるようにしたいと考えています。ですので、もし明日お時間があって、もっと学びたいとお考えでしたら、木曜日にいくつかのワークショップが用意されていますので、ぜひチェックすることをお勧めします。これをスキャンして、詳細を学ぶための完全なリストをご覧いただけます。そして、もし木曜日にお時間がない場合でも、いつでもこれらのいずれかをスキャンして、始めることができ、独自に素早く習得することができます。多くのチュートリアルを用意しており、幅広くTrainium上で開発していただけることを本当に楽しみにしています。
それでは、改めて、ここにお越しいただいた皆様全員に感謝申し上げます。本当にありがとうございます。そして、もしお時間があれば、ぜひアンケートにご記入ください。
※ こちらの記事は Amazon Bedrock を利用し、元動画の情報をできる限り維持しつつ自動で作成しています。
Discussion