tl;dr 变换器在 NLP 领域取得了最先进的性能,并且正在成为众多其他任务的流行选择。它们计算成本高昂,这已经成为它们广泛生产化的障碍。随着 PyTorch 1.12 的发布,BetterTransformer 实现了与 torch.nn.TransformerEncoder
兼容的快速路径,用于 Transformer 编码器推理,并且不需要模型作者修改他们的模型。BetterTransformer 的改进可以在许多常见执行场景中实现超过 2 倍的速度和吞吐量提升。要使用 BetterTransformer,请安装 PyTorch 1.12,并开始使用 PyTorch API 今天使用高质量、高性能的 Transformer 模型。
变换器编码器架构图(来自《Attention Is All You Need》)。在推理期间,整个模块将作为一个单一的 PyTorch 原生函数执行。
在本文中,我们分享了以下主题——性能改进、向后兼容性和利用 FastPath。以下了解更多这些主题。
性能提升
BetterTransformer 发布,包含针对 CPU 和 GPU 的加速原生实现的多头注意力(MultiHeadAttention)和 Transformer 编码器层(TransformerEncoderLayer)。这些快速路径已集成到标准的 PyTorch Transformer API 中,并将加速 TransformerEncoder、TransformerEncoderLayer 和 MultiHeadAttention nn.modules。这些新模块实现了两种优化类型:(1) 熔合内核将通常用于实现 Transformers 的多个单个操作器合并,以提供更高效的实现;(2) 利用输入的稀疏性来避免对填充标记执行不必要的操作。填充标记在许多用于自然语言处理的 Transformer 模型中经常占输入批次的很大一部分。
向后兼容性
优先进步地,无需对模型进行任何更改即可从 BetterTransformer 提供的性能提升中受益。为了从快速路径执行中受益,输入和操作条件必须满足某些访问条件(见下文)。虽然 Transformer API 的内部实现已更改,但 PyTorch 1.12 与之前版本中提供的 Transformer 模块保持严格兼容,使 PyTorch 用户能够在使用 BetterTransformer 改进的同时,继续使用之前 PyTorch 版本创建和训练的模型。
除了启用 PyTorch nn.Modules 之外,BetterTransformer 还为 PyTorch 库提供了改进。性能提升将通过两种不同的启用路径提供:
-
透明加速:当前使用 PyTorch nn.Modules(如 MultiHeadAttention)以及高级 Transformer 组件的用户将自动受益于新 nn.Modules 的性能提升。例如,在 torchvision 库中使用的视觉 Transformer(ViT)实现(代码链接)。
-
火炬文本库加速:作为本项目的一部分,我们对火炬文本进行了优化,使其基于 PyTorch 核心 API 构建,以利用 BetterTransformer 增强功能,同时保持与之前版本库和用之前火炬文本版本训练的模型的严格和透明兼容性。在火炬文本中使用 PyTorch Transformers 也确保火炬文本将受益于 PyTorch Transformer 实现预期的未来增强。
利用 Fastpath 的优势
BetterTransformer 是 PyTorch Transformer API 的 fastpath。fastpath 是对 CPU 和 GPU 上关键 Transformer 函数的本地、专用实现,适用于常见的 Transformer 用例。
要利用输入稀疏性(即填充)来加速您的模型(见图 2),在实例化 TransformerEncoder 时设置关键字参数 enable_nested_tensor=True
,并在推理时传入 src_key_padding_mask
参数(表示填充标记)。这要求填充掩码是连续的,这是典型情况。
目前,BetterTransformer 的加速功能仅适用于推理中使用的 transformer 编码器模型。要利用 fastpath 执行,模型必须由以下任何一种组件组成:TransformerEncoder、TransformerEncoderLayer 或 MultiheadAttention(MHA)。fastpath 执行也受一些条件的限制。最重要的是,模型必须在推理模式下执行,并且操作不收集梯度记录信息的输入张量(例如,使用 torch.no_grad 运行)。有关 nn.MultiHeadAttention 和 nn.TransformerEncoder 的完整条件列表,请参阅以下链接。如果不符合这些条件,控制流将切换到传统的 PyTorch 1.11 Transformer 实现,该实现具有相同的 API,但缺乏 fastpath 性能提升。
其他 transformer 模型(如解码器模型)如果使用 PyTorch MultiheadAttention 模块,将受益于 BetterTransformer 的 fastpath。计划未来的工作是将端到端的 BetterTransformer fastpath 扩展到基于 TransformerDecoder 的模型,以支持流行的 seq2seq 和解码器仅(例如,OPT)模型架构,以及训练。
加速
以下图表展示了使用 BERT-base 模型在小规模和大规模输入下所达到的性能:
图 1:PyTorch 1.12 改进与 BetterTransformer 快速路径执行
图 2:PyTorch 1.12 改进与 BetterTransformer 快速路径执行
启用 enable_nested_tensor=True 的稀疏优化
BetterTransformer 包括两种优化类型:(1)通过单个内核实现多个操作以提高效率的融合内核,以及(2)通过避免对填充标记进行不必要的处理来利用稀疏性。对于小输入大小,增强性能主要得益于融合内核实现,并且无论填充量如何,都显示出恒定的性能提升。然而,对于大输入,虽然融合内核仍然有益,但计算密集型处理限制了融合内核可能带来的好处,因为基准性能已经接近理论峰值。但是,随着填充量的增加,性能会显著提高,因为可以通过利用 NLP 工作负载中由填充引入的稀疏性来避免越来越多的计算。
未来工作
作为我们针对 PyTorch BetterTransformer 的持续工作的一部分,我们正在努力将 BetterTransformer 的改进扩展到 Transformer 解码器。我们的目标是超越推理,扩展到训练。
我们正在合作,使 BetterTransformer 能够在 FairSeq、MetaSeq 和 HuggingFace 等更多库上运行,以惠及所有基于 Transformer 的 PyTorch 模型。我们将在此博客系列中提供关于 BetterTransformer 加速在更大 PyTorch 生态系统中的进展的未来更新。
致谢:作者们想感谢林乔、阿吉特·马休斯、安德鲁·图洛奇、德米特罗·祖尔加科夫、娜塔莉亚·吉梅尔申、艾玛德·埃尔-哈拉蒂、马克·萨拉菲姆、阿丹·阿齐兹、吉塔·乔汗和哈米德·肖贾纳泽里在项目过程中的支持、贡献以及许多有益的建议,以及在准备此博客方面的帮助。