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 时的平均推理时间（毫秒）性能基准。
在下表中选择 Subset 以切换不同的 GPU，以及在 PyTorch nightly 2.1.0dev 和启用 reduce-overhead 模式的 torch.compile 上的基准测试。