AWS Trainium2でネイティブPyTorchを用いて線形アテンションモデルを訓練

i14およびAWS Generative AI Innovation Centerの研究チームは、AWS Trainium2アクセラレーターを使用してLaplace 3Bモデルの訓練と提供に成功しました。2026年10月に発表されたこの取り組みは、カスタムGPUカーネルに依存せずに、線形アテンションアーキテクチャを専用ハードウェア向けに最適化できることを実証しています。

この協業は、訓練時の二次計算量増加や推論時のメモリコスト増大という課題を抱える従来のトランスフォーマーのsoftmaxアテンションからの転換を示しています。TorchNeuron NativeによるネイティブPyTorchサポートを活用することで、チームはシーケンスごとの固定メモリフットプリントを維持しながら、著しいパフォーマンス指標を達成しました。

何が起きたか

大規模言語モデル（LLM）は通常、softmaxアテンションを使用するトランスフォーマーアーキテクチャに依存しています。推論時、これらのモデルは会話やドキュメントの長さに比例して線形に成長するキー・バリュー（KV）キャッシュを維持します。この成長により同時接続数が制限され、サービングコストが増加し、多くの場合、アプリケーションは過去のコンテキストを破棄せざるを得なくなります。これらの制限に対処するため、i14は線形アテンション原則に基づくモデルであるLaplace 3Bを開発しました。

標準的なトランスフォーマーとは異なり、Laplace 3Bはコンテキストの長さに関係なく、シーケンスあたり29.5 MBで一定となる再帰状態を使用します。チームはAWSと協力し、このアーキテクチャをTrainium2ハードウェアに適応させました。その際、ネイティブPyTorchコードをAWS Trainiumアクセラレーター上で直接実行可能にするベータ機能であるTorchNeuron Nativeを利用しました。このアプローチにより、モデルのGated Linear Attention（GLA）設計とMuonオプティマイザーの使用が保持されました。

結果として、大幅な効率向上が見られました。単一のtrn2.48xlargeインスタンスでは、モデルは約120,000トークン/秒で訓練されました。推論時には、1,280の並列シーケンス全体で11,544出力トークン/秒を提供しました。これらの数値は、ソフトウェアスタックが基盤アーキテクチャと正しく整合していれば、線形アテンションモデルが専用ハードウェア上で高いスループットを達成できることを示しています。

仕組み

線形アテンションモデルは、ペアごとのトークン比較ではなく、再帰状態を中心に計算を再編成します。従来のトランスフォーマーでは、訓練中のアテンションスコアがシーケンス長に対して二次的に増加します。Laplace 3Bはこれを、各トークンごとに固定サイズの行列状態を更新するメカニズムに置き換えます。これにより、訓練計算量はO(T^2)からO(T)に変化し、推論メモリは一定に保たれます。モデルはこの状態に履歴を圧縮するため、特定のワークロードに対する検索品質の慎重な評価が必要です。

このモデルをTrainium2に適応させるには、3つの具体的なエンジニアリング変更が必要でした。まず、チームはGPU最適化されたTritonカーネルを標準的なPyTorch演算子に置き換えました。GLAカーネルを128トークンの16チャンクで構成し、NeuronCoreオンチップバッファの128パーティションに合わせて整列させました。次に、Muonオプティマイザーの通信を再設計しました。gather-and-broadcast方式の代わりに2つのall-to-all集合通信を使用することで、ランクあたりのトラフィックを24分の1に削減し、スループットを1.63倍向上させました。最後に、コンパイラのグラフサイズおよびメモリ制限に適合させるため、逐次的なバックワードパスを静的にアンロールしました。

Trainium2アーキテクチャはこれらの最適化において重要な役割を果たしました。各NeuronCoreは、128パーティションを持つ28 MiBのソフトウェア管理型State Buffer（SBUF）を備えています。キャッシュとは異なり、コンパイラはすべてのデータ転送を事前にスケジュールするため、 eviction（追い出し）やmiss（ミス）が発生しません。各128トークンのチャンクを1つのパーティションにマッピングすることで、チームは測定されたスループットを14パーセント向上させました。モデルのチャンクサイズとハードウェアの物理パーティション間のこの整合性が、パフォーマンスを引き出す鍵となりました。

主要な詳細

• Laplace 3Bは、コンテキスト長に依存せず、シーケンスあたり29.5 MBの固定再帰状態を維持します。
• 1つのtrn2.48xlargeインスタンスでの訓練は、約120,000トークン/秒に達しました。
• 推論では、1,280の並列シーケンス全体で11,544出力トークン/秒を提供しました。
• 実装では、Trainium2上で標準的なPyTorch演算子を実行するためにTorchNeuron Nativeを使用しました。
• Muonオプティマイザーの通信トラフィックは、all-to-all集合通信を使用することで24分の1に削減されました。
• モデルは、NeuronCores上の128 SBUFパーティションに合わせて、128トークンの16チャンクを使用します。

なぜ重要なのか

AI製品を構築するエンジニアにとって、KVキャッシュのコスト増大は主要なボトルネックです。コンテキストが長くなるにつれてメモリ使用量が急増し、サーバーが処理できる同時ユーザー数が制限されます。Laplace 3Bのような線形アテンションモデルは、予測可能なサービングコストへの道を提供します。状態サイズがコンテキストとともに成長しないため、アプリケーションはメモリ割り当ての比例的な増加なしに長い履歴を保持できます。これは、永続的なメモリが不可欠なエージェント型ワークフローや大規模ドキュメント分析において特に価値があります。

さらに、この取り組みは、Trainium2のような専用ハードウェアが低レベルのカーネル書き直しを必要とせずに新しいアーキテクチャをサポートできることを実証しています。TorchNeuron Nativeを通じてネイティブPyTorchを使用することで、開発者は慣れ親しんだフレームワーク内に留まりながら、ハードウェア固有の最適化を活用できます。Muonオプティマイザーのためのall-to-all通信戦略の成功は、シャーディングされたシステムにおける分散訓練のための青写真も提供します。計算FLOPsが変わらない場合でも、通信オーバーヘッドを削減することで顕著なスループット向上が得られることを示しています。

実施可能なこと

• 長いコンテキストウィンドウと予測可能なメモリ使用量が必要なユースケースのために、線形アテンションモデルを検討してください。
• AWS Trainiumインスタンスでカスタムアーキテクチャを展開している場合は、TorchNeuron Nativeをテストしてください。
• 専用アクセラレーター上でスループットを最大化するために、データチャンクサイズをハードウェアバッファパーティションに合わせてください。
• ネットワークトラフィックのボトルネックを削減するために、分散オプティマイザーでall-to-all通信パターンを検討してください。
• softmaxアテンションから状態圧縮メカニズムへ切り替える際は、検索品質を監視してください。
• KVキャッシュの成長が並列度を制限していないかを確認するために、推論エンジンをプロファイリングしてください。