PyTorch torch.compile 性能优化完全指南:Profiling、基准测试与 CUDAGraph Trees 实战
2026/9/11 20:44:58 网站建设 项目流程

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 前后吞吐,并深入解析mode="reduce-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 混合精度:

  • TorchBench:PyTorch 官方维护的端到端基准套件
  • Huggingface:主流 Transformer 模型集合
  • TIMM:图像分类模型集合(仓库中的对应脚本为 benchmarks/dynamo/timm_models.py)

同时看板还覆盖 TorchInductor 的三种配置:

  • default:默认配置
  • with_cudagraphs:默认配置叠加 CUDA Graph
  • dynamic:默认配置叠加动态形状(dynamic_shapes)

1.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), device=device) 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 中可以看到:

  • CompiledFunctionCompiledFunctionBackward事件,对应 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 program

2.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), device=device) 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), device=device, requires_grad=True) 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 Region(PyTorch 2.2 引入):覆盖整个编译区域的 profiler 事件。图断裂几乎总是表现为嵌套的 "Torch-Compiled Region" 事件。PyTorch 2.5 起该事件还包含 frame ID(帧的唯一标识)与 frame compile ID(该帧被编译的次数)。若两个函数分别独立应用torch.compile,通常看到两个**相邻(非嵌套)**的 Torch-Compiled Region;若遇到图断裂(或disable()/被跳过的区域),则看到嵌套事件。
  • CompiledFunction(PyTorch 2.0 引入):当任何输入需要梯度时出现。每次图断裂都会打断一个 CompiledFunction 块、将其一分为二。CompiledFunction 只在涉及 Autograd(即图中部分输入张量requires_grad=True)时出现;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 内核

  1. CPU 侧事件以triton_前缀出现,信息量较少(只有内核名与启动),不如普通 aten 内核启动那样包含输入形状、类型等
  2. 内核启动显示为cuLaunchKernel而非 aten 算子常见的cudaLaunchKernel
  3. GPU 侧事件是否详细取决于 Inductor 配置unique_kernel_names

非 Inductor 生成的 Triton 内核:CPU 侧事件可能不出现(自动插入 profiler 事件的机制目前实现在 Inductor 层,绕过 Inductor 的 Triton 内核除非手动注解否则不会出现);内核启动同样为cuLaunchKernel;GPU 侧事件按作者命名。

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 中识别内核类别;默认关闭是为了提高编译缓存命中率0
TORCHINDUCTOR_BENCHMARK_KERNEL让 Inductor 的 codegen harness 为每个 Triton 内核生成独立的基准代码0
TORCHINDUCTOR_MAX_AUTOTUNE让 Inductor 的自动调优器尝试更多triton.Configs并挑选性能最佳者,会增加编译时间以换取性能提升0

这些配置在仓库中的落点包括 torch/_inductor/config.py 中的max_autotunebenchmark_kerneltriton.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_NAMES=1 TORCHINDUCTOR_BENCHMARK_KERNEL=1 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_AUTOTUNE=1 python /tmp/k.py

也可以临时添加更多 reduction 启发式规则后再次运行,观察对内核的影响。

四、CUDAGraph Trees:reduce-overhead 模式的核心机制

4.1 背景:CUDA Graph 与 PyTorch 集成

CUDA Graph(CUDA 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_trees=False时,不同图捕获之间不复用内存,可能导致显著的内存回退。即使模型没有图断裂,前向与反向也是两次独立的图捕获,内存池不共享——前向中保存的激活内存无法在反向中被回收。

4.2 CUDAGraph Trees 的工作原理

与 Graphed Callables 类似,CUDAGraph Trees 在所有图捕获间使用单一内存池;不同的是它不需要单一调用序列,而是为 CUDA Graph 捕获构建独立的"树"。官方文档给出的示意示例:

@torch.compile(mode="reduce-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, device="cuda")) # The second run does a CUDA Graph recording, and replays it foo(torch.arange(0, 10, device="cuda")) # Finally we hit the optimized, CUDA Graph replay path foo(torch.arange(0, 10, device="cuda"))

该函数有两条路径: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 4

4.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'被就地更新,但由于xx'位于不同 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(mode="reduce-overhead") def foo(x): return x + 1 @torch.compile(mode="reduce-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], device="cuda") # 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], device="cuda") tmp = foo(inp) # While `tmp` is a CUDAGraph Tree managed function's 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_graphs=True(torch/_inductor/config.py)跳过动态形状输入的函数,只对静态形状函数使用 cudagraph。

4.5 NCCL 支持

CUDAGraph Trees 支持包含 NCCL 算子的函数:Trees 按设备录制 CUDA Graph,而 NCCL 支持允许跨设备通信。示例:

@torch.compile(mode="reduce-overhead") def func(x): y = x * x y = torch.distributed.all_reduce(y, op=torch.distributed.ReduceOp.SUM) x = torch.nn.functional.silu(x) return x * y

4.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_scalar

4.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_partition=True(定义于 torch/_inductor/config.py)。

考虑如下例子:xy是 GPU 输入但y_cpu是 CPU 张量。无 partition 时该函数因 CPU 算子被整体跳过;启用 partition 后 CPU 算子被切分出去,剩余 GPU 算子被 cudagraph 化,得到两个独立的 CUDA Graph:

def 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(mode="reduce-overhead") def my_model(x): y = torch.matmul(x, x) return y x = torch.randn(10, 10, device="cuda") 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 性能调优路径

结合上述官方文档与仓库源码,一条可落地的调优路径为:

  1. 宏观对比:用 benchmarks/dynamo 下的基准脚本(--performance --amp --inference/--training --backend inductor)复现看板命令,先拿到 eager 与编译后的加速比基线;有 PR 性能顾虑时借助 nightly 看板工作流做回归对比。
  2. 整体剖析:用torch.profiler导出 chrome trace,检查图断裂(嵌套 Torch-Compiled Region / 分裂的 CompiledFunction)、CPU 启动开销(GPU 内核间大段空隙)与 GPU 利用率;小 batch 场景优先尝试 CUDA Graphs。
  3. 内核级拆解:设置TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1 TORCHINDUCTOR_BENCHMARK_KERNEL=1运行基准脚本,定位编译模块路径并执行python fwd.py -p,按内核类别与执行次数确定优化优先级;对最贵的内核用TORCHINDUCTOR_MAX_AUTOTUNE=1 python k.py单独微基准验证。
  4. 运行时优化:在mode="reduce-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),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询