Transformers 文档

TensorFlow模型的XLA集成

TensorFlow模型的XLA集成

加速线性代数,简称XLA,是一个用于加速TensorFlow模型运行时的编译器。来自官方文档

XLA(加速线性代数)是一个特定领域的线性代数编译器,可以在不更改源代码的情况下加速TensorFlow模型。

在 TensorFlow 中使用 XLA 非常简单——它已经包含在 tensorflow 库中,并且可以通过在任何图形创建函数(如 tf.function)中使用 jit_compile 参数来触发。当使用像 fit()predict() 这样的 Keras 方法时,您只需将 jit_compile 参数传递给 model.compile() 即可启用 XLA。然而,XLA 不仅限于这些方法——它还可以用于加速任何任意的 tf.function

🤗 Transformers 中的多个 TensorFlow 方法已被重写为与 XLA 兼容,包括针对 GPT2T5OPT 等模型的文本生成,以及针对 Whisper 等模型的语音处理。

虽然加速的确切程度很大程度上取决于模型,但对于🤗 Transformers中的TensorFlow文本生成模型,我们注意到加速了约100倍。本文将解释如何为这些模型使用XLA以获得最大性能。如果您有兴趣了解更多关于基准测试和我们XLA集成背后的设计理念,我们还将提供额外资源的链接。

使用XLA运行TF函数

让我们考虑以下在TensorFlow中的模型:

import tensorflow as tf

model = tf.keras.Sequential(
    [tf.keras.layers.Dense(10, input_shape=(10,), activation="relu"), tf.keras.layers.Dense(5, activation="softmax")]
)

上述模型接受的输入维度为 (10, )。我们可以像这样使用模型进行前向传递:

# Generate random inputs for the model.
batch_size = 16
input_vector_dim = 10
random_inputs = tf.random.normal((batch_size, input_vector_dim))

# Run a forward pass.
_ = model(random_inputs)

为了使用XLA编译的函数运行前向传递,我们需要执行以下操作:

xla_fn = tf.function(model, jit_compile=True)
_ = xla_fn(random_inputs)

默认的call()函数用于编译XLA图。但如果你想将任何其他模型函数编译到XLA中,也可以通过以下方式实现:

my_xla_fn = tf.function(model.my_xla_fn, jit_compile=True)

使用XLA从🤗 Transformers运行TF文本生成模型

要在🤗 Transformers中启用XLA加速生成,您需要安装最新版本的transformers。您可以通过运行以下命令来安装:

pip install transformers --upgrade

然后你可以运行以下代码:

import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

# Will error if the minimal version of Transformers is not installed.
from transformers.utils import check_min_version

check_min_version("4.21.0")


tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("openai-community/gpt2")
input_string = ["TensorFlow is"]

# One line to create an XLA generation function
xla_generate = tf.function(model.generate, jit_compile=True)

tokenized_input = tokenizer(input_string, return_tensors="tf")
generated_tokens = xla_generate(**tokenized_input, num_beams=2)

decoded_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
print(f"Generated -- {decoded_text}")
# Generated -- TensorFlow is an open-source, open-source, distributed-source application # framework for the

正如你所注意到的,在generate()上启用XLA只需一行代码。其余代码保持不变。然而,上述代码片段中有一些特定于XLA的注意事项。你需要了解这些,才能实现XLA带来的加速。我们将在下一节讨论这些内容。

需要注意的事项

当你第一次执行一个启用了XLA的函数(如上面的xla_generate())时,它会在内部尝试推断计算图,这个过程是耗时的。这个过程被称为“tracing”

你可能会注意到生成时间并不快。连续调用xla_generate()(或任何其他启用XLA的函数)时,如果函数的输入与最初构建计算图时的形状相同,就不需要再次推断计算图。虽然这对于具有固定输入形状的模式(例如,图像)来说不是问题,但如果你正在处理具有可变输入形状的模式(例如,文本),则必须注意。

为了确保 xla_generate() 始终使用相同的输入形状操作,你可以在调用分词器时指定 padding 参数。

import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("openai-community/gpt2")
input_string = ["TensorFlow is"]

xla_generate = tf.function(model.generate, jit_compile=True)

# Here, we call the tokenizer with padding options.
tokenized_input = tokenizer(input_string, pad_to_multiple_of=8, padding=True, return_tensors="tf")

generated_tokens = xla_generate(**tokenized_input, num_beams=2)
decoded_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
print(f"Generated -- {decoded_text}")

这样,您可以确保xla_generate()的输入始终会接收到与其跟踪时相同的形状,从而加快生成时间。您可以使用以下代码验证这一点:

import time
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("openai-community/gpt2")

xla_generate = tf.function(model.generate, jit_compile=True)

for input_string in ["TensorFlow is", "TensorFlow is a", "TFLite is a"]:
    tokenized_input = tokenizer(input_string, pad_to_multiple_of=8, padding=True, return_tensors="tf")
    start = time.time_ns()
    generated_tokens = xla_generate(**tokenized_input, num_beams=2)
    end = time.time_ns()
    print(f"Execution time -- {(end - start) / 1e6:.1f} ms\n")

在 Tesla T4 GPU 上,您可以期待如下输出:

Execution time -- 30819.6 ms

Execution time -- 79.0 ms

Execution time -- 78.9 ms

第一次调用xla_generate()由于追踪的原因会耗时较长,但后续的调用会快得多。请记住,生成选项的任何更改都会触发重新追踪,从而导致生成时间变慢。

我们没有涵盖🤗 Transformers在本文档中提供的所有文本生成选项。我们鼓励您阅读文档以了解高级用例。

附加资源

在这里,如果你想更深入地了解🤗 Transformers中的XLA以及一般情况,我们为你提供了一些额外的资源。

< > Update on GitHub