使用Optimum的BetterTransformer为常见的transformer结构模型进行推理加速

原创

已于 2024-07-03 14:37:42 修改 · 875 阅读

9 ·

CC 4.0 BY-SA版权

文章标签：

#transformer #深度学习 #人工智能

于 2024-07-01 17:38:04 首次发布

概述

🤗 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 的情况下可以预期加速。