概述
🤗 Optimum 提供了一个名为 BetterTransformer 的 API,这是标准 PyTorch Transformer API 的快速路径,能够通过稀疏性和融合内核(如 Flash Attention)在 CPU 和 GPU 上实现有趣的加速。目前,BetterTransformer 支持来自原生 nn.TransformerEncoderLayer
的快速路径,以及 torch.nn.functional.scaled_dot_product_attention
中的 Flash Attention 和 Memory-Efficient Attention。
快速入门
自 1.13 版本以来,PyTorch 发布 了其标准 Transformer API 的稳定版本,为基于 Transformer 的模型提供了开箱即用的性能改进。你可以在大多数消费类设备上获得有趣的加速,包括 CPU、旧版和新版的 NVIDIA GPU。现在,你可以在 🤗 Optimum 中与 Transformers 一起使用这个功能,并将其用于 Hugging Face 生态系统中的主要模型。
在 2.0 版本中,PyTorch 将本地缩放点积注意力运算符 (SDPA) 包含在 torch.nn.functional
中。这个函数包含了几种实现方式,可以根据输入和使用的硬件进行应用。有关更多信息,请参见官方文档和此博客文章中的基准测试。
我们在 🤗 Optimum 中提供了这些优化的开箱即用集成,因此你可以转换任何支持的 🤗 Transformers 模型,使其在相关情况下使用优化路径和 scaled_dot_product_attention
函数。
PyTorch 原生的 scaled_dot_product_attention
正在慢慢成为默认值,并在 🤗 Transformers 中集成。对于在 Transformers 中支持 SDPA 的模型,我们不再推荐使用 BetterTransformer,而是建议直接使用 Transformers 和 PyTorch 的最新版本进行注意力优化(Flash Attention 和 memory-efficient attention),通过 SDPA 实现。
PyTorch 原生的 scaled_dot_product_attention
运算符只能在未提供 attention_mask
的情况下调度到 Flash Attention。
因此,默认情况下,在训练模式中,BetterTransformer 集成 放弃了 mask 支持,仅可用于不需要批量训练填充 mask 的训练。例如,这种情况适用于掩码语言模型或因果语言模型。BetterTransformer 不适用于需要填充 mask 任务的模型微调。
在推理模式中,为了保证正确性,保留填充 mask,因此仅在 batch size = 1 的情况下可以预期加速。