用于推理的 torch.compile

torch.compile 将 PyTorch 代码编译为优化内核，从而显著加速推理。该功能依赖 TorchDynamo 将代码编译为图，并依赖 TorchInductor 进一步将图编译为优化内核。它是一个强大的优化工具，在许多情况下，只需添加一行代码即可使用。

使用 torch.compile 包装模型，以编译并返回优化后的模型。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model)

首次调用 torch.compile 时速度较慢，因为模型需要被编译。后续对已编译模型的调用会快得多，因为无需再次编译。

编译过程有多个参数可供自定义。其中两个较为重要的参数如下所示。完整的参数列表请参阅 torch.compile 的 文档。

模式

mode 参数提供了多种用于编译的性能选项。请尝试不同的模式，以找到最适合您使用场景的模式。

• default 是速度和内存之间的平衡选项。
• reduce-overhead 会略微增加内存占用，但能降低 Python 开销，从而可能获得更快的速度。
• max-autotune 提供最快的速度，但编译时间更长。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead")

Fullgraph

Fullgraph 会尝试将整个模型编译为单个图，以最大化性能。如果 torch.compile 遇到图中断（graph break），即无法将模型编译为单个图，则会抛出错误。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead", fullgraph=True)

基准测试

下表列出了在不同视觉任务中，针对同一图像，在不同 GPU 和批大小下，启用与禁用 torch.compile 时的平均推理时间（毫秒）性能基准对比。

在下方表格中选择 Subset 可在不同的 GPU 之间切换，并查看 PyTorch nightly 2.1.0dev 以及启用 reduce-overhead 模式的 torch.compile 的基准测试结果。

后续步骤

• 请参阅 用于训练的 torch.compile 指南，了解如何使用 torch.compile 进行训练。