xformers

xFormers - 加速Transformer研究的工具箱

xFormers是:

• 可定制的构建模块：可独立使用、无需样板代码的可定制构建模块。这些组件与领域无关，xFormers被视觉、自然语言处理等领域的研究人员使用。
• 以研究为先：xFormers包含尖端组件，这些组件在主流库如PyTorch中尚未提供。
• 注重效率：因为迭代速度很重要，组件尽可能快速和内存高效。xFormers包含自己的CUDA内核，但在相关时也会调用其他库。

安装xFormers

• （推荐，Linux）使用conda安装最新稳定版：需要使用conda安装PyTorch 2.4.0

conda install xformers -c xformers

• （推荐，Linux和Windows）使用pip安装最新稳定版：需要PyTorch 2.4.0

# cuda 11.8版本
pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu118
# cuda 12.1版本
pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu121

• 开发版二进制文件：

# 使用conda或pip，与稳定版要求相同
conda install xformers -c xformers/label/dev
pip install --pre -U xformers

• 从源代码安装：如果你想使用其他版本的PyTorch（包括每日构建版）

# （可选）使构建速度更快
pip install ninja
# 如果在不同GPU类型上运行和构建，请设置TORCH_CUDA_ARCH_LIST
pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers
# （这可能需要几十分钟）

基准测试

内存高效的多头注意力机制 设置：A100使用f16，测量前向+反向传播的总时间

注意，这是精确的注意力计算，而不是近似，只需调用xformers.ops.memory_efficient_attention

更多基准测试

xFormers提供了许多组件，更多基准测试可在BENCHMARKS.md中找到。

（可选）测试安装

以下命令将提供有关xFormers安装的信息，以及已构建/可用的内核：

python -m xformers.info

使用xFormers

Transformer的关键概念

让我们从Transformer架构的经典概述开始（插图来自Lin等人的"A Survey of Transformers"）

您将在此插图中找到主要的仓库边界：Transformer通常由一系列注意力机制、用于编码位置信息的嵌入、前馈块和残差路径（通常称为前层或后层归一化）组成。这些边界并不适用于所有模型，但我们发现在实践中，通过一些调整，它可以涵盖大多数最先进的模型。

因此，模型并非以单一文件的形式实现，这些文件通常难以处理和修改。上图中的大多数概念都对应一个抽象层次，当一个子块存在变体时，应该始终可以选择其中任何一个。您可以专注于特定的封装层次并根据需要进行修改。

仓库结构

├── ops                         # 函数操作符
    └ ...
├── components                  # 组件库，任何组件都可直接使用
│   ├── attention
│   │    └ ...                  # 所有支持的注意力机制
│   ├── feedforward             #
│   │    └ ...                  # 所有支持的前馈网络
│   ├── positional_embedding    #
│   │    └ ...                  # 所有支持的位置编码
│   ├── activations.py          #
│   └── multi_head_dispatch.py  # （可选）多头包装器
|
├── benchmarks
│     └ ...                     # 大量可用于测试各个部分的基准测试
└── triton
      └ ...                     # （可选）所有triton部分，需要triton + CUDA GPU

主要特性

1. 多种可互换的注意力机制
2. 优化的构建模块，超越PyTorch原语
   1. 内存高效的精确注意力 - 速度提高10倍
   2. 稀疏注意力
   3. 块稀疏注意力
   4. 融合softmax
   5. 融合线性层
   6. 融合层归一化
   7. 融合dropout(activation(x+bias))
   8. 融合SwiGLU
3. 基准测试和测试工具
   1. 微基准测试
   2. Transformer块基准测试
   3. LRA，支持SLURM
4. 程序化和可扫描友好的层和模型构建
   1. 兼容层次化Transformer，如Swin或Metaformer
5. 可黑客化
   1. 不使用单一CUDA内核，可组合的构建模块
   2. 使用Triton进行某些优化部分，显式、Python化且用户可访问
   3. 原生支持SquaredReLU（除了ReLU、LeakyReLU、GeLU等），可扩展激活函数

安装疑难解答

• NVCC和当前CUDA运行时匹配。根据你的设置，你可能可以通过module unload cuda; module load cuda/xx.x更改CUDA运行时，可能还包括nvcc
• 你使用的GCC版本与当前NVCC功能匹配
• TORCH_CUDA_ARCH_LIST环境变量设置为你想要支持的架构。建议设置（构建慢但全面）是export TORCH_CUDA_ARCH_LIST="6.0;6.1;6.2;7.0;7.2;7.5;8.0;8.6"
• 如果从源代码构建时内存溢出，可以通过MAX_JOBS减少ninja的并行度（例如MAX_JOBS=2）
• 如果在使用conda安装时遇到UnsatisfiableError，请确保你的conda环境中安装了PyTorch，并且你的设置（PyTorch版本、cuda版本、python版本、操作系统）与xFormers的现有二进制文件匹配

许可证

xFormers采用BSD风格的许可证，详见LICENSE文件。

引用xFormers

如果你在出版物中使用xFormers，请使用以下BibTeX条目进行引用。

@Misc{xFormers2022,
  author =       {Benjamin Lefaudeux and Francisco Massa and Diana Liskovich and Wenhan Xiong and Vittorio Caggiano and Sean Naren and Min Xu and Jieru Hu and Marta Tintore and Susan Zhang and Patrick Labatut and Daniel Haziza and Luca Wehrstedt and Jeremy Reizenstein and Grigory Sizov},
  title =        {xFormers: A modular and hackable Transformer modelling library},
  howpublished = {\url{https://github.com/facebookresearch/xformers}},
  year =         {2022}
}

致谢

xFormers使用了以下仓库，或是以接近原始形式使用，或是作为灵感来源：

• Sputnik
• GE-SpMM
• Triton
• LucidRain Reformer
• RevTorch
• Nystromformer
• FairScale
• Pytorch Image Models
• CUTLASS
• Flash-Attention