AWS Neuron が Flash Attention カーネルを導入し、高いパフォーマンスと長いシーケンス長を実現

投稿日: 2024年7月11日

本日、AWS は Neuron 2.19 のリリースを発表しました。これにより Flash Attention カーネルのサポートが導入され、長いシーケンス長でも高性能な LLM モデルのトレーニングと推論が可能になります。

AWS Neuron は、生成 AI 向けに構築された AWS Inferentia および Trainium ベースのインスタンス用の SDK です。Neuron は PyTorch などの一般的な ML フレームワークと統合されます。Neuron には、Trn1 および Inf2 インスタンスでの AI モデルの高性能トレーニングと推論をサポートするコンパイラ、ランタイム、ツール、ライブラリが含まれています。

今回のリリースでは、トレーニングと推論の両方に新機能とパフォーマンスの改善が追加され、PyTorch 2.1 と PyTorch 1.13 用の新しい Ubuntu 22 Neuron DLAMI が追加されました。Neuron 2.19 では、Flash Attention カーネルのサポートが追加され、大規模なシーケンス長 (8K 以上) のトレーニング、Llama3 モデルのトレーニング、およびインターリーブパイプライン並列処理が可能になり、トレーニングの効率とリソース使用率が向上しました。推論用として、このリリースでは Flash Attention カーネルのサポートが追加され、最大 32,000 のコンテキスト長の LLM 推論が可能になります。Neuron 2.19 ではさらに、Llama3 モデル推論のサポートが追加され、Mistral-7b-V0.2 モデルでの連続バッチ処理のベータサポートが追加されました。Neuron 2.19 では以下の新しいツールが導入されました: EKS の Neuron Node Problem Detector and Recovery プラグイン、および Kubernetes の強化された Neuron メトリクスモニタリングを可能にする Neuron Monitor for EKS です。

AWS Neuron SDK を使用して、Trn1 インスタンスと Inf2 インスタンスでモデルをトレーニングおよびデプロイできます。これらのインスタンスは、オンデマンドインスタンス、リザーブドインスタンス、スポットインスタンスとして、または Savings Plan の一部として AWS リージョンで利用できます。

Neuron 2.19 の機能のリストについては、Neuron リリースノートを参照してください。Neuron を使い始めるには、以下を参照してください。
AWS Neuron
Inf2 インスタンス
Trn1 インスタンス