pytorch-profiling-collection

Ascend NPU Profiling Collection Guide

This Skill applies to general scenarios that are not MindSpeed-LLM or MindSpeed-MM, using torch_npu.profiler to collect performance data.

Step 0: Confirm prerequisites

Reply to the user:

Before starting Profiling collection, please provide the following information:

Required:

• Training/inference script path — a runnable Python script (.py) or shell launch script (.sh, commonly used with deepspeed / torchrun launches)

Optional (with default values):

• Collection level: level0 / level1(recommended) / level2
• Collection steps (active): default 3
• CPU collection: enabled by default
• Memory collection: disabled by default
• Tensor Shape recording: enabled by default
• Stack collection (with_stack): disabled by default
• Output directory: default ./profiling_result
• Collection mode: training loop / single inference / specified code segment (default training loop)

When the user does not provide a script path: stop. Reply: "Profiling requires a runnable Python script; please prepare one before collecting."

User input and variable mapping

User says	Mapping result
"collect 3 steps", "active=5"	active=5
"level0", "operators only"	profiler_level=Level0
"level1"	profiler_level=Level1
"level2", "full collection"	profiler_level=Level2
"collect CPU"	activities include CPU
"no CPU", "npu only"	activities NPU only
"memory"	profile_memory=True
"stack", "with_stack"	with_stack=True
"no shape", "no dimensions"	record_shapes=False
"output to XXX", "output=XXX"	OUTPUT_DIR=XXX
"inference script", "no training loop"	Single inference mode (without schedule)
"code segment", "only profile a certain function"	Wrap the target code segment with start/stop
"skip first wait", "skip_first_wait"	skip_first_wait=1(skip the wait phase of the first loop)

Step 1: Display the configuration confirmation table

After collecting the information, display the confirmation table; execute only after the user confirms.

Configuration item	Value
Script path	(path provided by the user)
Script type	.py / .sh
Collection level	Level1
Collection steps (active)	3
Collection mode	Training loop / single inference / specified code segment
CPU collection	Enabled
Memory collection	Disabled
Stack collection	Disabled
Tensor Shape	Enabled
Output directory	./profiling_result

After confirmation, proceed to Step 2.

Step 2: Check the environment

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 unavailable → stop: "No NPU detected in the current environment. Please run in an environment with an Ascend NPU."
• torch_npu unavailable → Stop: "Please install torch_npu and try again."

Step 3: Create a Profiling script and run it

Modifying the user's original script is prohibited. Create a new script based on the original script (with a timestamp, e.g. your_script_profiling_20260101120000.py）。

3.0 Detect the script type

# 检测用户提供的脚本后缀
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 → go to branch A (inject code → python new.py）
• .sh → go to branch B (extract the referenced .py → inject code → modify the path in run.sh → bash new.sh）

---

Branch A: user provides a .py script

A1. Code injection template

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. Three collection methods

Method A — Training loop collection (default):

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

Method B — Single inference (no training loop):

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

Method C — Specified code section:

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

A3. Run the new script

python {PROFILING_SCRIPT_PATH}

Run succeeds → proceed to Step 4. If it fails, report the error and stop.

---

Branch B: user provides a .sh script

B1. Copy the shell script

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. Extract the referenced Python script path

From run.sh , find the main entry Python script path, in the following order of priority:

1. --module parameter(deepspeed style): extract the module name and convert it to a .py path
2. python / torchrun / deepspeed after the command .py file: extract the file path

# 尝试提取 --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. Inject profiling code into the extracted Python script

Using the injection template from branch A (A1 + A2), create a profiling version in the same directory as the original .py :

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. Modify the new shell script

Replace the original Python path referenced in NEW_SH with the injected path:

# 替换 --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. Run the new shell script

Keep the original environment variables and distributed launch method unchanged:

bash "$NEW_SH"

Run succeeds → proceed to Step 4. If it fails, report the error and stop.

Step 4: Verify and report the results

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

After verification passes, report to the user:

Profiling collection complete! Output directory: {OUTPUT_DIR} Collection configuration: Level{LEVEL}, {ACTIVE} steps, CPU={CPU}, Memory={MEMORY} Next step: use the profiling-analysis series of Skills for performance analysis

Troubleshooting

Symptom	Solution
Too few active steps	Confirm that the training step actually reaches the target number of steps
No trace file after collection	Check prof.stop() whether it is called, whether the path is writable, and whether the NPU is available
NPU out of memory	Reduce the batch size in the training script
Error "Profiler already started"	Check whether there are multiple prof.start()
Output directory is empty	Confirm active > 0, and that training has executed to the corresponding step

Prohibitions

• Modifying the user's original script is prohibited — a new script must be created
• Placing prof.start() / prof.stop() inside the loop is prohibited — it must be outside the loop
• Setting active=0 is prohibited — collect at least 1 step
• Do not collect all steps at a large scale — 3 to 5 steps are enough for analysis
• Do not continue execution without a script
• Do not use --profile and other MindSpeed-LLM parameters — this Skill is implemented by injecting Python code
• Do not install frameworks, download models, or convert weights
• Do not use .sh in the sed script to fully replace all .py with the same name without informing the user — if run.sh contains multiple .pywith the same name (such as eval/preprocess sub-scripts), they will all be rewritten together; you must first grep review the hit scope before replacing

Reference documentation

See references/torch-npu-profiler-config-reference.md

• Ascend PyTorch Profiler official documentation
• profiling-analysis series Skills