[Bug已解决] Profiler trace-only 导出与默认导出在 CUDA 下不一致DISABLED test绕过方案解决方案一、现象长什么样你用 PyTorch 的Profiler性能分析器抓 CUDA 训练的耗时并尝试把结果导出成 traceChrome trace / tensorboard结果遇到一个被官方DISABLED的测试DISABLED test_trace_only_export_matches_default_use_cuda_True (__main__.TestProfiler)即 pytorch/pytorch#186909。含义Profiler 的「trace-only 导出」模式与「默认完整导出」模式在 CUDA 下产出的结果不一致测试被禁用。 对你来说实际表现可能是用profiler.export_chrome_trace(x.json)和profiler.key_averages()两种口径看 CUDA 算子耗时数字对不上或启用profile_memory/with_stack/trace_onlyTrue这类选项后CUDA 事件的时间戳 / 算子归属和默认模式不同导致你基于 trace 做的性能判断失真。 本文聚焦PyTorch Profiler 的 trace 导出机制、为什么 CUDA 下两种导出会不一致、作为用户怎么拿到可信的 CUDA profile 数据。二、背景PyTorch Profiler 的两种导出PyTorch Profilertorch.profiler.profile抓取的记录可以以多种方式消费默认导出完整分析prof.key_averages()、prof.table()、export_chrome_trace()综合产出包含 CPU 算子、CUDA 内核、内存、栈等trace-only 导出只导出「trace 事件流」Chrome trace 格式侧重于时间线不附加 key_averages 的聚合。 两者的数据来源应是同一份底层记录但在 CUDA 下CUDA 事件的关联kernel 与 CPU 算子的对应、时间对齐需要特殊后处理。测试test_trace_only_export_matches_default_use_cuda_True想验证「trace-only 导出的 CUDA 事件」与「默认导出的 CUDA 事件」完全一致但当前不一致 → 禁用。 不一致的可能表现同一个 CUDA kernel 在 trace-only 里归属到不同 CPU 算子、或时间戳偏移、或被重复/漏记。三、为什么 CUDA 下两种导出会不一致CUDA 事件是异步的CPU 发起 kernel Launch 后立刻返回kernel 真正在 GPU 上跑的时间由 CUDA 事件event记录。Profiler 要把「CPU 侧的 launch」和「GPU 侧的执行」通过流stream和时间戳关联起来这叫GPU-CPU 事件关联correlation。 问题trace-only 导出可能没走完整的关系联逻辑它只想要时间线跳过了一些聚合/对齐步骤而默认导出走了完整关联。于是同一个 CUDA kernel 在两种路径下关联的 CPU 算子不同correlation 缺失 → 归到错误的父算子时间区间计算不同是否包含 launch 开销、是否对齐到 streamCUDA 图CUDA Graph捕获的 kernel 在 trace-only 下可能被不同处理。 本质trace-only 是一条「轻量导出捷径」它在 CUDA 关联这一步和完整路径有实现差异属于 Profiler 内部 bug。四、最小可运行演示带守卫下面演示 Profiler 的基本 CUDA 用法以及两种导出口径结构示意import torch import torch.nn as nn from torch.profiler import profile, ProfilerActivity def demo_profiler(): if not torch.cuda.is_available(): print([skip] 无 GPU仅说明 Profiler 用法) return model nn.Linear(64, 64).cuda() x torch.randn(8, 64, devicecuda) with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapesTrue, ) as prof: for _ in range(3): model(x).sum().backward() # 默认导出聚合表含 CUDA 内核 print(prof.key_averages().table(sort_bycuda_time_total, row_limit5)) # trace 导出 prof.export_chrome_trace(/tmp/trace.json) print(已导出 trace 到 /tmp/trace.json) if __name__ __main__: demo_profiler()要点key_averages()与export_chrome_trace在 CUDA 下理论应一致但本 bug 说它们有偏差。下文给「拿到可信数据」的做法。五、解决方案一用默认导出key_averages / table做权威分析既然 trace-only 与默认在 CUDA 下不一致而默认导出完整关联是更可信的口径做性能分析时以它为准from torch.profiler import profile, ProfilerActivity with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapesTrue, with_stackTrue, ) as prof: run_training_step() # 权威用 key_averages 看 CUDA 时间 print(prof.key_averages().table(sort_bycuda_time_total, row_limit10)) # 若要 trace 文件做可视化也用默认路径导出而非 trace_only 捷径 prof.export_chrome_trace(full_trace.json)经验法则性能数字看key_averages()可视化看export_chrome_trace二者都来自完整路径互相对得上别用trace_onlyTrue这条捷径拿 CUDA 数字。六、解决方案二避免 trace_only 捷径选项如果你之前用了「只导出 trace、不做聚合」的捷径开关如某些旧 API 的use_traceTrue或自定义导出改走完整路径# 不推荐触发不一致路径 # prof.export_chrome_trace(x.json, use_trace_onlyTrue) # 假设存在 # 推荐完整导出 prof.export_chrome_trace(x.json)如需把 trace 给 Chromechrome://tracing或 Perfetto 看完整路径的 trace 同样可用只是文件更大一点多了关联信息但时间线正确。七、解决方案三CUDA 事件关联失效时的兜底手动计时若你需要「某个算子到底在 GPU 上跑了多久」的精确数字且对 Profiler 的关联不放心用CUDA 事件Event手动计时兜底import torch def cuda_time_ms(fn, repeat20): if not torch.cuda.is_available(): return None # 预热 for _ in range(5): fn() torch.cuda.synchronize() start torch.cuda.Event(enable_timingTrue) end torch.cuda.Event(enable_timingTrue) start.record() for _ in range(repeat): fn() end.record() torch.cuda.synchronize() return start.elapsed_time(end) / repeat # 单次平均 ms def kernel(): x torch.randn(1024, 1024, devicecuda) return (x x).sum() print(手动 CUDA 计时 ms:, cuda_time_ms(kernel))torch.cuda.Event直接量 GPU 执行时间不依赖 Profiler 的 CPU-GPU 关联逻辑是最可靠的 CUDA 耗时来源可用来交叉验证 Profiler 数字。八、解决方案四升级并关注该测试启用#186909 是内部测试被禁用官方会修 Profiler 的关联逻辑。升级pip install --upgrade torch --index-url https://download.pytorch.org/whl/cu124判断修复该 DISABLED 测试被启用且通过即代表 trace-only 与默认导出在 CUDA 下一致届时两种口径可混用。修复前以默认导出 手动 Event 计时为准。九、排查清单发现 trace 与key_averages的 CUDA 数字对不上 → 确认 #186909trace-only 与默认在 CUDA 不一致。性能分析以key_averages()/table()为权威别用 trace_only 捷径拿数字。可视化用export_chrome_trace完整路径时间线正确。CUDA 精确耗时用torch.cuda.Event手动计时交叉验证。升级关注该 DISABLED 测试启用修复后两口径一致。通用Profiler 开activities[CPU, CUDA]否则看不到 GPU 内核。十、小结DISABLED test_trace_only_export_matches_default_use_cuda_True#186909的本质是PyTorch Profiler 的「trace-only 导出」走的是轻量捷径在 CUDA 的「CPU-GPU 事件关联」这一步和「默认完整导出」实现不同导致同一 CUDA kernel 在两种口径下归属/时间不一致测试被禁用。这是 Profiler 内部 bug不是你抓数据的方式错。 应对以默认导出为权威性能数字看key_averages()/table()可视化用export_chrome_trace完整路径二者一致别用 trace_only 捷径它绕过了完整关联CUDA 数字不可信手动交叉验证用torch.cuda.Event直接量 GPU 执行时间最可靠等升级该 DISABLED 测试启用即代表修复两口径可混用。 记住CUDA 事件异步Profiler 要把「CPU launch」关联到「GPU 执行」需专门后处理trace-only 捷径跳了这步所以 CUDA 数字会偏。要可靠就走完整路径或手动 Event 计时。