torch.compile

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

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

使用 TorchDynamo 和指定的后端优化给定的模型/函数。如果您正在编译一个 torch.nn.Module ，您还可以使用 torch.nn.Module.compile() 就地编译模块而不改变其结构。

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

参数:

• 模型（Callable）- 要优化的模块/函数

• fullgraph（布尔值）- 如果为 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.org/docs/main/torch.compiler_custom_backends.html#registering-custom-backends

• mode (str) –

  可以为“默认”、“减少开销”、“最大自动调整”或“最大自动调整-不使用 cudagraphs”

  • “默认”是默认模式，在性能和开销之间取得良好平衡

  • “减少开销”是一种减少 CUDA 图 Python 开销的模式，适用于小批量。减少开销可能会以增加内存使用为代价，因为我们将会缓存调用所需的工作空间内存，以便在后续运行中不必重新分配。减少开销不一定总是有效；目前，我们只为不修改输入的 CUDA 图减少开销；在其他情况下，CUDA 图可能不适用；使用 TORCH_LOG=perf_hints 进行调试。

  • “最大自动调整”是一种利用支持设备上的 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 可以减少 Python 与 CUDA 图的 overhead

  • trace.enabled 是最实用的调试标志，可以开启

  • trace.graph_diagram 可以显示融合后的图示

  • 对于电感器，可以通过调用 torch._inductor.list_options() 来查看它支持的完整配置列表

• 禁用（布尔值）- 将 torch.compile()转换为无操作以进行测试

示例：

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