Transformers 文档

torch.compile

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

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)

基准测试

请参阅下表,了解在不同 GPU 和批处理大小下,针对不同视觉任务,在同一图像上启用和禁用 torch.compile 时平均推理时间(毫秒)的性能基准比较。

在下表中选择**子集**以切换不同的 GPU,以及 PyTorch nightly 2.1.0dev 和启用 `reduce-overhead` 模式的 torch.compile 的基准测试。

< > 在 GitHub 上更新