torch.compiler

创建日期：2023 年 7 月 28 日 | 最后更新日期：2025 年 12 月 7 日

torch.compiler 是一个命名空间，一些内部编译器方法通过该命名空间暴露给用户。该命名空间中的主要函数和特性是 torch.compile。

torch.compile 是 PyTorch 2.x 中引入的一个 PyTorch 函数，旨在解决 PyTorch 中准确图捕获的问题，并最终使软件工程师能够更快地运行他们的 PyTorch 程序。torch.compile 是用 Python 编写的，它标志着 PyTorch 从 C++ 向 Python 的转变。

torch.compile 利用以下底层技术：

• TorchDynamo (torch._dynamo) 是一个内部 API，它使用称为 Frame Evaluation API 的 CPython 功能来安全地捕获 PyTorch 图。 PyTorch 用户可用的外部方法通过 torch.compiler 命名空间暴露。

• TorchInductor 是默认的 torch.compile 深度学习编译器，可为多个加速器和后端生成快速代码。您需要使用后端编译器才能通过 torch.compile 实现加速。对于 NVIDIA、AMD 和 Intel GPU，它利用 OpenAI Triton 作为关键构建块。

• AOT Autograd 不仅捕获用户级代码，还捕获反向传播，从而“提前”捕获反向传播。这使得使用 TorchInductor 加速前向和反向传播成为可能。

要更好地理解 torch.compile 在您的代码上的跟踪行为，或要了解有关 torch.compile 内部机制的更多信息，请参阅 torch.compile 编程模型。

注意

在某些情况下，术语 torch.compile、TorchDynamo、torch.compiler 在本文档中可能会互换使用。

警告

torch.compile 可能不支持最近发布的 Python 主要版本。

如果您尝试在不受支持的 Python 环境中使用 @torch.compile，您可能会遇到类似以下错误的错误：

RuntimeError: torch.compile is not supported on Python 3.xx.0+

请确保您当前的 Python 版本在 PyTorch 对 torch.compile 支持的范围内。

如果您在过新的 Python 版本上安装了 PyTorch，您将需要切换到较早的 Python 版本才能使用 torch.compile。

如上所述，为了更快地运行您的工作流，通过 TorchDynamo 的 torch.compile 需要一个后端来将捕获的图转换为快速的机器代码。不同的后端可以带来各种优化收益。默认后端称为 TorchInductor，也称为 _inductor_。TorchDynamo 拥有我们合作伙伴开发的受支持后端列表，可以通过运行 torch.compiler.list_backends() 来查看，每个后端都有其可选的依赖项。

一些最常用的后端包括：

训练和推理后端

后端	描述
torch.compile(m, backend="inductor")	使用 TorchInductor 后端。 阅读更多
torch.compile(m, backend="cudagraphs")	带 AOT Autograd 的 CUDA 图。 阅读更多
torch.compile(m, backend="ipex")	在 CPU 上使用 IPEX。 阅读更多

仅推理后端

后端	描述
torch.compile(m, backend="tensorrt")	使用 Torch-TensorRT 进行推理优化。需要在调用脚本中 import torch_tensorrt 来注册后端。 阅读更多
torch.compile(m, backend="ipex")	在 CPU 上使用 IPEX 进行推理。 阅读更多
torch.compile(m, backend="tvm")	使用 Apache TVM 进行推理优化。 阅读更多
torch.compile(m, backend="openvino")	使用 OpenVINO 进行推理优化。 阅读更多