Transformers 文档
torch.compile
加入 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 上更新