PyTorch Profiler 实现原理深度解析从 RecordFunction 到 Kineto 的 CPU/GPU 全链路采集架构【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorchtorch.profiler是 PyTorch 官方提供的性能剖析工具它能够记录模型执行过程中的算子耗时、内存分配、Python 调用栈与 GPU 事件并导出为 Chrome Trace / Perfetto 可读的 JSON 文件。本文以仓库中 torch/csrc/profiler/README.md 为核心骨架结合 torch/csrc/profiler/collection.h、torch/csrc/profiler/collection.cpp、torch/csrc/autograd/profiler_kineto.cpp 等源码深入剖析 Profiler 的底层实现事件从 CPU 算子回调、内存分配、自动求导、Python 栈到 GPU 活动是如何被采集、关联与导出的。读完本文你将理解 Profiler 的完整数据链路、各采集阶段的内部机制以及如何基于这些原理正确配置和使用剖析参数。代码库结构Profiler 前端、后端与 C 采集层的职责划分Profiler 的实现横跨 Python 与 C 两大层次README 给出了核心文件布局相对仓库根路径关键组成如下torch/ │ ├── profiler/ # 主 Python 包包含核心前端逻辑 │ ├── __init__.py # profiler 包的初始化文件 │ ├── profiler.py # 主 Profiler 前端类profile、schedule、supported_activities 等 │ └── _utils.py # FunctionEvent 工具函数 │ ├── autograd/ # Autograd 包 │ ├── __init__.py # autograd 包初始化文件 │ ├── profiler.py # 主 Profiler 后端类 │ └── profiler_utils.py # FunctionEvent 工具函数 │ ├── csrc/ # C 与 C 源码 │ └── profiler/ # Profiler C 源码 │ ├── collection.cpp # 主要采集逻辑 │ ├── collection.h # 采集定义 │ ├── kineto_client_interface.cpp # 从 kineto 调用 Profiler 的接口仅 on-demand 模式 │ ├── kineto_client_interface.h # Client 接口定义 │ ├── kineto_shim.cpp # 从 profiler 调用 kineto 的 shim │ ├── kineto_shim.h # shim 定义 │ ├── util.cpp # 处理 profiler 事件中参数的 utils │ ├── util.h # util 定义 │ └── README.md # 本文所讲解的文档 │ └── autograd/ # Autograd C 源码 │ ├── profiler_python.cpp # 主要的 Python 调用栈采集逻辑 │ ├── profiler_python.h # Python 调用栈采集定义 │ ├── profiler_kineto.cpp # 启动采集/kineto 的 Profiler 后端逻辑 │ └── profiler_kineto.h # 启动采集/kineto 的 Profiler 后端定义 │ └── ATen/ # ATen C 源码 │ ├── record_function.cpp # RecordFunction 采集逻辑 │ └── record_function.h # RecordFunction 定义从源码结构可以清晰地看到职责边界Python 前端torch/profiler/profiler.py负责向用户暴露profile()、schedule()、ProfilerActivity等 API并把用户的配置如record_shapes、with_stack翻译成底层ProfilerConfig。Python 后端torch/autograd/profiler.pytorch.autograd.profiler模块中的profile类历史上是旧式 Profiler 的入口现在与 Kineto 后端并存。C 采集层torch/csrc/profiler/实现线程级事件缓冲ThreadLocalSubqueue、事件树构建Result、Kineto shim 与时钟转换。Autograd C 层torch/csrc/autograd/profiler_kineto.cpp实现回调注册与前后端桥接profiler_python.cpp实现基于 CPython C API 的调用栈追踪。RecordFunctionCPU 侧事件插桩的基础设施RecordFunction是 Profiler 插桩 CPU 侧事件的核心机制其定义位于 aten/src/ATen/record_function.h。它是一个通用的函数调用插桩手段并非 Profiler 专有也可以用于其他通用场景例如 PyTorch 的 大规模部署特性。在 PyTorch 内部它已被安插在若干关键位置最典型的是dispatcher中环绕每一个算子调用见 aten/src/ATen/core/dispatch/Dispatcher.h。RecordFunction的核心设计是回调注册制用户或 PyTorch 自身可以注册回调每当执行路径遇到一个RecordFunction守卫时这些回调就会被触发。Profiler 正是利用这一机制来记录每个算子调用的开始与结束时间以及用户自定义的RecordFunction注释。从工程角度RecordFunction机制被刻意设计为低开销尤其是在没有注册任何回调时。这是因为算子调度路径是训练/推理的热路径任何额外开销都会被放大。尽管如此一旦启用回调依然会引入一定的性能损耗——这解释了为什么 Profiler 只在需要时开启采集。RecordFunction还提供了 Python 绑定with torch.profiler.record_function(...)。这是用户在代码中标记模块级事件的常用手段例如import torch with torch.profiler.profile(activities[torch.profiler.ProfilerActivity.CPU]) as prof: with torch.profiler.record_function(my_custom_module): x torch.randn(1024, 1024) y x.mm(x)自定义的USER_SCOPE事件与算子事件一样会被采集并出现在最终的 trace 中。Autograd 集成用序列号与前向线程 ID 关联前反向自动求导引擎负责自动计算梯度。Profiler 从 autograd 引擎记录两类关键信息用于在 trace 中把反向算子与触发它的前向算子关联起来序列号Sequence Number定义位于 aten/src/ATen/SequenceNumber.h。这是一个每线程唯一的索引分配给前向传播中的每次算子调用。当反向算子被触发时它会被赋予与其来源前向算子相同的序列号。利用这一点Profiler 能够完成前向/反向算子的匹配在 Chrome Trace 中该功能以fwd_bwd flow events的形式呈现。需要特别注意的是README 中的脚注只有输入张量需要梯度的算子调用才会被分配序列号。这意味着纯推理无梯度需求场景下序列号关联的价值有限而在训练场景中该机制是剖析反向传播热点的重要依据。前向线程 IDForward Thread IDAutograd 可以在多线程环境中使用。前向线程 ID 表示前向算子执行所在线程的 ID相关字段可见于 aten/src/ATen/record_function.h 的RecordFunction定义中。之所以需要它是因为上述序列号只在单个线程内唯一不同线程上可能产生相同的序列号必须用前向线程 ID 加以区分。在 torch/csrc/profiler/collection.h 的TorchOpBasicFields中可以看到这两个字段的直接体现struct TorchOpBasicFields { int64_t sequence_number_{0}; uint64_t forward_tid_{0}; ... };ExtraFieldsEventType::TorchOp继承TorchOpBasicFields在采集时随事件一起写入供后处理阶段构建前反向关联。Torch 算子采集流程回调注册、线程子队列与热路径优化本节描述auto-trace进程内、同步模式下 torch 算子的通用采集流程。关于 on-demand进程外、异步追踪的细节可参考 Libkineto 的 READMEKineto 作为第三方子模块被引入。回调的注册profiler_kineto.cpp当一次 trace 开始时autograd/profiler 后端会调用 torch/csrc/autograd/profiler_kineto.cpp 来准备、启动或停止采集。在 trace 启动时定义于该文件中的onFunctionEnter与onFunctionExit回调会被注册到RecordFunction机制中。源码中回调分为两组torch/csrc/autograd/profiler_kineto.cpponFunctionEnterGlobal/onFunctionExitGlobal全局回调用于KINETO_ONDEMAND或配置了profile_all_threads的KINETO会话onFunctionEnterTLS/onFunctionExitTLS线程本地回调仅对 trace 启动时存在的线程生效。回调的注册范围由ExperimentalConfig决定分为两种模式Global全局回调注册到执行期间的所有线程。Local本地回调仅注册到trace 开始时刻已存在的线程上。auto recordFunctionCallback at::RecordFunctionCallback(onFunctionEnterGlobal, onFunctionExitGlobal) .needsInputs(state_ptr-config().report_input_shapes) .scopes(scopes);needsInputs(config.report_input_shapes)表明只有当用户开启record_shapes时回调才需要算子输入信息否则采集开销更小。线程子队列与begin_op在onFunctionEnter内部Profiler 会为每个线程创建一个ThreadLocalSubqueue实例见 torch/csrc/profiler/collection.h确保每个 CPU 算子都与它实际执行的线程关联。当一个 torch 算子进入时Profiler 调用定义在collection.cpp的begin_op来记录必要信息。ThreadLocalSubqueue的设计值得关注内部使用AppendOnlyList只追加、块大小为 512作为事件存储避免每次算子都分配独立 vector每种事件类型对应独立的存储区torch_ops_算子事件、allocations_内存分配、backend_events_后端事件、vulkan_events_、ooms_OOM、py_calls_Python 调用算子事件的correlation id由EventBlock按块起始 ID 块内偏移计算无需额外分配torch/csrc/profiler/collection.cpp。README 特别强调begin_op被刻意设计得非常轻量因为它处于 profiling 的 hot path热路径上过高的额外开销会扭曲 profile 结果、降低其参考价值。因此回调期间只收集最必要的信息大部分逻辑都推迟到后处理阶段完成。这也是InputOutputEncoder存在的原因——它在采集期把算子输入的形状、dtype 编码进连续的AppendOnlyList避免为每个算子创建 vector而在后处理时才解码还原见 torch/csrc/profiler/collection.cpp。全局会话的并发安全对于全局回调profiler_kineto.cpp中还有一个GlobalCallbackSession协调机制torch/csrc/profiler/kineto_profiler.cpp由于全局回调会在任意线程上触发而disableProfiler()可能在另一个线程上销毁 profiler 状态因此每个回调通过enter()/exit()只在极短的临界区getGlobal()RecordQueue访问内持有 in-flight 计数teardown 时drain()关闭会话并等待所有 in-flight 回调退出避免 use-after-free。退出回调还会校验session_generation_丢弃跨会话悬空的退出事件。内存分配事件采集绕过 RecordFunction 的瞬时事件与有起止时刻的算子事件不同内存分配事件被表示为cpu_instant_event零持续时间。因此分配事件不经过RecordFunction而是由reportMemoryUsage直接调用emplace_allocation_event把事件入队到对应的ThreadLocalSubqueue。在 torch/csrc/autograd/profiler_kineto.cpp 中可以看到void reportMemoryUsage( void* ptr, int64_t alloc_size, size_t total_allocated, size_t total_reserved, c10::Device device) override { if (config_.profile_memory !config_.disabled()) { recordQueue.getSubqueue()-emplace_allocation_event( c10::getApproximateTime(), ptr, alloc_size, total_allocated, total_reserved, device.type(), device.index()); } }同样地OOM内存不足事件通过reportOutOfMemory→emplace_ooms_event进入队列。RawAllocation结构torch/csrc/profiler/collection.h记录了分配指针、大小、累计分配量/预留量以及设备信息并被static_assert(std::is_trivial_vRawAllocation)强制保持平凡类型以提升性能。在 Python 前端对应的开关是profile()的profile_memory参数torch/profiler/profiler.pytorch.profiler.profile( activities[torch.profiler.ProfilerActivity.CPU], profile_memoryTrue, # 跟踪张量内存分配/释放 )Kineto 集成多架构事件采集、预热与导出Kineto 是 Profiler 获取 GPU 及其他加速器事件的关键抽象层。它作为第三方子模块third_party/kineto/目录被引入通过与 CUPTI 等库交互来接收 GPU 与加速器事件再转发给前端 Profiler。预热warmup的原因Kineto 需要时间prepare也称 warmup这些第三方模块以避免初始化例程扭曲 profile 结果。理论上可以在作业启动时就完成预热但让 CUPTI 这类重量级库持续运行会带来显著的不必要开销因此预热被安排在正式 trace 之前的短暂窗口内完成。双向桥接与后处理如前所述profiler_kineto.cpp在后台调用合适的 profiler 阶段同时它还会调用kineto_shim.cpp后者触发 Kineto 中的对应例程。当一次 trace 完成后Kineto 收集的所有事件会被转发给 Profiler原因有两个合并所有数据并完成 Profiler 事件与 Kineto 事件之间的后处理例如把 Kineto 事件嵌入Result事件树见 torch/csrc/profiler/collection.h 的ExtraFieldsEventType::Kineto其中Flow结构用于将 Kineto 的 flow 事件映射进 profiler 树把这些事件作为FunctionEvents转发给 Python 前端供prof.key_averages()、prof.table()等 API 使用。文件导出集成的最后一步是文件导出。所有事件收集并后处理完成后它们可以导出为 JSON 文件供Perfetto或Chrome Tracer可视化。这一步通过调用 Kineto 的ActivityTraceInterface::save完成把所有事件信息写入磁盘。在 Python 侧tensorboard_trace_handler会自动生成 trace 目录from torch.profiler import profile, ProfilerActivity, tensorboard_trace_handler with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], on_trace_readytensorboard_trace_handler(./log/resnet18), ) as prof: model(x) # 输出的 JSON 文件位于 ./log/resnet18/ 目录下Python 栈追踪基于 CPython 剖析 API 的实现当在 profiler 中设置with_stackTrue时Python 栈追踪器会使用PythonTracerBase中定义的make函数生成具体实现位于profiler_python.cpptorch/csrc/autograd/profiler_python.cpp。使用PyEval_SetProfile追踪执行事件为了剖析调用栈实现使用PyEval_SetProfile来追踪并处理 Python 程序中的各种执行事件对应源码中recordPyCall/recordCCall的实现与PyTrace_*分支见 torch/csrc/autograd/profiler_python.cpp。它针对以下具体场景做出响应CPython 事件处理函数采集内容PyTrace_CALLrecordPyCall记录每次 Python 函数调用捕获后续分析所需的关键细节PyTrace_C_CALLrecordCCall记录对 C 函数的调用包括相关参数提供程序执行流的完整视图PyTrace_RETURN—记录 Python 函数的退出时间实现精确的函数执行时长测量PyTrace_C_RETURN与PyTrace_C_EXCEPTION—记录 C 函数的退出时间无论正常结束还是因异常结束确保所有执行路径都被统计在RecordQueue/ThreadLocalSubqueue中Python 调用事件被以(TraceKey, approx_time_t)对的形式存入py_calls_torch/csrc/profiler/collection.h后处理阶段通过匹配入口与出口来构建完整的调用事件。注意事项README 原文对于 Python 3.12.0–3.12.4CPython 中存在一个 bug需要使用sys.monitoring作为 workaround。这意味着不同 Python 小版本下 Python 栈追踪的底层实现路径可能不同。Python 前端的对应参数为with_stacktorch/profiler/profiler.pywith torch.profiler.profile( activities[torch.profiler.ProfilerActivity.CPU], with_stackTrue, # 记录算子的源码信息文件名与行号 with_modulesTrue, # 记录模块层级含函数名 ) as prof: ...需要说明的是README 与本仓库当前代码均提示with_modules已被标记为弃用deprecated未来的版本可能移除而旧式的with_flops参数在 eager 模式下已不起作用建议以with_stackTrue与record_shapesTrue组合来获得可归因的堆栈与形状信息。时钟对齐TSC 周期与 Unix 纳秒之间的转换在创建时间戳时Profiler 会根据系统环境选择最有效的时钟。大多数 Linux 系统的默认选择是TSC它以 CPU 周期cycles的形式记录时间。为了把这个时间转换为纳秒级 Unix 时间Profiler 会创建一个时钟转换器clock converter。如果 profiler 中包含了 Kineto这个转换器也会被传入 Kineto以确保两端时间轴对齐——这正是 trace 中 CPU 事件与 GPU 事件能够精确排布在统一时间线上的基础。在源码中可以看到这一机制的直接体现torch/csrc/autograd/profiler_kineto.cppauto converter clockConverter.makeConverter(); #ifdef USE_KINETO libkineto::get_time_converter() converter; #endif auto records_and_trace recordQueue.getRecords(std::move(converter), startTime, end_time);c10::ApproximateClockToUnixTimeConverterclockConverter的类型负责建立 TSC 计数与 Unix 时间的映射getRecords在把各线程子队列中的原始近似时间事件物化为Result时统一使用该转换器换算为纳秒时间戳。相关类型定义位于 c10/util/ApproximateClock.h。总结一次完整 trace 的端到端数据流综合以上各节一次torch.profiler.profile(...)调用的完整数据流可以归纳为启动Python 前端profile()解析参数并调用后端profiler_kineto.cpp创建KinetoThreadLocalState与RecordQueue注册onFunctionEnter/Exit回调全局或线程本地采集每个算子经 dispatcher 触发RecordFunction回调 →begin_op在对应线程的ThreadLocalSubqueue中写入轻量事件分配与 OOM 事件绕过RecordFunction直接入队Python 栈事件由PyEval_SetProfile驱动GPU 事件由 Kineto 通过 CUPTI 等收集结束finalizeTrace()停止队列、构建时钟转换器并同步给 KinetogetRecords()把原始事件物化为Result树与 Kineto 事件合并、完成后处理前反向匹配、栈/模块/形状解码、未完成事件补记结束时间导出通过ActivityTraceInterface::save写出 JSON供 Chrome Trace / Perfetto 可视化或转成FunctionEvents供 Python 侧key_averages()等 API 做表格统计。理解这条链路后你将能更有针对性地使用 Profiler用record_shapes观察算子输入形状、用with_stack定位调用来源、用profile_memory追踪显存分配并通过 Chrome Trace 中的 fwd_bwd flow 事件分析前反向传播的时间占比。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 PyTorch Profiler 实现原理深度解析：从 RecordFunction 到 Kineto 的 CPU/GPU 全链路采集架构
PyTorch Profiler 实现原理深度解析从 RecordFunction 到 Kineto 的 CPU/GPU 全链路采集架构【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorchtorch.profiler是 PyTorch 官方提供的性能剖析工具它能够记录模型执行过程中的算子耗时、内存分配、Python 调用栈与 GPU 事件并导出为 Chrome Trace / Perfetto 可读的 JSON 文件。本文以仓库中 torch/csrc/profiler/README.md 为核心骨架结合 torch/csrc/profiler/collection.h、torch/csrc/profiler/collection.cpp、torch/csrc/autograd/profiler_kineto.cpp 等源码深入剖析 Profiler 的底层实现事件从 CPU 算子回调、内存分配、自动求导、Python 栈到 GPU 活动是如何被采集、关联与导出的。读完本文你将理解 Profiler 的完整数据链路、各采集阶段的内部机制以及如何基于这些原理正确配置和使用剖析参数。代码库结构Profiler 前端、后端与 C 采集层的职责划分Profiler 的实现横跨 Python 与 C 两大层次README 给出了核心文件布局相对仓库根路径关键组成如下torch/ │ ├── profiler/ # 主 Python 包包含核心前端逻辑 │ ├── __init__.py # profiler 包的初始化文件 │ ├── profiler.py # 主 Profiler 前端类profile、schedule、supported_activities 等 │ └── _utils.py # FunctionEvent 工具函数 │ ├── autograd/ # Autograd 包 │ ├── __init__.py # autograd 包初始化文件 │ ├── profiler.py # 主 Profiler 后端类 │ └── profiler_utils.py # FunctionEvent 工具函数 │ ├── csrc/ # C 与 C 源码 │ └── profiler/ # Profiler C 源码 │ ├── collection.cpp # 主要采集逻辑 │ ├── collection.h # 采集定义 │ ├── kineto_client_interface.cpp # 从 kineto 调用 Profiler 的接口仅 on-demand 模式 │ ├── kineto_client_interface.h # Client 接口定义 │ ├── kineto_shim.cpp # 从 profiler 调用 kineto 的 shim │ ├── kineto_shim.h # shim 定义 │ ├── util.cpp # 处理 profiler 事件中参数的 utils │ ├── util.h # util 定义 │ └── README.md # 本文所讲解的文档 │ └── autograd/ # Autograd C 源码 │ ├── profiler_python.cpp # 主要的 Python 调用栈采集逻辑 │ ├── profiler_python.h # Python 调用栈采集定义 │ ├── profiler_kineto.cpp # 启动采集/kineto 的 Profiler 后端逻辑 │ └── profiler_kineto.h # 启动采集/kineto 的 Profiler 后端定义 │ └── ATen/ # ATen C 源码 │ ├── record_function.cpp # RecordFunction 采集逻辑 │ └── record_function.h # RecordFunction 定义从源码结构可以清晰地看到职责边界Python 前端torch/profiler/profiler.py负责向用户暴露profile()、schedule()、ProfilerActivity等 API并把用户的配置如record_shapes、with_stack翻译成底层ProfilerConfig。Python 后端torch/autograd/profiler.pytorch.autograd.profiler模块中的profile类历史上是旧式 Profiler 的入口现在与 Kineto 后端并存。C 采集层torch/csrc/profiler/实现线程级事件缓冲ThreadLocalSubqueue、事件树构建Result、Kineto shim 与时钟转换。Autograd C 层torch/csrc/autograd/profiler_kineto.cpp实现回调注册与前后端桥接profiler_python.cpp实现基于 CPython C API 的调用栈追踪。RecordFunctionCPU 侧事件插桩的基础设施RecordFunction是 Profiler 插桩 CPU 侧事件的核心机制其定义位于 aten/src/ATen/record_function.h。它是一个通用的函数调用插桩手段并非 Profiler 专有也可以用于其他通用场景例如 PyTorch 的 大规模部署特性。在 PyTorch 内部它已被安插在若干关键位置最典型的是dispatcher中环绕每一个算子调用见 aten/src/ATen/core/dispatch/Dispatcher.h。RecordFunction的核心设计是回调注册制用户或 PyTorch 自身可以注册回调每当执行路径遇到一个RecordFunction守卫时这些回调就会被触发。Profiler 正是利用这一机制来记录每个算子调用的开始与结束时间以及用户自定义的RecordFunction注释。从工程角度RecordFunction机制被刻意设计为低开销尤其是在没有注册任何回调时。这是因为算子调度路径是训练/推理的热路径任何额外开销都会被放大。尽管如此一旦启用回调依然会引入一定的性能损耗——这解释了为什么 Profiler 只在需要时开启采集。RecordFunction还提供了 Python 绑定with torch.profiler.record_function(...)。这是用户在代码中标记模块级事件的常用手段例如import torch with torch.profiler.profile(activities[torch.profiler.ProfilerActivity.CPU]) as prof: with torch.profiler.record_function(my_custom_module): x torch.randn(1024, 1024) y x.mm(x)自定义的USER_SCOPE事件与算子事件一样会被采集并出现在最终的 trace 中。Autograd 集成用序列号与前向线程 ID 关联前反向自动求导引擎负责自动计算梯度。Profiler 从 autograd 引擎记录两类关键信息用于在 trace 中把反向算子与触发它的前向算子关联起来序列号Sequence Number定义位于 aten/src/ATen/SequenceNumber.h。这是一个每线程唯一的索引分配给前向传播中的每次算子调用。当反向算子被触发时它会被赋予与其来源前向算子相同的序列号。利用这一点Profiler 能够完成前向/反向算子的匹配在 Chrome Trace 中该功能以fwd_bwd flow events的形式呈现。需要特别注意的是README 中的脚注只有输入张量需要梯度的算子调用才会被分配序列号。这意味着纯推理无梯度需求场景下序列号关联的价值有限而在训练场景中该机制是剖析反向传播热点的重要依据。前向线程 IDForward Thread IDAutograd 可以在多线程环境中使用。前向线程 ID 表示前向算子执行所在线程的 ID相关字段可见于 aten/src/ATen/record_function.h 的RecordFunction定义中。之所以需要它是因为上述序列号只在单个线程内唯一不同线程上可能产生相同的序列号必须用前向线程 ID 加以区分。在 torch/csrc/profiler/collection.h 的TorchOpBasicFields中可以看到这两个字段的直接体现struct TorchOpBasicFields { int64_t sequence_number_{0}; uint64_t forward_tid_{0}; ... };ExtraFieldsEventType::TorchOp继承TorchOpBasicFields在采集时随事件一起写入供后处理阶段构建前反向关联。Torch 算子采集流程回调注册、线程子队列与热路径优化本节描述auto-trace进程内、同步模式下 torch 算子的通用采集流程。关于 on-demand进程外、异步追踪的细节可参考 Libkineto 的 READMEKineto 作为第三方子模块被引入。回调的注册profiler_kineto.cpp当一次 trace 开始时autograd/profiler 后端会调用 torch/csrc/autograd/profiler_kineto.cpp 来准备、启动或停止采集。在 trace 启动时定义于该文件中的onFunctionEnter与onFunctionExit回调会被注册到RecordFunction机制中。源码中回调分为两组torch/csrc/autograd/profiler_kineto.cpponFunctionEnterGlobal/onFunctionExitGlobal全局回调用于KINETO_ONDEMAND或配置了profile_all_threads的KINETO会话onFunctionEnterTLS/onFunctionExitTLS线程本地回调仅对 trace 启动时存在的线程生效。回调的注册范围由ExperimentalConfig决定分为两种模式Global全局回调注册到执行期间的所有线程。Local本地回调仅注册到trace 开始时刻已存在的线程上。auto recordFunctionCallback at::RecordFunctionCallback(onFunctionEnterGlobal, onFunctionExitGlobal) .needsInputs(state_ptr-config().report_input_shapes) .scopes(scopes);needsInputs(config.report_input_shapes)表明只有当用户开启record_shapes时回调才需要算子输入信息否则采集开销更小。线程子队列与begin_op在onFunctionEnter内部Profiler 会为每个线程创建一个ThreadLocalSubqueue实例见 torch/csrc/profiler/collection.h确保每个 CPU 算子都与它实际执行的线程关联。当一个 torch 算子进入时Profiler 调用定义在collection.cpp的begin_op来记录必要信息。ThreadLocalSubqueue的设计值得关注内部使用AppendOnlyList只追加、块大小为 512作为事件存储避免每次算子都分配独立 vector每种事件类型对应独立的存储区torch_ops_算子事件、allocations_内存分配、backend_events_后端事件、vulkan_events_、ooms_OOM、py_calls_Python 调用算子事件的correlation id由EventBlock按块起始 ID 块内偏移计算无需额外分配torch/csrc/profiler/collection.cpp。README 特别强调begin_op被刻意设计得非常轻量因为它处于 profiling 的 hot path热路径上过高的额外开销会扭曲 profile 结果、降低其参考价值。因此回调期间只收集最必要的信息大部分逻辑都推迟到后处理阶段完成。这也是InputOutputEncoder存在的原因——它在采集期把算子输入的形状、dtype 编码进连续的AppendOnlyList避免为每个算子创建 vector而在后处理时才解码还原见 torch/csrc/profiler/collection.cpp。全局会话的并发安全对于全局回调profiler_kineto.cpp中还有一个GlobalCallbackSession协调机制torch/csrc/profiler/kineto_profiler.cpp由于全局回调会在任意线程上触发而disableProfiler()可能在另一个线程上销毁 profiler 状态因此每个回调通过enter()/exit()只在极短的临界区getGlobal()RecordQueue访问内持有 in-flight 计数teardown 时drain()关闭会话并等待所有 in-flight 回调退出避免 use-after-free。退出回调还会校验session_generation_丢弃跨会话悬空的退出事件。内存分配事件采集绕过 RecordFunction 的瞬时事件与有起止时刻的算子事件不同内存分配事件被表示为cpu_instant_event零持续时间。因此分配事件不经过RecordFunction而是由reportMemoryUsage直接调用emplace_allocation_event把事件入队到对应的ThreadLocalSubqueue。在 torch/csrc/autograd/profiler_kineto.cpp 中可以看到void reportMemoryUsage( void* ptr, int64_t alloc_size, size_t total_allocated, size_t total_reserved, c10::Device device) override { if (config_.profile_memory !config_.disabled()) { recordQueue.getSubqueue()-emplace_allocation_event( c10::getApproximateTime(), ptr, alloc_size, total_allocated, total_reserved, device.type(), device.index()); } }同样地OOM内存不足事件通过reportOutOfMemory→emplace_ooms_event进入队列。RawAllocation结构torch/csrc/profiler/collection.h记录了分配指针、大小、累计分配量/预留量以及设备信息并被static_assert(std::is_trivial_vRawAllocation)强制保持平凡类型以提升性能。在 Python 前端对应的开关是profile()的profile_memory参数torch/profiler/profiler.pytorch.profiler.profile( activities[torch.profiler.ProfilerActivity.CPU], profile_memoryTrue, # 跟踪张量内存分配/释放 )Kineto 集成多架构事件采集、预热与导出Kineto 是 Profiler 获取 GPU 及其他加速器事件的关键抽象层。它作为第三方子模块third_party/kineto/目录被引入通过与 CUPTI 等库交互来接收 GPU 与加速器事件再转发给前端 Profiler。预热warmup的原因Kineto 需要时间prepare也称 warmup这些第三方模块以避免初始化例程扭曲 profile 结果。理论上可以在作业启动时就完成预热但让 CUPTI 这类重量级库持续运行会带来显著的不必要开销因此预热被安排在正式 trace 之前的短暂窗口内完成。双向桥接与后处理如前所述profiler_kineto.cpp在后台调用合适的 profiler 阶段同时它还会调用kineto_shim.cpp后者触发 Kineto 中的对应例程。当一次 trace 完成后Kineto 收集的所有事件会被转发给 Profiler原因有两个合并所有数据并完成 Profiler 事件与 Kineto 事件之间的后处理例如把 Kineto 事件嵌入Result事件树见 torch/csrc/profiler/collection.h 的ExtraFieldsEventType::Kineto其中Flow结构用于将 Kineto 的 flow 事件映射进 profiler 树把这些事件作为FunctionEvents转发给 Python 前端供prof.key_averages()、prof.table()等 API 使用。文件导出集成的最后一步是文件导出。所有事件收集并后处理完成后它们可以导出为 JSON 文件供Perfetto或Chrome Tracer可视化。这一步通过调用 Kineto 的ActivityTraceInterface::save完成把所有事件信息写入磁盘。在 Python 侧tensorboard_trace_handler会自动生成 trace 目录from torch.profiler import profile, ProfilerActivity, tensorboard_trace_handler with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], on_trace_readytensorboard_trace_handler(./log/resnet18), ) as prof: model(x) # 输出的 JSON 文件位于 ./log/resnet18/ 目录下Python 栈追踪基于 CPython 剖析 API 的实现当在 profiler 中设置with_stackTrue时Python 栈追踪器会使用PythonTracerBase中定义的make函数生成具体实现位于profiler_python.cpptorch/csrc/autograd/profiler_python.cpp。使用PyEval_SetProfile追踪执行事件为了剖析调用栈实现使用PyEval_SetProfile来追踪并处理 Python 程序中的各种执行事件对应源码中recordPyCall/recordCCall的实现与PyTrace_*分支见 torch/csrc/autograd/profiler_python.cpp。它针对以下具体场景做出响应CPython 事件处理函数采集内容PyTrace_CALLrecordPyCall记录每次 Python 函数调用捕获后续分析所需的关键细节PyTrace_C_CALLrecordCCall记录对 C 函数的调用包括相关参数提供程序执行流的完整视图PyTrace_RETURN—记录 Python 函数的退出时间实现精确的函数执行时长测量PyTrace_C_RETURN与PyTrace_C_EXCEPTION—记录 C 函数的退出时间无论正常结束还是因异常结束确保所有执行路径都被统计在RecordQueue/ThreadLocalSubqueue中Python 调用事件被以(TraceKey, approx_time_t)对的形式存入py_calls_torch/csrc/profiler/collection.h后处理阶段通过匹配入口与出口来构建完整的调用事件。注意事项README 原文对于 Python 3.12.0–3.12.4CPython 中存在一个 bug需要使用sys.monitoring作为 workaround。这意味着不同 Python 小版本下 Python 栈追踪的底层实现路径可能不同。Python 前端的对应参数为with_stacktorch/profiler/profiler.pywith torch.profiler.profile( activities[torch.profiler.ProfilerActivity.CPU], with_stackTrue, # 记录算子的源码信息文件名与行号 with_modulesTrue, # 记录模块层级含函数名 ) as prof: ...需要说明的是README 与本仓库当前代码均提示with_modules已被标记为弃用deprecated未来的版本可能移除而旧式的with_flops参数在 eager 模式下已不起作用建议以with_stackTrue与record_shapesTrue组合来获得可归因的堆栈与形状信息。时钟对齐TSC 周期与 Unix 纳秒之间的转换在创建时间戳时Profiler 会根据系统环境选择最有效的时钟。大多数 Linux 系统的默认选择是TSC它以 CPU 周期cycles的形式记录时间。为了把这个时间转换为纳秒级 Unix 时间Profiler 会创建一个时钟转换器clock converter。如果 profiler 中包含了 Kineto这个转换器也会被传入 Kineto以确保两端时间轴对齐——这正是 trace 中 CPU 事件与 GPU 事件能够精确排布在统一时间线上的基础。在源码中可以看到这一机制的直接体现torch/csrc/autograd/profiler_kineto.cppauto converter clockConverter.makeConverter(); #ifdef USE_KINETO libkineto::get_time_converter() converter; #endif auto records_and_trace recordQueue.getRecords(std::move(converter), startTime, end_time);c10::ApproximateClockToUnixTimeConverterclockConverter的类型负责建立 TSC 计数与 Unix 时间的映射getRecords在把各线程子队列中的原始近似时间事件物化为Result时统一使用该转换器换算为纳秒时间戳。相关类型定义位于 c10/util/ApproximateClock.h。总结一次完整 trace 的端到端数据流综合以上各节一次torch.profiler.profile(...)调用的完整数据流可以归纳为启动Python 前端profile()解析参数并调用后端profiler_kineto.cpp创建KinetoThreadLocalState与RecordQueue注册onFunctionEnter/Exit回调全局或线程本地采集每个算子经 dispatcher 触发RecordFunction回调 →begin_op在对应线程的ThreadLocalSubqueue中写入轻量事件分配与 OOM 事件绕过RecordFunction直接入队Python 栈事件由PyEval_SetProfile驱动GPU 事件由 Kineto 通过 CUPTI 等收集结束finalizeTrace()停止队列、构建时钟转换器并同步给 KinetogetRecords()把原始事件物化为Result树与 Kineto 事件合并、完成后处理前反向匹配、栈/模块/形状解码、未完成事件补记结束时间导出通过ActivityTraceInterface::save写出 JSON供 Chrome Trace / Perfetto 可视化或转成FunctionEvents供 Python 侧key_averages()等 API 做表格统计。理解这条链路后你将能更有针对性地使用 Profiler用record_shapes观察算子输入形状、用with_stack定位调用来源、用profile_memory追踪显存分配并通过 Chrome Trace 中的 fwd_bwd flow 事件分析前反向传播的时间占比。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考