评价此页
★ ★ ★ ★ ★

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][源码]#
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 或模板进行矩阵乘法（在支持的设备上）以及 Triton 进行卷积（在 GPU 上）的模式。它在 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 与编译一起保存。这是一个不安全的特性，并且对于作为数据类型的 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() 变成一个无操作，用于测试

示例

@torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
def foo(x):
    return torch.sin(x) + torch.cos(x)
在此页面上
显示源代码
PyTorch 库