torch_npu 性能分析模块技术解析
·
背景概述
在深度学习模型的训练与推理过程中，性能分析（Profiling）是优化算力利用率、定位性能瓶颈的关键环节。torch_npu 是昇腾 NPU 上的 PyTorch 适配框架，其性能分析模块（Profiler）在 PyTorch Profiler 接口规范基础上进行扩展，提供算子执行、内存使用、通信操作等多维度的性能数据采集、解析与可视化能力。本文从模块架构、核心流程、状态机调度与工程实践等角度系统解析该模块的内部实现，为在昇腾 NPU 上进行性能调优的开发者提供参考。
1. 模块概述
torch_npu Profiler 模块是一款针对华为昇腾（Ascend）NPU 设备深度优化的性能分析工具。该模块在 PyTorch 原生 Profiler 功能基础上进行了扩展，为开发者提供涵盖 NPU 算子执行、内存使用、通信操作等多维度的全面性能分析能力。
1.1 主要功能
- 算子性能采集：精确采集 CPU 端与 NPU 端算子的执行耗时，为性能瓶颈定位提供数据基础
- 内存分析：全链路跟踪内存的分配、释放及使用情况，辅助识别内存泄漏与过度占用问题
- Python 调用栈追踪：完整记录 Python 函数调用关系，帮助开发者理解代码执行路径
- 通信分析：深入分析分布式训练场景下的通信操作效率，识别通信瓶颈
- MSTX 标记：支持用户自定义性能标记点，实现细粒度的性能观测
- TensorBoard 集成：生成兼容 TensorBoard 的可视化 trace 文件，提升数据分析效率
- 多种导出格式：支持 JSON（Chrome Trace）、数据库（DB）等多种数据导出格式，满足不同分析场景需求
1.2 核心职责
| 组件 | 职责 |
|---|---|
| ProfilerMgr | 管理 Profiler 生命周期，控制数据采集开关 |
| NpuProfilerThreadLocalState | 管理线程本地的 Profiler 状态 |
| PythonTracer | 追踪 Python 函数调用 |
| DataDumper | 将采集数据写入文件 |
| ProfilingParser | 解析原始数据并生成可视化文件 |
2. 架构设计
2.1 Profiler 整体架构
+-----------------------------------------------------------------------------------+
| 用户层 (User Layer) |
| +-------------+ +------------------+ +-------------------+ +---------------+ |
| | profile() | | schedule() | | tensorboard_ | | analyse() | |
| | context | | ProfilerAction | | trace_handler() | | | |
| +-------------+ +------------------+ +-------------------+ +---------------+ |
+-----------------------------------------------------------------------------------+
|
v
+-----------------------------------------------------------------------------------+
| Python 接口层 (Python Interface) |
| +-------------------+ +--------------------+ +---------------------------+ |
| | _KinetoProfile | | _ProfInterface | | _ExperimentalConfig | |
| | profile | | | | | |
| +-------------------+ +--------------------+ +---------------------------+ |
| +-------------------+ +--------------------+ +---------------------------+ |
| | ProfActionController| | ProfPathCreator | | ProfGCDetector | |
| +-------------------+ +--------------------+ +---------------------------+ |
+-----------------------------------------------------------------------------------+
|
v
+-----------------------------------------------------------------------------------+
| C++ 核心层 (C++ Core Layer) |
| +-------------------+ +--------------------+ +---------------------------+ |
| | ProfilerMgr | | NpuProfilerThread | | PythonTracer | |
| | (单例模式) | | LocalState | | | |
| +-------------------+ +--------------------+ +---------------------------+ |
| +-------------------+ +--------------------+ +---------------------------+ |
| | MstxMgr | | FeatureMgr | | DataDumper | |
| | (MSTX标记管理) | | (特性管理) | | (数据转储) | |
| +-------------------+ +--------------------+ +---------------------------+ |
+-----------------------------------------------------------------------------------+
|
v
+-----------------------------------------------------------------------------------+
| 底层接口层 (Low-level Interface) |
| +-------------------+ +--------------------+ +---------------------------+ |
| | ACL Profiler API | | RecordFunction | | PyEval_SetProfile | |
| | (aclprof*) | | Callback | | (Python C API) | |
| +-------------------+ +--------------------+ +---------------------------+ |
+-----------------------------------------------------------------------------------+
|
v
+-----------------------------------------------------------------------------------+
| 数据解析层 (Data Parsing Layer) |
| +-------------------+ +--------------------+ +---------------------------+ |
| | ProfilingParser | | FwkFileParser | | CANNFileParser | |
| +-------------------+ +--------------------+ +---------------------------+ |
| +-------------------+ +--------------------+ +---------------------------+ |
| | TreeBuilder | | TraceEventManager | | BinaryDecoder | |
| +-------------------+ +--------------------+ +---------------------------+ |
+-----------------------------------------------------------------------------------+
|
v
+-----------------------------------------------------------------------------------+
| 可视化层 (Visualization Layer) |
| +-------------------+ +--------------------+ +---------------------------+ |
| | TraceViewParser | | OperatorViewParser | | MemoryViewParser | |
| +-------------------+ +--------------------+ +---------------------------+ |
| +-------------------+ +--------------------+ +---------------------------+ |
| | KernelViewParser | | StackViewParser | | CommunicationParser | |
| +-------------------+ +--------------------+ +---------------------------+ |
+-----------------------------------------------------------------------------------+
2.2 数据采集层设计
数据采集层负责在运行时收集各类性能数据，是整个 Profiler 系统的数据基础：
+-----------------------------------------------------------------------------------+
| 数据采集层 |
+-----------------------------------------------------------------------------------+
| |
| +------------------------+ +------------------------+ |
| | CPU 数据采集 | | NPU 数据采集 | |
| +------------------------+ +------------------------+ |
| | | | | |
| | +------------------+ | | +------------------+ | |
| | | RecordFunction | | | | ACL Profiler | | |
| | | Callback | | | | (aclprofStart) | | |
| | +------------------+ | | +------------------+ | |
| | | | | | | |
| | v | | v | |
| | +------------------+ | | +------------------+ | |
| | | OpRangeData | | | | NPU Trace Data | | |
| | | (算子范围数据) | | | | (NPU追踪数据) | | |
| | +------------------+ | | +------------------+ | |
| | | | | |
| +------------------------+ +------------------------+ |
| |
| +------------------------+ +------------------------+ |
| | Python 追踪 | | 内存数据采集 | |
| +------------------------+ +------------------------+ |
| | | | | |
| | +------------------+ | | +------------------+ | |
| | | PyEval_SetProfile| | | | MemoryAllocator | | |
| | | PythonTracer | | | | Hooks | | |
| | +------------------+ | | +------------------+ | |
| | | | | | | |
| | v | | v | |
| | +------------------+ | | +------------------+ | |
| | | TraceEvent | | | | MemoryData | | |
| | | (追踪事件) | | | | (内存数据) | | |
| | +------------------+ | | +------------------+ | |
| | | | | |
| +------------------------+ +------------------------+ |
| |
+-----------------------------------------------------------------------------------+
|
v
+------------------------+
| RingBuffer |
| (环形缓冲区) |
+------------------------+
|
v
+------------------------+
| DataDumper |
| (数据转储器) |
+------------------------+
|
v
+------------------------+
| Binary Files |
| (二进制文件) |
+------------------------+
2.3 数据解析层设计
数据解析层负责将采集到的原始二进制数据转换为结构化对象，为上层可视化提供统一的数据模型：
+-----------------------------------------------------------------------------------+
| 数据解析层 |
+-----------------------------------------------------------------------------------+
| |
| 输入文件 |
| +----------------+ +----------------+ +----------------+ |
| | torch.op_range | | torch.op_mark | | torch.memory | |
| | (算子范围) | | (标记点) | | (内存使用) | |
| +----------------+ +----------------+ +----------------+ |
| +----------------+ +----------------+ +----------------+ |
| | python_tracer | | CANN profiler | | gc_record | |
| | _func/hash | | data | | (GC记录) | |
| +----------------+ +----------------+ +----------------+ |
| | | | |
| v v v |
| +------------------------------------------------------------------+ |
| | BinaryDecoder / TLVDecoder | |
| | (二进制解码器) | |
| +------------------------------------------------------------------+ |
| | |
| v |
| +------------------------------------------------------------------+ |
| | Bean 对象层 | |
| | +---------------+ +---------------+ +---------------+ | |
| | | TorchOpBean | | OpMarkBean | | MemoryUseBean | | |
| | +---------------+ +---------------+ +---------------+ | |
| | +---------------+ +---------------+ +---------------+ | |
| | | PythonTracer | | GCRecordBean | | NodeInfoBean | | |
| | | FuncBean | | | | | | |
| | +---------------+ +---------------+ +---------------+ | |
| +------------------------------------------------------------------+ |
| | |
| v |
| +------------------------------------------------------------------+ |
| | TreeBuilder | |
| | (构建算子调用树) | |
| +------------------------------------------------------------------+ |
| |
+-----------------------------------------------------------------------------------+
2.4 可视化层设计
可视化层将解析后的数据转换为多种视图格式，支持不同维度的性能分析需求：
+-----------------------------------------------------------------------------------+
| 可视化层 |
+-----------------------------------------------------------------------------------+
| |
| +------------------------------------------------------------------+ |
| | ParserConfig (解析器配置) | |
| | - COMMON_CONFIG: 通用配置 | |
| | - ONLY_FWK_CONFIG: 仅框架配置 | |
| | - LEVEL_NONE_CONFIG: Level_none配置 | |
| +------------------------------------------------------------------+ |
| | |
| v |
| +------------------------------------------------------------------+ |
| | ConcurrentTasksManager (并发任务管理器) | |
| +------------------------------------------------------------------+ |
| | |
| +------------------------+------------------------+ |
| | | | |
| v v v |
| +----------------+ +----------------+ +----------------+ |
| | TraceView | | OperatorView | | MemoryView | |
| | Parser | | Parser | | Parser | |
| | (时间线视图) | | (算子视图) | | (内存视图) | |
| +----------------+ +----------------+ +----------------+ |
| | | | |
| v v v |
| +----------------+ +----------------+ +----------------+ |
| | trace_view.json| | operator_*.csv | | memory_*.csv | |
| +----------------+ +----------------+ +----------------+ |
| |
| +------------------------+------------------------+ |
| | | | |
| v v v |
| +----------------+ +----------------+ +----------------+ |
| | KernelView | | StackView | | Communication | |
| | Parser | | Parser | | Parser | |
| | (Kernel视图) | | (调用栈视图) | | (通信视图) | |
| +----------------+ +----------------+ +----------------+ |
| | | | |
| v v v |
| +----------------+ +----------------+ +----------------+ |
| | kernel_*.csv | | stack_*.txt | | comm_*.json | |
| +----------------+ +----------------+ +----------------+ |
| |
+-----------------------------------------------------------------------------------+
|
v
+-----------------------------------------------------------------------------------+
| 输出 (Output) |
| +----------------+ +----------------+ +----------------+ |
| | Chrome Trace | | CSV Files | | SQLite DB | |
| | JSON | | | | | |
| +----------------+ +----------------+ +----------------+ |
| | | | |
| +------------------------+------------------------+ |
| | |
| v |
| +----------------+ |
| | TensorBoard | |
| | 可视化 | |
| +----------------+ |
+-----------------------------------------------------------------------------------+
3. 整体流程
3.1 Profiler 启动流程
Profiler 的启动遵循一套完整的初始化流程，确保各组件按序就绪：
用户调用 profile().__enter__()
|
v
+---------------------------+
| ProfPathCreator.init() |
| 初始化输出路径 |
+---------------------------+
|
v
+---------------------------+
| _ProfInterface.init_trace()|
| 1. 创建 profiling 目录 |
| 2. 调用 _init_profiler() |
+---------------------------+
|
v
+---------------------------+
| initNpuProfiler() |
| (C++ 层) |
| 1. 验证路径有效性 |
| 2. ProfilerMgr::Init() |
| 3. AclProfilingInit() |
+---------------------------+
|
v
+---------------------------+
| 根据 schedule 决定动作 |
| - WARMUP: warmup_trace() |
| - RECORD: start_trace() |
+---------------------------+
|
v (如果是 start_trace)
+---------------------------+
| startNpuProfiler() |
| (C++ 层) |
| 1. 创建 ThreadLocalState |
| 2. ProfilerMgr::Start() |
| 3. 注册 RecordFunction |
| Callback |
| 4. 启动 PythonTracer |
+---------------------------+
|
v
+---------------------------+
| ProfilerMgr::Start() |
| 1. EnableMsProfiler() |
| 2. StartDataReceiver() |
| 3. 设置 report_enable_ |
+---------------------------+
3.2 数据采集流程
运行时数据采集是 Profiler 的核心工作环节，通过多通道并行采集确保数据的完整性和实时性：
+-----------------------------------------------------------------------------------+
| 运行时数据采集 |
+-----------------------------------------------------------------------------------+
CPU 算子数据采集:
+---------------------------+
| PyTorch 执行算子 |
| 触发 RecordFunction |
+---------------------------+
|
v
+---------------------------+
| 进入回调 (enter callback) |
| 1. 获取 ThreadLocalState |
| 2. 创建 NpuObserverContext|
| 3. 记录开始时间戳 |
| 4. 解析输入形状/类型 |
| 5. 记录调用栈 |
+---------------------------+
|
v
+---------------------------+
| 退出回调 (exit callback) |
| 1. 记录结束时间戳 |
| 2. ProfilerMgr::Upload() |
+---------------------------+
|
v
+---------------------------+
| DataDumper::Report() |
| 1. 数据入队 RingBuffer |
| 2. 后台线程异步写文件 |
+---------------------------+
4. 核心数据采集流程
4.1 Python 追踪数据采集
Python 层追踪用于记录 Python 函数的调用关系与执行时间，其采集流程如下：
+---------------------------+
| Python 函数调用 |
| 触发 PyTrace_CALL |
+---------------------------+
|
v
+---------------------------+
| PythonTracer::pyProfileFn |
| 1. 生成 hash id |
| 2. 记录 TraceEvent |
+---------------------------+
|
v
+---------------------------+
| 达到阈值时批量上报 |
| reportTraceData() |
+---------------------------+
当 Python 函数被调用时，CPython 解释器触发 PyTrace_CALL 事件，PythonTracer::pyProfileFn 回调函数随即被调用。该函数为每个调用生成唯一的 hash id，并记录对应的 TraceEvent。当事件累积到一定阈值后，系统会批量上报数据，以降低频繁 I/O 带来的性能开销。
4.2 内存数据采集
内存数据采集关注 NPU 设备上的内存分配与释放行为，其流程如下：
+---------------------------+
| 内存分配/释放操作 |
| CachingAllocator 调用 |
+---------------------------+
|
v
+---------------------------+
| reportMemoryDataToNpuProfiler|
| 1. 创建 MemoryData |
| 2. ProfilerMgr:: |
| UploadWithLock() |
+---------------------------+
当 CachingAllocator 执行内存分配或释放操作时，会触发 reportMemoryDataToNpuProfiler 回调。该回调构造 MemoryData 对象，并通过 ProfilerMgr::UploadWithLock() 方法将数据安全地提交至 Profiler 管理器，确保多线程环境下的数据一致性。
5. 数据导出流程
当用户调用 profile().__exit__() 或 step() 方法时，会触发 RECORD_AND_SAVE 动作，从而启动完整的数据导出流程：
用户调用 profile().__exit__() 或 step() 触发 RECORD_AND_SAVE
|
v
+---------------------------+
| ProfActionController. |
| transit_action() |
+---------------------------+
|
v
+---------------------------+
| _ProfInterface.stop_trace()|
| 1. _stop_profiler() |
| 2. _disable_event_record()|
+---------------------------+
|
v
+---------------------------+
| stopNpuProfiler() |
| (C++ 层) |
| 1. 移除 RecordFunction |
| Callback |
| 2. 停止 PythonTracer |
| 3. ProfilerMgr::Stop() |
+---------------------------+
|
v
+---------------------------+
| ProfilerMgr::Stop() |
| 1. npuSynchronizeDevice() |
| 2. StopDataReceiver() |
| 3. AclProfilingStop() |
| 4. 销毁 profConfig |
+---------------------------+
|
v
+---------------------------+
| _ProfInterface. |
| finalize_trace() |
| 1. _finalize_profiler() |
| 2. _dump_profiler_info() |
| 3. _dump_metadata() |
+---------------------------+
|
v
+---------------------------+
| on_trace_ready 回调 |
| tensorboard_trace_handler |
+---------------------------+
|
v
+---------------------------+
| NpuProfiler.analyse() |
| 1. 获取 profiler 路径列表 |
| 2. 创建多进程池 |
| 3. 提交解析任务 |
+---------------------------+
|
v
+---------------------------+
| ProfilingParser. |
| analyse_profiling_data() |
| 1. 加载配置 |
| 2. 运行各类解析器 |
| 3. 数据简化处理 |
+---------------------------+
整个导出流程分为三个阶段：
- 停止采集阶段：通过
stopNpuProfiler()移除回调、停止 Python Tracer，并调用ProfilerMgr::Stop()完成设备同步、停止数据接收器及 ACL Profiling 会话。 - 数据收尾阶段：
finalize_trace()负责将 Profiler 信息与元数据落盘。 - 数据解析阶段：通过
on_trace_ready回调触发NpuProfiler.analyse()，利用多进程池并行执行ProfilingParser中的各类解析器，最终生成可视化所需的中间文件。
6. 与 TensorBoard 集成流程
TensorBoard 集成是 Profiler 数据可视化的核心出口，其完整流程如下：
+---------------------------+
| tensorboard_trace_handler |
| 初始化输出路径 |
+---------------------------+
|
v
+---------------------------+
| profile 完成数据采集 |
| 触发 on_trace_ready |
+---------------------------+
|
v
+---------------------------+
| handler_fn(prof_inst) |
| 调用 prof_inst.analyse() |
+---------------------------+
|
v
+---------------------------+
| NpuProfiler.analyse() |
| analysis_type = |
| TENSORBOARD_TRACE_HANDLER |
+---------------------------+
|
v
+---------------------------+
| ProfilingParser 执行 |
| ParserConfig.COMMON_CONFIG|
| 中的所有解析器 |
+---------------------------+
|
+---> TraceViewParser --> trace_view.json
|
+---> OperatorViewParser --> operator_details.csv
|
+---> KernelViewParser --> kernel_details.csv
|
+---> MemoryViewParser --> memory_record.csv
|
+---> CommunicationParser --> communication.json
|
v
+---------------------------+
| 输出到 ASCEND_PROFILER_ |
| OUTPUT 目录 |
+---------------------------+
|
v
+---------------------------+
| TensorBoard 加载 |
| trace_view.json |
| tensorboard --logdir=xxx |
+---------------------------+
当 tensorboard_trace_handler 被注册为 on_trace_ready 回调后，Profiler 会在数据采集完成后自动调用 NpuProfiler.analyse()，并指定分析类型为 TENSORBOARD_TRACE_HANDLER。随后，ProfilingParser 会依次执行 ParserConfig.COMMON_CONFIG 中注册的各类解析器，分别生成 trace 视图、算子详情、内核详情、内存记录及通信数据等文件，最终输出至指定目录供 TensorBoard 加载展示。
7. UML 类图设计
7.1 C++ 核心类图
C++ 核心层是 Profiler 的引擎，负责生命周期管理、数据采集与上报。核心类之间的关系如下：
关键设计说明：
ProfilerMgr采用单例模式，是 Profiler 的全局控制中心，负责协调数据接收、状态管理及资源释放。NpuProfilerThreadLocalState保存线程本地的 Profiler 配置与回调句柄，确保多线程环境下的隔离性。DataDumper与TraceDataDumper分别负责通用性能数据与 Python 追踪数据的异步落盘。BaseReportData作为数据基类，OpRangeData与MemoryData分别承载算子范围与内存事件信息。
7.2 Python 接口类图
Python 层封装了用户交互接口，负责将用户调用映射至 C++ 核心层：
关键设计说明：
profile类继承自_KinetoProfile，是用户最常使用的入口，支持上下文管理器与显式start/stop两种模式。Schedule类实现了灵活的采集调度策略，支持skip_first/wait/warmup/active/repeat等参数组合。ProfActionController负责根据调度结果驱动 Profiler 状态机迁移。NpuProfiler.analyse()为静态方法，负责启动数据解析流程。
7.3 数据解析类图
数据解析层将原始 Profiling 数据转换为可视化文件，各类解析器协同工作：
关键设计说明：
BaseParser为所有解析器提供统一的接口与公共属性。TraceViewParser是核心解析器，依赖FwkFileParser与CANNFileParser分别读取框架侧与 CANN 侧数据，并通过TreeBuilder构建算子树、TraceEventManager生成 Trace 事件。- 各视图解析器（Operator/Kernel/Memory/Stack/Communication）独立运行，输出对应的 CSV 或 JSON 文件。
8. 关键代码实现细节
8.1 ProfilerActivity 的实现
ProfilerActivity 枚举定义了 Profiler 可采集的活动类型，涵盖 CPU 端与 NPU 端算子追踪。
C++ 定义（npu_profiler.h）：
enum class NpuActivityType {
NONE = 0,
CPU, // CPU 端算子追踪
NPU, // NPU 端算子追踪
};
Python 绑定（init.cpp）：
py::enum_(m, "ProfilerActivity")
.value("CPU", NpuActivityType::CPU)
.value("NPU", NpuActivityType::NPU);
活动查询接口（profiler_interface.py）：
@no_exception_func(set())
def supported_activities():
return _supported_npu_activities()
8.2 NPU 钩子的实现
NPU Profiler 通过 RecordFunction 回调机制捕获 CPU 端算子执行事件。
回调注册（npu_profiler.cpp）：
static void registerCallback(const std::unordered_set &scopes)
{
auto registeration_state_ptr = NpuProfilerThreadLocalState::getTLS();
TORCH_INTERNAL_ASSERT(registeration_state_ptr, "Expected profiler state set");
auto handle = at::addThreadLocalCallback(
at::RecordFunctionCallback(
// 进入回调 - 记录算子开始
[](const at::RecordFunction &fn) -> std::unique_ptr {
auto state_ptr = NpuProfilerThreadLocalState::getTLS();
if (!state_ptr) {
return nullptr;
}
const auto &config = state_ptr->config();
auto ctx_ptr = state_ptr->newOpEvent(fn);
auto &data_ptr = ctx_ptr->data_;
// 记录输入形状和类型
if ((C10_UNLIKELY(config.report_input_shapes))) {
parseInputShapesAndDtypes(fn, data_ptr.get(), report_details);
}
// 记录调用栈
if (C10_UNLIKELY(config.with_stack && fn.scope() != at::RecordScope::BACKWARD_FUNCTION)) {
auto cs = torch::profiler::impl::prepareCallstack(torch::jit::currentCallstack());
data_ptr->stack = torch::profiler::impl::callstackStr(cs);
}
return ctx_ptr;
},
// 退出回调 - 记录算子结束
[](const at::RecordFunction &fn, at::ObserverContext *ctx_ptr) {
auto state_ptr = NpuProfilerThreadLocalState::getTLS();
if (!state_ptr) {
return;
}
auto *npu_ctx_ptr = static_cast(ctx_ptr);
auto data_ptr = std::move(npu_ctx_ptr->data_);
data_ptr->end_ns = static_cast(Utils::GetClockTime());
data_ptr->end_thread_id = Utils::GetTid();
// 上报数据
if (ProfilerMgr::GetInstance()->ReportEnable().load(std::memory_order_relaxed)) {
ProfilerMgr::GetInstance()->Upload(std::move(data_ptr));
}
}
)
.needsInputs(registeration_state_ptr->config().report_input_shapes)
.scopes(scopes)
);
registeration_state_ptr->setCallbackHandle(handle);
}
8.3 日志记录机制
Profiler 采用异步数据转储机制，最小化对训练性能的影响。
DataDumper 类（data_dumper.h）：
class DataDumper : public Thread {
public:
void Init(const std::string &path, size_t capacity);
void Report(std::unique_ptr data);
void Start();
void Stop();
private:
std::string path_;
std::atomic start_;
std::atomic init_;
RingBuffer> data_chunk_buf_; // 环形缓冲区
std::map fd_map_;
};
数据上报流程（profiler_mgr.cpp）：
void ProfilerMgr::Upload(std::unique_ptr data)
{
if (enable_warmup_.load()) {
return; // warmup 期间不记录数据
}
dataReceiver_.Report(std::move(data));
}
void ProfilerMgr::UploadWithLock(std::unique_ptr data)
{
if (enable_warmup_.load()) {
return;
}
std::lock_guard lock(reportDataMutex_);
dataReceiverWithLock_.Report(std::move(data));
}
8.4 Python Tracer 实现
Python Tracer 利用 CPython 的 Profile 机制追踪 Python 函数调用。
核心实现（profiler_python.cpp）：
class PythonTracer final {
public:
static int pyProfileFn(
PyObject* obj,
PyFrameObject* frame,
int what,
PyObject* arg);
private:
void start(size_t max_threads = max_py_threads);
void stop();
void recordPyCall(TraceContext* ctx, PyFrameObject* frame);
void recordCCall(TraceContext* ctx, PyFrameObject* frame, PyObject* arg);
void recordReturn(TraceContext* ctx, PyFrameObject* frame, TraceTag tag);
std::atomic active_{false};
PyInterpreterState* interpreter_{nullptr};
std::unordered_map py_call_cache_;
AppendOnlyList events_;
};
// Profile 回调函数
int PythonTracer::pyProfileFn(
PyObject* obj,
PyFrameObject* frame,
int what,
PyObject* arg)
{
auto ctx = reinterpret_cast(obj);
auto active_tracer = ctx->thread_local_result_->active_tracer_;
switch (what) {
case PyTrace_CALL:
active_tracer->recordPyCall(ctx, frame);
break;
case PyTrace_C_CALL:
active_tracer->recordCCall(ctx, frame, arg);
break;
case PyTrace_RETURN:
active_tracer->recordReturn(ctx, frame, TraceTag::kPy_Return);
break;
case PyTrace_C_RETURN:
active_tracer->recordReturn(ctx, frame, TraceTag::kC_Return);
break;
}
return 0;
}
8.5 调度器实现
调度器控制 Profiler 的采集周期，支持灵活的采样策略。
文件路径：scheduler.py
class ProfilerAction(Enum):
NONE = 0 # 不采集
WARMUP = 1 # 预热阶段
RECORD = 2 # 正常记录
RECORD_AND_SAVE = 3 # 记录并保存
class Schedule:
"""
调度器控制 Profiler 的采集周期:
skip_first -> wait -> warmup -> active (repeat)
"""
def __init__(self, wait: int, active: int, warmup: int = 0,
repeat: int = 0, skip_first: int = 0):
self.wait = wait
self.active = active
self.warmup = warmup
self.repeat = repeat
self.skip_first = skip_first
def __call__(self, step: int) -> ProfilerAction:
if step < self.skip_first:
return ProfilerAction.NONE
step -= self.skip_first
num_steps = self.wait + self.warmup + self.active
if self.repeat > 0 and step / num_steps >= self.repeat:
return ProfilerAction.NONE
mod_step = step % num_steps
if mod_step < self.wait:
return ProfilerAction.NONE
elif mod_step < self.wait + self.warmup:
return ProfilerAction.WARMUP
else:
return (ProfilerAction.RECORD
if mod_step < num_steps - 1
else ProfilerAction.RECORD_AND_SAVE)
8.6 状态转换控制器
状态转换控制器负责根据调度结果驱动 Profiler 状态机迁移。
文件路径：_profiler_action_controller.py
class ProfActionController:
def __init__(self, prof, prof_inst, on_trace_ready):
self.prof = prof
self.prof_inst = prof_inst
self.action_map = self._init_action_map()
self.on_trace_ready = on_trace_ready
9. 状态机调度与工程实践
9.1 状态机调度机制
在 Profiler 的调度系统中，状态机是核心的调度控制机制。通过定义明确的动作状态和转换规则，状态机能够精确控制 Profiler 在不同阶段的启停行为，确保数据采集的准确性和完整性。
9.1.1 状态定义
Profiler 的调度状态机定义了以下核心状态：
| 状态 | 说明 |
|---|---|
NONE | 空闲状态，Profiler 未执行任何采集动作 |
WARMUP | 预热阶段，执行初始化但不记录数据 |
RECORD | 正式采集阶段，记录性能数据 |
RECORD_AND_SAVE | 采集并保存数据阶段 |
9.1.2 状态转换实现
状态转换通过动作映射表驱动，每个状态对对应一组待执行的动作：
def transit_action(self, prev_action, current_action):
action_list = self.action_map.get((prev_action, current_action))
if action_list:
for action in action_list:
action()
def _init_action_map(self):
return {
# NONE -> WARMUP: 初始化并预热
(ProfilerAction.NONE, ProfilerAction.WARMUP): [
self.prof_inst.init_trace,
self.prof_inst.warmup_trace
],
# WARMUP -> RECORD: 开始正式记录
(ProfilerAction.WARMUP, ProfilerAction.RECORD): [
self.prof_inst.start_trace
],
# RECORD_AND_SAVE -> NONE: 停止并保存
(ProfilerAction.RECORD_AND_SAVE, ProfilerAction.NONE): [
self.prof_inst.stop_trace,
self.prof_inst.finalize_trace,
self._trace_ready
],
# ... 其他状态转换
}
9.2 关键文件路径汇总
| 功能模块 | 文件路径 |
|---|---|
| C++ 核心实现 | — |
| Profiler 主入口 | torch_npu/csrc/profiler/npu_profiler.cpp |
| Profiler 头文件 | torch_npu/csrc/profiler/npu_profiler.h |
| Profiler 管理器 | torch_npu/csrc/profiler/profiler_mgr.cpp |
| Python Tracer | torch_npu/csrc/profiler/profiler_python.cpp |
| MSTX 管理器 | torch_npu/csrc/profiler/mstx_mgr.cpp |
| 特性管理器 | torch_npu/csrc/profiler/feature_mgr.cpp |
| Python 绑定 | torch_npu/csrc/profiler/init.cpp |
| 数据转储器 | torch_npu/csrc/toolkit/profiler/inc/data_dumper.h |
| 数据上报结构 | torch_npu/csrc/toolkit/profiler/inc/data_reporter.h |
| Python 接口 | — |
| 模块入口 | torch_npu/profiler/__init__.py |
| Profile 类 | torch_npu/profiler/profiler.py |
| 调度器 | torch_npu/profiler/scheduler.py |
| 实验配置 | torch_npu/profiler/experimental_config.py |
| Profiler 接口 | torch_npu/profiler/profiler_interface.py |
| 动作控制器 | torch_npu/profiler/_profiler_action_controller.py |
| 路径创建器 | torch_npu/profiler/_profiler_path_creator.py |
| 数据解析 | — |
| 解析入口 | torch_npu/profiler/analysis/_npu_profiler.py |
| 解析器主类 | torch_npu/profiler/analysis/_profiling_parser.py |
| 框架文件解析 | torch_npu/profiler/analysis/prof_parse/_fwk_file_parser.py |
| Trace 视图解析 | torch_npu/profiler/analysis/prof_view/_trace_view_parser.py |
| 常量定义 | torch_npu/profiler/analysis/prof_common_func/_constant.py |
9.3 特性与优化方向
9.3.1 核心优势
| 优势 | 说明 |
|---|---|
| 完整的 PyTorch 兼容性 | 采用与 PyTorch Profiler 相同的 API 设计，用户学习成本低 |
| 多级别追踪支持 | 支持 Level0/Level1/Level2 等不同精度的追踪，用户可按需选择 |
| 异步数据采集 | 使用 RingBuffer + 后台线程异步写入，最小化对训练性能的影响 |
| 灵活的调度机制 | 支持 skip_first/wait/warmup/active/repeat 等多种调度参数 |
| 完善的 Python 追踪 | 支持 Python 函数调用栈追踪，便于定位性能瓶颈 |
| TensorBoard 集成 | 原生支持 TensorBoard 可视化，便于分析 |
| 多进程解析 | 支持多进程并行解析，提高大规模数据解析效率 |
| 丰富的视图输出 | 支持 Trace View、Operator View、Memory View 等多种视图 |
| GC 检测 | 支持 Python GC 事件检测，帮助分析内存问题 |
| MSTX 标记 | 支持用户自定义标记点和区间，便于精细化分析 |
9.3.2 后续演进方向
性能优化
- 采样机制：引入采样率参数，允许用户只记录部分算子；实现基于时间窗口的采样策略
- 增量式数据转储：支持边采集边写入的流式转储，进一步降低数据落盘对内存的占用
- 压缩存储：对重复的字符串（如算子名、调用栈）进行去重存储；支持 gzip 压缩输出
功能增强
- 实时监控：提供实时 profiling 数据流接口；支持 Prometheus/Grafana 集成
- 智能分析：自动识别性能瓶颈；提供优化建议
- 分布式支持增强：自动聚合多 rank 数据；支持集群级别的性能分析
- 更细粒度的控制：支持按算子名过滤；支持按模块过滤
易用性改进
- 错误处理与提示：进一步完善参数校验与错误提示，提升使用友好性
- 文档完善：提供更多使用示例；增加 FAQ 文档
- CLI 工具：提供命令行分析工具；支持离线数据查看
代码质量
- 单测与集成：持续完善 C++ 层单元测试与集成测试，保障模块长期质量
- 代码重构：持续统一错误码定义、精简重复代码，提升可维护性
9.4 Python API 使用指南
9.4.1 基础用法
import torch
import torch_npu
from torch_npu.profiler import profile, ProfilerActivity
# 基础用法 - 使用 context manager
model = MyModel().npu()
inputs = torch.randn(32, 3, 224, 224).npu()
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.NPU],
record_shapes=True,
profile_memory=True,
with_stack=True,
) as prof:
for _ in range(10):
output = model(inputs)
# 导出 Chrome Trace 格式
prof.export_chrome_trace("trace.json")
9.4.2 使用调度器控制采集周期
import torch
import torch_npu
from torch_npu.profiler import (
profile,
ProfilerActivity,
schedule,
tensorboard_trace_handler
)
model = MyModel().npu()
inputs = torch.randn(32, 3, 224, 224).npu()
# 使用调度器: 跳过前5步,等待1步,预热1步,采集3步,重复2次
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.NPU],
schedule=schedule(
skip_first=5,
wait=1,
warmup=1,
active=3,
repeat=2
),
on_trace_ready=tensorboard_trace_handler(
dir_name="./profiler_logs",
worker_name="worker0"
),
record_shapes=True,
profile_memory=True,
) as prof:
for step in range(20):
output = model(inputs)
prof.step() # 必须调用 step() 来触发调度
9.4.3 使用实验配置
import torch
import torch_npu
from torch_npu.profiler import (
profile,
ProfilerActivity,
_ExperimentalConfig,
ProfilerLevel,
AiCMetrics,
ExportType
)
# 配置高级选项
experimental_config = _ExperimentalConfig(
profiler_level=ProfilerLevel.Level1, # 使用 Level1 追踪级别
aic_metrics=AiCMetrics.PipeUtilization, # AI Core 流水线利用率
l2_cache=True, # 开启 L2 缓存统计
mstx=True, # 开启 MSTX 标记
data_simplification=True, # 数据简化
export_type=[ExportType.Text, ExportType.Db], # 同时导出文本和数据库
gc_detect_threshold=0.1, # GC 检测阈值(秒)
)
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.NPU],
experimental_config=experimental_config,
record_shapes=True,
profile_memory=True,
with_stack=True,
with_modules=True,
) as prof:
for _ in range(10):
output = model(inputs)
9.4.4 子线程 Profiling
import threading
import torch
import torch_npu
from torch_npu.profiler import profile, ProfilerActivity
def worker_thread():
# 在子线程中启用 profiler
profile.enable_profiler_in_child_thread(
record_shapes=True,
profile_memory=True,
with_stack=True
)
try:
# 执行计算
x = torch.randn(100, 100).npu()
y = torch.matmul(x, x)
finally:
# 禁用 profiler
profile.disable_profiler_in_child_thread()
# 主线程开启 profiler
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.NPU],
) as prof:
# 创建并启动子线程
thread = threading.Thread(target=worker_thread)
thread.start()
thread.join()
9.4.5 使用 MSTX 标记
import torch
import torch_npu
from torch_npu.profiler import profile, ProfilerActivity, _ExperimentalConfig
# MSTX 标记需要开启 mstx 选项
experimental_config = _ExperimentalConfig(mstx=True)
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.NPU],
experimental_config=experimental_config,
) as prof:
# 使用 record_function 添加标记
with torch.profiler.record_function("my_custom_region"):
x = torch.randn(1000, 1000).npu()
y = torch.matmul(x, x)
# 另一个标记
with torch.profiler.record_function("another_region"):
z = torch.relu(y)
9.4.6 导出内存时间线
import torch
import torch_npu
from torch_npu.profiler import profile, ProfilerActivity
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.NPU],
record_shapes=True,
profile_memory=True,
with_stack=True, # 内存分析需要调用栈
) as prof:
for _ in range(10):
x = torch.randn(1000, 1000).npu()
y = torch.matmul(x, x)
del x, y
torch_npu.npu.empty_cache()
# 导出内存时间线
prof.export_memory_timeline("memory_timeline.html", device="npu:0")
9.4.7 离线数据分析
from torch_npu.profiler import analyse
# 离线分析已采集的 profiling 数据
analyse(
profiler_path="./profiler_logs/worker0_20240101120000_ascend_pt",
max_process_number=8, # 使用 8 个进程并行解析
export_type=["text", "db"] # 导出文本和数据库格式
)
9.4.8 完整训练循环示例
import torch
import torch.nn as nn
import torch.optim as optim
import torch_npu
from torch_npu.profiler import (
profile,
ProfilerActivity,
schedule,
tensorboard_trace_handler,
_ExperimentalConfig,
ProfilerLevel
)
# 模型定义
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(784, 256)
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = torch.relu(self.fc1(x))
return self.fc2(x)
# 初始化
device = torch.device("npu:0")
model = SimpleModel().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
# 配置 Profiler
experimental_config = _ExperimentalConfig(
profiler_level=ProfilerLevel.Level1,
data_simplification=True,
)
# 训练循环
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.NPU],
schedule=schedule(skip_first=10, wait=5, warmup=1, active=3, repeat=2),
on_trace_ready=tensorboard_trace_handler("./tb_logs"),
record_shapes=True,
profile_memory=True,
with_stack=True,
experimental_config=experimental_config,
) as prof:
for epoch in range(3):
for batch_idx in range(100):
# 模拟数据
data = torch.randn(32, 784).to(device)
target = torch.randint(0, 10, (32,)).to(device)
# 前向传播
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
# 反向传播
loss.backward()
optimizer.step()
# 重要: 调用 step() 触发调度器
prof.step()
print("Profiling 完成! 使用 TensorBoard 查看结果:")
print("tensorboard --logdir=./tb_logs")
10. 总结
本文系统解析了 torch_npu 性能分析模块（Profiler）的内部实现，涵盖模块架构与整体流程、核心数据采集与导出链路、TensorBoard 集成方案、关键类图设计与代码实现，以及状态机调度机制与工程实践。通过分层架构、环形缓冲与异步写入、多进程解析等机制，该模块在保证性能数据完整性的同时，尽可能降低了对训练任务的干扰，并为开发者提供从基础用法、调度控制、实验配置到离线分析的一整套使用路径。开发者可结合本文的分析思路与 API 示例，针对具体业务场景在昇腾 NPU 上开展性能分析与优化。
附录 A: 数据文件格式
A.1 二进制文件结构
torch.op_range 文件采用 TLV (Type-Length-Value) 编码：
+------------------+------------------+------------------+
| Type (2 bytes) | Length (4 bytes) | Value (N bytes) |
+------------------+------------------+------------------+
A.2 输出目录结构
profiler_output/
├── FRAMEWORK/ # 框架层数据
│ ├── torch.op_range # 算子范围数据
│ ├── torch.op_mark # 标记点数据
│ ├── torch.memory_usage # 内存使用数据
│ ├── torch.python_tracer_func # Python 函数追踪
│ └── torch.python_tracer_hash # Python 函数哈希表
├── PROF_*/ # CANN 层数据
│ └── ...
├── ASCEND_PROFILER_OUTPUT/ # 解析输出
│ ├── trace_view.json # Chrome Trace 格式
│ ├── operator_details.csv # 算子详情
│ ├── kernel_details.csv # Kernel 详情
│ ├── memory_record.csv # 内存记录
│ └── ...
├── profiler_info.json # Profiler 配置信息
└── profiler_metadata.json # 元数据
附录 B: 常见问题与解决方案
Q1: 如何控制 Profiling 采集带来的额外开销？
A：性能分析工具在采集数据时会产生一定的额外开销，这是 Profiling 机制的通用特性，可通过以下方式将其控制在合理范围：
- 减少 profiling 范围，使用 schedule 限制采集步数
- 按需关闭 with_stack 和 Python Tracer
- 按需选择较低的追踪级别（如 Level0）
Q2: 无法生成 trace_view.json？
A：
- 检查输出目录是否可写
- 确认 profiling 正常结束（调用了 stop）
- 查看是否在 daemon 进程中
Q3: 子线程数据没有被采集？
A：
需要在子线程中显式调用：
profile.enable_profiler_in_child_thread()
# ... 执行代码 ...
profile.disable_profiler_in_child_thread()
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。
更多推荐
0
0- 0
已为社区贡献6条内容
所有评论(0)