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() 原地编译模块，而无需更改其结构。

具体来说，对于在已编译区域内执行的每个帧，我们将尝试编译它并在代码对象上缓存编译结果以供将来使用。如果先前的编译结果不适用于后续调用（这称为“保护失败”），则单个帧可能会被编译多次，您可以使用 TORCH_LOGS=guards 调试这些情况。最多可以将多个编译结果与一个帧关联 torch._dynamo.config.cache_size_limit，默认为 8；此时，我们将回退到急切模式。请注意，编译缓存是针对代码对象的，而不是帧；如果您动态创建函数的多个副本，它们将共享相同的代码缓存。

参数

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

• fullgraph (bool) – 如果为 False（默认值），则 torch.compile 会尝试发现它将优化的函数中的可编译区域。如果为 True，则要求整个函数能够捕获到单个图中。如果这是不可能的（即，如果存在图中断），则这将引发错误。

• dynamic (bool 或 None) – 使用动态形状跟踪。当此值为 True 时，我们将预先尝试生成尽可能动态的内核，以避免在大小更改时重新编译。这可能并不总是有效，因为某些操作/优化将强制专业化；使用 TORCH_LOGS=dynamic 调试过度专业化。当此值为 False 时，我们将永远不会生成动态内核，我们将始终进行专业化。默认情况下（None），我们将自动检测是否发生了动态性，并在重新编译时编译更动态的内核。

• backend (str 或 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 将逐点操作融合到模板中。需要也设置 max_autotune

  • max_autotune 将进行分析以选择最佳的 matmul 配置

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

  • shape_padding 将矩阵形状填充以更好地对齐 GPU 上的加载，尤其是在张量核心方面

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

  • trace.enabled 是要开启的最有用的调试标志

  • trace.graph_diagram 将在融合后向您显示图的图片

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

• disable (bool) – 将 torch.compile() 转换为测试的无操作

示例

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