太长不看：使用 TensorFlow 在 🤗 transformers 上进行的文本生成现在可以使用 XLA 编译。它比以前快了100倍，甚至比 PyTorch 还要快 — 在下面的 colab 上查看吧！
文本生成
随着大型语言模型的质量提高，我们对这些模型能做什么的期望也在增加。特别是自 OpenAI 的 GPT-2 发布以来，带有文本生成功能的模型一直备受关注。而且有充分的理由 — 这些模型可以用于摘要、翻译，甚至在某些语言任务上展示出了零样本学习的能力。本博客文章将展示如何在 TensorFlow 中充分利用这项技术。
🤗 transformers 库从自然语言处理模型开始，所以文本生成对我们来说至关重要。这是 Hugging Face 民主化努力的一部分，以确保它易于访问、易于控制和高效。关于不同类型的文本生成有一篇先前的博客文章。然而，在下面有一个核心功能的快速回顾 — 如果您熟悉我们的 generate 函数并且想要直接跳入 TensorFlow 的特定性，请随意跳过。
让我们从基础知识开始。文本生成可以是确定性的或随机的，这取决于 do_sample 标志。默认情况下，它设置为 False ，导致输出是确定性的，也被称为贪婪解码。当它设置为 True 时，也被称为采样，输出将是随机的，但您仍然可以通过 seed 参数获得可重现的结果（与无状态的 TensorFlow 随机数生成中的格式相同）。作为一个经验法则，如果您希望从模型中获得事实信息，您需要确定性生成，如果您希望获得更有创造性的输出，则需要随机生成。
# 需要 transformers >= 4.21.0;
# 根据您的硬件，采样输出可能会有所不同。
from transformers import AutoTokenizer, TFAutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
model.config.pad_token_id = model.config.eos_token_id
inputs = tokenizer(["TensorFlow 是"], return_tensors="tf")
generated = model.generate(**inputs, do_sample=True, seed=(42, 0))
print("采样输出：", tokenizer.decode(generated[0]))
# > 采样输出：TensorFlow 是一个很棒的学习平台，用于学习有关数据结构和数据科学中的结构。根据目标应用程序，可能需要更长的输出。您可以使用 max_new_tokens 控制生成输出的长度，但请记住，更长的生成将需要更多的资源。
generated = model.generate(
**inputs, do_sample=True, seed=(42, 0), max_new_tokens=5
)
print("限制为 5 个新令牌：", tokenizer.decode(generated[0]))
# > 限制为 5 个新令牌：TensorFlow 是一个很棒的学习平台
generated = model.generate(
**inputs, do_sample=True, seed=(42, 0), max_new_tokens=30
)
print("限制为 30 个新令牌：", tokenizer.decode(generated[0]))
# > 限制为 30 个新令牌：TensorFlow 是一个很棒的学习平台，用于学习有关数据结构和数据科学中的结构。采样有一些可以用来控制随机性的旋钮。最重要的是 temperature ，它设置输出的整体熵 — 小于 1.0 的值将优先采样具有更高可能性的令牌，而大于 1.0 的值则相反。将其设置为 0.0 将行为减少到贪婪解码，而非常大的值则近似于均匀采样。
generated = model.generate(
**inputs, do_sample=True, seed=(42, 0), temperature=0.7
)
print("温度为 0.7：", tokenizer.decode(generated[0]))
# > 温度为 0.7：TensorFlow 是一个很棒的方式来做这样的事情........
generated = model.generate(
**inputs, do_sample=True, seed=(42, 0), temperature=1.5
)
print("温度为 1.5：", tokenizer.decode(generated[0]))
# > 温度为 1.5：TensorFlow 正在开发用于 Cython 和 Bamboo 的双模式。
# 在 Bamboo 上......与采样不同，贪婪解码会在每次生成的迭代中选择最可能的标记。然而，它经常会产生次优的输出。您可以通过使用num_beams参数来提高结果的质量。当它大于1时，它会触发束搜索，不断探索高概率序列。这种探索会以额外的资源和计算时间为代价。
generated = model.generate(**inputs, num_beams=2)
print("束搜索输出:", tokenizer.decode(generated[0]))
# > 束搜索输出: TensorFlow是一个开源的、开源的、
# 分布式的应用框架 for the最后，在运行采样或束搜索时，您可以使用num_return_sequences来返回多个序列。对于采样来说，它等同于从相同的输入提示多次生成，而对于束搜索来说，它返回得分最高的生成束按降序排列。
generated = model.generate(**inputs, num_beams=2, num_return_sequences=2)
print(
"所有生成的假设:",
"\n".join(tokenizer.decode(out) for out in generated)
)
# > 所有生成的假设: TensorFlow是一个开源的、开源的、
# 分布式的应用框架 for the
# > TensorFlow是一个开源的、开源的、分布式的应用框架，它允许正如您所见，文本生成的基本控制非常简单。然而，上面的示例中没有涵盖许多选项，建议阅读文档了解高级用例。不幸的是，当您使用TensorFlow运行generate时，您可能会注意到它执行起来需要一些时间。如果您的目标应用程序期望低延迟或大量的输入提示，使用TensorFlow进行文本生成似乎是一项昂贵的努力。😬
不要担心，本博客的其余部分旨在证明一行代码可以带来巨大的改进。如果您宁愿直接行动，colab提供了一个可以调试的交互式示例！
TensorFlow和XLA
XLA，即加速线性代数，是一个最初用于加速TensorFlow模型的编译器。现在，它也是JAX背后的编译器，并且甚至可以与PyTorch一起使用。尽管对一些人来说，“编译器”这个词听起来有些令人生畏，但是XLA在TensorFlow中的使用非常简单–它被打包在tensorflow库中，并且可以通过在任何创建图形的函数中添加jit_compile参数来触发。
对于熟悉TensorFlow 1 🧓的人来说，TensorFlow图的概念很自然，因为它是唯一的操作模式。首先，您以声明方式定义操作以创建图形。然后，您可以将输入通过图形并观察输出。快速、高效，但很难调试。TensorFlow 2引入了Eager Execution和以命令式编码模型的能力–TensorFlow团队在他们的博客文章中更详细地解释了这种差异。
Hugging Face使用Eager Execution来编写他们的TensorFlow模型。透明度是一项核心价值，能够在任何时候检查模型内部非常有益。然而，这也意味着某些模型的使用不能从图模式的性能优势中受益（例如当调用model(args)时）。
幸运的是，TensorFlow团队考虑到了像我们这样的用户🥳！将包含TensorFlow代码的函数包装在tf.function中将尝试在调用包装函数时将其转换为图形。如果您正在训练模型，调用model.compile()（不带run_eagerly=True）正是这种包装，以便在调用model.fit()时从图模式中受益。由于tf.function可以在包含TensorFlow代码的任何函数中使用，这意味着可以将它应用于超出模型推断范围的函数，从而创建一个优化的单一图形。
现在您已经知道如何创建TensorFlow图形，并使用XLA进行编译非常简单–只需将jit_compile=True作为参数添加到上述函数（tf.function和tf.keras.Model.compile）中即可。假设一切顺利（下面会有更多详细信息），并且您正在使用GPU或TPU，您会注意到第一次调用需要一些时间，但是后续调用会快得多。这是一个执行模型推断和一些后处理的简单示例函数：
# 注意：执行时间在很大程度上取决于硬件 -- 这里使用了3090。
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
inputs = tokenizer(["TensorFlow 是"], return_tensors="tf")
def most_likely_next_token(inputs):
model_output = model(inputs)
return tf.argmax(model_output.logits[:, -1, :], axis=-1)
print("调用普通函数，使用 TensorFlow 代码...")
most_likely_next_token(inputs)
# > 执行时间 -- 48.8 毫秒一行代码即可从上面的函数创建一个启用 XLA 加速的函数。
xla_most_likely_next_token = tf.function(most_likely_next_token, jit_compile=True)
print("调用 XLA 函数...（第一次调用将会很慢）")
xla_most_likely_next_token(inputs)
# > 执行时间 -- 3951.0 毫秒
print("调用 XLA 函数...（第二次调用将会很快）")
xla_most_likely_next_token(inputs)
# > 执行时间 -- 1.6 毫秒使用 TensorFlow 和 XLA 进行文本生成
与任何优化过程一样，没有免费的午餐 — XLA 也不例外。从文本生成用户的角度来看，只需要记住一个技术方面的细节。在不深入细节的情况下，XLA 在调用时对 tf.function 进行即时编译（JIT），这依赖于多态性。
当以这种方式编译函数时，XLA 会跟踪每个张量的形状和类型，以及每个非张量函数输入的数据。该函数被编译为二进制文件，每次以相同的张量形状和类型（具有任何张量数据）及相同的非张量参数调用时，都可以重用已编译的函数。相反，如果使用不同形状或类型的输入张量调用函数，或者使用不同的非张量参数，那么将进行一次新的昂贵的编译步骤。以下是一个简单示例：
# 注意：执行时间在很大程度上取决于硬件 -- 这里使用了3090。
import tensorflow as tf
@tf.function(jit_compile=True)
def max_plus_constant(tensor, scalar):
return tf.math.reduce_max(tensor) + scalar
# 慢：第一次调用时会触发 XLA 编译
max_plus_constant(tf.constant([0, 0, 0]), 1)
# > 执行时间 -- 520.4 毫秒
# 快：不是首次使用具有相同张量形状、张量类型和完全相同的非张量参数调用
max_plus_constant(tf.constant([1000, 0, -10]), 1)
# > 执行时间 -- 0.6 毫秒
# 慢：不同的张量类型
max_plus_constant(tf.constant([0, 0, 0], dtype=tf.int64), 1)
# > 执行时间 -- 27.1 毫秒
# 慢：不同的张量形状
max_plus_constant(tf.constant([0, 0, 0, 0]), 1)
# > 执行时间 -- 25.5 毫秒
# 慢：不同的非张量参数
max_plus_constant(tf.constant([0, 0, 0]), 2)
# > 执行时间 -- 24.9 毫秒在实践中，对于文本生成来说，这意味着输入应该被填充为某个长度的倍数（以便具有有限数量的可能形状），并且使用不同的选项在第一次使用时会很慢。让我们看看当你不加考虑地调用 XLA 进行生成时会发生什么。
# 注意：执行时间在很大程度上取决于硬件 -- 这里使用了3090。
import time
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM
# 注意新增的参数 `padding_side="left"` -- 仅解码器模型（可以使用 TFAutoModelForCausalLM 实例化）应该进行左填充，因为它们从输入提示开始生成。
tokenizer = AutoTokenizer.from_pretrained(
"gpt2", padding_side="left", pad_token="</s>"
)
model = TFAutoModelForCausalLM.from_pretrained("gpt2")
model.config.pad_token_id = model.config.eos_token_id
input_1 = ["TensorFlow 是"]
input_2 = ["TensorFlow 是一个"]
# 一行代码创建一个启用 XLA 生成的函数
xla_generate = tf.function(model.generate, jit_compile=True)
# 调用 XLA 生成，不进行填充
tokenized_input_1 = tokenizer(input_1, return_tensors="tf") # 长度 = 4
tokenized_input_2 = tokenizer(input_2, return_tensors="tf") # 长度 = 5
print(f"`tokenized_input_1` 形状 = {tokenized_input_1.input_ids.shape}")
print(f"`tokenized_input_2` 形状 = {tokenized_input_2.input_ids.shape}")
print("使用 tokenized_input_1 调用 XLA 生成...")
print("(首次调用将会很慢)")
start = time.time_ns()
xla_generate(**tokenized_input_1)
end = time.time_ns()
print(f"执行时间 -- {(end - start) / 1e6:.1f} 毫秒\n")
# > 执行时间 -- 9565.1 毫秒
print("使用 tokenized_input_2 调用 XLA 生成...")
print("(具有不同的长度 = 将触发重新跟踪)")
start = time.time_ns()
xla_generate(**tokenized_input_2)
end = time.time_ns()
print(f"执行时间 -- {(end - start) / 1e6:.1f} 毫秒\n")
# > 执行时间 -- 6815.0 毫秒哦不，这太慢了！通过填充来控制不同形状的组合是一种解决方案，就像上面提到的那样。Tokenizer类有一个pad_to_multiple_of参数，可以用于在接受任意输入长度和限制追踪之间取得平衡。
padding_kwargs = {"pad_to_multiple_of": 8, "padding": True}
tokenized_input_1_with_padding = tokenizer(
input_1, return_tensors="tf", **padding_kwargs
) # length = 8
tokenized_input_2_with_padding = tokenizer(
input_2, return_tensors="tf", **padding_kwargs
) # length = 8
print(
"`tokenized_input_1_with_padding` shape = ",
f"{tokenized_input_1_with_padding.input_ids.shape}"
)
print(
"`tokenized_input_2_with_padding` shape = ",
f"{tokenized_input_2_with_padding.input_ids.shape}"
)
print("使用tokenized_input_1_with_padding调用XLA生成...")
print("(慢，第一次使用这个长度)")
start = time.time_ns()
xla_generate(**tokenized_input_1_with_padding)
end = time.time_ns()
print(f"执行时间 -- {(end - start) / 1e6:.1f} 毫秒\n")
# > 执行时间 -- 6815.4 毫秒
print("使用tokenized_input_2_with_padding调用XLA生成...")
print("(会很快！)")
start = time.time_ns()
xla_generate(**tokenized_input_2_with_padding)
end = time.time_ns()
print(f"执行时间 -- {(end - start) / 1e6:.1f} 毫秒\n")
# > 执行时间 -- 19.3 毫秒这好多了，以后使用这种方式连续进行生成调用将比以前快上几个数量级。请记住，尝试任何新的生成选项，都会触发追踪。
print("使用相同的输入，但使用新选项调用XLA生成...")
print("(再次慢)")
start = time.time_ns()
xla_generate(**tokenized_input_1_with_padding, num_beams=2)
end = time.time_ns()
print(f"执行时间 -- {(end - start) / 1e6:.1f} 毫秒\n")
# > 执行时间 -- 9644.2 毫秒从开发者的角度来看，依赖于XLA意味着需要注意一些额外的细微差别。XLA在数据结构的大小预先已知的情况下才能发挥作用，例如在模型训练中。另一方面，当无法确定数据结构的维度或使用某些动态切片时，XLA无法进行编译。现代文本生成的实现是自回归的，其自然行为是扩展张量并在进行操作时突然中断 — 换句话说，默认情况下不适合XLA。我们已经重写了整个TensorFlow文本生成代码库，以向量化操作并使用具有填充的固定大小结构。我们还修改了NLP模型，以在存在填充结构时正确使用它们的位置嵌入。对于TensorFlow文本生成用户来说，结果应该是看不见的，除了可以使用XLA编译。
基准测试和结论
上面你看到了你可以将TensorFlow函数转换为图形，并使用XLA编译加速它们。目前的文本生成形式只是一个自回归函数，它在模型的前向传递和一些后处理之间交替进行，每次迭代生成一个标记。通过XLA编译，整个过程得到了优化，从而实现更快的执行。但是速度快多少呢？下面的Gradio演示包含了对Hugging Face的文本生成在两个主要ML框架TensorFlow和PyTorch上使用多个GPU模型的一些基准测试。
如果你探索结果，两个结论很快就能看到：
- 正如本博文一直在介绍的，使用XLA的TensorFlow文本生成要快得多。在某些情况下，速度提升超过100倍，这真正展示了编译图形的强大之处🚀
- 在绝大多数情况下，使用XLA的TensorFlow文本生成是最快的选择，有些情况下快了多达9倍，推翻了PyTorch是进行严肃NLP任务的首选框架的错误观点💪
尝试一下这个colab，享受使用XLA加速的文本生成的强大力量吧！