自定义后端¶
概述¶
torch.compile 提供了一种简单的方法，使用户能够定义自定义后端。
一个后端函数具有契约
(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]) -> Callable。
后端函数可以通过TorchDynamo调用，它是用于跟踪FX图的组件torch.compile,
在跟踪一个FX图之后，预期返回一个与跟踪后的FX图等效的编译函数。
返回的可调用对象应与传递给后端的原始torch.fx.GraphModule的forward函数具有相同的契约:
(*args: torch.Tensor) -> List[torch.Tensor]。
为了使TorchDynamo调用你的后端，请将你的后端函数作为backend参数传递给
torch.compile。例如,
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(...):
...
如下是更多示例。
注册自定义后端¶
You can register your backend using the register_backend decorator, for example,
from torch._dynamo.optimizations import register_backend
@register_backend
def my_compiler(gm, example_inputs):
...
除了 register_backend 装饰器之外，如果你的后端位于另一个 Python 包中，你也可以通过 Python 包的入口点来注册你的后端，这为一个包为另一个包注册插件提供了一种方式。
提示
您可以在 entry_points 的
Python 包装文档 中了解更多。
要通过 entry_points 注册你的后端，你可以将你的后端函数添加到你包中的 torch_dynamo_backends 入口点组的
setup.py 文件中，例如：
...
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.optimizations.training.aot_autograd 包裹你的后端，并像之前一样使用 torch.compile 和 backend 参数。
被 aot_autograd 包裹的后端函数应与之前保持相同的契约。
后端函数通过 aot_autograd 的 fw_compiler（前向编译器）
或 bw_compiler（反向编译器）的 kwargs 传递。如果未指定 bw_compiler，
则反向编译函数默认使用前向编译函数。
一个需要注意的地方是，AOTAutograd 要求后端返回的编译函数必须被“封装”。这可以通过将编译函数用 functorch.compile.make_boxed_func 包裹来实现。
例如，
from torch._dynamo.optimizations.training 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)
示例¶
调试后端¶
如果你想更好地理解编译过程中的情况，可以创建一个自定义编译器，本节中称为后端（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,),) {}
最后两个图的顺序是不确定的，这取决于即时编译器首先遇到哪一个。
快速后端¶
集成一个提供卓越性能的自定义后端也非常简单，我们将使用 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 包含许多后端，可以在
backends.py
或 torch._dynamo.list_backends() 中找到。你可以使用以下代码将这些后端组合在一起：
from torch._dynamo.optimizations import BACKENDS
def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
try:
trt_compiled = BACKENDS["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 = BACKENDS["inductor"](gm, example_inputs)
if inductor_compiled is not None:
return inductor_compiled
except Exception:
pass
return gm.forward