torch.compile

torch.compile(model: Callable[[_InputT], _RetT], *, fullgraph: bool = False, dynamic: Optional[bool] = None, backend: Union[str, Callable] = 'inductor', mode: Optional[str] = None, options: Optional[dict[str, Union[str, int, bool]]] = None, disable: bool = False) → Callable[[_InputT], _RetT]

torch.compile(model: None = None, *, fullgraph: bool = False, dynamic: Optional[bool] = None, backend: Union[str, Callable] = 'inductor', mode: Optional[str] = None, options: Optional[dict[str, Union[str, int, bool]]] = None, disable: bool = False) → Callable[[Callable[[_InputT], _RetT]], Callable[[_InputT], _RetT]]

使用 TorchDynamo 和指定的后端优化给定的模型/函数。如果要编译 torch.nn.Module，还可以使用 torch.nn.Module.compile() 就地（inplace）编译模块，而无需改变其结构。

具体而言，对于在编译区域内执行的每个帧（frame），我们将尝试对其进行编译，并将编译结果缓存到代码对象（code object）上以备将来使用。如果之前的编译结果不适用于后续调用（这被称为“守卫失败”或“guard failure”），单个帧可能会被编译多次。您可以使用 TORCH_LOGS=guards 来调试这些情况。一个帧最多可以关联 torch._dynamo.config.recompile_limit 个编译结果，默认值为 8；达到此限制后，我们将回退到 eager 模式。请注意，编译缓存是基于代码对象（code object）而非帧（frame）的；如果您动态创建函数的多个副本，它们将共享相同的代码缓存。

参数

• model (Callable) – 要优化的模块/函数

• fullgraph (bool) – 如果为 False（默认值），torch.compile 会尝试在函数中发现可编译的区域进行优化。如果为 True，则要求整个函数必须能被捕获到一个单独的计算图中。如果无法做到（即存在图断点graph breaks），则会引发错误。

• dynamic (bool or None) – 使用动态形状追踪。当此参数为 True 时，我们将提前尝试生成尽可能动态的内核，以避免在形状变化时重新编译。但这并非总是有效，因为某些操作/优化会强制进行特化；使用 TORCH_LOGS=dynamic 来调试过度特化问题。当此参数为 False 时，我们将永远不会生成动态内核，而是始终进行特化。默认值（None）下，我们自动检测是否发生了动态行为，并在重新编译时编译更动态的内核。

• backend (str or Callable) –

  要使用的后端

  • “inductor” 是默认后端，它在性能和开销之间取得了很好的平衡

  • 非实验性的内建后端可以使用 torch._dynamo.list_backends() 查看

  • 实验性或调试用的内建后端可以使用 torch._dynamo.list_backends(None) 查看

  • 注册外部自定义后端：https://pytorch.ac.cn/docs/main/torch.compiler_custom_backends.html#registering-custom-backends

• mode (str) –

  可以是 “default”、“reduce-overhead”、“max-autotune” 或 “max-autotune-no-cudagraphs” 之一

  • “default” 是默认模式，它在性能和开销之间取得了很好的平衡

  • “reduce-overhead” 模式使用 CUDA 图减少 Python 的开销，对小批量数据很有用。减少开销可能会以增加内存使用为代价，因为我们将缓存调用所需的 workspace 内存，以便后续运行时不必重新分配。减少开销并非始终有效；目前，我们只对不修改输入的纯 CUDA 图减少开销。在其他一些情况下，CUDA 图不适用；可以使用 TORCH_LOG=perf_hints 进行调试。

  • “max-autotune” 模式在支持的设备上利用 Triton 或基于模板的矩阵乘法，并在 GPU 上利用基于 Triton 的卷积。它在 GPU 上默认启用 CUDA 图。

  • “max-autotune-no-cudagraphs” 模式类似于 “max-autotune”，但不使用 CUDA 图

  • 要查看每种模式设置的具体配置，可以调用 torch._inductor.list_mode_options()

• options (dict) –

  传递给后端的可选项字典。一些值得尝试的重要选项包括：

  • epilogue_fusion 将逐点（pointwise）操作融合到模板中。需要同时设置 max_autotune

  • max_autotune 将通过 profile 选择最佳的 matmul 配置

  • fallback_random 在调试精度问题时很有用

  • shape_padding 填充矩阵形状，以更好地对齐 GPU 上的加载，特别是对于张量核心（tensor cores）

  • triton.cudagraphs 使用 CUDA 图减少 Python 的开销

  • trace.enabled 是最有用的调试标志之一

  • trace.graph_diagram 将显示融合后的计算图图像

  • 对于 inductor 后端，可以通过调用 torch._inductor.list_options() 查看其支持的完整配置列表

• disable (bool) – 将 torch.compile() 变成空操作（no-op）用于测试

示例

@torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
def foo(x):
    return torch.sin(x) + torch.cos(x)