PyTorch torch.compile 性能优化完全指南Profiling、基准测试与 CUDAGraph Trees 实战【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读torch.compile是 PyTorch 2.x 中默认开启的即时编译器它将 Dynamo 图捕获、TorchInductor 代码生成与 CUDA Graph 优化融为一体。本文将基于 PyTorch 官方性能文档系统讲解如何用torch.profiler定位编译区域与图断裂、用 TorchInductor 的 GPU 剖析工具拆解每个 Triton 内核的开销、借助 nightly 性能看板对比 PR 前后吞吐并深入解析modereduce-overhead背后的 CUDAGraph Trees 记忆池复用机制与其限制。读完本文你将掌握一套从宏观基准对比到单内核微基准再到CUDA Graph 级优化的完整性能调优方法论。一、PyTorch 2.0 性能看板宏观基准对比的入口1.1 看板测量什么PyTorch 团队通过 HUD Performance Dashboard 每夜跟踪torch.compile的性能。官方文档描述其运行环境为 12 个 GCP A100 节点每个节点含一块 40GB A100 GPU 和一块 6 核 2.2GHz Intel Xeon CPU对应 CI 工作流文件为.github/workflows/inductor-perf-test-nightly.yml。看板针对三个基准套件测量推理inference与训练training性能全部使用 AMP 混合精度TorchBenchPyTorch 官方维护的端到端基准套件Huggingface主流 Transformer 模型集合TIMM图像分类模型集合仓库中的对应脚本为 benchmarks/dynamo/timm_models.py同时看板还覆盖 TorchInductor 的三种配置default默认配置with_cudagraphs默认配置叠加 CUDA Graphdynamic默认配置叠加动态形状dynamic_shapes1.2 三个关键指标看板除通过率外报告三个核心性能指标指标对比基准含义Geometric mean speedup几何平均加速比PyTorch eager 模式越大越好衡量编译后整体提速Mean compilation time平均编译时间—衡量第一次调用torch.compile的编译开销Peak memory footprint compression ratio峰值内存压缩比PyTorch eager 模式越大越好衡量编译后内存占用改善看板上每个单独性能数字都可以点击进入该基准套件下所有测试的详细数值视图并可通过顶部的下拉列表切换表格与图表默认展示过去 7 天 TorchBench 的 AMP 训练性能趋势。1.3 在合并前检查 PR 对 TorchInductor 性能的影响如果担心某个 PR 会改变 TorchInductor 的生成代码质量可以在 CI 的 Actions 页面手动触发inductor-perf-test-nightly.yml工作流在Run workflow时选择 PR 分支提交。跑完后在看板 UI 中选中对应分支名和 commit ID 即可查看结果。注意这是一次昂贵的 CI 运行官方文档明确提示应谨慎使用该功能。1.4 本地复现看板命令完整 dashboard 运行所用的确切命令行可从任意近期 CI 运行日志中找到。官方文档给出的示例为python benchmarks/dynamo/huggingface.py --performance --cold-start-latency --inference --amp --backend inductor --disable-cudagraphs --device cuda只要本机有能跑 PyTorch 2.x 的 GPU 即可本地执行。python benchmarks/dynamo/huggingface.py -h会给出基准脚本的详细选项说明。类似的基准脚本还有 benchmarks/dynamo/timm_models.py、benchmarks/dynamo/torchbench.py 等统一位于 benchmarks/dynamo 目录。二、用 torch.profiler 理解 torch.compile 性能2.1 torch.profiler 的定位与示例torch.profiler适合在内核级粒度理解程序性能例如展示图断裂位置与资源利用率帮助判断下一步该往哪个方向排查。若需要更底层的内核级分析可结合 Nvidia Nsight Compute、AMD Omnitrace、Intel VTune Profiler 或下文将介绍的 Inductor 剖析工具。官方文档给出的 resnet18 剖析示例包含三个关键写法加入 warm-up 运行等待编译完成同时预热 CUDA caching allocator用torch.profiler.profile()上下文包住感兴趣的部分用prof.export_chrome_trace(trace.json)导出剖析产物import torch from torchvision.models import resnet18 device cuda # or cpu, xpu, etc. model resnet18().to(device) inputs [torch.randn((5, 3, 224, 224), devicedevice) for _ in range(10)] model_c torch.compile(model) def fwd_bwd(inp): out model_c(inp) out.sum().backward() # warm up fwd_bwd(inputs[0]) with torch.profiler.profile() as prof: for i in range(1, 4): fwd_bwd(inputs[i]) prof.step() prof.export_chrome_trace(trace.json)查看 chrome trace在 Chrome 浏览器打开chrome://tracing并加载 json 文件用w/s键缩放、a/d键左右平移按?可查看快捷键帮助。在 trace 中可以看到CompiledFunction与CompiledFunctionBackward事件对应 Dynamo 编译的区域上方为 CPU 事件下方为 GPU 事件2.2 CPU 与加速器事件之间的 Flow加速器上的每个内核都是由 CPU 代码启动的而除少数例外加速器内核是异步启动的因此 profiler 可以绘制 CPU 事件与加速器事件之间的流flow连接标明是哪个 CPU 事件启动了该内核。查看方式点击某个 GPU kernel再点击 ac2g或者用顶部 Flow events 下拉菜单打开全部 flow。2.3 绕开 CUDA Graph 剖析问题启用 CUDA Graph 后某些 CUDA 配置driver 版本低于 525.85.12或 CUDA 12会在剖析工具与 CUDA Graph 之间产生冲突。修复方法是在程序顶部加一个空的剖析上下文import torch torch.profiler._utils._init_for_cuda_graphs() # ... rest of program2.4 理解编译时间剖析首次调用要弄清编译为什么慢可以剖析torch.compile程序的第一次调用。需注意编译类负载与典型 PyTorch 负载差异很大trace 失真比普通剖析更明显且 trace 文件可能很大——超过 1GB 的 trace 用 chrome tracing 工具很难打开。官方示例同时展示了用torch.profiler.record_function标记 warmup compile 与 resnet18 compile 两个阶段的做法第一次调用必须在剖析期间发生才能捕获编译过程同时要先加一个小的 warm-up 编译以初始化需要惰性初始化的系统。import torch from torchvision.models import resnet18 # user can switch between cuda and xpu device cuda model resnet18().to(device) inputs [torch.randn((5, 3, 224, 224), devicedevice) for _ in range(10)] model_c torch.compile(model) def fwd_bwd(inp): out model_c(inp) out.sum().backward() def warmup_compile(): def fn(x): return x.sin().relu() x torch.rand((2, 2), devicedevice, requires_gradTrue) fn_c torch.compile(fn) out fn_c(x) out.sum().backward() with torch.profiler.profile() as prof: with torch.profiler.record_function(warmup compile): warmup_compile() with torch.profiler.record_function(resnet18 compile): fwd_bwd(inputs[0]) prof.export_chrome_trace(trace_compile.json)非图形化替代方案torch._dynamo.utils.compile_times()提供大致相同的信息。它不会展示各编译步骤何时发生但能展示每个步骤耗时多少且不受剖析开销影响。2.5 用 Torch-Compiled Region 与 CompiledFunction 定位图断裂尽管已有日志工具用于识别图断裂profiler 提供了一种快速的可视化方法关注两个事件Torch-Compiled RegionPyTorch 2.2 引入覆盖整个编译区域的 profiler 事件。图断裂几乎总是表现为嵌套的 Torch-Compiled Region 事件。PyTorch 2.5 起该事件还包含 frame ID帧的唯一标识与 frame compile ID该帧被编译的次数。若两个函数分别独立应用torch.compile通常看到两个**相邻非嵌套**的 Torch-Compiled Region若遇到图断裂或disable()/被跳过的区域则看到嵌套事件。CompiledFunctionPyTorch 2.0 引入当任何输入需要梯度时出现。每次图断裂都会打断一个 CompiledFunction 块、将其一分为二。CompiledFunction 只在涉及 Autograd即图中部分输入张量requires_gradTrue时出现trace 中它通常与反向传播的CompiledFunctionBackward成对出现若反向函数被调用两者之间会出现 fwd-bwd link。若你的图不需要梯度且不包含 Torch-Compiled Region 事件可以借助是否存在 Inductor 生成的 Triton 内核这一线索来判断torch.compile是否真正生效。官方文档还给出了一个使用torch._dynamo.graph_break()强制打断四个连续 Sequential 模块的合成示例ModelWithBreaks并在 3 次循环中分别对 10 个输入做fwd_bwd最终导出trace_break.json观察嵌套区域与多个 CompiledFunction 事件。2.6 内核事件三层结构与启动开销当一个算子被启动时通常看到三类事件CPU 侧事件、内核启动GPU kernel 场景事件、GPU 侧事件。Inductor 生成的 Triton 内核CPU 侧事件以triton_前缀出现信息量较少只有内核名与启动不如普通 aten 内核启动那样包含输入形状、类型等内核启动显示为cuLaunchKernel而非 aten 算子常见的cudaLaunchKernelGPU 侧事件是否详细取决于 Inductor 配置unique_kernel_names非 Inductor 生成的 Triton 内核CPU 侧事件可能不出现自动插入 profiler 事件的机制目前实现在 Inductor 层绕过 Inductor 的 Triton 内核除非手动注解否则不会出现内核启动同样为cuLaunchKernelGPU 侧事件按作者命名。Inductor 生成的 CPU 内核CPU 侧事件不会出现尚未添加剖析支持也没有内核启动与 GPU 侧事件。非 Triton 内核aten 内核或自定义算子有时 Inductor 会回退到原始算子实现此时会看到对 aten 算子的调用。启动开销问题GPU 利用率不佳的快速识别方法是看 GPU 上内核之间是否有大段空隙。这通常源于 CPU 开销——两次内核启动之间 CPU 耗时大于 GPU 处理内核的耗时在小 batch size 下更常见。当启动开销是主要矛盾时启用 CUDA Graphs 往往能显著改善。三、TorchInductor GPU 剖析拆解到单个 Triton 内核当模型没有达到预期速度时需要检查单个内核。通常 GPU 时间占比最大的内核最值得关注之后可以单独运行该内核并检查其性能。PyTorch 为此提供了一整套工具核心步骤是先整体拆解 GPU 时间再对单个内核做微基准。3.1 相关环境变量环境变量作用默认值TORCHINDUCTOR_UNIQUE_KERNEL_NAMES为 Triton 内核生成更有意义的名称如triton_poi_fused_cat_155poi表示 pointwise后跟原始 ATen 算子名便于在 trace 中识别内核类别默认关闭是为了提高编译缓存命中率0TORCHINDUCTOR_BENCHMARK_KERNEL让 Inductor 的 codegen harness 为每个 Triton 内核生成独立的基准代码0TORCHINDUCTOR_MAX_AUTOTUNE让 Inductor 的自动调优器尝试更多triton.Configs并挑选性能最佳者会增加编译时间以换取性能提升0这些配置在仓库中的落点包括 torch/_inductor/config.py 中的max_autotune、benchmark_kernel与triton.unique_kernel_names后者通过TORCHINDUCTOR_UNIQUE_KERNEL_NAMES环境变量控制以及 torch/_inductor/codegen/triton.py 中按config.benchmark_kernel拼接imports_for_benchmark_kernel()的实现。3.2 把模型 GPU 时间拆解到单个内核以 mixnet_l 为例步骤 1运行基准脚本TORCHINDUCTOR_UNIQUE_KERNEL_NAMES1 TORCHINDUCTOR_BENCHMARK_KERNEL1 python -u benchmarks/dynamo/timm_models.py --backend inductor --amp --performance --dashboard --only mixnet_l --disable-cudagraphs --training官方文档特别提示该工具依赖内核名判断其类别因此开启TORCHINDUCTOR_UNIQUE_KERNEL_NAMES至关重要。步骤 2在输出日志中查找编译模块路径**Compiled module path: /tmp/torchinductor_shunting/qz/cqz7hvhood7y3psp7fy6msjxsxyli7qiwiybizdwtjw6ffyq5wwd.py**每个编译模块一行。若没有额外图断裂日志中应看到两行——一个对应前向图一个对应反向图。步骤 3对单个编译模块剖析把前向图对应的模块重命名为fwd.py直接运行python fwd.py -p输出中有三处值得关注Chrome trace 文件日志中查找Chrome trace for the profile is written to /tmp/compiled_module_profile.json用chrome://tracing加载即可交互式查看。GPU 忙碌时间占比日志行形如Percent of time when GPU is busy: 102.88%。可能看到大于 100% 的值原因是 PyTorch 在开启剖析时用内核执行时间、关闭剖析时用 wall time剖析会对内核执行时间造成轻微失真但总体影响不大。例如densenet121小 batch 下前向图GPU 忙碌占比仅 32.69%说明模型有大量 CPU 开销——这与启用 cudagraphs 能大幅改善 densenet121 性能的事实一致。按内核类别拆解 GPU 时间mixnet_l 示例中pointwise 内核占 28.58%、reduction 内核占 13.85%、persistent reduction 内核占 3.89%其余为 mm/conv 的 cutlass/cudnn 内核占 56.57%。该信息位于每个内核类别报告的最后一行汇总中。还可以放大到某一类内核如 reduction看到按执行时间排序的明细表及每个内核的执行次数——这有助于判断优化优先级仅占 0.1% 的内核即使优化到极致收益也有限占 2% 的内核提速 2 倍能带来 1% 的整体收益才值得投入。3.3 单独基准某个 Triton 内核假设想仔细研究最贵的 reduction 内核triton_red_fused__native_batch_norm_legit_functional_16占前向图整体 wall time 的 2.19%。在fwd.py中查找该内核名可找到注释行# kernel path: /tmp/torchinductor_shunting/jk/cjk2vm3446xrk7rth7hr6pun7xxo3dnzubwcn6ydrpifal4eykrz.py将其重命名为k.py——它是一个独立的 Python 模块包含内核代码与其基准。直接运行python k.py会报告该内核的执行时间与带宽。还可以验证 max-autotune 是否对它有帮助TORCHINDUCTOR_MAX_AUTOTUNE1 python /tmp/k.py也可以临时添加更多 reduction 启发式规则后再次运行观察对内核的影响。四、CUDAGraph Treesreduce-overhead 模式的核心机制4.1 背景CUDA Graph 与 PyTorch 集成CUDA GraphCUDA 10 引入允许把一系列 CUDA 内核定义并封装为单个单元操作图通过一次 CPU 操作启动多个 GPU 操作从而降低启动开销。它对高 CPU 开销或小计算量的模型提速尤其明显但也有严格限制不支持任意控制流不过torch.cond()表达的控制流可以捕获进 CUDA Graph触发 host 到 device 同步的内核如.item()会报错所有输入参数固定为录制时的值CUDA 内存地址固定但该地址上的内存值可以变化不允许必需的 CPU 算子或 CPU 副作用PyTorch 的 torch.cuda.CUDAGraph 封装处理了与 caching allocator 交互的几个棘手问题CachingAllocator 为所有新分配使用独立内存池录制期间内存的分配、记账与释放与 eager 运行完全一致重放时只调用内核allocator 不发生变化录制之后 allocator 并不知道用户程序中哪些内存正被使用。eager 分配与 cudagraph 分配使用独立内存池若两者都有大量内存分配可能增加程序总内存。Make Graphed Callablestorch.cuda.make_graphed_callables是让一系列 callable 共享单个内存池的抽象利用录制时 caching allocator 精确记账内存的事实安全地共享内存每次调用时输出保持为存活内存防止一个 callable 覆盖另一个的存活内存它只能按单一顺序调用第一次运行的内存地址会被固化。旧版 TorchDynamo CUDA Graph 集成的问题cudagraph_treesFalse时不同图捕获之间不复用内存可能导致显著的内存回退。即使模型没有图断裂前向与反向也是两次独立的图捕获内存池不共享——前向中保存的激活内存无法在反向中被回收。4.2 CUDAGraph Trees 的工作原理与 Graphed Callables 类似CUDAGraph Trees 在所有图捕获间使用单一内存池不同的是它不需要单一调用序列而是为 CUDA Graph 捕获构建独立的树。官方文档给出的示意示例torch.compile(modereduce-overhead) def foo(x): # GRAPH 1 y x * x * x # graph break triggered here if y.sum() 0: # GRAPH 2 z y ** y else: # GRAPH 3 z (y.abs() ** y.abs()) torch._dynamo.graph_break() # GRAPH 4 return z * torch.rand_like(z) # the first run warms up each graph, which does things like CuBlas or Triton benchmarking foo(torch.arange(0, 10, devicecuda)) # The second run does a CUDA Graph recording, and replays it foo(torch.arange(0, 10, devicecuda)) # Finally we hit the optimized, CUDA Graph replay path foo(torch.arange(0, 10, devicecuda))该函数有两条路径1 → 2 → 4 或 1 → 3 → 4。系统通过构建一条 CUDA Graph 录制磁带如 1 → 2 → 4在多次录制间共享全部内存并施加不变量保证内存始终位于录制时的位置、用户程序中不存在可能被覆盖的存活张量CUDA Graph 的既有约束依然适用必须以相同参数静态尺寸、地址等调用相同内核录制与重放之间必须观察到相同的内存模式某图的一个输出张量在录制期间于另一个图之后死亡重放时也必须如此CUDA 池中的存活内存在两次录制间强制形成依赖这些录制只能以单一顺序调用1 → 2 → 4所有内存共享单一内存池因此相对 eager没有额外内存开销。新路径Graph 3出现时Graph 1 被重放后遇到尚未录制的 Graph 3。由于重放时私有内存池不更新y不会反映在 allocator 中若不处理会被覆盖。为了在重放其他图之后复用同一内存池系统把内存池checkpoint 回 Graph 1 结束时的状态使存活张量反映到 caching allocator 中即可安全运行新图。此时先命中已录制的CUDAGraph.replay()路径Graph 1再遇到 Graph 3——同样需要先 warm up 一次再录制warm-up 运行时内存地址未固定Graph 4 也会回退到 Inductor 的非 cudagraph 调用。第二次遇到 Graph 3 时即可录制随后因输入内存地址变化再次录制 Graph 4由此形成一棵 CUDA Graph 录制树1 / \ 2 3 \ \ 4 44.3 输入突变Input Mutation支持输入突变函数指对输入张量进行就地写入的函数例如def foo(x, y): # mutates input x x.add_(1) return x y输入突变对 CUDAGraph Trees 是挑战CUDA Graph 要求静态内存地址Trees 可能为每个输入张量x分配一个静态地址x执行时先把x复制到x再重放已录制的图。对输入突变函数x被就地更新但由于x与x位于不同 CUDA 内存地址该更新不会反映到输入张量x上。深入分析后发现输入有三类来自 eager 的输入假定每次执行的张量地址都变化因 cudagraph 冻结内存地址必须在录制与执行前复制到静态地址张量参数与缓冲区假定并运行时检查每次执行地址相同无需复制内容CUDAGraph Trees 先前的输出cudagraph 输出地址固定若先运行 CUDAGraph1 再运行 CUDAGraph2从 1 流入 2 的输入地址固定像参数与缓冲区一样无需复制运行时若检查到不稳定则重新录制因此 CUDAGraph Trees 支持对参数、缓冲区及先前输出张量的突变对 eager 输入的突变Trees 会不带 CUDA Graph 运行该函数并输出skipping due to mutated inputs日志。官方示例演示了启用输入突变支持的方式import torch torch.compile(modereduce-overhead) def foo(x): return x 1 torch.compile(modereduce-overhead) def mut(x): return x.add_(2) # Enable input mutation support torch._inductor.config.triton.cudagraph_support_input_mutation True for i in range(3): torch.compiler.cudagraph_mark_step_begin() inp torch.rand([4], devicecuda) # CUDAGraph is applied since foo does not mutate inp tmp foo(inp) # Although mut mutates tmp, which is an output of a CUDAGraph # managed function. So CUDAGraph is still applied. mut(tmp) torch.compiler.cudagraph_mark_step_begin() inp torch.rand([4], devicecuda) tmp foo(inp) # While tmp is a CUDAGraph Tree managed functions output, tmp.clone() # is not. So CUDAGraph is not applied to mut and there is a log # skipping cudagraphs due to mutated inputs mut(tmp.clone())仓库中该配置项位于 torch/_inductor/config.py 的triton.cudagraph_support_input_mutation默认值not is_fbcode()。若函数要突变 eager 输入官方建议重写函数以避免输入突变。4.4 动态形状支持动态形状指输入张量在不同调用间形状不同。由于 CUDA Graph 要求固定张量地址CUDAGraph Trees 会为输入张量的每种唯一形状重新录制CUDA Graph导致单个 inductor 图对应多个 CUDA Graph。形状有限如推理中的 batch size时重录是划算的若形状频繁变化甚至每次调用都变重录可能不划算。官方文档还提到直到 CUDA 12.4 与 Driver 550 之前NVIDIA 每个内核启动在 CUDA Graph 中占用 64KB 设备内存大量重录时该内存成本可观。建议对形状频繁变化的函数把输入张量 padding 到少数固定形状以继续享受 CUDA Graph 收益同时可设置torch._inductor.config.triton.cudagraph_skip_dynamic_graphsTruetorch/_inductor/config.py跳过动态形状输入的函数只对静态形状函数使用 cudagraph。4.5 NCCL 支持CUDAGraph Trees 支持包含 NCCL 算子的函数Trees 按设备录制 CUDA Graph而 NCCL 支持允许跨设备通信。示例torch.compile(modereduce-overhead) def func(x): y x * x y torch.distributed.all_reduce(y, optorch.distributed.ReduceOp.SUM) x torch.nn.functional.silu(x) return x * y4.6 跳过 CUDA Graph 的原因由于 CUDA Graph 要求静态输入地址且不支持 CPU 算子Trees 会检查函数是否满足要求并在必要时跳过。常见原因输入突变跳过就地突变 eager 输入的函数突变参数/缓冲区/先前输出仍支持CPU 算子跳过含 CPU 算子的函数建议拆成多个函数并对仅含 GPU 算子的部分应用多设备算子含多设备算子的函数被跳过目前按设备应用跨设备通信请用 NCCL 等受支持库自由 unbacked 符号通常出现在动态形状场景Trees 目前为每种唯一输入形状录制一张图CUDAGraph 不安全的自定义算子可能包含 cudagraph 不安全算子导致跳过不兼容算子官方给出完整的在确定性算法开启时同样不兼容的不兼容算子清单aten._fused_moving_avg_obs_fq_helper.default aten._fused_moving_avg_obs_fq_helper_functional.default aten.multinomial.default fbgemm.dense_to_jagged.default fbgemm.jagged_to_padded_dense.default run_and_save_rng_state run_with_rng_state aten._local_scalar_dense aten._assert_scalar4.7 CUDAGraph 不安全自定义算子自定义算子默认被视为对 CUDA Graph 安全但某些算子可能包含 CPU 算子等不支持操作。由于编译器把自定义算子当作黑盒用户必须显式打上torch._C.Tag.cudagraph_unsafe标签来标记其不安全torch.library.custom_op( mylib::modify, mutates_args(), tags(torch._C.Tag.cudagraph_unsafe,), ) def modify(pic: torch.Tensor) - torch.Tensor: pic1 pic 1 pic1_cpu (pic1.cpu() 1) * 2 return pic1_cpu.cuda() pic modify.register_fake def _(pic): return torch.empty_like(pic)含此类算子的函数会被 CUDA Graph 跳过除非启用 CUDAGraph partition。4.8 CUDAGraph Partition自动切分不兼容算子CUDAGraph partition 是一种编译器解决方案自动把不支持的算子切分出去、重排算子以减少分区数量并对每个分区单独应用 CUDA Graph。启用方式为torch._inductor.config.graph_partitionTrue定义于 torch/_inductor/config.py。考虑如下例子x、y是 GPU 输入但y_cpu是 CPU 张量。无 partition 时该函数因 CPU 算子被整体跳过启用 partition 后 CPU 算子被切分出去剩余 GPU 算子被 cudagraph 化得到两个独立的 CUDA Graphdef f(x, y): x1 x 1 y1 y 1 y_cpu y1.cpu() 1 z x y return x1 y1 z y_cpu.cuda()目前 partition 支持切分以下类型的算子非 GPU 算子如 CPU 张量上的计算设备复制算子设备间数据传输如例子中的y1.cpu()控制流算子尚未被 CUDA Graph 支持会被切分CUDAGraph 不安全自定义算子带torch._C.Tag.cudagraph_unsafe标签的算子Unbacked Symints见动态形状支持一节4.9 限制与跨迭代存活的张量由于 CUDA Graph 固定内存地址它对上一次调用遗留的存活张量处理不佳。官方示例import torch torch.compile(modereduce-overhead) def my_model(x): y torch.matmul(x, x) return y x torch.randn(10, 10, devicecuda) y1 my_model(x) y2 my_model(x) print(y1) # RuntimeError: Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run.在旧版独立 CUDA Graph 实现中第二次调用的输出会覆盖第一次的输出。CUDAGraph Trees 既不想在迭代间引入意外依赖导致错过热路径也不想过早释放上一次调用的内存。其启发式规则是推理时每次调用torch.compile都开启新迭代训练时只要没有未调用的 pending backward 也如此。若启发式判断错误可用torch.compiler.cudagraph_mark_step_begin()手动标记新迭代开始其实现位于 torch/compiler/init.py 并委托给torch._inductor.cudagraph_trees.mark_step_begin或在开始下一次运行前在torch.compile之外克隆上一次迭代的张量。若用户可见输出必须跨迭代存活可设置torch._inductor.config.triton.cudagraph_trees_generation_cloning user_visible。开启该 opt-in 行为后Trees 会在开始新一代之前把存活的用户可见输出存储从 CUDA Graph 内存池克隆出去这不适用于梯度、保存的激活或其他内部张量。4.10 两种方案对比Footguns独立 CudaGraphCUDAGraph Trees内存可能增加每次图编译新尺寸等时若同时运行非 cudagraph 内存时录制时机任何一次新的图调用程序中任何新的唯一路径出现时重新录制Footguns一次图的调用会覆盖上一次调用无法在模型的不同运行之间持久化内存——一次训练循环或一次推理运行五、总结一条完整的 torch.compile 性能调优路径结合上述官方文档与仓库源码一条可落地的调优路径为宏观对比用 benchmarks/dynamo 下的基准脚本--performance --amp --inference/--training --backend inductor复现看板命令先拿到 eager 与编译后的加速比基线有 PR 性能顾虑时借助 nightly 看板工作流做回归对比。整体剖析用torch.profiler导出 chrome trace检查图断裂嵌套 Torch-Compiled Region / 分裂的 CompiledFunction、CPU 启动开销GPU 内核间大段空隙与 GPU 利用率小 batch 场景优先尝试 CUDA Graphs。内核级拆解设置TORCHINDUCTOR_UNIQUE_KERNEL_NAMES1 TORCHINDUCTOR_BENCHMARK_KERNEL1运行基准脚本定位编译模块路径并执行python fwd.py -p按内核类别与执行次数确定优化优先级对最贵的内核用TORCHINDUCTOR_MAX_AUTOTUNE1 python k.py单独微基准验证。运行时优化在modereduce-overhead下理解并善用 CUDAGraph Trees——通过cudagraph_support_input_mutation支持参数/缓冲区突变、用cudagraph_skip_dynamic_graphs规避频繁重录、用graph_partition自动切分 CPU 算子、必要时用cudagraph_mark_step_begin()修正迭代边界启发式。通过这一路径你可以把模型跑得慢从模糊的直觉逐步收敛为可量化、可定位、可验证的具体内核与配置问题。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 PyTorch torch.compile 性能优化完全指南：Profiling、基准测试与 CUDAGraph Trees 实战
PyTorch torch.compile 性能优化完全指南Profiling、基准测试与 CUDAGraph Trees 实战【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读torch.compile是 PyTorch 2.x 中默认开启的即时编译器它将 Dynamo 图捕获、TorchInductor 代码生成与 CUDA Graph 优化融为一体。本文将基于 PyTorch 官方性能文档系统讲解如何用torch.profiler定位编译区域与图断裂、用 TorchInductor 的 GPU 剖析工具拆解每个 Triton 内核的开销、借助 nightly 性能看板对比 PR 前后吞吐并深入解析modereduce-overhead背后的 CUDAGraph Trees 记忆池复用机制与其限制。读完本文你将掌握一套从宏观基准对比到单内核微基准再到CUDA Graph 级优化的完整性能调优方法论。一、PyTorch 2.0 性能看板宏观基准对比的入口1.1 看板测量什么PyTorch 团队通过 HUD Performance Dashboard 每夜跟踪torch.compile的性能。官方文档描述其运行环境为 12 个 GCP A100 节点每个节点含一块 40GB A100 GPU 和一块 6 核 2.2GHz Intel Xeon CPU对应 CI 工作流文件为.github/workflows/inductor-perf-test-nightly.yml。看板针对三个基准套件测量推理inference与训练training性能全部使用 AMP 混合精度TorchBenchPyTorch 官方维护的端到端基准套件Huggingface主流 Transformer 模型集合TIMM图像分类模型集合仓库中的对应脚本为 benchmarks/dynamo/timm_models.py同时看板还覆盖 TorchInductor 的三种配置default默认配置with_cudagraphs默认配置叠加 CUDA Graphdynamic默认配置叠加动态形状dynamic_shapes1.2 三个关键指标看板除通过率外报告三个核心性能指标指标对比基准含义Geometric mean speedup几何平均加速比PyTorch eager 模式越大越好衡量编译后整体提速Mean compilation time平均编译时间—衡量第一次调用torch.compile的编译开销Peak memory footprint compression ratio峰值内存压缩比PyTorch eager 模式越大越好衡量编译后内存占用改善看板上每个单独性能数字都可以点击进入该基准套件下所有测试的详细数值视图并可通过顶部的下拉列表切换表格与图表默认展示过去 7 天 TorchBench 的 AMP 训练性能趋势。1.3 在合并前检查 PR 对 TorchInductor 性能的影响如果担心某个 PR 会改变 TorchInductor 的生成代码质量可以在 CI 的 Actions 页面手动触发inductor-perf-test-nightly.yml工作流在Run workflow时选择 PR 分支提交。跑完后在看板 UI 中选中对应分支名和 commit ID 即可查看结果。注意这是一次昂贵的 CI 运行官方文档明确提示应谨慎使用该功能。1.4 本地复现看板命令完整 dashboard 运行所用的确切命令行可从任意近期 CI 运行日志中找到。官方文档给出的示例为python benchmarks/dynamo/huggingface.py --performance --cold-start-latency --inference --amp --backend inductor --disable-cudagraphs --device cuda只要本机有能跑 PyTorch 2.x 的 GPU 即可本地执行。python benchmarks/dynamo/huggingface.py -h会给出基准脚本的详细选项说明。类似的基准脚本还有 benchmarks/dynamo/timm_models.py、benchmarks/dynamo/torchbench.py 等统一位于 benchmarks/dynamo 目录。二、用 torch.profiler 理解 torch.compile 性能2.1 torch.profiler 的定位与示例torch.profiler适合在内核级粒度理解程序性能例如展示图断裂位置与资源利用率帮助判断下一步该往哪个方向排查。若需要更底层的内核级分析可结合 Nvidia Nsight Compute、AMD Omnitrace、Intel VTune Profiler 或下文将介绍的 Inductor 剖析工具。官方文档给出的 resnet18 剖析示例包含三个关键写法加入 warm-up 运行等待编译完成同时预热 CUDA caching allocator用torch.profiler.profile()上下文包住感兴趣的部分用prof.export_chrome_trace(trace.json)导出剖析产物import torch from torchvision.models import resnet18 device cuda # or cpu, xpu, etc. model resnet18().to(device) inputs [torch.randn((5, 3, 224, 224), devicedevice) for _ in range(10)] model_c torch.compile(model) def fwd_bwd(inp): out model_c(inp) out.sum().backward() # warm up fwd_bwd(inputs[0]) with torch.profiler.profile() as prof: for i in range(1, 4): fwd_bwd(inputs[i]) prof.step() prof.export_chrome_trace(trace.json)查看 chrome trace在 Chrome 浏览器打开chrome://tracing并加载 json 文件用w/s键缩放、a/d键左右平移按?可查看快捷键帮助。在 trace 中可以看到CompiledFunction与CompiledFunctionBackward事件对应 Dynamo 编译的区域上方为 CPU 事件下方为 GPU 事件2.2 CPU 与加速器事件之间的 Flow加速器上的每个内核都是由 CPU 代码启动的而除少数例外加速器内核是异步启动的因此 profiler 可以绘制 CPU 事件与加速器事件之间的流flow连接标明是哪个 CPU 事件启动了该内核。查看方式点击某个 GPU kernel再点击 ac2g或者用顶部 Flow events 下拉菜单打开全部 flow。2.3 绕开 CUDA Graph 剖析问题启用 CUDA Graph 后某些 CUDA 配置driver 版本低于 525.85.12或 CUDA 12会在剖析工具与 CUDA Graph 之间产生冲突。修复方法是在程序顶部加一个空的剖析上下文import torch torch.profiler._utils._init_for_cuda_graphs() # ... rest of program2.4 理解编译时间剖析首次调用要弄清编译为什么慢可以剖析torch.compile程序的第一次调用。需注意编译类负载与典型 PyTorch 负载差异很大trace 失真比普通剖析更明显且 trace 文件可能很大——超过 1GB 的 trace 用 chrome tracing 工具很难打开。官方示例同时展示了用torch.profiler.record_function标记 warmup compile 与 resnet18 compile 两个阶段的做法第一次调用必须在剖析期间发生才能捕获编译过程同时要先加一个小的 warm-up 编译以初始化需要惰性初始化的系统。import torch from torchvision.models import resnet18 # user can switch between cuda and xpu device cuda model resnet18().to(device) inputs [torch.randn((5, 3, 224, 224), devicedevice) for _ in range(10)] model_c torch.compile(model) def fwd_bwd(inp): out model_c(inp) out.sum().backward() def warmup_compile(): def fn(x): return x.sin().relu() x torch.rand((2, 2), devicedevice, requires_gradTrue) fn_c torch.compile(fn) out fn_c(x) out.sum().backward() with torch.profiler.profile() as prof: with torch.profiler.record_function(warmup compile): warmup_compile() with torch.profiler.record_function(resnet18 compile): fwd_bwd(inputs[0]) prof.export_chrome_trace(trace_compile.json)非图形化替代方案torch._dynamo.utils.compile_times()提供大致相同的信息。它不会展示各编译步骤何时发生但能展示每个步骤耗时多少且不受剖析开销影响。2.5 用 Torch-Compiled Region 与 CompiledFunction 定位图断裂尽管已有日志工具用于识别图断裂profiler 提供了一种快速的可视化方法关注两个事件Torch-Compiled RegionPyTorch 2.2 引入覆盖整个编译区域的 profiler 事件。图断裂几乎总是表现为嵌套的 Torch-Compiled Region 事件。PyTorch 2.5 起该事件还包含 frame ID帧的唯一标识与 frame compile ID该帧被编译的次数。若两个函数分别独立应用torch.compile通常看到两个**相邻非嵌套**的 Torch-Compiled Region若遇到图断裂或disable()/被跳过的区域则看到嵌套事件。CompiledFunctionPyTorch 2.0 引入当任何输入需要梯度时出现。每次图断裂都会打断一个 CompiledFunction 块、将其一分为二。CompiledFunction 只在涉及 Autograd即图中部分输入张量requires_gradTrue时出现trace 中它通常与反向传播的CompiledFunctionBackward成对出现若反向函数被调用两者之间会出现 fwd-bwd link。若你的图不需要梯度且不包含 Torch-Compiled Region 事件可以借助是否存在 Inductor 生成的 Triton 内核这一线索来判断torch.compile是否真正生效。官方文档还给出了一个使用torch._dynamo.graph_break()强制打断四个连续 Sequential 模块的合成示例ModelWithBreaks并在 3 次循环中分别对 10 个输入做fwd_bwd最终导出trace_break.json观察嵌套区域与多个 CompiledFunction 事件。2.6 内核事件三层结构与启动开销当一个算子被启动时通常看到三类事件CPU 侧事件、内核启动GPU kernel 场景事件、GPU 侧事件。Inductor 生成的 Triton 内核CPU 侧事件以triton_前缀出现信息量较少只有内核名与启动不如普通 aten 内核启动那样包含输入形状、类型等内核启动显示为cuLaunchKernel而非 aten 算子常见的cudaLaunchKernelGPU 侧事件是否详细取决于 Inductor 配置unique_kernel_names非 Inductor 生成的 Triton 内核CPU 侧事件可能不出现自动插入 profiler 事件的机制目前实现在 Inductor 层绕过 Inductor 的 Triton 内核除非手动注解否则不会出现内核启动同样为cuLaunchKernelGPU 侧事件按作者命名。Inductor 生成的 CPU 内核CPU 侧事件不会出现尚未添加剖析支持也没有内核启动与 GPU 侧事件。非 Triton 内核aten 内核或自定义算子有时 Inductor 会回退到原始算子实现此时会看到对 aten 算子的调用。启动开销问题GPU 利用率不佳的快速识别方法是看 GPU 上内核之间是否有大段空隙。这通常源于 CPU 开销——两次内核启动之间 CPU 耗时大于 GPU 处理内核的耗时在小 batch size 下更常见。当启动开销是主要矛盾时启用 CUDA Graphs 往往能显著改善。三、TorchInductor GPU 剖析拆解到单个 Triton 内核当模型没有达到预期速度时需要检查单个内核。通常 GPU 时间占比最大的内核最值得关注之后可以单独运行该内核并检查其性能。PyTorch 为此提供了一整套工具核心步骤是先整体拆解 GPU 时间再对单个内核做微基准。3.1 相关环境变量环境变量作用默认值TORCHINDUCTOR_UNIQUE_KERNEL_NAMES为 Triton 内核生成更有意义的名称如triton_poi_fused_cat_155poi表示 pointwise后跟原始 ATen 算子名便于在 trace 中识别内核类别默认关闭是为了提高编译缓存命中率0TORCHINDUCTOR_BENCHMARK_KERNEL让 Inductor 的 codegen harness 为每个 Triton 内核生成独立的基准代码0TORCHINDUCTOR_MAX_AUTOTUNE让 Inductor 的自动调优器尝试更多triton.Configs并挑选性能最佳者会增加编译时间以换取性能提升0这些配置在仓库中的落点包括 torch/_inductor/config.py 中的max_autotune、benchmark_kernel与triton.unique_kernel_names后者通过TORCHINDUCTOR_UNIQUE_KERNEL_NAMES环境变量控制以及 torch/_inductor/codegen/triton.py 中按config.benchmark_kernel拼接imports_for_benchmark_kernel()的实现。3.2 把模型 GPU 时间拆解到单个内核以 mixnet_l 为例步骤 1运行基准脚本TORCHINDUCTOR_UNIQUE_KERNEL_NAMES1 TORCHINDUCTOR_BENCHMARK_KERNEL1 python -u benchmarks/dynamo/timm_models.py --backend inductor --amp --performance --dashboard --only mixnet_l --disable-cudagraphs --training官方文档特别提示该工具依赖内核名判断其类别因此开启TORCHINDUCTOR_UNIQUE_KERNEL_NAMES至关重要。步骤 2在输出日志中查找编译模块路径**Compiled module path: /tmp/torchinductor_shunting/qz/cqz7hvhood7y3psp7fy6msjxsxyli7qiwiybizdwtjw6ffyq5wwd.py**每个编译模块一行。若没有额外图断裂日志中应看到两行——一个对应前向图一个对应反向图。步骤 3对单个编译模块剖析把前向图对应的模块重命名为fwd.py直接运行python fwd.py -p输出中有三处值得关注Chrome trace 文件日志中查找Chrome trace for the profile is written to /tmp/compiled_module_profile.json用chrome://tracing加载即可交互式查看。GPU 忙碌时间占比日志行形如Percent of time when GPU is busy: 102.88%。可能看到大于 100% 的值原因是 PyTorch 在开启剖析时用内核执行时间、关闭剖析时用 wall time剖析会对内核执行时间造成轻微失真但总体影响不大。例如densenet121小 batch 下前向图GPU 忙碌占比仅 32.69%说明模型有大量 CPU 开销——这与启用 cudagraphs 能大幅改善 densenet121 性能的事实一致。按内核类别拆解 GPU 时间mixnet_l 示例中pointwise 内核占 28.58%、reduction 内核占 13.85%、persistent reduction 内核占 3.89%其余为 mm/conv 的 cutlass/cudnn 内核占 56.57%。该信息位于每个内核类别报告的最后一行汇总中。还可以放大到某一类内核如 reduction看到按执行时间排序的明细表及每个内核的执行次数——这有助于判断优化优先级仅占 0.1% 的内核即使优化到极致收益也有限占 2% 的内核提速 2 倍能带来 1% 的整体收益才值得投入。3.3 单独基准某个 Triton 内核假设想仔细研究最贵的 reduction 内核triton_red_fused__native_batch_norm_legit_functional_16占前向图整体 wall time 的 2.19%。在fwd.py中查找该内核名可找到注释行# kernel path: /tmp/torchinductor_shunting/jk/cjk2vm3446xrk7rth7hr6pun7xxo3dnzubwcn6ydrpifal4eykrz.py将其重命名为k.py——它是一个独立的 Python 模块包含内核代码与其基准。直接运行python k.py会报告该内核的执行时间与带宽。还可以验证 max-autotune 是否对它有帮助TORCHINDUCTOR_MAX_AUTOTUNE1 python /tmp/k.py也可以临时添加更多 reduction 启发式规则后再次运行观察对内核的影响。四、CUDAGraph Treesreduce-overhead 模式的核心机制4.1 背景CUDA Graph 与 PyTorch 集成CUDA GraphCUDA 10 引入允许把一系列 CUDA 内核定义并封装为单个单元操作图通过一次 CPU 操作启动多个 GPU 操作从而降低启动开销。它对高 CPU 开销或小计算量的模型提速尤其明显但也有严格限制不支持任意控制流不过torch.cond()表达的控制流可以捕获进 CUDA Graph触发 host 到 device 同步的内核如.item()会报错所有输入参数固定为录制时的值CUDA 内存地址固定但该地址上的内存值可以变化不允许必需的 CPU 算子或 CPU 副作用PyTorch 的 torch.cuda.CUDAGraph 封装处理了与 caching allocator 交互的几个棘手问题CachingAllocator 为所有新分配使用独立内存池录制期间内存的分配、记账与释放与 eager 运行完全一致重放时只调用内核allocator 不发生变化录制之后 allocator 并不知道用户程序中哪些内存正被使用。eager 分配与 cudagraph 分配使用独立内存池若两者都有大量内存分配可能增加程序总内存。Make Graphed Callablestorch.cuda.make_graphed_callables是让一系列 callable 共享单个内存池的抽象利用录制时 caching allocator 精确记账内存的事实安全地共享内存每次调用时输出保持为存活内存防止一个 callable 覆盖另一个的存活内存它只能按单一顺序调用第一次运行的内存地址会被固化。旧版 TorchDynamo CUDA Graph 集成的问题cudagraph_treesFalse时不同图捕获之间不复用内存可能导致显著的内存回退。即使模型没有图断裂前向与反向也是两次独立的图捕获内存池不共享——前向中保存的激活内存无法在反向中被回收。4.2 CUDAGraph Trees 的工作原理与 Graphed Callables 类似CUDAGraph Trees 在所有图捕获间使用单一内存池不同的是它不需要单一调用序列而是为 CUDA Graph 捕获构建独立的树。官方文档给出的示意示例torch.compile(modereduce-overhead) def foo(x): # GRAPH 1 y x * x * x # graph break triggered here if y.sum() 0: # GRAPH 2 z y ** y else: # GRAPH 3 z (y.abs() ** y.abs()) torch._dynamo.graph_break() # GRAPH 4 return z * torch.rand_like(z) # the first run warms up each graph, which does things like CuBlas or Triton benchmarking foo(torch.arange(0, 10, devicecuda)) # The second run does a CUDA Graph recording, and replays it foo(torch.arange(0, 10, devicecuda)) # Finally we hit the optimized, CUDA Graph replay path foo(torch.arange(0, 10, devicecuda))该函数有两条路径1 → 2 → 4 或 1 → 3 → 4。系统通过构建一条 CUDA Graph 录制磁带如 1 → 2 → 4在多次录制间共享全部内存并施加不变量保证内存始终位于录制时的位置、用户程序中不存在可能被覆盖的存活张量CUDA Graph 的既有约束依然适用必须以相同参数静态尺寸、地址等调用相同内核录制与重放之间必须观察到相同的内存模式某图的一个输出张量在录制期间于另一个图之后死亡重放时也必须如此CUDA 池中的存活内存在两次录制间强制形成依赖这些录制只能以单一顺序调用1 → 2 → 4所有内存共享单一内存池因此相对 eager没有额外内存开销。新路径Graph 3出现时Graph 1 被重放后遇到尚未录制的 Graph 3。由于重放时私有内存池不更新y不会反映在 allocator 中若不处理会被覆盖。为了在重放其他图之后复用同一内存池系统把内存池checkpoint 回 Graph 1 结束时的状态使存活张量反映到 caching allocator 中即可安全运行新图。此时先命中已录制的CUDAGraph.replay()路径Graph 1再遇到 Graph 3——同样需要先 warm up 一次再录制warm-up 运行时内存地址未固定Graph 4 也会回退到 Inductor 的非 cudagraph 调用。第二次遇到 Graph 3 时即可录制随后因输入内存地址变化再次录制 Graph 4由此形成一棵 CUDA Graph 录制树1 / \ 2 3 \ \ 4 44.3 输入突变Input Mutation支持输入突变函数指对输入张量进行就地写入的函数例如def foo(x, y): # mutates input x x.add_(1) return x y输入突变对 CUDAGraph Trees 是挑战CUDA Graph 要求静态内存地址Trees 可能为每个输入张量x分配一个静态地址x执行时先把x复制到x再重放已录制的图。对输入突变函数x被就地更新但由于x与x位于不同 CUDA 内存地址该更新不会反映到输入张量x上。深入分析后发现输入有三类来自 eager 的输入假定每次执行的张量地址都变化因 cudagraph 冻结内存地址必须在录制与执行前复制到静态地址张量参数与缓冲区假定并运行时检查每次执行地址相同无需复制内容CUDAGraph Trees 先前的输出cudagraph 输出地址固定若先运行 CUDAGraph1 再运行 CUDAGraph2从 1 流入 2 的输入地址固定像参数与缓冲区一样无需复制运行时若检查到不稳定则重新录制因此 CUDAGraph Trees 支持对参数、缓冲区及先前输出张量的突变对 eager 输入的突变Trees 会不带 CUDA Graph 运行该函数并输出skipping due to mutated inputs日志。官方示例演示了启用输入突变支持的方式import torch torch.compile(modereduce-overhead) def foo(x): return x 1 torch.compile(modereduce-overhead) def mut(x): return x.add_(2) # Enable input mutation support torch._inductor.config.triton.cudagraph_support_input_mutation True for i in range(3): torch.compiler.cudagraph_mark_step_begin() inp torch.rand([4], devicecuda) # CUDAGraph is applied since foo does not mutate inp tmp foo(inp) # Although mut mutates tmp, which is an output of a CUDAGraph # managed function. So CUDAGraph is still applied. mut(tmp) torch.compiler.cudagraph_mark_step_begin() inp torch.rand([4], devicecuda) tmp foo(inp) # While tmp is a CUDAGraph Tree managed functions output, tmp.clone() # is not. So CUDAGraph is not applied to mut and there is a log # skipping cudagraphs due to mutated inputs mut(tmp.clone())仓库中该配置项位于 torch/_inductor/config.py 的triton.cudagraph_support_input_mutation默认值not is_fbcode()。若函数要突变 eager 输入官方建议重写函数以避免输入突变。4.4 动态形状支持动态形状指输入张量在不同调用间形状不同。由于 CUDA Graph 要求固定张量地址CUDAGraph Trees 会为输入张量的每种唯一形状重新录制CUDA Graph导致单个 inductor 图对应多个 CUDA Graph。形状有限如推理中的 batch size时重录是划算的若形状频繁变化甚至每次调用都变重录可能不划算。官方文档还提到直到 CUDA 12.4 与 Driver 550 之前NVIDIA 每个内核启动在 CUDA Graph 中占用 64KB 设备内存大量重录时该内存成本可观。建议对形状频繁变化的函数把输入张量 padding 到少数固定形状以继续享受 CUDA Graph 收益同时可设置torch._inductor.config.triton.cudagraph_skip_dynamic_graphsTruetorch/_inductor/config.py跳过动态形状输入的函数只对静态形状函数使用 cudagraph。4.5 NCCL 支持CUDAGraph Trees 支持包含 NCCL 算子的函数Trees 按设备录制 CUDA Graph而 NCCL 支持允许跨设备通信。示例torch.compile(modereduce-overhead) def func(x): y x * x y torch.distributed.all_reduce(y, optorch.distributed.ReduceOp.SUM) x torch.nn.functional.silu(x) return x * y4.6 跳过 CUDA Graph 的原因由于 CUDA Graph 要求静态输入地址且不支持 CPU 算子Trees 会检查函数是否满足要求并在必要时跳过。常见原因输入突变跳过就地突变 eager 输入的函数突变参数/缓冲区/先前输出仍支持CPU 算子跳过含 CPU 算子的函数建议拆成多个函数并对仅含 GPU 算子的部分应用多设备算子含多设备算子的函数被跳过目前按设备应用跨设备通信请用 NCCL 等受支持库自由 unbacked 符号通常出现在动态形状场景Trees 目前为每种唯一输入形状录制一张图CUDAGraph 不安全的自定义算子可能包含 cudagraph 不安全算子导致跳过不兼容算子官方给出完整的在确定性算法开启时同样不兼容的不兼容算子清单aten._fused_moving_avg_obs_fq_helper.default aten._fused_moving_avg_obs_fq_helper_functional.default aten.multinomial.default fbgemm.dense_to_jagged.default fbgemm.jagged_to_padded_dense.default run_and_save_rng_state run_with_rng_state aten._local_scalar_dense aten._assert_scalar4.7 CUDAGraph 不安全自定义算子自定义算子默认被视为对 CUDA Graph 安全但某些算子可能包含 CPU 算子等不支持操作。由于编译器把自定义算子当作黑盒用户必须显式打上torch._C.Tag.cudagraph_unsafe标签来标记其不安全torch.library.custom_op( mylib::modify, mutates_args(), tags(torch._C.Tag.cudagraph_unsafe,), ) def modify(pic: torch.Tensor) - torch.Tensor: pic1 pic 1 pic1_cpu (pic1.cpu() 1) * 2 return pic1_cpu.cuda() pic modify.register_fake def _(pic): return torch.empty_like(pic)含此类算子的函数会被 CUDA Graph 跳过除非启用 CUDAGraph partition。4.8 CUDAGraph Partition自动切分不兼容算子CUDAGraph partition 是一种编译器解决方案自动把不支持的算子切分出去、重排算子以减少分区数量并对每个分区单独应用 CUDA Graph。启用方式为torch._inductor.config.graph_partitionTrue定义于 torch/_inductor/config.py。考虑如下例子x、y是 GPU 输入但y_cpu是 CPU 张量。无 partition 时该函数因 CPU 算子被整体跳过启用 partition 后 CPU 算子被切分出去剩余 GPU 算子被 cudagraph 化得到两个独立的 CUDA Graphdef f(x, y): x1 x 1 y1 y 1 y_cpu y1.cpu() 1 z x y return x1 y1 z y_cpu.cuda()目前 partition 支持切分以下类型的算子非 GPU 算子如 CPU 张量上的计算设备复制算子设备间数据传输如例子中的y1.cpu()控制流算子尚未被 CUDA Graph 支持会被切分CUDAGraph 不安全自定义算子带torch._C.Tag.cudagraph_unsafe标签的算子Unbacked Symints见动态形状支持一节4.9 限制与跨迭代存活的张量由于 CUDA Graph 固定内存地址它对上一次调用遗留的存活张量处理不佳。官方示例import torch torch.compile(modereduce-overhead) def my_model(x): y torch.matmul(x, x) return y x torch.randn(10, 10, devicecuda) y1 my_model(x) y2 my_model(x) print(y1) # RuntimeError: Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run.在旧版独立 CUDA Graph 实现中第二次调用的输出会覆盖第一次的输出。CUDAGraph Trees 既不想在迭代间引入意外依赖导致错过热路径也不想过早释放上一次调用的内存。其启发式规则是推理时每次调用torch.compile都开启新迭代训练时只要没有未调用的 pending backward 也如此。若启发式判断错误可用torch.compiler.cudagraph_mark_step_begin()手动标记新迭代开始其实现位于 torch/compiler/init.py 并委托给torch._inductor.cudagraph_trees.mark_step_begin或在开始下一次运行前在torch.compile之外克隆上一次迭代的张量。若用户可见输出必须跨迭代存活可设置torch._inductor.config.triton.cudagraph_trees_generation_cloning user_visible。开启该 opt-in 行为后Trees 会在开始新一代之前把存活的用户可见输出存储从 CUDA Graph 内存池克隆出去这不适用于梯度、保存的激活或其他内部张量。4.10 两种方案对比Footguns独立 CudaGraphCUDAGraph Trees内存可能增加每次图编译新尺寸等时若同时运行非 cudagraph 内存时录制时机任何一次新的图调用程序中任何新的唯一路径出现时重新录制Footguns一次图的调用会覆盖上一次调用无法在模型的不同运行之间持久化内存——一次训练循环或一次推理运行五、总结一条完整的 torch.compile 性能调优路径结合上述官方文档与仓库源码一条可落地的调优路径为宏观对比用 benchmarks/dynamo 下的基准脚本--performance --amp --inference/--training --backend inductor复现看板命令先拿到 eager 与编译后的加速比基线有 PR 性能顾虑时借助 nightly 看板工作流做回归对比。整体剖析用torch.profiler导出 chrome trace检查图断裂嵌套 Torch-Compiled Region / 分裂的 CompiledFunction、CPU 启动开销GPU 内核间大段空隙与 GPU 利用率小 batch 场景优先尝试 CUDA Graphs。内核级拆解设置TORCHINDUCTOR_UNIQUE_KERNEL_NAMES1 TORCHINDUCTOR_BENCHMARK_KERNEL1运行基准脚本定位编译模块路径并执行python fwd.py -p按内核类别与执行次数确定优化优先级对最贵的内核用TORCHINDUCTOR_MAX_AUTOTUNE1 python k.py单独微基准验证。运行时优化在modereduce-overhead下理解并善用 CUDAGraph Trees——通过cudagraph_support_input_mutation支持参数/缓冲区突变、用cudagraph_skip_dynamic_graphs规避频繁重录、用graph_partition自动切分 CPU 算子、必要时用cudagraph_mark_step_begin()修正迭代边界启发式。通过这一路径你可以把模型跑得慢从模糊的直觉逐步收敛为可量化、可定位、可验证的具体内核与配置问题。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考