Build on Trainium

Build on Trainium は、 AWS Trainium での次世代のイノベーションと開発を支援するための AI 研究と大学教育に焦点を当てた 1 億 1000 万 USD の投資プログラムです。AWS Trainium は、最先端の AI のアイデアとアプリケーションを進化させるために独自に設計された AI シストリックアレイチップです。Build on Trainium は、Trainium に関する新しい AI 研究に資金を提供し、主要な学術チームに投資して、新しいモデルアーキテクチャ、ML ライブラリ、最適化、大規模分散システムなどの重要な分野でイノベーションを起こしています。この複数年にわたる取り組みにより、学術コミュニティが Trainium 周辺のオープンソースコミュニティを活用し、投資し、貢献するよう促すことで、AI の未来の基礎を築いています。これらの利点と、 Neuron Software Development Kit (SDK)、および最近リリースされた Neuron Kernel Interface (NKI) を組み合わせることで、Trainium のお客様はクラウドで大規模なイノベーションを実現できるようになりました。

最大 40,000 個の Trainium チップを搭載した専用の Trainium リサーチクラスターを作成しました。これは、 Amazon EC2 UltraClusters を使用してペタビット単位の単一のノンブロッキングスケールネットワークに接続された Amazon EC2 Trn1 インスタンスを介して利用できるようになります。研究チームと学生は、Amazon EC2 Capacity Blocks for ML を使用してセルフマネージド型のキャパシティブロック予約を通じてこのチップにアクセスできます。

幅広い研究コミュニティを対象に Amazon Research Awards (ARA) による提案募集 (CFP) を複数回実施しています。選ばれた提案には AWS Trainium クレジットと Trainium リサーチクラスターへのアクセス権が付与されます。Build on Trainium は、人気のあるオープンソースの ML ライブラリとフレームワークを活用し、ML デベロッパーコミュニティのリソースを強化するオープンソースに貢献する研究提案を歓迎します。

Neuron Kernel Interface (NKI) は、AWS AI チップ、Trainium および Inferentia 用の新しいプログラミングインターフェイスです。NKI では、AWS Trainium や Inferentia で利用できるハードウェアプリミティブや命令に直接アクセスできるため、研究者はコンピューティングカーネルを構築およびチューニングして最適なパフォーマンスを得ることができます。Python ベースのプログラミング環境で、一般的に使用されている Triton のような構文とタイルレベルのセマンティクスを採用しています。研究者は NKI を使用して、新しい機能、最適化、科学的な革新によって深層学習モデルを強化できます。詳細については、 NKI のドキュメントページを参照してください。

MIT のデバイス実現ラボでは、AWS Trainium を使用して医療 AI 研究の限界に挑戦しています。当社の 3D 超音波セグメンテーションと音速推定モデルは、かつてないほど高速かつ効率的にトレーニングを行い、最先端の精度を実現しながら実験時間を半分以上短縮します。AWS Trainium のおかげで、従来の GPU システムでは不可能だった方法で研究を拡大することができました。AWS Trainium (trn.32xlarge) で 3D 完全畳み込みニューラルネットワークをトレーニングすることで、NVIDIA A100 インスタンスと比較して 50% 高いスループットと低コストで最先端のパフォーマンスを実現しました。32 ノードの Trainium クラスターを使用して、180 を超えるアブレーション実験を効率的に実施し、総トレーニング時間を数か月から数週間に短縮し、MIT での医療 AI イノベーションを加速させました。将来的には、Trainium を使用して、デジタル超音波ワークフローを操作および自動化できる AI エージェントモデルのトレーニングを行う予定です。これにより、臨床医の時間を大幅に節約し、患者により良いケアを提供できるようになります。」

「CMU Catalyst 研究グループは、ML システムの最適化に取り組んでいます。私たちのプロジェクトは、さまざまな ML システム間での最適化をより簡単にすることを目的としています。Trainium は、低レベルの制御と Neuron Kernel Interface (NKI) を介したアクセス可能なプログラミングインターフェイスの両方を提供する点でユニークです。

Build on Trainium プログラムを通じた AWS の支援を受けて、私たちの研究者は重要なカーネルである FlashAttention（フラッシュアテンション）の高度な最適化を検討することができました。私たちが最も驚いたのは、イテレーションの速さでした。公開されている NKI、Neuron プロファイラー、およびアーキテクチャドキュメントを使用して、わずか 1 週間で以前の最先端技術に加えて有意義な改善を実現しました。強力なツールと明確なハードウェアの洞察を組み合わせることで、私たちのチームは洗練された低レベルの最適化を利用できるようになりました。

AWS Trainium と Neuron Kernel Interface (NKI) は、私たちのような研究者がより迅速にイノベーションを起こすことを可能にし、通常はハードウェア固有の最適化作業を遅らせる障壁を取り除いてくれます。」

 

「Build on Trainium プログラムを通じて、彼のチームは AWS Neuron の新しい NKI オープンソースコンパイラスタックにフルアクセスできるようになりました。これには、Trainium ISA や正確なスケジューリングとメモリ割り当てのための API への直接の可視性も含まれます。このレベルの可視性と制御により、学生は最適化の機会をより簡単に分析し、パフォーマンスの高い実装をより効果的に見つけることができます。」

カリフォルニア大学バークレー校、Associate Professor of Computer Science、Christopher Fletcher 氏
 

「AWS Trainium と Inferentia へのアクセスは、大規模で効率的な AI システムに関する研究と教育を進めるのに役立ちました。これらのプラットフォームを使用して、専門家によるトレーニングと推論の最適化、新しいアクセラレータアーキテクチャのスケーラビリティ、効率、移植性を向上させる新しいランタイムおよびスケジューリング手法のプロトタイプを作成します。UIUCの研究者は、Neuron Developerスタックを活用することで、AIワークロードの効率と移植性を高める新しいランタイムおよびスケジューリング技術を開発しています。チームは、Neuron Developer stackのオープン性に特に感銘を受けました。これにより、これらのプラットフォームはランタイムリサーチにとって価値があり、従来のGPUアーキテクチャを超えるスパース性、メモリ階層、通信効率の革新が可能になります。」

「AWS Trainium と Build on Trainium プログラムを活用することで、学生と私は量子回路シミュレーションを大幅に加速することができました。このプロジェクトには、高性能シミュレーターを共同で構築した強力な学生グループが集まり、これまで不可能だった規模でのより深い実験と実践的な学習が可能になりました。」

「UTS の研究チームは、カスタム Neuron NKI カーネルを開発することにより、ツリーリングウォーターマーキングアルゴリズムの統合を検討しています。Build on Trainium プログラムを通じてオープンソースの Neuron スタックにアクセスできたことは、大きな変革をもたらしました。これにより、Trainium アーキテクチャをこれまでにないほど可視化し、ハードウェアレベルで直接作業できるようになります。Trainium にアクセスできたおかげで、私たちのチームはウォーターマーキングのワークロードを大幅に加速でき、イテレーションサイクルが短縮され、より複雑なモデルやテクニックを探求できるようになりました。このようなアクセスの深さにより、研究者は新しいアイデアのプロトタイプを作成したり、低レベルの最適化を実験したり、ウォーターマーキングシステムが最新の AI アクセラレータで実現できることの限界を押し広げたりすることができます。」