meta-pytorch/attention-gym
Helpful tools and examples for working with flex-attention
何を解決するか
Attention Gym は、新しいアテンションメカニズムの実装と実験を行うためのプレイグラウンドとビルディングブロックのライブラリを提供します。カスタムアテンションマスクやスコアの変更を簡単に作成できるようにし、通常は手動で構築するのが複雑なスパースアテンションや線形アテンションの効率的な実装を提供します。
仕組み
このプロジェクトは3つの主要なワークフローで動作します：
- FlexAttention 統合：PyTorch の FlexAttention API に直接渡せる
mask_mod および score_mod 関数を提供し、アテンションの計算方法をカスタマイズできます。
- 専用カーネル：
gather_attn を通じたスパースアテンション用の専用 API とカーネル、および GDN と KDA パスを含む線形アテンション用のカーネルを提供します。
- 実用的な例：ページ化アテンション、リングアテンション、CUDA Graphs をカバーするベンチマークと実世界のワークロードのセットを提供し、これらのビルディングブロックを本番環境に近いシナリオでどのように適用するかを示します。
対象ユーザー
PyTorch 2.5+ を使って新しいアテンション変種を構築している、またはアテンションのパフォーマンスを最適化したい AI リサーチャーおよび開発者向けに設計されています。
特徴
- PyTorch FlexAttention との統合により、マスクやスコアの変更が簡単です。
- スパースアテンションパターンと効率的な線形アテンションカーネルのサポート。
- GDN および KDA（チャンク、再帰、デコードパス）用の専用プリミティブ。
- 変動スパースアテンション、コンパイル、プロファイリングをカバーする包括的な例。
関連
- プロジェクト
lucidrains/perceiver-pytorchPerceiverおよびPerceiver IOアーキテクチャ（および小さな実験的バリアント）を実装したPyTorchライブラリです。すぐに使えるPerceiver、PerceiverIO、PerceiverLMクラスを提供し、フーリエ位置エンコーディングをサポートしており、pip install perceiver-pytorchでインストール可能です。このリポジトリは、大規模なマルチモーダル入力を処理するための最先端のTransformerモデルに焦点を当てたAI/MLプロジェクトです。
- プロジェクト
Shmuma/ptanPTAN (PyTorch AgentNet) は、PyTorchベースのエージェント向けに、経験バッファ、学習ループ、Ignite統合などの再利用可能な強化学習用ビルディングブロックを提供する小型のオープンソースライブラリです。OpenAI-Gym環境に対応しています。
- プロジェクト
ixaxaar/pytorch-dncPyTorchライブラリで、Differentiable Neural Computers (DNC)、Sparse DNC (SDNC)、Sparse Access Memory (SAM) を実装。インストール手順、APIの使い方、デバッグサポート、および例題トレーニングタスク（コピー、加算、argmax）を含む。
- プロジェクト
meta-pytorch/torchtunetorchtuneは、PyTorchベースのライブラリで、Llama、Gemma、Qwen、Phiなど多くの現代的なLLM（大規模言語モデル）のファインチューニング、RL-HF、蒸留、量子化対応学習、推論のための即時実行可能なレシピと設定を提供します。分散学習、メモリ節約テクニック、デバイスサポートをシンプルなCLI（`tune`）とYAMLファイルの背後で抽象化し、Hugging Face、torch-ao、FSDP2、ログ記録ツールと統合されています。2025年に開発が停止しましたが、LLMの実験に実用的でオープンソースなツールキットとして残っています。
- プロジェクト
webis-de/small-textscikit-learn、PyTorch、Hugging Face transformersをサポートする、テキスト分類用の即時利用可能なアクティブラーニングループを提供するPythonライブラリ。多数のクエリ戦略、初期化方法、停止基準を備えています。