torch.compile

torch.compile(model, *, fullgraph=False, dynamic=None, backend="inductor", options=None, disable=False) -> Callable[[_InputT], _RetT]

torch.compile(fullgraph=False, dynamic=None, backend="inductor", options=None, disable=False) -> Callable[[Callable[[_InputT], _RetT]], Callable[[_InputT], _RetT]]

使用TorchDynamo和指定后端来优化给定的模型或函数。如果你正在编译torch.nn.Module，你也可以直接使用torch.nn.Module.compile()方法来就地编译模块而不改变其结构。

具体来说，对于编译区域内的每个执行帧，我们将尝试对其进行编译，并将编译结果缓存到代码对象中以供将来使用。如果之前的编译结果不适用于后续调用（这称为“guard failure”），同一个帧可能会被多次重新编译。你可以使用 TORCH_LOGS=guards 来调试这些情况。每个帧最多可以关联多个编译结果，上限为 torch._dynamo.config.cache_size_limit（默认值为 8）。超过这个限制后，我们将回退到即时执行模式（eager mode）。需要注意的是，编译缓存是按代码对象而非帧进行的；如果你动态创建了多个函数副本，它们将共享同一个代码缓存。

参数

• model (Callable) – 待优化的模块或函数

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

• dynamic (bool 或 None) – 使用动态形状跟踪。当此参数为 True 时，我们将尝试生成尽可能动态的内核，以避免在大小变化时重新编译。然而，并非所有操作和优化都能支持这种动态性；某些情况下会强制进行专门化处理。使用 TORCH_LOGS=dynamic 可以调试过度专门化的问题。当此参数为 False 时，我们永远不会生成动态内核，始终进行专门化处理。默认情况下（None），我们将自动检测是否发生了动态变化，并在重新编译时生成更动态的内核。

• 后端 (str 或 Callable) –

  所需的后端

  • “inductor” 是默认后端，在性能和开销之间提供了一个很好的平衡。

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

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

  • 要注册一个外部自定义后端，请参阅：https://pytorch.org/docs/main/torch.compiler_custom_backends.html#registering-custom-backends

• 模式 (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 会进行性能分析，以选择最佳的矩阵乘法配置。

  • fallback_random 在调试准确性问题时非常有用

  • shape_padding 用于填充矩阵形状，以在 GPU 上更好地对齐加载，特别是在使用张量核心时。

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

  • trace.enabled 是最实用的调试标志，开启它可以帮助有效调试。

  • trace.graph_diagram，这将展示融合后图形的图像

  • 你可以通过调用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)