feifeibear/long-context-attention

何を解決するか

YunChang は、長文コンテキスト LLM のトレーニングと推論に使われる2つの一般的な分散注意機構である DeepSpeed-Ulysses と Ring-Attention の制限を克服する Unified Sequence Parallelism (USP) を実装しています。Ulysses はアテンションヘッド数に制限があり、GQA や MQA には適していません。一方 Ring-Attention は計算・通信の効率が低く、大規模展開時にデッドロックのリスクがあります。

動作方法

USP は DeepSpeed-Ulysses と Ring-Attention の長所を組み合わせたハイブリッドなシーケンス並列アプローチです。ユーザーは Ulysses の度合いと Ring の度合いを設定して、シーケンスを GPU 間で分割できます。このプロジェクトは FlashAttention (v2 および v3)、Torch Efficient など、さまざまなアテンションバックエンドと統合されており、因果マスクの最適化のためのロードバランシング戦略（例："zigzag" および "stripe"）も提供しています。

対象ユーザー

長文コンテキストの生成AIモデルをトレーニングまたはデプロイする研究者やエンジニアで、複数の GPU にわたるシーケンス長をスケーリングする必要がある人、特に GQA/MQA アーキテクチャや異種ネットワークハードウェアを使用している人向けです。

特徴

• ハイブリッドアプローチ：Ulysses と Ring アテンションを統合し、ヘッド数の制限と通信の非効率性を克服します。
• 広範なハードウェア対応：NVIDIA GPU（H100、A100 など）および NPU に対応。
• 柔軟なバックエンド：FlashAttention v2/v3、SageAttention、FlashInfer をサポート。
• 統合性：NVIDIA の TransformerEngine で既に採用されており、Megatron-LM フレームワークとも互換性があります。