PyTorch Profiler 使用指南从性能剖析到执行追踪的完整实践【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读torch.profiler是 PyTorch 内置的性能剖析工具用于在训练和推理过程中采集性能指标帮助开发者定位开销最大的算子、查看输入形状与调用栈、研究设备GPU/XPU 等上的内核活动并可视化执行轨迹。本文以仓库中的 profiler 官方文档 为骨架结合 torch/profiler 模块的源码与 测试用例系统讲解 profiler 的核心 API、调度机制、Trace 导出、内存与执行追踪能力以及 ITT 仪器化 API 的用法让读者能够直接在自己的训练/推理代码中落地一套完整的性能分析方案。模块定位与整体架构PyTorch Profiler 是 PyTorch 中用于采集训练与推理性能指标的工具其核心 API 以上下文管理器形式提供可帮助开发者回答三类问题哪些模型算子最耗时、算子的输入形状与调用栈是什么、设备上的内核活动及执行轨迹如何分布。从 torch/profiler/init.py 的模块文档可以看到profiler 建立在底层torch._C._autograd的 Kineto 支持之上且明确说明早期位于torch.autograd模块中的旧版 API 被视为遗留实现将来会被弃用——因此新代码应统一使用torch.profiler。模块的顶层导出见init.py包括profile性能剖析上下文管理器核心入口schedule、supported_activities、tensorboard_trace_handler调度、活动集合查询与 TensorBoard 导出辅助函数ProfilerAction、ProfilerActivity调度动作与可剖析活动的枚举kineto_available、DeviceType、record_function、ExecutionTraceObserver底层可用性标志、设备类型、手工打点与执行追踪观察者。从源码结构看profiler 的实现由以下几部分构成文件职责torch/profiler/profiler.pyprofile、schedule、ProfilerAction、tensorboard_trace_handler、ExecutionTraceObserver等核心实现torch/profiler/init.py模块入口与公共导出torch/profiler/itt.pyIntel Instrumentation and Tracing TechnologyITTAPI 封装torch/profiler/python_tracer.pyPython 调用跟踪torch/profiler/_memory_profiler.py内存时间线分析torch/profiler/_chrome_trace_export.pyChrome Trace JSON 导出核心 API 总览本文档对应的 API 参考docs/source/profiler.md围绕以下对象展开torch.profiler.profile剖析上下文管理器profile的所有成员方法均可用torch.profiler.ProfilerAction调度动作枚举torch.profiler.ProfilerActivity可剖析活动枚举schedule、supported_activities、tensorboard_trace_handler三个模块级函数。ProfilerActivity可剖析的活动类型ProfilerActivity枚举定义了可采集的活动组从 torch/_C/_profiler.pyi 的类型声明可见其取值CPUCPU 上的算子执行Torch 算子事件CUDACUDA 设备上的内核与内存事件依赖 CUPTI 库XPUIntel XPU 设备活动MTIAMeta MTIA 设备活动HPUHabana HPU 设备活动PrivateUse1通过torch._C._get_privateuse1_backend_name注册的私有后端设备。在 profiler.py 中_KinetoProfile会根据activities中是否包含上述枚举自动推断use_devicecuda/xpu/mtia/hpu或私有后端名该设备将作为 trace 中记录活动的主设备。值得注意的新特性activities参数不仅接受裸的ProfilerActivity枚举还接受dict[ProfilerActivity, list[str]]形式的活动子集过滤例如{ProfilerActivity.CUDA: [GPU_MEMCPY, CUDA_RUNTIME]}只采集指定的 CUDA 活动类型空列表{ProfilerActivity.CUDA: []}则表示该组什么都不采集。解析逻辑见 profiler.py同一活动组不能重复出现否则会抛出ValueError。supported_activities()查询当前环境支持的活动supported_activities()返回当前环境中可用的剖析活动集合实现为对底层torch.autograd._supported_activities()的透传见 profiler.py。其 docstring 特别强调了一个兼容性细节profiler 使用 CUPTI 库追踪设备上的 CUDA 内核若 CUDA 已启用但 CUPTI 不可用传入ProfilerActivity.CUDA会退回到遗留的 CUDA 剖析代码路径此时 CUDA 时间仍会出现在 profiler 表格输出中但不会出现在 JSON trace 中。在测试用例 test_profiler.py 中supported_activities()被用来动态选择非 CPU 设备并构造profile(activitiessupported_activities())这是编写跨设备兼容测试的推荐写法。profile 上下文管理器最常用的性能分析入口profile是文档 API Reference 的核心对象。它是一个上下文管理器基本用法如下与文档及源码 docstring 中的示例一致import torch with torch.profiler.profile( activities[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ] ) as p: # 需要分析的代码 code_to_profile() print(p.key_averages().table(sort_byself_cuda_time_total, row_limit-1))profile 的完整参数说明结合 profile 的 docstring 与构造签名profiler.py各参数含义如下参数默认值作用activitiesNone自动选择 CPU 可用设备采集的活动集合支持枚举或枚举→活动名列表字典scheduleNone默认持续采集接收 step 整数、返回ProfilerAction的可调用对象控制各 step 的动作on_trace_readyNone每个剖析周期结束时schedule返回RECORD_AND_SAVE时被调用的回调接收profile实例典型用途是导出 trace 或打印汇总record_shapesFalse记录算子输入形状profile_memoryFalse追踪张量内存的分配/释放with_stackFalse记录算子的源文件与行号with_flopsFalse用公式估算特定算子矩阵乘法与 2D 卷积的 FLOPswith_modulesFalse记录算子对应的模块层级已弃用仅对 TorchScript 模型生效eager 模式下无效建议改用with_stackTrueexperimental_configNone透传给 Kineto 的实验选项集合不保证向后兼容execution_trace_observerNonePyTorch Execution Trace 观察者对象其start()/stop()与 profiler 的时间窗口保持一致acc_eventsFalse跨多个剖析周期累积FunctionEventscustom_trace_id_callbackNone用户提供的 trace ID 生成器每个周期调用一次可通过get_trace_id()获取post_processing_timeout_sNone剖析后处理超时秒事件解析超时后返回部分结果适合超大规模 trace使用注意事项源码 docstring 明确提示了两类开销profiler.py开启 shape 与 stack 记录会带来额外开销record_shapesTrue时profiler 会临时持有张量引用可能阻碍依赖引用计数的优化并引入额外张量拷贝。因此在生产训练中建议仅在定位问题的短窗口内开启record_shapes/with_stack。另外profile构造时会校验activities非空profiler.pyactivitiesNone时会自动回退到supported_activities()如果用户传入的调度在第一步就返回ProfilerAction.DEVICE_STOPPED会直接抛出ValueError——该动作只能由 profiler 内部在设备采集提前停止如 CUPTI 缓冲溢出时设置用户调度不得返回它。schedule 调度器长时训练的分段剖析torch.profiler.schedule返回一个可调用对象用作profile的schedule参数用于在长时间训练任务中分阶段开启/关闭采集并在不同迭代产生多条 trace。参数与行为torch.profiler.schedule( wait1, warmup1, active2, repeat1, )waitint每个周期内先等待不采集的 step 数warmupint预热 step 数。预热阶段不记录事件用于让 CUDA 上下文初始化、算子 autotuning 等稳定下来避免污染测量结果activeint实际记录事件的 step 数最后一个 active step 会触发RECORD_AND_SAVE并调用on_trace_readyrepeatint周期重复次数0表示一直重复直到剖析结束skip_firstint最开始跳过的 step 数这些 step 返回NONEskip_first_waitint非零时第一个周期跳过wait阶段。调度状态机与参数校验见 schedule 实现wait、warmup、repeat、skip_first必须 ≥ 0active必须 0否则抛出AssertionErrorwarmup0时会发出警告Profiler wont be using warmup, this can skew profiler results——预热是保证测量准确性的重要步骤。一个完整的训练循环示例文档与源码 docstring 给出的标准示例profiler.pydef trace_handler(prof): print( prof.key_averages().table(sort_byself_cuda_time_total, row_limit-1) ) # prof.export_chrome_trace(/tmp/test_trace_ str(prof.step_num) .json) with torch.profiler.profile( activities[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], scheduletorch.profiler.schedule(wait1, warmup1, active2, repeat1), on_trace_readytrace_handler, ) as p: for iter in range(N): code_iteration_to_profile(iter) p.step() # 通知 profiler 进入下一个迭代对于wait1, warmup1, active2, repeat1profiler 跳过第一个 step第二个 step 预热第三、四个 step 记录事件之后 trace 就绪on_trace_ready被调用周期随后重复。ProfilerAction 枚举ProfilerAction定义于 profiler.py枚举值含义NONE不执行任何操作跳过/等待阶段WARMUP预热初始化而不记录事件RECORD记录事件RECORD_AND_SAVE记录事件并保存/导出触发on_trace_readyDEVICE_STOPPED设备采集提前停止如 CUPTI 缓冲溢出仅由 profiler 内部设置在内部profile通过一张前动作→后动作的状态转移表action_map驱动prepare_trace、start_trace、stop_trace、_trace_ready的调用。例如(NONE, RECORD)会执行prepare_trace与start_trace而(RECORD_AND_SAVE, RECORD)会依次执行stop_trace、_trace_ready、prepare_trace、start_trace从而实现周期的无缝衔接。非法转移如RECORD直接跳到NONE会发出警告并尽力恢复。从源码结构可以看出这套状态机保证了即使设备采集中途异常退出DEVICE_STOPPEDprofiler 也能在下个周期边界恢复不会拖垮训练进程。结果查看key_averages 表格与 Chrome Tracekey_averages()算子级汇总profile实例的key_averages方法profiler.py按算子名聚合事件返回torch.autograd.profiler_util.EventList并支持按输入形状group_by_input_shapeTrue、调用栈group_by_stack_n、重载名group_by_overload_name分组或包含 Python 函数include_python_functionsTrue。调用.table(sort_byself_cuda_time_total, row_limit-1)可打印按CUDA 自身耗时排序的完整表格。需要注意使用 shape/stack 分组前必须在创建 profiler 时设置record_shapes/with_stack。events()方法则返回未聚合的原始FunctionEvent列表当experimental_config.trace_onlyTrue时events()不可用需改用export_chrome_trace()。export_chrome_trace()导出 Chrome JSON 轨迹prof.export_chrome_trace(trace.json)导出的 JSON 文件可在 Chrome 的chrome://tracing或 Perfetto 中可视化。export_chrome_trace的实现细节profiler.py包括路径以.gz结尾时自动输出 gzip 压缩的 JSON传入use_python_exportTrue或cuda_graph_annotations/graph_lanes参数时切换到 Python 导出器若启用 Kineto则只会导出调度中最后一个周期的 tracecuda_graph_annotations可将 CUDA 图内核注解烘焙进 tracegraph_lanes控制被图化事件是否移动到展示通道none保持默认流布局all将其移动到注解命名的通道并记录original_stream。此外还可通过add_metadata(key, value)、add_metadata_json(key, value)与preset_metadata_json(key, value)向 trace 中注入用户自定义元数据后者用于 profiler 尚未启动时预设、启动后自动写入见 profiler.py。export_stacks()调用栈火焰图数据prof.export_stacks(stacks.txt, metricself_cpu_time_total)将栈跟踪保存到文件metric支持self_cpu_time_total或self_cuda_time_total见 profiler.py配合火焰图工具可快速定位热点调用路径。toggle_collection_dynamic()运行中动态开关采集在剖析过程中可按需开/关指定活动的采集profiler.pywith torch.profiler.profile( activities[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ] ) as p: code_to_profile_0() # 关闭 CUDA 活动采集 p.toggle_collection_dynamic(False, [torch.profiler.ProfilerActivity.CUDA]) code_to_profile_1() # 重新开启 CUDA 活动采集 p.toggle_collection_dynamic(True, [torch.profiler.ProfilerActivity.CUDA]) code_to_profile_2()tensorboard_trace_handler与 TensorBoard 集成tensorboard_trace_handler(dir_name, worker_nameNone, use_gzipFalse, use_python_exportFalse)生成一个可直接用作on_trace_ready的回调将 trace 文件输出到指定目录之后该目录可直接作为 TensorBoard 的 logdir 使用with torch.profiler.profile( activities[torch.profiler.ProfilerActivity.CPU], on_trace_readytorch.profiler.tensorboard_trace_handler(./log), ) as p: ...tensorboard --logdir ./log实现细节profiler.pyworker_name在分布式场景下应为每个 worker 唯一默认为{hostname}_{pid}文件名使用纳秒级时间戳{worker_name}.{time.time_ns()}.pt.trace.json避免同名冲突use_gzipTrue时输出追加.gz后缀的压缩文件use_python_exportTrue时走 Python 导出路径。仓库测试 test_profiler.py 中大量使用profile(activitiessupported_activities())与schedule组合验证 trace 导出与 step 计数行为可作为集成参考。进阶能力内存剖析与执行追踪内存时间线导出已弃用建议迁移启用record_shapesTrue, profile_memoryTrue, with_stackTrue后可通过export_memory_timeline(path, device)导出内存事件时间线profiler.py输出格式由path后缀决定.html嵌入 PNG 的内存时间线图.json/.json.gz[times, [sizes by category]]形式的绘图点.raw.json.gz原始内存事件点(timestamp, action, numbytes, category)其中action取PREEXISTING、CREATE、INCREMENT_VERSION、DESTROYcategory来自torch.profiler._memory_profiler.Category。注意该方法在源码中已标记为deprecated官方建议改用torch.cuda.memory._record_memory_history与torch.cuda.memory._export_memory_snapshot。ExecutionTraceObserverExecution Trace 执行追踪ExecutionTraceObserver提供基于图的 AI/ML 负载表示对应论文《PyTorch Execution Traces》可用于重放基准、模拟器与仿真器。它与 profiler 同窗口采集传入execution_trace_observer后观察者的start()/stop()会在与 profiler 相同的窗口内被调用。文档给出的示例profiler.pywith torch.profiler.profile( ... execution_trace_observer( ExecutionTraceObserver().register_callback(./execution_trace.json) ), ) as p: for iter in range(N): code_iteration_to_profile(iter) p.step()ExecutionTraceObserver的核心行为profiler.py每个进程只能有一个实例重复调用register_callback()不会重复注册register_callback(output_file_path)把观察者挂到 record function 回调上输出路径以.gz结尾时先写临时未压缩文件、注销时再压缩start()/stop()控制实际采集unregister_callback()负责收尾保存 Triton 生成的内核文件、压缩输出、移除观察者set_extra_resource_collection(True)会额外收集生成的内核、索引张量数据等重放 Execution Trace 所需的资源存放在输出文件同目录的{文件名}_resources子目录也可通过环境变量ENABLE_PYTORCH_EXECUTION_TRACE1配合ENABLE_PYTORCH_EXECUTION_TRACE_EXTRAS1自动启用无需改代码见 build_execution_trace_obs_from_env分布式场景下启动时会以## process_group:init ##节点记录进程组配置profiler.py。仓库中的 test/profiler/test_execution_trace.py 覆盖了该功能的完整行为文档注释亦指引读者参考test_execution_trace_with_kineto()用例。ITT APIIntel 仪器化与追踪技术文档第三部分docs/source/profiler.md介绍torch.profiler.itt模块它封装了 Intel 的 Instrumentation and Tracing TechnologyITTAPI用于在 Intel 平台上配合 VTune Profiler 等工具进行打点分析。API 一览函数作用torch.profiler.itt.is_available()检查当前构建是否可用 ITT 功能torch.profiler.itt.mark(msg)标记某一时刻发生的瞬时事件torch.profiler.itt.range_push(msg)将一个范围压入嵌套范围栈返回该范围的零基深度torch.profiler.itt.range_pop()从嵌套范围栈弹出最近的范围返回结束范围的零基深度torch.profiler.itt.range(msg)上下文管理器/装饰器进入作用域时 push、退出时 pop用法示例import torch.profiler.itt as itt if itt.is_available(): itt.mark(epoch_start) itt.range_push(forward_pass) # ... 前向计算 ... itt.range_pop() # 或者使用上下文管理器支持 str.format 参数 with itt.range(layer_{}, 3): # ... 某个层的计算 ... pass实现要点torch/profiler/itt.py底层通过torch._C._itt提供rangePush、rangePop、mark等 C 绑定若当前构建未安装 ITT会回退到桩实现is_available()返回False其余函数抛出RuntimeErrorrange_push(msg)要求 ASCII 消息返回新范围在嵌套栈中的零基深度range_pop()返回被结束范围的深度range是contextmanager实现支持msg.format(*args, **kwargs)格式化。分布式与元数据辅助能力_KinetoProfile在启动 trace 时start_trace会自动完成多项元数据注入根据profile_memory、with_stack、record_shapes、with_modules、with_flops写入对应的元数据标记若 Kineto 可用且torch.distributed已初始化自动注入distributedInfobackend、rank、world_size、进程组配置NCCL 后端还含 NCCL 版本并支持把preset_metadata合并进 trace当检测到 Inductor 的 CUDA Graphs 配置torch._inductor.config.triton.cudagraphs且 CUDA 版本 12.6 时会设置DISABLE_CUPTI_LAZY_REINIT1与TEARDOWN_CUPTI0规避 CUPTI 与 CUDA Graph 的已知兼容性问题见 profiler.py。profile还支持通过set_custom_trace_id_callback(callback)动态更换 trace ID 生成器并用get_trace_id()获取当前周期的 trace IDprofiler.py默认使用随机 UUID。调试建议与最佳实践长训练任务务必使用schedule默认调度会持续采集所有事件带来显著开销wait/warmup/active组合既能跳过启动阶段又能通过 warmup 稳定 CUDA 状态保证测量数据可信。按需开启record_shapes/with_stack二者是定位 shape 相关性能问题与调用来源的关键但会引入额外开销应在问题定位窗口内使用。用supported_activities()做设备兼容在混合设备或自定义后端环境下优先基于supported_activities()动态构造activities避免在无 CUPTI 或目标设备不可用时报错并注意 CUPTI 不可用时 CUDA 时间不会出现在 JSON trace 中。合理选择导出目标本地可视化用export_chrome_trace团队分享/分布式训练用tensorboard_trace_handler(dir_name)tensorboard --logdir超大 trace 可考虑.gz压缩输出与post_processing_timeout_s超时保护。复杂模型结合 Execution Trace需要重放、模拟或仿真时使用ExecutionTraceObserver输出.et.json轨迹并开启额外资源收集或通过ENABLE_PYTORCH_EXECUTION_TRACE1环境变量一键开启。Intel 平台利用 ITT 打点配合 VTune 等工具时使用itt.range/itt.mark标注关键代码段注意消息须为 ASCII 且需要 ITT 构建支持先检查is_available()。参考文件索引本文档主体docs/source/profiler.md核心实现torch/profiler/profiler.py、torch/profiler/init.pyITT 封装torch/profiler/itt.py类型声明torch/_C/_profiler.pyi测试用例test/profiler/test_profiler.py、test/profiler/test_execution_trace.py相关模块torch/profiler/_memory_profiler.py、torch/profiler/_chrome_trace_export.py【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考