自定义后端¶
概述¶
torch.compile
提供了一种定义自定义后端的直接方法。
后端函数的契约是 (gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]) -> Callable
。
后端函数可以由 TorchDynamo(torch.compile
的图跟踪组件)在跟踪 FX 图后调用,并期望返回一个已编译的函数,该函数等效于跟踪的 FX 图。返回的可调用对象应具有与传递给后端的原始 torch.fx.GraphModule
的 forward
函数相同的契约:(*args: torch.Tensor) -> List[torch.Tensor]
。
为了让 TorchDynamo 调用您的后端,请将您的后端函数作为 backend
kwarg 传递给 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(...):
...
请参阅下方以获取更多示例。
注册自定义后端¶
您可以使用 register_backend
装饰器注册您的后端,例如,
from torch._dynamo import register_backend
@register_backend
def my_compiler(gm, example_inputs):
...
除了 register_backend
装饰器,如果您的后端在另一个 Python 包中,您还可以通过 Python 包的入口点 (entry_points
) 注册您的后端,这提供了一种为一个包注册插件到另一个包的方式。
提示
您可以在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")
。对于与精简器一起使用是必需的。精简器生成的任何代码都必须调用注册您的后端函数的代码,通常是通过
import
语句。
AOTAutograd 后的自定义后端¶
可以定义由 AOTAutograd 而不是 TorchDynamo 调用的自定义后端。这主要有两个原因:
用户可以定义支持模型训练的后端,因为 AOTAutograd 可以生成用于编译的反向图。
AOTAutograd 生成由核心 Aten 算子组成的 FX 图。因此,自定义后端只需支持核心 Aten 算子集,这是一个比整个 torch/Aten 算子集小得多的算子集。
使用 torch._dynamo.backends.common.aot_autograd
包装您的后端,并像以前一样使用 torch.compile
和 backend
kwarg。由 aot_autograd
包装的后端函数应具有与之前相同的契约。
后端函数通过 fw_compiler
(前向编译器)或 bw_compiler
(反向编译器)kwarg 传递给 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