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；此时我们将回退到 eager 模式。请注意，编译缓存是每个代码对象，而不是帧；如果您动态创建函数的多个副本，它们将共享相同的代码缓存。

参数

• 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 开销的模式，适用于小批量。降低开销可能会以增加内存使用量为代价，因为我们将缓存调用所需的工作区内存，以便不必在后续运行中重新分配它。不保证降低开销有效；目前，我们仅降低不改变输入的纯 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，它将减少 Python 和 CUDA 图的开销

  • 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)