SandAI-org/MagiAttention

何を解決するか

MagiAttentionは、分散学習における超長文脈および異種マスクパターンに対するアテンション機構のスケーラビリティの課題に対処します。線形スケーラビリティと高いパフォーマンスを提供し、通常のコンテキスト並列（CP）学習設定で発生する計算および通信オーバーヘッドを低減することを目指しています。

動作方法

以下の主要なイノベーションを通じて分散アテンション機構を実装しています：

• Flex-Flash-Attention (FFA)：一般化されたアテンションマスク定式化（AttnSlice）と、多様なマスクタイプのコンパクトな表現と効率的な分散パーティショニングを可能にするカスタムカーネル。
• 計算負荷のバランス：細粒度のチャンクレベルシャーディング戦略とディスパッチソルバーを用いて、CPランク間で作業負荷を均等に分配します。
• ゼロ冗長通信：従来のリング型P2P通信を、前方および後方パスで冗長通信量を排除する新しいGroupCastおよびGroupReduceプリミティブに置き換えます。
• 適応的マルチステージオーバーラップ：計算と通信のスケジューリングにより、レイテンシを隠蔽しGPUの利用効率を最大化します。

対象ユーザー

超長文脈を持つ大規模モデルを学習する研究者や開発者向けです。たとえば、自己回帰型動画生成モデル（例：Magi-1）や、Megatron-LM、PyTorch FSDP、HuggingFace Transformersなどのフレームワークと統合するケースに適しています。

主な特徴

• 線形スケーラビリティ：H100およびB200 GPU上で分散学習環境でほぼ線形のスケーラビリティを実現。
• 広範なハードウェア対応：Hopper GPU上でFlash-Attention 3と同等のパフォーマンスを発揮し、Flash-Attention 4によるBlackwell GPUへの早期対応を実現。
• フレームワーク統合：Megatron-LM、PyTorch FSDP、HuggingFace Transformersとの容易な統合をサポート。
• 高度なマスク処理：多様かつ重複するアテンションマスクパターンをネイティブでサポート。
• アテンションシンク：分散学習可能なアテンションシンク機構の拡張を含む。