使用 Better Transformer 进行快速 Transformer 推理¶

作者: Michael Gschwind

本教程介绍了 Better Transformer (BT) 作为 PyTorch 1.12 版本的一部分。在本教程中，我们将展示如何使用 Better Transformer 和 torchtext 进行生产推理。Better Transformer 是一种可用于生产的快速路径，可通过在 CPU 和 GPU 上提供高性能来加速 Transformer 模型的部署。对于基于 PyTorch 核心 nn.module 或 torchtext 的模型，快速路径功能可以透明地工作。

可以使用 Better Transformer 快速路径执行加速的模型是使用以下 PyTorch 核心 torch.nn.module 类 TransformerEncoder、TransformerEncoderLayer 和 MultiHeadAttention 的模型。此外，torchtext 已更新为使用核心库模块以受益于快速路径加速。（将来可能会在快速路径执行中启用其他模块。）

Better Transformer 提供两种类型的加速

针对 CPU 和 GPU 的原生多头注意力 (MHA) 实现，以提高整体执行效率。
利用 NLP 推理中的稀疏性。由于输入长度可变，输入标记可能包含大量填充标记，可以跳过对这些标记的处理，从而实现显著的加速。

快速路径执行受某些条件约束。最重要的是，模型必须在推理模式下执行，并对不收集梯度带信息的输入张量进行操作（例如，使用 torch.no_grad 运行）。

要在 Google Colab 中遵循此示例，请点击此处。

本教程中的 Better Transformer 功能¶

加载预训练模型（在 PyTorch 1.12 版本之前创建，不使用 Better Transformer）
在 CPU 上运行和基准测试推理，使用和不使用 BT 快速路径（仅限原生 MHA）
在（可配置）DEVICE 上运行和基准测试推理，使用和不使用 BT 快速路径（仅限原生 MHA）
启用稀疏性支持
在（可配置）DEVICE 上运行和基准测试推理，使用和不使用 BT 快速路径（原生 MHA + 稀疏性）

其他信息¶

有关 Better Transformer 的更多信息，请参阅 PyTorch.Org 博客更快 Transformer 推理的 Better Transformer。

设置

1.1 加载预训练模型

我们按照 torchtext.models 中的说明，从预定义的 torchtext 模型中下载 XLM-R 模型。我们还设置 DEVICE 以在加速器上执行测试。（根据您的环境适当地启用 GPU 执行。）

import torch
import torch.nn as nn

print(f"torch version: {torch.__version__}")

DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

print(f"torch cuda available: {torch.cuda.is_available()}")

import torch, torchtext
from torchtext.models import RobertaClassificationHead
from torchtext.functional import to_tensor
xlmr_large = torchtext.models.XLMR_LARGE_ENCODER
classifier_head = torchtext.models.RobertaClassificationHead(num_classes=2, input_dim = 1024)
model = xlmr_large.get_model(head=classifier_head)
transform = xlmr_large.transform()

1.2 数据集设置

我们设置了两种类型的输入：一个小输入批次和一个带有稀疏性的大的输入批次。

small_input_batch = [
               "Hello world",
               "How are you!"
]
big_input_batch = [
               "Hello world",
               "How are you!",
               """`Well, Prince, so Genoa and Lucca are now just family estates of the
Buonapartes. But I warn you, if you don't tell me that this means war,
if you still try to defend the infamies and horrors perpetrated by
that Antichrist- I really believe he is Antichrist- I will have
nothing more to do with you and you are no longer my friend, no longer
my 'faithful slave,' as you call yourself! But how do you do? I see
I have frightened you- sit down and tell me all the news.`

It was in July, 1805, and the speaker was the well-known Anna
Pavlovna Scherer, maid of honor and favorite of the Empress Marya
Fedorovna. With these words she greeted Prince Vasili Kuragin, a man
of high rank and importance, who was the first to arrive at her
reception. Anna Pavlovna had had a cough for some days. She was, as
she said, suffering from la grippe; grippe being then a new word in
St. Petersburg, used only by the elite."""
]

接下来，我们选择小或大输入批次，预处理输入并测试模型。

input_batch=big_input_batch

model_input = to_tensor(transform(input_batch), padding_value=1)
output = model(model_input)
output.shape

最后，我们设置基准迭代次数

ITERATIONS=10

执行

2.1 在 CPU 上运行和基准测试推理，使用和不使用 BT 快速路径（仅限原生 MHA）

我们在 CPU 上运行模型，并收集性能信息

第一次运行使用传统的（“慢路径”）执行。
第二次运行通过使用 model.eval() 将模型置于推理模式并使用 torch.no_grad() 禁用梯度收集来启用 BT 快速路径执行。

当模型在 CPU 上执行时，您可以看到性能提升（幅度将取决于 CPU 模型）。请注意，快速路径性能分析显示大部分执行时间都在原生 TransformerEncoderLayer 实现 aten::_transformer_encoder_layer_fwd 中。

print("slow path:")
print("==========")
with torch.autograd.profiler.profile(use_cuda=False) as prof:
  for i in range(ITERATIONS):
    output = model(model_input)
print(prof)

model.eval()

print("fast path:")
print("==========")
with torch.autograd.profiler.profile(use_cuda=False) as prof:
  with torch.no_grad():
    for i in range(ITERATIONS):
      output = model(model_input)
print(prof)

2.2 在（可配置）DEVICE 上运行和基准测试推理，使用和不使用 BT 快速路径（仅限原生 MHA）

我们检查 BT 稀疏性设置

model.encoder.transformer.layers.enable_nested_tensor

我们禁用 BT 稀疏性

model.encoder.transformer.layers.enable_nested_tensor=False

我们在 DEVICE 上运行模型，并收集在 DEVICE 上原生 MHA 执行的性能信息

第一次运行使用传统的（“慢路径”）执行。
第二次运行通过使用 model.eval() 将模型置于推理模式并使用 torch.no_grad() 禁用梯度收集来启用 BT 快速路径执行。

在 GPU 上执行时，您应该会看到显著的加速，特别是对于小输入批次设置

model.to(DEVICE)
model_input = model_input.to(DEVICE)

print("slow path:")
print("==========")
with torch.autograd.profiler.profile(use_cuda=True) as prof:
  for i in range(ITERATIONS):
    output = model(model_input)
print(prof)

model.eval()

print("fast path:")
print("==========")
with torch.autograd.profiler.profile(use_cuda=True) as prof:
  with torch.no_grad():
    for i in range(ITERATIONS):
      output = model(model_input)
print(prof)

2.3 在（可配置）DEVICE 上运行和基准测试推理，使用和不使用 BT 快速路径（原生 MHA + 稀疏性）

我们启用稀疏性支持

model.encoder.transformer.layers.enable_nested_tensor = True

我们在 DEVICE 上运行模型，并收集在 DEVICE 上原生 MHA 和稀疏性支持执行的性能信息

第一次运行使用传统的（“慢路径”）执行。
第二次运行通过使用 model.eval() 将模型置于推理模式并使用 torch.no_grad() 禁用梯度收集来启用 BT 快速路径执行。

在 GPU 上执行时，您应该会看到显著的加速，特别是对于包含稀疏性的大的输入批次设置

model.to(DEVICE)
model_input = model_input.to(DEVICE)

print("slow path:")
print("==========")
with torch.autograd.profiler.profile(use_cuda=True) as prof:
  for i in range(ITERATIONS):
    output = model(model_input)
print(prof)

model.eval()

print("fast path:")
print("==========")
with torch.autograd.profiler.profile(use_cuda=True) as prof:
  with torch.no_grad():
    for i in range(ITERATIONS):
      output = model(model_input)
print(prof)

总结¶

在本教程中，我们介绍了使用 PyTorch 核心 Better Transformer 对 Transformer Encoder 模型的支持，在 torchtext 中使用 Better Transformer 快速路径执行进行快速 Transformer 推理。我们演示了在 BT 快速路径执行可用之前训练的模型中使用 Better Transformer。我们演示并基准测试了两种 BT 快速路径执行模式的使用，即原生 MHA 执行和 BT 稀疏性加速。

使用 Better Transformer 进行快速 Transformer 推理¶

本教程中的 Better Transformer 功能¶

其他信息¶

总结¶

文档

教程

资源