PyTorch 编译体系实战指南：深入解析 torch.compiler 与 torch.compile

PyTorch 编译体系实战指南深入解析 torch.compiler 与 torch.compile【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读torch.compiler是 PyTorch 2.x 引入的一个命名空间它把 PyTorch 内部编译器的关键方法尤其是torch.compile以公开 API 的形式暴露给用户。本文以 docs/source/user_guide/torch_compiler/torch.compiler.md 为骨架结合仓库内 TorchDynamo、TorchInductor 等源码与配套文档系统讲解torch.compile的三大底层技术TorchDynamo、TorchInductor、AOT Autograd、上手示例、后端选择、torch.compiler.precompile预编译 API、以及图中断graph break、守卫guards、重编译recompilation、动态形状dynamic shapes等编程模型核心概念。读完本文你将能够用一行代码加速自己的推理/训练程序并能系统排查为什么没有提速的问题。一、torch.compiler 是什么torch.compiler是一个命名空间将部分内部编译器方法暴露给用户使用。该命名空间的核心函数与核心特性就是torch.compile见 docs/source/user_guide/torch_compiler/torch.compiler.md。torch.compile是 PyTorch 2.x 引入的函数目标是解决 PyTorch 中精确图捕获accurate graph capturing这一难题最终让软件工程师能够更快地运行 PyTorch 程序。它完全用 Python 编写标志着 PyTorch 从 C 主导转向 Python 主导的一次转变。需要特别注意的是在官方文档中torch.compile、TorchDynamo、torch.compiler这几个术语有时会被混用同时torch.compiler.compile与torch.compile两种写法等价都可以在你的代码中正常工作见 docs/source/user_guide/torch_compiler/torch.compiler_fine_grain_apis.md。三大底层技术torch.compile依赖以下三项底层技术协同工作技术作用说明TorchDynamotorch._dynamo图捕获前端使用 CPython 的 Frame Evaluation API 安全地捕获 PyTorch 计算图对外暴露的方法通过torch.compiler命名空间提供给用户TorchInductor默认深度学习编译器为多种加速器和后端生成快速代码NVIDIA、AMD 和 Intel GPU 上以 OpenAI Triton 为核心构建模块AOT Autograd前向 反向捕获不仅捕获用户级代码还捕获反向传播实现提前ahead-of-time捕获反向传播从而让 TorchInductor 同时加速前向和反向要理解torch.compile在你自己代码上的追踪行为或深入了解其内部机制可以参考torch.compile编程模型。torch.compiler.precompileAOT 预编译 APItorch.compiler还包含一个提前ahead-of-timeAPItorch.compiler.precompile。它捕获整个计算fn(*example_inputs)——其中模型作为example_inputs之一传入例如precompile(lambda model, x: model(x), model, x)——并将其降级为一段自包含、可运行的 Python 源码字符串加上一个加速缓存。随后用torch.compiler.precompile.load重新加载该产物由于没有烘焙任何权重运行时需要再次传入模型。详见 API 参考 与 torch.compiler API reference。核心用法示例python_code, cache torch.compiler.precompile(lambda m, x: m(x), model, x) f torch.compiler.precompile.load(python_code, cache) out f(model, x) # 运行时再次传入模型关键语义依据 docs/source/torch.compiler_api.md不烘焙权重nn.Module参数parameters/buffers会被提升为图输入运行时把模型再传一次即可默认make_fx追踪器捕获是非严格non-strict的控制流针对示例输入特化形状是静态的每个尺寸都被烘焙进去。唯一例外是在调用前用torch._dynamo.decorators.mark_unbacked显式标记为 unbacked 的张量维度仅 inductor 后端支持该维度会被捕获为 unbacked symint一份产物即可服务任意运行时尺寸后端inductor默认经由 AOTAutograd Inductor 降级eager保留捕获的 ATen 图布局灵活、无内核但形状仍针对示例特化反向传播语义如果fn执行 backward产物会重跑整个前向和反向并把参数梯度散布到运行时模型的parameters()的.grad字段上累加方式p.grad g与 eager.backward()完全一致——所以请保留你惯用的zero_grad()/optimizer.step()循环安全警告load会以exec执行python_code请只加载你自己生成或完全信任的(python_code, cache)对就像对待任何即将运行的代码一样。二、快速上手从三角函数融合到真实模型2.1 第一个 torch.compile 示例下面是一个简单的torch.compile推理示例演示了torch.cos()与torch.sin()这两个 pointwise逐元素算子的使用见 docs/source/user_guide/torch_compiler/torch.compiler_get_started.md。这个例子本身可能看不出显著性能提升但能帮助你建立直觉。import torch def fn(x): a torch.cos(x) b torch.sin(a) return b new_fn torch.compile(fn, backendinductor) input_tensor torch.randn(10000).to(devicecuda:0) a new_fn(input_tensor)注意运行此脚本需要机器上至少有一块 GPU。如果没有 GPU可以删除代码中的.to(devicecuda:0)它将改在 CPU 上运行也可以把 device 设置为xpu:0在 Intel® GPU 上运行。2.2 Inductor 的核心优化融合fusion更常见的 pointwise 算子是torch.relu()。在 eager 模式下pointwise 算子效率不高每个算子都需要从内存读取张量、做修改、再写回。而 Inductor 执行的最重要优化就是融合。以上面cos/sin为例eager 模式2 次读x、a 2 次写a、b融合后1 次读x 1 次写b。这一点对新一代 GPU 尤为关键——瓶颈往往在内存带宽向 GPU 送数据的速度而非计算能力。此外 Inductor 还自动支持CUDA graphs可消除从 Python 程序逐个 launch kernel 的开销。2.3 查看生成的 Triton 内核TorchDynamo 支持许多后端TorchInductor 通过生成 Triton 内核来工作。把上面的例子保存为example.py然后运行TORCH_COMPILE_DEBUG1 python example.py脚本执行时终端会打印DEBUG消息在日志接近末尾处可以看到一个包含torchinductor_your_username的文件夹路径。该文件夹中的output_code.py包含生成的 kernel 代码大致如下pointwise(size_hints[16384], filename__file__, triton_meta{signature: {in_ptr0: *fp32, out_ptr0: *fp32, xnumel: i32}, device: 0, constants: {}, mutated_arg_names: [], configs: [AttrsDescriptor(divisible_by_16(0, 1, 2), equal_to_1())]}) triton.jit def triton_(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr): xnumel 10000 xoffset tl.program_id(0) * XBLOCK xindex xoffset tl.arange(0, XBLOCK)[:] xmask xindex xnumel x0 xindex tmp0 tl.load(in_ptr0 (x0), xmask, other0.0) tmp1 tl.cos(tmp0) tmp2 tl.sin(tmp1) tl.store(out_ptr0 (x0 tl.zeros([XBLOCK], tl.int32)), tmp2, xmask)注意以上代码只是示例具体生成的代码会因硬件而异。可以验证cos与sin确实发生了融合两个操作出现在同一个 Triton kernel 内临时变量保存在访问极快的寄存器中。由于生成的代码本身是 Python 写的即使没写过多少 CUDA kernel 也相当容易读懂。2.4 真实模型ResNet50、BERT 与 TIMMResNet50PyTorch Hubimport torch model torch.hub.load(pytorch/vision:v0.10.0, resnet50, pretrainedTrue) opt_model torch.compile(model, backendinductor) opt_model(torch.randn(1,3,64,64))HuggingFace BERT设计目标之一就是让 TorchDynamo/TorchInductor 对任何用户自定义模型开箱即用import torch from transformers import BertTokenizer, BertModel tokenizer BertTokenizer.from_pretrained(bert-base-uncased) model BertModel.from_pretrained(bert-base-uncased).to(devicecuda:0) model torch.compile(model, backendinductor) # 这是唯一修改的一行代码 text Replace me by any text youd like. encoded_input tokenizer(text, return_tensorspt).to(devicecuda:0) output model(**encoded_input)如果去掉模型和encoded_input上的to(devicecuda:0)Triton 会生成针对 CPU 优化的 C 内核。BERT 的内核比三角函数示例复杂但同样可以快速浏览来理解 PyTorch 的工作方式。TIMMresnext101_32x8dimport timm import torch model timm.create_model(resnext101_32x8d, pretrainedTrue, num_classes2) opt_model torch.compile(model, backendinductor) opt_model(torch.randn(64,3,7,7))三、后端Backend体系与选择要真正跑得更快torch.compile经由 TorchDynamo 需要一个后端把捕获的图转换成快速的机器码。不同的后端带来不同程度的优化收益。默认后端是 TorchInductor也叫inductor。TorchDynamo 维护着一个由合作伙伴开发的受支持后端列表可通过torch.compiler.list_backends()查看——每个后端都带各自的可选依赖。在源码层面该函数在 torch/compiler/init.py 中定义为list_backends(exclude_tags(debug, experimental))返回可直接传给torch.compile(..., backendname)的合法字符串后端注册表位于 torch/_dynamo/backends/registry.py。3.1 训练与推理后端后端描述torch.compile(m, backendinductor)使用 TorchInductor 后端默认torch.compile(m, backendcudagraphs)CUDA graphs AOT Autogradtorch.compile(m, backendipex)在 CPU 上使用 IPEX3.2 仅推理后端后端描述torch.compile(m, backendtensorrt)使用 Torch-TensorRT 做推理优化需要在调用脚本中import torch_tensorrt注册后端torch.compile(m, backendipex)在 CPU 上用 IPEX 做推理torch.compile(m, backendtvm)使用 Apache TVM 做推理优化torch.compile(m, backendopenvino)使用 OpenVINO 做推理优化在 REPL 里运行torch.compiler.list_backends()可以看到所有可用后端可尝试用cudagraphs作为下一个实验对象。四、Python 版本兼容性警告torch.compile可能不支持最近发布的大版本 Python。如果在不受支持的 Python 环境中尝试使用torch.compile可能会遇到类似下面的错误RuntimeError: torch.compile is not supported on Python 3.xx.0请确保当前 Python 版本处于 PyTorch 对torch.compile支持的范围之内。如果 PyTorch 安装在了过新的 Python 版本上你需要切换到更早的 Python 版本才能使用torch.compile。五、编程模型理解 torch.compile 的追踪行为要预测编译器在你代码上的行为、并在遇到问题时系统地自我排障需要理解torch.compile的编程模型。它做两件事一是澄清torch.compile的部分内部行为让你能更好地预测编译器行为二是提供对torch.compile进行更细粒度控制的方法见 docs/source/user_guide/torch_compiler/compile/programming_model.md。5.1 Dynamo 追踪Dynamotorch.compile的前端是一个自定义的 Python 字节码解释器它在保留 Python 全部灵活性的同时为 PyTorch 程序实现图编译。给定一个待编译函数Dynamo 解释 Python 字节码把 PyTorch 操作序列提取为一个或多个 FX 图可由后端进一步优化并同时产出FX 图接收原始输入加上函数所需的额外输入Python 字节码可作为f的即插即用替代品负责获取额外输入并传给图同时包含无法优化的 Python 副作用如列表 appendguards守卫图与字节码保持有效的条件列表。除非另有指定Dynamo 生成的图会针对输入张量的形状进行特化。5.2 图中断Graph BreaksDynamo 追踪你的代码尝试把 PyTorch 代码捕获到单一计算图FX graph中但这并非总是可行。遇到无法追踪的代码时就会发生graph break。在默认torch.compile设置下一次图中断的流程是编译目前已确定下来的 FX 图 → 用普通 Python 运行不受支持的代码 → 在不受支持代码之后用新的 FX 图恢复追踪。图中断是 Dynamo 的一个特性它允许 Dynamo 运行任意 Python 代码并切分出一个个可单独优化的功能性子图。但它也可能导致torch.compile出现意外的慢——如果你没有获得预期的加速建议检查并消除图中断。图中断可能发生在依赖数据的 if 语句、许多 Python 内置函数、C 函数等。图中断示例因调用不受支持的torch.savetorch.compile def f(x): y x ** 2 / 2 torch.save(y, foo.pt) # torch.save 是不受支持的操作 z y ** 3 / 6 return z x torch.randn(3) print(f(x))其语义等价于def compiled_f_semantics(x): y torch.compile(g, fullgraphTrue)(x) torch.save(y, foo.pt) z torch.compile(h, fullgraphTrue)(x) return z def g(x): return x ** 2 / 2 def h(x): return y ** 3 / 65.3 Guards守卫torch.compile在追踪代码时会基于某些运行时值做出假设并据此生成guards——对这些假设的运行时检查。在后续调用已编译函数时guards 会被执行以决定能否复用已编译的代码。常见的运行时检查包括常量值、类型和对象 ID。由于 guards 每次调用都会运行它们会带来每次调用的开销。如果这个开销对你的模型影响显著可以参考 Reducing Guard Overhead。下面的例子展示了生成的 guards——TENSOR_MATCHguard 检查输入的 type、device、dtype、shape 等torch.compile def fn(x): return x 1 print(fn(torch.ones(3, 3)))从 Dynamo 内部视角看见 docs/source/user_guide/torch_compiler/torch.compiler_dynamo_overview.mdcheck_tensor这个 guard 会检查以下torch.Tensor属性张量的 Python 类是否张量子类化等dtypedevicerequires_graddispatch key应用线程局部 include/exclude 后ndimsizesstrides5.4 重编译Recompilations如果 guards 对所有已编译的代码实例都失败torch.compile就必须重新编译该函数——这意味着需要再次追踪原始代码。下面的例子中由于检查张量形状的 guard 失败重编译是必要的。重编译会增加整体编译时间相关优化可参考 Dealing with Recompilations 与 Reducing Compile Time。torch.compile def fn(x): return x 1 print(fn(torch.ones(3, 3))) print(fn(torch.ones(4, 4))) # 形状变了 → 触发重编译5.5 动态形状Dynamic Shapestorch.compile最初假设张量形状是静态/恒定的并基于这些假设设置 guards。使用dynamic shapes可以让torch.compile生成能接受不同形状张量输入的编译代码——从而避免每次形状变化都重新编译。默认情况下torch.compile(dynamicNone)开启了自动动态形状如果编译因形状不匹配而失败会用动态形状重试编译。动态形状也可以完全启用dynamicTrue或禁用dynamicFalse。torch.compile(dynamicTrue) def fn(x): return x 1 print(fn(torch.ones(3, 3))) print(fn(torch.ones(4, 4))) # 不再需要重新编译六、细粒度追踪控制 APItorch.compile会对整个用户模型执行 TorchDynamo 追踪。但模型代码的一小部分可能无法被编译器处理此时你可能想在这一小部分上禁用编译同时对模型其余部分继续编译。下表汇总了用于定义跳过编译的代码区域的 API见 docs/source/user_guide/torch_compiler/torch.compiler_fine_grain_apis.mdAPI描述何时使用torch.compiler.disable在被装饰函数及其递归调用的函数上禁用 Dynamo模型一小部分无法被torch.compile处理时用于快速解阻塞torch._dynamo.disallow_in_graph禁止被标记的算子进入 TorchDynamo 图TorchDynamo 产生图中断并在 eager不编译模式下运行该算子适合算子级处理自定义算子如torch.ops.fbgemm.*导致torch.compile出问题时既可用于调试也可用于解阻塞torch.compile.allow_in_graph被注解的可调用对象原样放入 TorchDynamo 图对 Dynamo 而言是黑盒。注意 AOT Autograd 仍会追踪它因此allow_in_graph只是 Dynamo 层面的概念适用于含 hooks、autograd.Function等已知 Dynamo 难支持特性的模型部分但每次使用都必须仔细审查不能有图中断、不能有闭包torch._dynamo.graph_break添加一个图中断图中断前后的代码都经过 TorchDynamo部署时很少有用如果你觉得需要它大概率应该用disable或disallow_in_graphtorch.compiler.is_compiling指示某图是否作为torch.compile()或torch.export()的一部分被执行/追踪—torch.compiler.is_dynamo_compiling指示某图是否经由 TorchDynamo 追踪比is_compiling()更严格只有使用 TorchDynamo 时才为 True—torch.compiler.is_exporting指示某图是否经由 export 追踪比is_compiling()更严格只有使用torch.export时才为 True—6.1 torch.compiler.disabletorch.compiler.disable在装饰的函数帧及其递归调用的所有函数帧上禁用编译。TorchDynamo 会拦截每个 Python 函数帧的执行假设函数fn调用a_fn和b_fn而a_fn又调用aa_fn、ab_fn。eager 模式下这些函数帧原样运行使用torch.compile时 TorchDynamo 会拦截其中每一个函数帧。如果a_fn给torch.compile带来麻烦且不是模型的关键部分就可以在a_fn上加torch.compiler.disableTorchDynamo 将不再查看从a_fn调用发起的帧恢复原始 Python 行为。也可以用非装饰器语法运行时打补丁跳过编译但官方建议尽量避免这种风格因为你必须确保原函数的所有使用者都在使用打过补丁的版本。6.2 torch._dynamo.disallow_in_graphtorch._dynamo.disallow_in_graph禁止算子而非函数出现在 TorchDynamo 提取的图中。典型场景模型能被 TorchDynamo 提取图但下游编译器失败如缺少 meta kernel或某个算子的 Autograd dispatch key 设置不正确。此时可把该算子标记为disallow_in_graphTorchDynamo 会图中断并用 eager 模式运行它。两个注意点需要找到对应的Dynamo 层算子而非 ATen 层算子详见下文 Limitationstorch._dynamo.disallow_in_graph是全局标志如果你在比较不同后端编译器切换到另一个编译器时可能需要为该被禁止算子调用allow_in_graph。6.3 torch.compiler.allow_in_graph当相关函数帧含 hooks、autograd.Function等已知难以被 TorchDynamo 支持的特性且你确信下游组件如 AOTAutograd能安全追踪该函数时可用torch.compiler.allow_in_graph。被装饰函数会被 TorchDynamo 视为黑盒原样放进生成的图。警告allow_in_graph会完全跳过 TorchDynamo 对被装饰函数的检查包括图中断处理、闭包处理等所有安全性检查。AOTAutograd 等下游组件依赖 TorchDynamo 处理复杂 Python 特性而allow_in_graph绕过了 TorchDynamo可能导致 soundness 问题和难以调试的问题。请谨慎使用。6.4 限制Limitations所有这些 API 都作用于 TorchDynamo 层面因此它们只能看到 TorchDynamo 所见的内容这可能导致令人困惑的场景。例如torch._dynamo.disallow_in_graph对 ATen 算子不起作用因为 ATen 算子对 AOT Autograd 是可见的——torch._dynamo.disallow_in_graph(torch.ops.aten.add)在上面的例子中不会生效。七、Dynamo 内部机制速览TorchDynamo简称 Dynamo是一个 Python 层面的即时JIT编译器旨在让未经修改的 PyTorch 程序跑得更快。它通过 CPython 的 Frame Evaluation APIPEP 523在字节码执行前动态修改 Python 字节码把 PyTorch 操作序列重写提取为 FX 图再用可定制的后端编译。Dynamo 设计上混合 Python 执行与编译后端兼顾可用性与性能。想观察 Dynamo 到底在做什么可以运行TORCH_LOGSdynamo,guards,bytecode日志会展示FX 图 → 原始字节码 → Dynamo 修改后的字节码 → 反编译源码 → guards。在修改后的字节码中__compiled_fn_0是编译后图的入口__resume_at_30_1、__resume_at_38_2是图中断后恢复执行的续接函数。若想获得人类可读的源码而非字节码可安装并使用 depyf 反编译钩子import depyf depyf.install()以上机制对用户的实际意义在于torch.compile先把用户函数中的计算图提取出来编译成优化函数再组装回一个新函数——它在功能上等价于用户代码但计算速度更快。八、继续深入torch.compiler主题下还有大量延伸文档可供查阅入门后续torch.compiler_get_started.md 的 Next Steps训练教程、API 参考、细粒度追踪核心概念core_concepts.md编程模型、Dynamo 概览、nn.Module 处理、backward 处理性能performance.md故障排查troubleshooting_faqs.mdAPI 与配置api_reference.md完整 API 文档、配置项、细粒度编译器控制深入主题动态形状手册compile/programming_model.md 下的一系列dynamic_shapes_*.md、torch.compiler_custom_backends.md自定义后端、torch.compiler_aot_inductor.mdAOT Inductor等从源码层面继续探索可关注 torch/compiler/init.py命名空间导出与list_backends、torch/_dynamoTorchDynamo 前端以及 torch/_inductorTorchInductor 默认后端等目录。结语torch.compiler命名空间把 PyTorch 2.x 的编译能力浓缩为一行torch.compile(model)TorchDynamo 负责安全图捕获AOT Autograd 把反向也纳入捕获TorchInductor 默认把图编译为 GPU 上的 Triton 内核或 CPU 上的 C 代码。理解 graph break、guards、recompilation 与 dynamic shapes 这四个编程模型概念是排查为什么没提速的关键细粒度 APIdisable/disallow_in_graph/allow_in_graph则让你在无法整体编译时精准解阻塞torch.compiler.precompile更进一步把整个计算提前降级为可分发、可重载的自包含产物。掌握这些能力即可在自己的 PyTorch 项目中系统性地获得编译加速。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考