PyTorch 自定义编译后端Custom Backends完整指南从 TorchDynamo 到 AOTAutograd 的接入实战【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorchtorch.compile为用户提供了一条直接定义自定义编译后端backend的通道一个后端本质上是一个以torch.fx.GraphModule和示例输入为参数、返回等价可调用对象的函数。本文以 torch.compiler_custom_backends.md 为骨架结合当前仓库中的注册表实现registry.py、AOTAutograd 封装common.py与后端初始化钩子eval_frame.py源码系统讲解后端的契约、注册方式、AOTAutograd 训练后端接入、eager 初始化钩子以及 Debug/Speedy/Composable 三类典型后端示例。读完本文你将能够独立编写、注册并组合自己的torch.compile后端。后端契约一个函数连接 Dynamo 与编译产物torch.compile的图追踪组件 TorchDynamo 在完成字节码分析并抽取出一张 FX 图之后会调用用户提供的后端函数。后端函数必须满足如下契约(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]) - Callable其中gm是 Dynamo 从用户代码中抽取出的 FXGraphModuleexample_inputs是用于推导形状等信息的示例输入张量列表在 Dynamo 内部通常为 FakeTensor返回值是一个“已编译函数”其行为必须与传入的 FX 图等价。返回的可调用对象契约与原torch.fx.GraphModule的forward一致(*args: torch.Tensor) - List[torch.Tensor]在 registry.py 中这两类签名被形式化定义为CompiledFn与CompilerFn两个类型class CompiledFn(Protocol): def __call__(self, *args: torch.Tensor) - tuple[torch.Tensor, ...]: ... CompilerFn Callable[[fx.GraphModule, list[torch.Tensor]], CompiledFn]要让 TorchDynamo 调用你的后端只需把后端函数作为backend关键字参数传给torch.compile既支持函数式调用也支持装饰器形式import torch def my_custom_backend(gm, example_inputs): return gm.forward def f(...): ... f_opt torch.compile(f, backendmy_custom_backend) torch.compile(backendmy_custom_backend) def g(...): ...最朴素的后端就是直接返回gm.forward即“不优化、仅捕获”的 eager 语义后端这也是后续所有示例的起点。注册自定义后端装饰器与 Entry Points 双通道使用register_backend装饰器你可以用torch._dynamo.register_backend装饰器将后端注册进全局注册表from torch._dynamo import register_backend register_backend def my_compiler(gm, example_inputs): ...从源码看register_backend 的核心逻辑如下若不传name则默认使用compiler_fn.__name__作为后端名tags参数可对后端打标签分类仓库预置了register_debug_backendtags(debug,)与register_experimental_backendtags(experimental,)两个便捷变体见 registry.py注册后后端函数会同时被写入_BACKENDS、_COMPILER_FNS、_BACKEND_TAGS三个全局字典若名字重复会抛出AssertionError(duplicate name: ...)。注册带来的直接收益是你可以用字符串代替函数本身来调用后端例如torch.compile(model, backendmy_compiler)。lookup_backendregistry.py会把字符串展开为真正的函数先在_BACKENDS中查找未命中则触发_lazy_import()与_discover_entrypoint_backends()仍找不到时会借助difflib.get_close_matches给出相近名字建议并抛出InvalidBackend。通过 Python 包 Entry Points 注册如果你的后端位于独立的 Python 包中可以通过 setuptools 的 entry points 机制注册插件这也是“一个包为另一个包注册插件”的标准做法。在包的setup.py中加入torch_dynamo_backends分组即可setup( ... entry_points{ torch_dynamo_backends: [ my_compiler your_module.submodule:my_compiler, ], }, ... )规则前是后端名称后是后端函数所在的模块路径与函数名。包安装后 entry point 即注入当前 Python 环境调用torch.compile(model, backendmy_compiler)时PyTorch 会先搜索register_backend注册的同名后端未命中再扫描 entry points 注册的后端。该机制的底层实现在 registry.py 的_discover_entrypoint_backends通过importlib.metadata.entry_points(grouptorch_dynamo_backends)枚举该分组的全部 entry point存入_BACKENDS字典lookup_backend在首次按名访问时通过entry_point.load()惰性加载并完成注册。注册有两个明确的用途允许以字符串形式传入torch.compile供 minifier最小化复现工具使用——minifier 生成的任何代码都必须调用注册你后端函数的代码通常是通过一条import语句实现。可用后端一览仓库内置大量后端可用torch._dynamo.list_backends()列出。其实现registry.py默认排除debug与experimental两类标签的后端返回按字母序排序的名字列表def list_backends(exclude_tags(debug, experimental)) - list[str]: ... return sorted(backends)面向训练的 AOTAutograd 后端TorchDynamo 之外你还可以定义由 AOTAutograd 调用的自定义后端。这有两个核心价值支持训练AOTAutograd 能生成用于编译的反向图因此这类后端天然支持模型训练更小的算子集合AOTAutograd 产出的 FX 图仅由 core Aten 算子 构成后端只需支持远小于完整 torch/Aten 算子集的 core Aten opset实现成本显著降低。接入方式是用torch._dynamo.backends.common.aot_autograd包装你的后端再照常通过backend关键字传给torch.compile。被包装的后端函数契约与之前完全一致。后端通过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_compilermy_compiler) # bw_compilermy_compiler model_opt torch.compile(model, backendmy_backend)从源码看AotAutograd 的实现要点包括构造时把所有关键字参数存入self.kwargs运行期转发给torch._functorch.aot_autograd.aot_module_simplified见 common.py若示例输入含 list/tuple/dict 等结构会先经flatten_graph_inputs展平common.py反向编译器会被双重disable包装既阻止 Dynamo 追踪 bw_compiler 函数本身也阻止追踪其生成的 backward passcommon.pybw_compiler缺省时回退到fw_compilerinference_compiler同理common.py可通过decompositions关键字传入分解表也支持返回表的零参 thunk以规避循环导入问题并在运行期解析为具体表common.py。其余可用关键字参数定义于 AotAutogradKwargs还包括partition_fn前反向图切分策略如min_cut_rematerialization_partition、keep_inference_input_mutations、ignore_shape_env、disable_functionalization、pre_grad_passes、compile_region_name等完整列表可查看该 TypedDict 定义。Eager 后端初始化_dynamo_backend_init钩子有些后端需要在torch.compile()时刻执行 eager 初始化例如加载原生库或初始化设备上下文。此时可以给后端定义一个_dynamo_backend_init属性——一个无参可调用对象在后端被解析时任何一次调用发生之前触发def my_backend(gm, example_inputs): return gm.forward def my_backend_init(): load_native_libs() # 在 compile() 时刻运行先于任何调用 my_backend._dynamo_backend_init my_backend_init torch.compile(backendmy_backend) def fn(x): return x 1该钩子的触发点在 eval_frame.py 的_maybe_fire_backend_initget_compiler_fn在lookup_backend解析后端之后、wrap_backend_debug包装之前调用它实现为“每次解析都触发”。由于属性通过getattr从后端对象读取因此无论是实例属性还是可经 MRO 解析的类方法都能命中。钩子行为要点生效范围广无论后端是直接传入、经register_backend按名注册还是通过torch.compiler.set_stance(force_backend...)强制指定钩子都会触发。torch._TorchCompileWrapper与AotAutograd都通过property把该属性转发给它们所包装的后端见 torch/init.py 与 common.py。AOTAutograd 场景使用aot_autograd(fw_compiler...)时把钩子设在内部的fw_compiler上即可——AotAutograd在触发时刻而非构造时刻读取它因此钩子可以在aot_autograd()构造之前或之后设置。只有fw_compiler上的钩子会被查询设在bw_compiler或inference_compiler上的钩子会被忽略。每次解析都会触发正常路径与fullgraphTrue路径都会触发且发生在任何调用之前因此环境损坏时torch.compile()会快速失败fail fast。后端被重复解析如set_stance(force_backend...)或compiled_autograd重建路径时每次解析触发一次需要一次性初始化的后端应在钩子内部自行去重import functools functools.cache # 每个进程仅执行一次 def my_backend_init(): load_native_libs() my_backend._dynamo_backend_init my_backend_init异常传播若钩子抛出异常异常会从torch.compile()中直接向外传播在torch.compiler.set_stance(force_backend...)场景下解析发生在首次调用时因此钩子触发以及失败从该次调用处浮出。实战示例三类典型自定义后端Debugging Backend打印 Dynamo 抽取的 FX 图想了解编译过程中发生了什么可以写一个打印 FX 图并返回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 # 返回一个 python callable torch.compile(backendmy_compiler) def fn(x, y): a torch.cos(x) b torch.sin(y) return a b fn(torch.randn(10), torch.randn(10))运行输出如下表格化的 FX IRmy_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.Modulefrom 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 # 返回一个 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, backendmy_compiler) optimized_mod(torch.randn(10))再看一个含控制流的例子它直观展示了 Dynamo 对条件分支的图切分能力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 # 返回一个 python callable torch.compile(backendmy_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))运行会依次输出三张子图其中if b.sum() 0的分支被拆成独立图最后两张图的输出顺序取决于 JIT 编译器先遇到哪一个是非确定性的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,),) {}Speedy Backend接入真实推理优化器接入一个性能更优的后端同样简单下面把torch.jit.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(backendoptimize_for_inference_compiler) def code_to_accelerate(): ...注意这里把 FX 图torch.jit.script成 TorchScript 模块再走 JIT 推理优化展示的是一条把 Dynamo 图转交给既有优化器管线的通用集成模式。Composable Backends后端组合与优雅降级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 # 第一个后端失败尝试其他后端... try: inductor_compiled lookup_backend(inductor)(gm, example_inputs) if inductor_compiled is not None: return inductor_compiled except Exception: pass return gm.forward该示例先用lookup_backend(tensorrt)尝试 TensorRT 后端失败抛异常或返回None则回退到inductor最终兜底返回gm.forward保证可用性。结合 registry.py 的实现可知lookup_backend对字符串会完成注册表查找、entry point 惰性加载与名字建议等全套解析逻辑因此既可按名取用内置后端也可取用你自己注册的后端使“组合后端”天然可扩展。小结与进一步阅读自定义后端是torch.compile开放生态的关键接口核心是(gm, example_inputs) - callable这一简洁契约register_backend与torch_dynamo_backendsentry points 提供了函数级与包级两种注册通道aot_autograd包装把后端扩展到了前向/反向图同时编译的训练场景并以 core Aten opset 大幅降低实现成本_dynamo_backend_init钩子则让后端能在编译时刻完成 eager 环境准备并快速失败。三者配合即可构建从调试、推理优化到多后端降级组合的完整后端体系。仓库中可供继续深挖的相关实现与测试后端注册表全量实现torch/_dynamo/backends/registry.pyAOTAutograd 后端封装torch/_dynamo/backends/common.py后端解析与初始化钩子触发点torch/_dynamo/eval_frame.py后端初始化钩子转发属性torch/init.py后端注册与 entry points 的测试用例test/dynamo/test_backends.py相关 IR 概念core Aten IRtorch.compiler_ir.md【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考