xFormers — Transformer 加速库

xFormers - 加速 Transformer 研究的工具箱

xFormers 是：
- 可定制的构建模块：独立/可定制的构建模块，无需样板代码即可使用。这些组件与领域无关，已被视觉、NLP 等领域的研究人员广泛使用。
- 研究优先：xFormers 包含前沿组件，这些组件在 PyTorch 等主流库中尚不可用。
- 为效率而生：考虑到迭代速度至关重要，组件设计得尽可能快速且内存高效。xFormers 包含自己的 CUDA 内核，但在适当时也会调度到其他库。

安装 xFormers

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

# [linux & win] cuda 12.6 版本
pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu126
# [linux & win] cuda 12.8 版本
pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu128
# [linux & win] cuda 13.0 版本
pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu130
# [仅限 linux] (实验性) rocm 7.1 版本
pip3 install -U xformers --index-url https://download.pytorch.org/whl/rocm7.1

• 开发版二进制包：

# 与稳定版要求相同
pip install --pre -U xformers

• 从源码安装：例如，如果你想与其他版本的 PyTorch（包括 nightly 版本）一起使用

# (可选) 使构建过程更快
pip install ninja
# 如果在不同类型 GPU 上运行和构建，请设置 TORCH_CUDA_ARCH_LIST
# 注意：必须先安装 pytorch！
pip install -v --no-build-isolation -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers
# (这可能需要几十分钟)

性能基准

内存高效的多头注意力 (MHA)

配置：A100 GPU，f16 精度，测量一次前向+后向传递的总时间

请注意，这是精确注意力，而非近似，只需调用 xformers.ops.memory_efficient_attention 即可。

更多基准测试

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

（可选）测试安装

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

python -m xformers.info

使用 xFormers

主要特性

1. 超越 PyTorch 原语的优化构建模块
2. 内存高效的精确注意力 - 速度提升高达 10 倍
3. 稀疏注意力
4. 块稀疏注意力
5. 融合 Softmax
6. 融合线性层
7. 融合层归一化
8. 融合 Dropout(activation(x+bias))
9. 融合 SwiGLU

安装故障排除

• 确保 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"
• 如果从源码构建时内存不足 (OOM)，可以通过 MAX_JOBS 减少 ninja 的并行度（例如 MAX_JOBS=2）
• 如果在 Windows 上遇到错误信息 Filename longer than 260 characters，请确保在操作系统级别启用了长路径支持，并执行命令 git config --global core.longpaths true

许可证

xFormers 采用 BSD 风格许可证，详见 LICENSE 文件。
它包含来自 triton-lang/kernels 仓库的代码。

引用 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