评价此页
★ ★ ★ ★ ★

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, Callable]]] = None, disable: bool = False) → Callable[[_InputT], _RetT][source]#
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, Callable]]] = None, disable: bool = False) → Callable[[Callable[[_InputT], _RetT]], Callable[[_InputT], _RetT]]

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

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

参数

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

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

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

  • backend (str 或 Callable) –

    要使用的后端

  • 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 将进行剖析以选择最佳的矩阵乘法配置

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

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

    • triton.cudagraphs 将通过 CUDA 图减少 Python 的开销

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

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

    • guard_filter_fn 控制哪些 dynamo guard 会随编译一起保存。这是一个不安全的功能，并且对于 dynamo guard 作为数据类型没有向后兼容性保证。对于稳定的辅助函数，请参阅 torch.compiler 中的文档，例如：- torch.compiler.skip_guard_on_inbuilt_nn_modules_unsafe - torch.compiler.skip_guard_on_all_nn_modules_unsafe - torch.compiler.keep_tensor_guards_unsafe

    • 对于 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)
在此页面上
显示源代码
PyTorch 库