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/pytorch

torch.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.py）：torch.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 的调用栈追踪。

RecordFunction：CPU 侧事件插桩的基础设施

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 中的脚注）：只有输入张量需要梯度的算子调用才会被分配序列号。这意味着纯推理（无梯度需求）场景下序列号关联的价值有限，而在训练场景中该机制是剖析反向传播热点的重要依据。

前向线程 ID（Forward Thread ID）

Autograd 可以在多线程环境中使用。前向线程 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}; ... };

ExtraFields<EventType::TorchOp>继承TorchOpBasicFields，在采集时随事件一起写入，供后处理阶段构建前反向关联。

Torch 算子采集流程：回调注册、线程子队列与热路径优化

本节描述auto-trace（进程内、同步）模式下 torch 算子的通用采集流程。关于 on-demand（进程外、异步）追踪的细节，可参考 Libkineto 的 README（Kineto 作为第三方子模块被引入）。

回调的注册：profiler_kineto.cpp

当一次 trace 开始时，autograd/profiler 后端会调用 torch/csrc/autograd/profiler_kineto.cpp 来准备、启动或停止采集。在 trace 启动时，定义于该文件中的onFunctionEnter与onFunctionExit回调会被注册到RecordFunction机制中。

源码中回调分为两组（torch/csrc/autograd/profiler_kineto.cpp）：

• onFunctionEnterGlobal/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_v<RawAllocation>)强制保持平凡类型以提升性能。

在 Python 前端，对应的开关是profile()的profile_memory参数（torch/profiler/profiler.py）：

torch.profiler.profile( activities=[torch.profiler.ProfilerActivity.CPU], profile_memory=True, # 跟踪张量内存分配/释放 )

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，原因有两个：

1. 合并所有数据并完成 Profiler 事件与 Kineto 事件之间的后处理（例如把 Kineto 事件嵌入Result事件树，见 torch/csrc/profiler/collection.h 的ExtraFields<EventType::Kineto>，其中Flow结构用于将 Kineto 的 flow 事件映射进 profiler 树）；
2. 把这些事件作为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_ready=tensorboard_trace_handler("./log/resnet18"), ) as prof: model(x) # 输出的 JSON 文件位于 ./log/resnet18/ 目录下

Python 栈追踪：基于 CPython 剖析 API 的实现

当在 profiler 中设置with_stack=True时，Python 栈追踪器会使用PythonTracerBase中定义的make函数生成，具体实现位于profiler_python.cpp（torch/csrc/autograd/profiler_python.cpp）。

使用PyEval_SetProfile追踪执行事件

为了剖析调用栈，实现使用PyEval_SetProfile来追踪并处理 Python 程序中的各种执行事件（对应源码中recordPyCall/recordCCall的实现与PyTrace_*分支，见 torch/csrc/autograd/profiler_python.cpp）。它针对以下具体场景做出响应：

在RecordQueue/ThreadLocalSubqueue中，Python 调用事件被以(TraceKey, approx_time_t)对的形式存入py_calls_（torch/csrc/profiler/collection.h），后处理阶段通过匹配入口与出口来构建完整的调用事件。

注意事项（README 原文）：对于 Python 3.12.0–3.12.4，CPython 中存在一个 bug，需要使用sys.monitoring作为 workaround。这意味着不同 Python 小版本下 Python 栈追踪的底层实现路径可能不同。

Python 前端的对应参数为with_stack（torch/profiler/profiler.py）：

with torch.profiler.profile( activities=[torch.profiler.ProfilerActivity.CPU], with_stack=True, # 记录算子的源码信息（文件名与行号） with_modules=True, # 记录模块层级（含函数名） ) as prof: ...

需要说明的是，README 与本仓库当前代码均提示with_modules已被标记为弃用（deprecated），未来的版本可能移除；而旧式的with_flops参数在 eager 模式下已不起作用，建议以with_stack=True与record_shapes=True组合来获得可归因的堆栈与形状信息。

时钟对齐：TSC 周期与 Unix 纳秒之间的转换

在创建时间戳时，Profiler 会根据系统环境选择最有效的时钟。大多数 Linux 系统的默认选择是TSC，它以 CPU 周期（cycles）的形式记录时间。为了把这个时间转换为纳秒级 Unix 时间，Profiler 会创建一个时钟转换器（clock converter）。

如果 profiler 中包含了 Kineto，这个转换器也会被传入 Kineto，以确保两端时间轴对齐——这正是 trace 中 CPU 事件与 GPU 事件能够精确排布在统一时间线上的基础。

在源码中可以看到这一机制的直接体现（torch/csrc/autograd/profiler_kineto.cpp）：

auto 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::ApproximateClockToUnixTimeConverter（clockConverter的类型）负责建立 TSC 计数与 Unix 时间的映射；getRecords在把各线程子队列中的原始近似时间事件物化为Result时，统一使用该转换器换算为纳秒时间戳。相关类型定义位于 c10/util/ApproximateClock.h。

总结：一次完整 trace 的端到端数据流

综合以上各节，一次torch.profiler.profile(...)调用的完整数据流可以归纳为：

1. 启动：Python 前端profile()解析参数并调用后端，profiler_kineto.cpp创建KinetoThreadLocalState与RecordQueue，注册onFunctionEnter/Exit回调（全局或线程本地）；
2. 采集：每个算子经 dispatcher 触发RecordFunction回调 →begin_op在对应线程的ThreadLocalSubqueue中写入轻量事件；分配与 OOM 事件绕过RecordFunction直接入队；Python 栈事件由PyEval_SetProfile驱动；GPU 事件由 Kineto 通过 CUPTI 等收集；
3. 结束：finalizeTrace()停止队列、构建时钟转换器（并同步给 Kineto），getRecords()把原始事件物化为Result树，与 Kineto 事件合并、完成后处理（前反向匹配、栈/模块/形状解码、未完成事件补记结束时间）；
4. 导出：通过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