AWS Trainiumは、Amazon Web Services（AWS）がクラウドでの機械学習トレーニングおよび推論タスクを高速化するために開発した、カスタムアプリケーション固有集積回路（ASIC）のファミリーである。2020年11月に発表され、2022年に初めて利用可能となったTrainiumチップは、大規模なAIワークロード、特に深層学習モデル（大規模言語モデルやトランスフォーマーアーキテクチャなど）に関わるものに対して、汎用グラフィックス処理ユニット（GPU）よりもコスト効率が高く、エネルギー効率に優れた代替手段を提供するように設計されている。これらのチップはAWSのElastic Compute Cloud（EC2）インスタンスに統合され、ニューラルネットワーク操作の特定の計算パターンに最適化された専用ハードウェアを顧客に提供する。
Trainiumの開発は、AWSがAIハードウェアの提供を多様化し、Nvidiaなどの外部チップサプライヤーへの依存を減らすための、より広範な戦略の一部である。独自のシリコンを設計することで、AWSはプラットフォーム上で実行される主要なAIワークロードに対して、ドルあたりおよびワットあたりのパフォーマンスを最適化することを目指している。Trainiumは、推論に特化した別のチップファミリーであるAWS Inferentiaによって補完され、これらが一体となってAWSのカスタムAIアクセラレーションポートフォリオの中核を形成する。Trainiumファミリーは、2023年に発表され2024年に広範な展開が始まった第2世代のTrainium2や、2024年後半に発表された第3世代のTrainium3など、その後の世代で進化を続けている。
ハードウェアアーキテクチャ
Trainiumチップは、ニューラルネットワークトレーニングの中核操作である高スループットの行列乗算とデータ移動を重視した、専用設計のアーキテクチャ上に構築されている。各Trainiumチップには複数のテンソルコアが含まれており、FP32、FP16、BF16、FP8などのデータ型をサポートする混合精度計算を実行できる。第1世代のTrainiumチップは、TSMCによる7ナノメートルプロセスで製造され、FP16およびBF16操作で最大200テラフロップスのパフォーマンスを提供する。また、大容量のオンチップメモリと外部メモリへの高帯域幅インターフェースを備え、大規模なモデルパラメータとアクティベーションを効率的に処理できる。
第2世代のTrainium2は、パフォーマンスとメモリ容量を大幅に向上させている。5ナノメートルプロセスで製造され、各Trainium2チップは最大400テラフロップスのFP16/BF16パフォーマンスを提供し、32ギガバイトの高帯域幅メモリ（HBM）を搭載する。ラックレベルのシステムである単一のTrainium2 UltraServerには、高速ファブリックで相互接続された64個のTrainium2チップが含まれ、合計25.6ペタフロップスの計算能力と6テラバイトのメモリを提供する。この設計により、生成AIアプリケーションで使用されるような、数千億のパラメータを持つモデルのトレーニングが可能になる。
2024年12月に発表されたTrainium3は、3ナノメートルプロセスで製造される予定であり、トレーニングと推論の両方の効率に焦点を当て、Trainium2の約2倍のパフォーマンスを提供すると見込まれている。AWSはTrainium3が2025年に利用可能となり、一部の顧客には早期アクセスが提供されると述べている。
EC2インスタンスと可用性
AWSは、EC2サービスを通じてTrn1およびTrn2インスタンスファミリーとしてTrainiumベースの計算能力を提供している。2022年10月に開始された第1世代のTrn1インスタンスには、1つのTrainiumチップを搭載したtrn1.2xlargeと、16個のTrainiumチップを搭載したtrn1.32xlargeの2つのサイズがある。trn1.32xlargeは最大3.2ペタフロップスのFP16/BF16パフォーマンスと512ギガバイトのメモリを提供し、複数インスタンス間での分散トレーニングのためにElastic Fabric Adapter（EFA）を介した高速ネットワーキングをサポートする。
2024年12月に開始された第2世代のTrn2インスタンスは、Trainium2チップに基づいている。trn2.48xlargeインスタンスには16個のTrainium2チップが含まれ、6.4ペタフロップスのパフォーマンスと1.5テラバイトのメモリを提供する。AWSはまた、4つのtrn2.48xlargeインスタンスを64個のTrainium2チップを備えた単一の論理システムに統合するTrn2 UltraServerも導入し、最大1兆パラメータのモデルのトレーニングを可能にしている。これらのインスタンスは、米国東部（バージニア北部）、米国西部（オレゴン）、欧州（アイルランド）などの一部のAWSリージョンで利用可能であり、拡張が計画されている。
ソフトウェアスタックと統合
Trainiumハードウェアは、コンパイラ、ランタイム、プロファイリングツールを含むソフトウェア開発キットであるAWS Neuron SDKによってサポートされている。Neuron SDKは、PyTorchやTensorFlowなどの一般的なフレームワークで書かれたモデルを、Trainium上で実行される最適化された実行可能ファイルに変換する。また、自動モデル並列化をサポートしており、ユーザーの介入を最小限に抑えながら、大規模なモデルを複数のチップおよびインスタンスに分割できる。AWSはまた、TrainiumをSageMakerプラットフォームと統合し、顧客がGPUベースのトレーニングに使用されるのと同じ管理インターフェースを通じてTrainiumインスタンスでトレーニングジョブを開始できるようにしている。
さらに、AWSはOpenAIおよびAnthropicとの協力関係を構築し、それぞれのモデルをTrainium向けに最適化している。Amazonから多額の投資を受けているAnthropicは、主要な初期採用者であり、Trainium2 UltraServerを使用してClaudeファミリーのモデルをトレーニングおよび展開している。OpenAIも2024年に、一部のフロンティアモデルのトレーニングにTrainium2を使用すると発表し、カスタムチップの能力に対する注目すべき支持を示した。
パフォーマンスとコスト効率
AWSはTrainiumを、Nvidia A100やH100 GPUを搭載した同等のGPUベースのEC2インスタンスと比較して、トレーニングジョブあたり最大50%低いコストを提供すると宣伝している。MLPerfなどの独立したベンチマークでは、Trainium2が標準的なトレーニングタスク、特にトランスフォーマーベースのモデルで競争力のあるパフォーマンスを達成することが示されている。推論に関しては、Trainium2は低レイテンシで高スループットを提供するように設計されており、チャットボットやコード生成などのリアルタイムアプリケーションに適している。
ただし、TrainiumのエコシステムはNvidiaのCUDAプラットフォームほど成熟していない。一部の機械学習フレームワークやライブラリのサポートは限定的であり、スパースアテンションやカスタムカーネルなどの特定の高度な機能には、追加のエンジニアリング作業が必要になる場合がある。AWSはNeuron SDKの機能拡張に積極的に投資しており、顧客がピークに近い利用率を達成できるようにするための一連のパフォーマンス最適化ガイドもリリースしている。
採用とエコシステム
AnthropicやOpenAIに加えて、他のいくつかの組織もAIワークロードにTrainiumを採用している。GPUインフラストラクチャを専門とするクラウドサービスプロバイダーであるCoreWeaveは、2024年にTrainiumベースのインスタンスを顧客に提供すると発表し、AWS自身のクラウドを超えてチップの可用性を拡大した。AIチップ分野の競合であるCerebras Systemsも、Trainiumと並んで自社のハードウェアを提供するためにAWSと提携しているが、両製品は異なる市場セグメントを対象としている。
研究コミュニティでは、Trainiumは学術環境での大規模モデルのトレーニングに使用されている。例えば、スタンフォードAIラボやバークレーAIリサーチの研究者は、モデルスケーリングと効率に関する実験でTrainium2を使用した研究を発表している。AWSはまた、オープンソースのAIプロジェクトに取り組む学術機関に無料の計算クレジットを提供するTrainium研究プログラムも設立している。
競争環境
Trainiumは、Google CloudのTensor Processing Units（TPU）、Microsoft AzureのMaiaチップ、Groqの言語処理ユニットなど、他のカスタムAIアクセラレータと直接競合する。TPUは数年前から利用可能であり、成熟したソフトウェアスタックを備えているが、Trainiumの利点はAWSの広範なクラウドサービスとの緊密な統合と、積極的な価格設定にある。Trainium3の登場は、特にAI計算への需要が供給を上回り続ける中で、競争を激化させると予想される。
AWSのカスタムシリコンへの投資は、代替のGPUおよびアクセラレータオプションを提供するAMDやIntel、およびデータセンターで勢いを増しているArm Holdingsベースの設計に対しても、AWSを位置づけるものである。Trainiumの成功は、AWSが堅牢なソフトウェアエコシステムを維持し、長期的に明確なパフォーマンスとコストの利点を示す能力に依存するだろう。
今後の方向性
今後を見据えて、AWSはTrainiumがAnthropicとの協力で構築された大規模スーパーコンピューターであるProject Rainierで中心的な役割を果たすと示唆している。2024年に発表されたProject Rainierは、数十万個のTrainium2チップで構成され、2025年の運用開始が計画されており、世界最大級のAIトレーニングクラスターになると予想されている。このプロジェクトは、成長するAIインフラストラクチャ市場を捉えるための戦略的資産としてのカスタムシリコンへのAWSのコミットメントを強調している。
さらに、AWSはエッジおよびオンプレミス展開でのTrainiumの使用を検討しているが、具体的な製品は発表されていない。同社はまた、現在のUltraServerフォームファクターを超えてTrainiumシステムをスケールさせるための高度なパッケージングおよび相互接続技術にも投資しており、潜在的に数兆パラメータのモデルのトレーニングを可能にする。
結論
AWS Trainiumは、AI向けカスタムシリコンに対するAmazonの重要な賭けを表している。クラウドプラットフォームと緊密に統合された専用設計のチップを提供することで、AWSはパフォーマンス、コスト、スケーラビリティの魅力的な組み合わせを顧客に提供することを目指している。ソフトウェアの成熟度とエコシステムの採用には課題が残るものの、Trainiumファミリーの急速な進化と、AnthropicやOpenAIなどの主要なAIラボからの支援は、Trainiumが今後数年間AIハードウェアの状況において主要なプレーヤーであり続けることを示唆している。