摘要 Transformer 在自然语言处理(NLP)领域取得了最先进的性能,并正日益受到其他诸多任务的青睐。然而,由于它们计算成本高昂,这已成为其广泛产品化的一大障碍。随 PyTorch 1.12 发布,BetterTransformer 为 Transformer 编码器推理实现了 torch.nn.TransformerEncoder 的向后兼容快速路径,且无需模型作者修改现有模型。在许多常见的执行场景中,BetterTransformer 可带来超过 2 倍的加速和吞吐量提升。如需使用 BetterTransformer,请安装 PyTorch 1.12,即刻开始使用 PyTorch API 体验高质量、高性能的 Transformer 模型。

Transformer 编码器架构示意图(来自“Attention Is All You Need”)。在推理过程中,整个模块将作为一个原生的 PyTorch 函数执行。
在本篇博客中,我们将分享以下主题:性能改进、向后兼容性以及如何利用快速路径(FastPath)。请在下方了解更多相关信息。
性能改进
BetterTransformer 发布时,针对 CPU 和 GPU 实现了 MultiHeadAttention 和 TransformerEncoderLayer 的加速原生实现。这些快速路径已集成到标准的 PyTorch Transformer API 中,并将加速 TransformerEncoder、TransformerEncoderLayer 和 MultiHeadAttention nn.modules。这些新模块实现了两类优化:(1) 融合算子(fused kernels)结合了通常用于实现 Transformer 的多个独立算子,从而提供了更高效的实现;(2) 利用输入中的稀疏性,避免对填充(padding)标记执行不必要的计算。在许多用于自然语言处理的 Transformer 模型中,填充标记通常占据了输入批次(batch)的很大一部分。
向后兼容性
值得庆幸的是,无需对模型进行任何更改,即可享受 BetterTransformer 带来的性能提升。 要从快速路径执行中获益,输入和运行条件必须满足某些访问要求(见下文)。虽然 Transformer API 的内部实现已发生变化,但 PyTorch 1.12 严格保持了与之前版本中内置的 Transformer 模块的兼容性,使 PyTorch 用户既能使用旧版本 PyTorch 创建和训练的模型,又能享受到 BetterTransformer 带来的性能改进。
除了启用 PyTorch nn.Modules,BetterTransformer 还为 PyTorch 库提供了改进。性能优势将通过两种不同的实现途径提供:
- 透明加速: PyTorch nn.Modules(如
MultiHeadAttention)以及更高级 Transformer 组件的现有用户,将自动受益于新模块的性能提升。一个例子是 torchvision 库中使用的视觉 Transformer (ViT) 实现(代码链接)。 - Torchtext 库加速: 作为该项目的一部分,我们优化了 Torchtext,使其构建在 PyTorch 核心 API 之上,从而在保持与先前库版本及先前版本训练的模型严格且透明的兼容性的同时,受益于 BetterTransformer 的增强功能。在 Torchtext 中使用 PyTorch Transformer 也确保了 Torchtext 将从 PyTorch Transformer 实现未来的预期改进中获益。
利用快速路径
BetterTransformer 是 PyTorch Transformer API 的快速路径。该快速路径是针对 CPU 和 GPU 的核心 Transformer 函数的一种原生、专用实现,适用于常见的 Transformer 用例。
为了利用输入稀疏性(即填充)来加速您的模型(见图 2),请在实例化 TransformerEncoder 时设置关键字参数 enable_nested_tensor=True,并在推理过程中传入 src_key_padding_mask 参数(用于标记填充位)。这要求填充掩码必须是连续的,这也是通常情况下的做法。
目前,BetterTransformer 的加速仅适用于推理阶段的 Transformer 编码器模型。要从快速路径执行中获益,模型必须由以下组件中的任意一个构成:TransformerEncoder、TransformerEncoderLayer 或 MultiHeadAttention (MHA)。快速路径执行还需满足一些标准。最重要的是,模型必须在推理模式下执行,并在不收集梯度信息(例如运行在 torch.no_grad 下)的输入张量上运行。完整的条件列表可以分别在这些链接中找到:nn.MultiHeadAttention 和 nn.TransformerEncoder。如果不满足这些标准,控制流将转向传统的 PyTorch 1.11 Transformer 实现,该实现具有相同的 API,但缺乏快速路径带来的性能提升。
使用 PyTorch MultiheadAttention 模块的其他 Transformer 模型(例如解码器模型)也将受益于 BetterTransformer 快速路径。未来计划的工作是将端到端的 BetterTransformer 快速路径扩展到基于 TransformerDecoder 的模型,以支持流行的 seq2seq 和仅解码器(例如 OPT)模型架构,并扩展到训练任务中。
加速效果
下图显示了 BERT-base 模型在小规模和大规模输入下获得的性能表现:

图 1:使用 BetterTransformer 快速路径执行的 PyTorch 1.12 改进

图 2:使用 BetterTransformer 快速路径执行的 PyTorch 1.12 改进
通过 enable_nested_tensor=True 启用稀疏性优化
BetterTransformer 包含两类优化:(1) 融合算子,在单个内核中更高效地实现多个操作;(2) 利用稀疏性,避免对填充标记进行不必要的处理。小输入规模的性能增强主要受益于融合算子实现,无论填充量多少,都能显示出恒定的性能提升。虽然大输入仍然受益于融合算子,但繁重的计算处理限制了融合算子可能带来的好处,因为基准性能已经接近理论峰值。然而,随着我们增加填充量,性能会显著提升,因为通过利用 NLP 工作负载中填充引入的稀疏性,可以避免大量的计算。
未来工作
作为我们正在进行的 PyTorch BetterTransformer 工作的一部分,我们正致力于将 BetterTransformer 的改进扩展到 Transformer 解码器。我们的目标是将范围从推理扩展到训练。
我们正在建立合作伙伴关系,以便在 FairSeq、MetaSeq 和 HuggingFace 等其他库上启用 BetterTransformer,从而造福所有基于 Transformer 的 PyTorch 模型。作为本系列博客的一部分,我们将提供有关 BetterTransformer 加速在更广泛的 PyTorch 生态系统中进展的后续更新。
致谢:作者衷心感谢 Lin Qiao、Ajit Mathews、Andrew Tulloch、Dmytro Dzhulgakov、Natalia Gimelshein、Emad El-Haraty、Mark Saroufim、Adnan Aziz、Geeta Chauhan 和 Hamid Shojanazeri 在本项目过程中及博客准备阶段提供的支持、贡献和许多有益建议。