oam-tools msprof 实战：用 torch_npu.profiler API 在 PyTorch 训练/推理代码内插桩采集 step 级性能数据

【免费下载链接】oam-tools 本项目为开发者提供故障定位工具，包含故障信息收集，软硬件信息展示，AI core error报错分析等能力，提升故障问题定位效率，文档可在昇腾社区搜索“故障处理简介”（选择社区版）。 项目地址: https://gitcode.com/cann/oam-tools

本篇技术指南讲解 CANN oam-tools 仓库中 msprof 性能采集体验 Demo 的第四种方式——在 PyTorch 代码中通过 torch_npu.profiler.profile 进行白盒插桩采集。与 msprof 命令行（CLI）黑盒旁路采集不同，API 方式可以精确圈定要采集的 step 区间、拿到 aten op 级耗时与 Python 调用栈，并额外产出 CLI 没有的 step_trace_time.csv（每 step 的 Computing/Free 拆分）。读完本文，你将掌握 profile/schedule/tensorboard_trace_handler 的完整用法，学会把 profiler 插桩迁移到自己的模型与训练循环，并能基于 op_statistic.csv、step_trace_time.csv、api_statistic.csv 定位"算力瓶颈"还是"host 下发瓶颈"。

一、为什么需要 API 插桩：白盒 vs 黑盒

在 msprof 性能采集体验 Demo 中，共演示了四种性能数据采集方式：

#	方式	特点	适用场景
01	msprof 命令行	黑盒，旁路拉起应用，不改代码	有个能跑的程序，想快速看一眼
02	AscendC API	自定义算子核函数直调	写了算子，想看硬件利用率
03	pyACL API	加载 .om 离线模型	要部署离线模型
04	torch_npu.profiler API	白盒，代码内插桩	在跑 PyTorch，想要 step 级数据

本实验（04_pyTorch）属于第四种。它的核心思路是：在训练/推理循环里直接调用 torch_npu.profiler.profile 插桩，实现白盒采集。相对 CLI（见 01_cmdline 的说明），API 方式具备三个 CLI 不具备的能力：

1. 精确圈定采集区间：通过 schedule(wait/warmup/active) 只采集第 N~M 个 step，不采集预热阶段；
2. 拿到 aten op 级耗时与 Python 调用栈：CLI 是旁路黑盒，看不到 PyTorch aten op 的存在；API 插桩则能记录每个算子及其 Python 栈；
3. 额外产出 step_trace_time.csv：给出每个 step 的 Computing（计算）/ Free（等待 host 下发）占比，这是 CLI 没有的。

二、实验环境与文件结构

Demo 总 README（experiment/task-book/msprof_experience_demo/README.md）给出的本实验运行环境为：

• Atlas A2 训练系列（910B3）
• CANN 9.1.0
• torch_npu 2.7.1
• Python 3.12

运行前需先 source <CANN路径>/set_env.sh（脚本会自动定位，找不到会提示）。注意：API 方式要求当前环境已安装 torch 与 torch_npu，并且只有具备 NPU 真实计算负载的程序才能采集到有意义的性能数据。

04_pyTorch 目录下仅两个文件：

文件	作用
src/model_with_profiler.py	同一个 TinyMLP + profiler 插桩
run.sh	采集脚本，跑完打印 kernel Top + step 拆分

实验负载是一个最简 TinyMLP：3 层 Linear(1024,1024) + GELU 再加一层 Linear(1024,1024)，输入 [32, 1024]，与 01_cmdline 使用同一模型，方便横向对比两种采集方式的结果是否自洽。

三、插桩代码逐段解析

完整代码见 src/model_with_profiler.py，下面拆解关键部分。

3.1 导入 profiler 组件

from torch_npu.profiler import (
    AiCMetrics,          # AI Core 硬件指标口径
    ProfilerActivity,    # 采集侧：CPU / NPU
    ProfilerLevel,       # 采集层级：Level0 / Level1 / Level2
    profile,             # profiler 上下文管理器
    schedule,            # step 采集窗口调度
    tensorboard_trace_handler,  # 输出目录处理器
)

3.2 构造 profiler：四个可调"旋钮"

源码把采集配置集中封装在 make_profiler(out_dir) 中，方便替换：

level = ProfilerLevel.Level1                      # 采集层级
pmu = AiCMetrics.PipeUtilization                  # AI Core 指标口径
acts = [ProfilerActivity.CPU, ProfilerActivity.NPU]  # 采集两侧

exp_config_cls = getattr(torch_npu.profiler, "_ExperimentalConfig")
trace_cfg = exp_config_cls(
    profiler_level=level,
    aic_metrics=pmu,
    data_simplification=False,   # 保留原始 trace；生产场景可设 True 省空间
)
window = schedule(wait=0, warmup=WARMUP_STEPS, active=ACTIVE_STEPS, repeat=1)

return profile(
    activities=acts,
    schedule=window,
    on_trace_ready=tensorboard_trace_handler(out_dir),
    experimental_config=trace_cfg,
    record_shapes=True,   # 记录算子 input shape
    with_stack=True,      # 记录 Python 调用栈
)

各参数的实战含义：

• ProfilerLevel：采集层级。Level1 包含 AI Core PMU 硬件指标；想更轻量可降为 Level0（仅软件栈耗时）。
• AiCMetrics：AI Core 指标口径。示例用 PipeUtilization 看各计算流水线占用情况，是判断算子是否打满硬件的关键口径。
• ProfilerActivity.CPU + NPU：两侧都采，才能看清"host 下发 vs NPU 执行"之间的 gap。
• _ExperimentalConfig：torch_npu 仅以下划线前缀的 _ExperimentalConfig 暴露采集配置（profiler_level、aic_metrics、data_simplification 等）。源码用 getattr 按名取用，避免直接书写受保护成员。
• schedule(wait=0, warmup=3, active=5, repeat=1)：采集窗口调度，含义见下一节。
• record_shapes=True / with_stack=True：分别记录算子输入 shape 与 Python 调用栈，便于在时间线上定位算子来源。

3.3 训练循环：with prof 块 + prof.step()

total_steps = WARMUP_STEPS + ACTIVE_STEPS   # 3 + 5 = 8
profiler = make_profiler(out_dir)
with torch.no_grad(), profiler as prof:
    for _ in range(total_steps):
        model(x)
        torch.npu.synchronize()
        prof.step()   # 通知 profiler 进入下一 step

WARMUP_STEPS = 3、ACTIVE_STEPS = 5（源码顶部常量定义）。schedule(wait=0, warmup=3, active=5, repeat=1) 表示：前 3 个 step 为预热不采集，接下来 5 个 step 为活跃采集窗口，整体只重复 1 轮。由于实验是推理场景（model.eval() + torch.no_grad()），循环体是前向推理；如果是训练循环，则应把 forward + backward + optimizer.step() 放进 with prof: 块内，并在每个训练 step 末尾调用 prof.step()。

3.4 输出目录处理

on_trace_ready=tensorboard_trace_handler(out_dir) 把采集结果写入指定目录（默认为 ./prof_out，可用命令行参数覆盖）。采集完成后数据落在 <out_dir>/*_ascend_pt/ASCEND_PROFILER_OUTPUT/ 下，其中就包含本实验重点关注的 step_trace_time.csv、op_statistic.csv、api_statistic.csv。

四、运行方法

直接执行采集脚本即可，默认使用 device 7：

bash run.sh        # 默认 device 7
bash run.sh 5      # 指定卡号

run.sh 的核心逻辑（set -e 保证出错即停）：

DEV=${1:-7}
HERE=$(cd "$(dirname "$0")" && pwd)
OUT="$HERE/prof_out"
rm -rf "$OUT" && mkdir -p "$OUT"     # 幂等：先清输出目录再重采

ASCEND_VISIBLE_DEVICES=$DEV python3 "$HERE/src/model_with_profiler.py" "$OUT"

# 跑完自动打印 step 拆分（API 独有）
cat "$OUT"/*_ascend_pt/ASCEND_PROFILER_OUTPUT/step_trace_time.csv 2>/dev/null

脚本幂等设计：每次运行先清空再重建 prof_out，避免旧数据干扰；采集完成后自动 cat 出 step_trace_time.csv，让你第一时间看到 step 级拆分。

五、如何用到你自己的模型

Demo 作者在代码里用 ← 注释标出了所有需要替换的位置，迁移成本极低：

1. 替换模型与输入：把 build_model() 换成你的模型结构，把 x = torch.randn(BATCH, HIDDEN, device=device) 换成你模型的真实输入（两处均在源码中标注 ← 换成你的模型 / ← 换成你模型的输入）。
2. 改造训练循环：若是训练场景，把 forward + backward + optimizer.step() 放进 with prof: 块内，每个 step 末尾调 prof.step()，用 schedule(wait/warmup/active) 圈定要采集的 step（warmup 用来跳过预热、等算子自动调优稳定，active 才是真正采集的窗口）。
3. 采集配置不用动：make_profiler 那一段和 run.sh 都可以原样保留，直接 bash run.sh 7 即可。

六、实测结果与解读（910B3 示例）

以下是 Demo 在 910B3 上实际跑出的示例数据，用于说明每个输出文件"怎么读"。

6.1 算子聚合：op_statistic.csv

OP Type	Core Type	Count	Total(us)	Ratio	CLI 对照
MatMulV2	AI_CORE	20	136.50	78.30%	77.88%
Gelu	AI_VECTOR_CORE	15	37.82	21.70%	21.45%

两个观察点：

• 与 CLI 高度一致：CLI（01_cmdline）实测 MatMulV2 占 77.88%、Gelu 占 21.45%，与 API 方式的 78.30% / 21.70% 基本吻合，验证了两种采集方式的数据自洽性。
• 热点明确：MatMul 占绝对主导，符合全连接网络的预期。op_summary_*.csv 中含完整 AI Core PMU 指标，可进一步定位单算子卡在哪个计算单元。

6.2 step 拆分：step_trace_time.csv（API 独有）

Step	Computing(us)	Free(us)	解读
3	35.64	2214.72	Computing : Free ≈ 1 : 62
5	34.40	2113.74	≈ 1 : 61
7	35.52	1384.98	≈ 1 : 39

关键洞察：NPU 实际计算每个 step 仅约 35us，但 Free（等待 host 下发）高达约 2000us，是典型的 host bound——TinyMLP 模型太小，host 下发开销完全盖过了计算。这正是 API 模式相对 CLI 的核心价值：CLI 只告诉你"MatMul 占 78%"，API 进一步告诉你"整个 step 99% 时间在等 host"。

换成你的模型后，重点看 step_trace_time.csv 的 Computing:Free——它直接回答"是算力不够（Computing 高）还是 host 下发拖后腿（Free 高）"。

6.3 host 下发 Top：api_statistic.csv

API Name	Time(us)	Count	Avg(us)
aclnnAddmm	259.56	20	12.98
aclrtLaunchKernelWithHostArgs	229.12	35	6.55
aclnnGelu	200.51	15	13.37

api_statistic.csv 给出 host 侧 API 调用的耗时分布：aclnnAddmm（对应 MatMul 的 aclnn 接口）、aclrrtLaunchKernelWithHostArgs（kernel 下发）、aclnnGelu 位居前列。结合 6.2 的 Free 数据，可以进一步把"等 host"归因到具体是"算子下发"还是"其他同步开销"，为针对性优化（如算子融合、减少 host 侧同步、加大 batch 提升计算占比）提供直接依据。

七、与 msprof CLI 方式的互补关系

本实验与仓库中的 msprof 命令行采集形成互补。CLI 方式（参考 msprof 采集通用命令 的命令格式 msprof [options] <app> 与 examples/msprof/run.sh 的 --output 等参数用法）不改一行模型代码，作为父进程拉起应用旁路采集，适合快速看全局热点；而 API 方式牺牲了"零侵入"，换来 step 级视角与 Python 栈。实践中的建议是：先用 CLI 快速确认热点算子，再用 API 插桩深入分析 step 内"计算 vs 等待"的时间构成，两者结合即可形成完整的性能分析闭环。

八、小结

通过 04_pyTorch 这个 Demo，你可以快速掌握 torch_npu.profiler API 采集的完整套路：

1. 插桩三要素：profile(activities, schedule, on_trace_ready) + 每个 step 末尾 prof.step() + schedule(wait/warmup/active) 圈定采集窗口；
2. 迁移路径清晰：只替换 build_model() 与输入、把训练循环包进 with prof: 块，采集配置与脚本可复用；
3. 三个核心产物：op_statistic.csv（热点算子）、step_trace_time.csv（计算 vs 等待拆分，API 独有）、api_statistic.csv（host 下发 API 耗时）；
4. 判断口径明确：Computing:Free 比值直接给出"算力 bound 还是 host bound"的结论。

如果要在真实模型上做性能分析，可以从 model_with_profiler.py 出发，按第五节的三步完成迁移，并用第六节的读数方法定位瓶颈。

【免费下载链接】oam-tools 本项目为开发者提供故障定位工具，包含故障信息收集，软硬件信息展示，AI core error报错分析等能力，提升故障问题定位效率，文档可在昇腾社区搜索“故障处理简介”（选择社区版）。 项目地址: https://gitcode.com/cann/oam-tools