meta-pytorch/attention-gym

Helpful tools and examples for working with flex-attention

解决的问题

Attention Gym 提供了一个用于实现和实验新型注意力机制的 Playground 和构建模块库。它简化了自定义注意力掩码和分数修改的流程，并提供了高效实现稀疏注意力和线性注意力变体的方案，这些原本从零开始构建会非常复杂。

工作原理

该项目通过三种主要工作流运行：

  1. FlexAttention 集成：提供 mask_mod 和 score_mod 函数，可直接传递给 PyTorch 的 FlexAttention API，以自定义注意力的计算方式。
  2. 专用内核：包含针对稀疏注意力（通过 gather_attn）和线性注意力（包括 GDN 和 KDA 路径）的专用 API 和内核。
  3. 实用示例：提供一系列基准测试和真实工作负载，涵盖分页注意力、环形注意力和 CUDA Graphs，展示如何在类似生产环境的场景中应用这些构建模块。

适用人群

专为使用 PyTorch 2.5+ 构建新注意力变体或希望优化注意力性能的 AI 研究人员和开发者设计。

主要亮点

  • 与 PyTorch FlexAttention 集成，轻松实现掩码和分数修改。
  • 支持稀疏注意力模式和高效的线性注意力内核。
  • 针对 GDN 和 KDA（分块、递归、解码路径）的专用原语。
  • 包含可变稀疏注意力、编译和性能分析的全面示例。

相关

  • 项目
    lucidrains/perceiver-pytorch

    一个实现 Perceiver 和 Perceiver IO 架构（以及一个小型实验性变体）的 PyTorch 库。它提供开箱即用的 Perceiver、PerceiverIO 和 PerceiverLM 类，支持傅立叶位置编码，并可通过 pip install perceiver-pytorch 安装。此仓库是一个专注于处理大型多模态输入的前沿 Transformer 模型 AI/ML 项目。

  • 项目
    Shmuma/ptan

    PTAN (PyTorch AgentNet) 是一个小型开源库，为基于 PyTorch 的智能体提供可重用的强化学习构建模块（如经验回放池、训练循环以及 Ignite 集成），适用于 OpenAI-Gym 环境。

  • 项目
    ixaxaar/pytorch-dnc

    一个 PyTorch 库，实现可微神经计算机（DNC）、稀疏 DNC（SDNC）和稀疏访问记忆（SAM），包含安装说明、API 使用方法、调试支持和示例训练任务（复制、加法、argmax）。

  • 项目
    meta-pytorch/torchtune

    torchtune 是一个基于 PyTorch 的库，提供多种现代大语言模型（Llama、Gemma、Qwen、Phi 等）的微调、RL-HF、蒸馏、量化感知训练和推理的即用型配方与配置。它通过简单的 CLI（`tune`）和 YAML 文件抽象了分布式训练、内存优化技巧和设备支持，与 Hugging Face、torch-ao、FSDP2 和日志工具集成。尽管 2025 年开发已停止，但它仍是 LLM 实验的实用、开源工具包。

  • 项目
    webis-de/small-text

    一个支持 scikit-learn、PyTorch 和 Hugging Face transformers 的 Python 库，提供即用型文本分类主动学习循环，包含多种查询策略、初始化方法和停止条件。