pytorch-profiling-collection

name: pytorch-profiling-collection description: 使用 torch_npu.profiler 在非 MindSpeed-LLM 与非 MindSpeed-MM 的训练/推理脚本中采集 Ascend NPU Profiling 数据。覆盖 level0/level1/level2 采集级别、训练循环/单次推理/指定代码段三种方式，支持 .py 和 .sh 脚本路径。当用户需要性能采集、性能分析、查看算子耗时、定位训练瓶颈时使用。触发关键词：profiling、性能采集、性能分析、算子耗时、瓶颈定位、torch_npu.profiler。

Ascend NPU Profiling 采集指南

本 Skill 适用于非 MindSpeed-LLM 和非 MindSpeed-MM 的通用场景，使用 torch_npu.profiler 采集性能数据。

第 0 步：确认前置条件

回复用户：

开始 Profiling 采集前，请提供以下信息：

必要：

• 训练/推理脚本路径 — 可运行的 Python 脚本（.py）或 shell 启动脚本（.sh，常见于 deepspeed / torchrun 启动）

可选（有默认值）：

• 采集级别：level0 / level1（推荐）/ level2
• 采集步数（active）：默认 3
• CPU 采集：默认开启
• 内存采集：默认关闭
• Tensor Shape 记录：默认开启
• 堆栈采集（with_stack）：默认关闭
• 输出目录：默认 ./profiling_result
• 采集方式：训练循环 / 单次推理 / 指定代码段（默认训练循环）

用户未提供脚本路径时：停止。回复："Profiling 需要可运行的 Python 脚本，请先准备好再进行采集。"

用户输入与变量映射

第 1 步：展示配置确认表

收集完信息后，展示确认表，用户确认后再执行。

确认后进入第 2 步。

第 2 步：检查环境

source /usr/local/Ascend/ascend-toolkit/set_env.sh 2>/dev/null || source /usr/local/Ascend/cann/set_env.sh 2>/dev/null
npu-smi info
python -c "import torch; import torch_npu; print(f'torch_npu: {torch_npu.__version__}'); print(f'NPU available: {torch.npu.is_available()}')"

• NPU 不可用 → 停止："当前环境未检测到 NPU，请在有 Ascend NPU 的环境中运行。"
• torch_npu 不可用 → 停止："请安装 torch_npu 后再试。"

第 3 步：创建 Profiling 脚本并执行

禁止修改用户的原始脚本。 基于原始脚本创建新脚本（带时间戳，如 your_script_profiling_20260101120000.py）。

3.0 检测脚本类型

# 检测用户提供的脚本后缀
SCRIPT_PATH="{USER_SCRIPT_PATH}"
if [[ "$SCRIPT_PATH" == *.py ]]; then
    SCRIPT_TYPE="py"
elif [[ "$SCRIPT_PATH" == *.sh ]]; then
    SCRIPT_TYPE="sh"
else
    echo "错误：不支持的文件类型，请提供 .py 或 .sh 脚本"
    exit 1
fi

• .py → 走分支 A（注入代码 → python new.py）
• .sh → 走分支 B（提取引用的 .py → 注入代码 → 修改 run.sh 中的路径 → bash new.sh）

分支 A：用户提供 .py 脚本

A1. 代码注入模板

import torch
import torch_npu

experimental_config = torch_npu.profiler._ExperimentalConfig(
    aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
    profiler_level=torch_npu.profiler.ProfilerLevel.{LEVEL},
    mstx=False, l2_cache=False, op_attr=False, data_simplification=False,
    record_op_args=False, gc_detect_threshold=None, host_sys=[], sys_io=False, sys_interconnection=False
)

prof = torch_npu.profiler.profile(
    activities=[{ACTIVITIES}],
    schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active={ACTIVE}, repeat=1, skip_first=1, skip_first_wait=0),
    on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("{OUTPUT_DIR}"),
    record_shapes={RECORD_SHAPES},
    profile_memory={PROFILE_MEMORY},
    with_stack={WITH_STACK},
    with_modules=False, with_flops=False,
    experimental_config=experimental_config
)

A2. 三种采集方式

方式 A — 训练循环采集（默认）：

prof.start()
for step in range(total_steps):
    # 原始训练逻辑
    prof.step()
prof.stop()

方式 B — 单次推理（无训练循环）：

prof.start()
# 原始推理逻辑
prof.stop()

方式 C — 指定代码段：

prof.start()
# 目标代码段
prof.stop()

A3. 运行新脚本

python {PROFILING_SCRIPT_PATH}

运行成功 → 进入第 4 步。失败则报告错误并停止。

分支 B：用户提供 .sh 脚本

B1. 复制 shell 脚本

TIMESTAMP=$(date +%Y%m%d%H%M%S)
SCRIPT_DIR=$(dirname "$SCRIPT_PATH")
SCRIPT_BASENAME=$(basename "$SCRIPT_PATH" .sh)

NEW_SH="${SCRIPT_DIR}/${SCRIPT_BASENAME}_profiling_${TIMESTAMP}.sh"
cp "$SCRIPT_PATH" "$NEW_SH"

B2. 提取被引用的 Python 脚本路径

从 run.sh 中找到主入口 Python 脚本路径，优先级如下：

1. --module 参数（deepspeed 风格）：提取模块名，转成 .py 路径
2. python / torchrun / deepspeed 命令后的 .py 文件：提取文件路径

# 尝试提取 --module 参数（deepspeed 风格）
MODULE=$(grep -oP '(?<=--module )[^ ]+' "$SCRIPT_PATH" 2>/dev/null || true)
if [[ -n "$MODULE" ]]; then
    # 将模块名转换为文件路径，如 align_anything.trainers.xxx → align_anything/trainers/xxx.py
    ORIG_PY_PATH=$(echo "$MODULE" | tr '.' '/').py
    # 如果相对路径，基于脚本所在目录查找
    if [[ ! -f "$ORIG_PY_PATH" && -f "${SCRIPT_DIR}/${ORIG_PY_PATH}" ]]; then
        ORIG_PY_PATH="${SCRIPT_DIR}/${ORIG_PY_PATH}"
    fi
else
    # 降级：提取 python/torchrun/deepspeed 后的第一个 .py 文件
    ORIG_PY_PATH=$(grep -oP '(?<=python |torchrun |deepspeed )\S+\.py' "$SCRIPT_PATH" 2>/dev/null | head -1)
    if [[ -n "$ORIG_PY_PATH" && ! -f "$ORIG_PY_PATH" && -f "${SCRIPT_DIR}/${ORIG_PY_PATH}" ]]; then
        ORIG_PY_PATH="${SCRIPT_DIR}/${ORIG_PY_PATH}"
    fi
fi

if [[ -z "$ORIG_PY_PATH" || ! -f "$ORIG_PY_PATH" ]]; then
    echo "错误：无法从 $SCRIPT_PATH 中提取有效的 Python 脚本路径"
    exit 1
fi

B3. 对提取的 Python 脚本注入 profiling 代码

使用分支 A 的注入模板（A1 + A2），在原始 .py 同级目录创建 profiling 版本：

ORIG_PY_DIR=$(dirname "$ORIG_PY_PATH")
ORIG_PY_BASENAME=$(basename "$ORIG_PY_PATH" .py)
NEW_PY="${ORIG_PY_DIR}/${ORIG_PY_BASENAME}_profiling_${TIMESTAMP}.py"

# 注入 profiling 代码到新 .py 文件（复用 A1/A2 的注入逻辑）
# 注入完成后得到 NEW_PY

B4. 修改新的 shell 脚本

将 NEW_SH 中引用的原 Python 路径替换为注入后的路径：

# 替换 --module 参数（deepspeed 风格）
if [[ -n "$MODULE" ]]; then
    NEW_MODULE="${ORIG_PY_BASENAME}_profiling_${TIMESTAMP}"
    sed -i "s|--module ${MODULE}|--module ${NEW_MODULE}|g" "$NEW_SH"
else
    # 替换 python/torchrun/deepspeed 后的 .py 文件路径
    ORIG_PY_REL=$(basename "$ORIG_PY_PATH")
    NEW_PY_REL=$(basename "$NEW_PY")
    sed -i "s|${ORIG_PY_REL}|${NEW_PY_REL}|g" "$NEW_SH"
fi

B5. 运行新的 shell 脚本

保持原环境变量和分布式启动方式不变：

bash "$NEW_SH"

运行成功 → 进入第 4 步。失败则报告错误并停止。

第 4 步：验证并报告结果

# 验证输出目录
ls {OUTPUT_DIR}/*.pt.json 2>/dev/null
ls {OUTPUT_DIR}/worker_0_*/trace_*.json 2>/dev/null

通过验证后向用户报告：

Profiling 采集完成！ 输出目录：{OUTPUT_DIR} 采集配置：Level{LEVEL}, {ACTIVE} steps, CPU={CPU}, Memory={MEMORY} 下一步：使用 profiling-analysis 系列 Skill 进行性能分析

故障排查

禁止事项

• 禁止修改用户的原始脚本 — 必须创建新脚本
• 禁止将 prof.start() / prof.stop() 放在循环内部 — 必须在循环外
• 禁止设置 active=0 — 至少采集 1 步
• 禁止大规模采集所有 step — 3~5 步足够分析
• 禁止没有脚本继续执行
• 禁止使用 --profile 等 MindSpeed-LLM 参数 — 本 Skill 通过注入 Python 代码实现
• 禁止安装框架、下载模型或转换权重
• 禁止在 .sh 脚本里用 sed 全量替换同名 .py 而不告知用户 — run.sh 中若有多个同名 .py（如 eval/preprocess 子脚本）会被一并改写，必须先 grep 复核命中范围再替换

参考文档

详见 references/torch-npu-profiler-config-reference.md

• 昇腾 PyTorch Profiler 官方文档
• profiling-analysis 系列 Skill