xformers - 可破解和优化的Transformers 构建块,支持可组合的结构。

Created at: 2021-10-14 02:08:50
Language: Python
License: NOASSERTION

使用 conda 安装 下载 许可证 在科拉布中打开
圈子CI 科考夫 黑
欢迎公关


xFormers - 加速变压器研究的工具箱

xFormers 是:

  • 可自定义的构建基块:无需样板代码即可使用的独立/可自定义构建基块。这些组件与领域无关,xFormers被视觉,NLP等研究人员使用。
  • 首先研究:xFormers包含前沿组件,这些组件在Pytorch等主流库中尚不可用。
  • 构建时考虑效率:由于迭代速度很重要,因此组件尽可能快速且节省内存。xFormers包含自己的CUDA内核,但在相关时会发送到其他库。

安装 xFormer

  • (推荐)使用二进制文件:我们为 Linux 和最新的 PyTorch 版本提供二进制文件。在 conda 中安装 pytorch 后,使用 conda 安装 xFormers:
conda install xformers -c xformers/label/dev
  • 从源代码:或者,如果没有可用的二进制文件(例如对于 Windows),你也可以从源代码安装:
# (Optional) Makes the build much faster
pip install ninja
# Set TORCH_CUDA_ARCH_LIST if running and building on different GPU types
pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers
# (this can take dozens of minutes)
  • pip 轮:pip 上没有可用的更新包,请从 conda 或源代码安装

结果

内存效率高 MHA ViTS 基准测试 设置:f16 上的 A100,测量前进+后退传递的总时间

请注意,这是精确的注意,而不是近似值,只需调用xformers.ops.memory_efficient_attention

更多基准

xFormers提供了许多组件,BENCHMARKS.md 中提供了更多基准测试。

(可选)测试安装

此命令将提供有关 xForers 安装的信息,以及构建/可用的内核:

python -m xformers.info

使用 xFormer

变压器关键概念

让我们从变压器架构的经典概述开始(插图来自Lin等人,“变压器调查”)

你将在此图中找到关键存储库边界:转换器通常由一组注意力机制、用于编码某些位置信息的嵌入、前馈块和残差路径(通常称为层前或后层范数)组成。这些边界并不适用于所有模型,但我们在实践中发现,给定一些住宿,它可以捕获大部分最先进的技术。

因此,模型不会在整体式文件中实现,这些文件通常处理和修改起来很复杂。上图中存在的大多数概念都对应于抽象级别,当给定子块存在变体时,应该始终可以选择其中任何一个。你可以专注于给定的封装级别,并根据需要对其进行修改。

回购映射

├── ops                         # Functional operators
    └ ...
├── components                  # Parts zoo, any of which can be used directly
│   ├── attention
│   │    └ ...                  # all the supported attentions
│   ├── feedforward             #
│   │    └ ...                  # all the supported feedforwards
│   ├── positional_embedding    #
│   │    └ ...                  # all the supported positional embeddings
│   ├── activations.py          #
│   └── multi_head_dispatch.py  # (optional) multihead wrap
|
├── benchmarks
│     └ ...                     # A lot of benchmarks that you can use to test some parts
└── triton
      └ ...                     # (optional) all the triton parts, requires triton + CUDA gpu
注意力机制

前馈机制

位置嵌入

剩余路径

初始 化

这是完全可选的,并且仅在通过 xFormers 生成完整模型时发生,而不是在单独拾取零件时发生。

基本上有两种初始化机制公开,但用户可以在事后自由初始化权重,因为他/她认为合适。

  • 部件可以公开一个方法,该方法定义了合理的默认值
    init_weights()
  • xFormers 支持特定的 init 方案,这些方案可以优先于 init_weights()

如果使用第二个代码路径(通过模型工厂构造模型),我们会检查所有权重是否已初始化,如果不是这种情况,则可能会出错(如果你设置

xformers.factory.weight_init.__assert_if_not_initialized = True
)

支持的初始化方案包括:

指定 init 方案的一种方法是将字段设置为匹配的枚举值。这可以很容易地扩展,请随时提交 PR !

config.weight_init

主要特点

  1. 多种注意力机制,可互换
  2. 优化的构建块,超越 PyTorch 原语
    1. 记忆效率高的精确注意力 - 速度提高 10 倍
    2. 注意力稀少
    3. 块稀疏注意力
    4. 熔融软最大值
    5. 熔融线性层
    6. 熔融层规范
    7. 熔断压差(激活(x+偏置))
    8. 融合的斯威格鲁
  3. 基准测试和测试工具
    1. 微型工作台标记
    2. 变压器块基准
    3. LRA,支持SLURM
  4. 编程和扫描友好的层和模型构建
    1. 与分层变形金刚兼容,如 Swin 或 Metaformer
  5. 可破解
    1. 不使用整体式 CUDA 内核、可组合构建块
    2. Triton 用于一些优化的部分,显式、pythonic 和用户可访问的部件
    3. 对SquaredReLU的原生支持(在ReLU,LeakyReLU,GeLU之上),可扩展的激活

安装疑难解答

  • NVCC 和当前的 CUDA 运行时匹配。根据你的设置,你可能能够使用 更改 CUDA 运行时,也可能
    module unload cuda; module load cuda/xx.x
    nvcc
  • 你使用的 GCC 版本与当前的 NVCC 功能匹配
  • env 变量设置为要支持的架构。建议的设置(构建缓慢但全面)是
    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
    MAX_JOBS=2
    )
  • 如果你在使用 conda 安装时遇到 UnsatfiableError,请确保你在 conda 环境中安装了 pytorch,并且你的设置(pytorch 版本、cuda 版本、python 版本、OS )与 xForers 的现有二进制文件匹配

许可证

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},
  title =        {xFormers: A modular and hackable Transformer modelling library},
  howpublished = {\url{https://github.com/facebookresearch/xformers}},
  year =         {2022}
}

捐赠

以下存储库在 xFormers 中使用,要么接近原始形式,要么作为灵感来源: