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 遇到图中断（意味着它无法将模型编译成单个图），它将引发错误。

from transformers import AutoModelForCausalLM

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

基准测试

有关使用 torch.compile 启用和禁用在各种 GPU 和批处理大小下对同一图像进行不同视觉任务的平均推理时间（以毫秒为单位）的性能基准测试，请参阅下表。

在下表中选择 **Subset** 以在不同的 GPU 之间切换，以及在 PyTorch nightly 2.1.0dev 和使用 reduce-overhead 模式启用的 torch.compile 上的基准测试。