评价此页
★ ★ ★ ★ ★

自定义后端#

创建于: 2025 年 6 月 10 日 | 最后更新于: 2025 年 6 月 10 日

概述#

torch.compile 提供了一种直接的方法，让用户能够定义自定义后端。

后端函数具有如下约定：(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]) -> Callable。

torch.compile 的图跟踪组件 TorchDynamo 可以在跟踪完 FX 图后调用后端函数，并期望后端返回一个与被跟踪的 FX 图等效的编译后的函数。返回的可调用对象应该具有与传给后端的 torch.fx.GraphModule 的 forward 函数相同的约定：(*args: torch.Tensor) -> List[torch.Tensor]。

为了让 TorchDynamo 调用你的后端，请在 torch.compile 中将你的后端函数作为 backend 关键字参数传递。例如，

import torch

def my_custom_backend(gm, example_inputs):
    return gm.forward

def f(...):
    ...

f_opt = torch.compile(f, backend=my_custom_backend)

@torch.compile(backend=my_custom_backend)
def g(...):
    ...

更多示例请参见下文。

注册自定义后端#

你可以使用 register_backend 装饰器来注册你的后端，例如，

from torch._dynamo import register_backend

@register_backend
def my_compiler(gm, example_inputs):
    ...

除了 register_backend 装饰器之外，如果你的后端在另一个 Python 包中，你也可以通过 Python 包的入口点来注册你的后端，这为包提供了一种为另一个包注册插件的方式。

提示

你可以在 Python 打包文档 中了解更多关于 entry_points 的信息。

要通过 entry_points 注册你的后端，你可以在你的包的 setup.py 文件中将你的后端函数添加到 torch_dynamo_backends 入口点组，如下所示：

...
setup(
    ...
    'torch_dynamo_backends': [
        'my_compiler = your_module.submodule:my_compiler',
    ]
    ...
)

请将 = 前面的 my_compiler 替换为你后端的名称，并将 = 后面的部分替换为你后端函数的模块和函数名。入口点将在安装包后添加到你的 Python 环境中。当你调用 torch.compile(model, backend="my_compiler") 时，PyTorch 将首先搜索通过 register_backend 注册的名为 my_compiler 的后端。如果未找到，它将继续在通过 entry_points 注册的所有后端中进行搜索。

注册有两个目的：

  • 你可以将包含你的后端函数名称的字符串传递给 torch.compile，而不是直接传递函数本身，例如 torch.compile(model, backend="my_compiler")。

  • 这对于使用 minifier 是必需的。minifier 生成的任何代码都必须调用注册你的后端函数的代码，通常是通过 `import` 语句。

AOTAutograd 之后的自定义后端#

可以定义由 AOTAutograd 而不是 TorchDynamo 调用的自定义后端。这主要有两个原因：

  • 用户可以定义支持模型训练的后端，因为 AOTAutograd 可以生成向后图进行编译。

  • AOTAutograd 生成由 核心 Aten 算子 组成的 FX 图。因此，自定义后端只需支持核心 Aten 算子集，这是一个比整个 torch/Aten 算子集小得多的算子集。

将你的后端包装在 torch._dynamo.backends.common.aot_autograd 中，并像以前一样使用带有 backend 关键字参数的 torch.compile。由 aot_autograd 包装的后端函数应具有与之前相同的约定。

后端函数通过 fw_compiler（前向编译器）或 bw_compiler（后向编译器）关键字参数传递给 aot_autograd。如果未指定 bw_compiler，则后向编译函数默认为前向编译函数。

一个需要注意的地方是，AOTAutograd 要求后端返回的编译函数是“boxed”的。这可以通过使用 functorch.compile.make_boxed_func 来包装编译后的函数来实现。

例如，

from torch._dynamo.backends.common import aot_autograd
from functorch.compile import make_boxed_func

def my_compiler(gm, example_inputs):
    return make_boxed_func(gm.forward)

my_backend = aot_autograd(fw_compiler=my_compiler)  # bw_compiler=my_compiler

model_opt = torch.compile(model, backend=my_backend)

示例#

调试后端#

如果你想更好地理解编译过程中发生了什么，你可以创建一个自定义编译器（在本节中称为后端），它将打印出从 Dynamo 的字节码分析中提取的 FX GraphModule，并返回一个 forward() 可调用对象。

例如

from typing import List
import torch
def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    print("my_compiler() called with FX graph:")
    gm.graph.print_tabular()
    return gm.forward  # return a python callable
@torch.compile(backend=my_compiler)
def fn(x, y):
    a = torch.cos(x)
    b = torch.sin(y)
    return a + b
fn(torch.randn(10), torch.randn(10))

运行上面的示例会产生以下输出：

my_compiler() called with FX graph:
opcode         name    target                                                  args        kwargs
-------------  ------  ------------------------------------------------------  ----------  --------
placeholder    x       x                                                       ()          {}
placeholder    y       y                                                       ()          {}
call_function  cos     <built-in method cos of type object at 0x7f1a894649a8>  (x,)        {}
call_function  sin     <built-in method sin of type object at 0x7f1a894649a8>  (y,)        {}
call_function  add     <built-in function add>                                 (cos, sin)  {}
output         output  output                                                  ((add,),)   {}

这对于 torch.nn.Module 也同样有效，如下所示：

from typing import List
import torch
def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    print("my_compiler() called with FX graph:")
    gm.graph.print_tabular()
    return gm.forward  # return a python callable
class MockModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.relu = torch.nn.ReLU()
    def forward(self, x):
        return self.relu(torch.cos(x))
mod = MockModule()
optimized_mod = torch.compile(mod, backend=my_compiler)
optimized_mod(torch.randn(10))

让我们再看一个包含控制流的例子：

from typing import List
import torch
def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    print("my_compiler() called with FX graph:")
    gm.graph.print_tabular()
    return gm.forward  # return a python callable
@torch.compile(backend=my_compiler)
def toy_example(a, b):
    x = a / (torch.abs(a) + 1)
    if b.sum() < 0:
        b = b * -1
    return x * b
for _ in range(100):
    toy_example(torch.randn(10), torch.randn(10))

运行此示例会产生以下输出：

my_compiler() called with FX graph:
opcode         name     target                                                  args              kwargs
-------------  -------  ------------------------------------------------------  ----------------  --------
placeholder    a        a                                                       ()                {}
placeholder    b        b                                                       ()                {}
call_function  abs_1    <built-in method abs of type object at 0x7f8d259298a0>  (a,)              {}
call_function  add      <built-in function add>                                 (abs_1, 1)        {}
call_function  truediv  <built-in function truediv>                             (a, add)          {}
call_method    sum_1    sum                                                     (b,)              {}
call_function  lt       <built-in function lt>                                  (sum_1, 0)        {}
output         output   output                                                  ((truediv, lt),)  {}

my_compiler() called with FX graph:
opcode         name    target                   args         kwargs
-------------  ------  -----------------------  -----------  --------
placeholder    b       b                        ()           {}
placeholder    x       x                        ()           {}
call_function  mul     <built-in function mul>  (b, -1)      {}
call_function  mul_1   <built-in function mul>  (x, mul)     {}
output         output  output                   ((mul_1,),)  {}

my_compiler() called with FX graph:
opcode         name    target                   args       kwargs
-------------  ------  -----------------------  ---------  --------
placeholder    b       b                        ()         {}
placeholder    x       x                        ()         {}
call_function  mul     <built-in function mul>  (x, b)     {}
output         output  output                   ((mul,),)  {}

The order of the last two graphs is nondeterministic depending
on which one is encountered first by the just-in-time compiler.

快速后端#

集成一个提供更优性能的自定义后端也很容易，我们将集成一个实际的后端，并结合 optimize_for_inference。

def optimize_for_inference_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    scripted = torch.jit.script(gm)
    return torch.jit.optimize_for_inference(scripted)

然后，你可以使用以下方法来优化任何现有代码：

@torch.compile(backend=optimize_for_inference_compiler)
def code_to_accelerate():
    ...

可组合后端#

TorchDynamo 包含许多后端，可以使用 torch._dynamo.list_backends() 列出。你可以使用以下代码组合这些后端：

from torch._dynamo import lookup_backend
def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    try:
        trt_compiled = lookup_backend("tensorrt")(gm, example_inputs)
        if trt_compiled is not None:
            return trt_compiled
    except Exception:
        pass
    # first backend failed, try something else...
    try:
        inductor_compiled = lookup_backend("inductor")(gm, example_inputs)
        if inductor_compiled is not None:
            return inductor_compiled
    except Exception:
        pass
    return gm.forward
在此页面上
在 GitHub 上编辑
显示源代码
PyTorch 库