mindspeed-mm-train-profiler

MindSpeed-MM Training Profiling Data Collection (Ascend NPU)

Reply in the user's language.
Do not translate the following: commands, file paths, environment variables, package names, error messages.

Step 0: Collect the training script and confirm the configuration

Reply to the user:

Before starting Profiling collection, please provide the training script path (.sh) — a runnable training script with the model, data, and weights already configured.

Without a training script, collection is not possible.

When the user provides no information: Stop. Reply: "Profiling requires a runnable training script. Please prepare one before collection."

After the user provides the script: Read the script contents, automatically detect the training engine, and locate the configuration file.

Detect the training engine

Search target	Purpose
`MM_TOOL` variable	Directly locate tools.json (Megatron engine)
`pretrain_sora.py` / `pretrain_transformers.py`	Confirm it is an MM training entry point
YAML path (the first non-option argument after torchrun)	Read the YAML for further determination
In the YAML, `MM_TOOL_PATH` field	Locate tools.json (Megatron engine)
In the YAML, `tools.profile` structure	Directly in the YAML configuration (standalone FSDP2 engine)

If the script does not reference mindspeed_mm，Stop and recommend using pytorch-profiling-collection generic collection skill。

Select the Profiling configuration item by item

In order, use the question tool to select configuration items one by one, giving a default value for each item:

① Collection type: Use the question tool, with options 静态（推荐，固定起止 step） / 动态（运行中通过文件控制）, static selected by default.

② Collection level: Use the question tool, with options level0（基础） / level1（推荐，含 AICore 利用率） / level2（全量）, level1 selected by default.

③ Collection start/end step (static mode only): Use the question tool, with options 10~12（推荐） / 5~8 / 自定义, by default 10~12is selected. If Custom is selected, then use the question tool to let the user input or select the start/end step.

④ Collection card number: Use the question tool, with options 仅 rank 0（推荐） / 全部 rank / 自定义（如 0,1,2）, by default 仅 rank 0。

⑤ Collection content (multiple selections allowed): Use the question tool, list all options at once, set multiple: true, allowing the user to select multiple items simultaneously. The options are:

CPU 事件(selected by default)
内存占用
算子调用栈
Tensor 输入形状
数据简化模式

⑥ AI Core metric type: Use the question tool, with options PipeUtilization（默认） / ArithmeticUtilization, the former selected by default.

⑦ Online parsing: Use the question tool, with options 开启（推荐，训练完自动解析） / 关闭（需手动跑离线解析）, enabled selected by default.

⑧ Number of training steps: Use the question tool to let the user input the total number of training steps. The options are:20 / 100 / 1000 / 10000, defaults to 100. Users can also select "Custom" to enter the number of steps themselves. After confirmation, change the --train-iters of the script copy to this value.

Display the complete configuration confirmation table

After collecting all answers, display the final configuration table and wait for user confirmation:

Configuration item Value

Run script <路径>

Training engine Megatron / standalone FSDP2

Configuration file <路径>

Collection type static / dynamic

Collection level level1

Collection steps step N ~ M

Collection card number rank 0

Collection content CPU events, memory usage, operator call stack, Tensor shapes, data simplification (multiple selection)

AI Core metrics PipeUtilization

Online parsing Enable/Disable

Training steps User-specified (e.g., 100)

Configuration will begin after confirmation. If changes are needed, please specify.

Configuration item	Value
Run script	`<路径>`
Training engine	Megatron / standalone FSDP2
Configuration file	`<路径>`
Collection type	static / dynamic
Collection level	level1
Collection steps	step N ~ M
Collection card number	rank 0
Collection content	CPU events, memory usage, operator call stack, Tensor shapes, data simplification (multiple selection)
AI Core metrics	PipeUtilization
Online parsing	Enable/Disable
Training steps	User-specified (e.g., 100)

Do not download models, convert weights, or create training scripts.

User input and variable mapping

question tool return value	Variable	Default value
Training script (.sh)	`TRAIN_SCRIPT`	(required)
"Static" / "static"	`PROFILE_TYPE=static`	`static`
"Dynamic" / "dynamic"	`PROFILE_TYPE=dynamic`	—
"Level0/1/2"	`PROFILE_LEVEL`	`level1`
"step N to M"	`START_STEP=N END_STEP=M`	`10, 12`
rank 0 / all / "0,1,2"	`PROFILE_RANKS`	`[0]`(rank 0 only)
"CPU events"	`WITH_CPU=true`	`true`
"Memory usage"	`WITH_MEMORY=true`	`false`
"Operator call stack"	`WITH_STACK=true`	`false`
"Tensor input shapes"	`RECORD_SHAPES=true`	`false`
"Data simplification mode"	`DATA_SIMPLIFICATION=true`	`false`
"PipeUtilization" / "ArithmeticUtilization"	`AIC_METRICS_TYPE`	`PipeUtilization`
"Enable" / "Disable"	`ANALYSE_FLAG=true/false`	`true`
"All ranks"	`PROFILE_RANKS=[-1]`	`[0]`
"20" / "100" / "1000" / "10000" / custom	`TRAIN_ITERS=用户指定值`	`100`

Step 1: 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 2: Configure and execute Profiling

Do not modify the user's original script or configuration file. First create a copy of the script, then create a copy of the configuration file to modify.

Read the original script and check whether the key paths it references exist (weight directory, data file, configuration file).
If any path does not exist,stop and report:
A path referenced in the training script does not exist:
- <缺失路径> — not found
Please check the path configuration in the script and ensure the training environment is ready before performing Profiling collection.

Based on the engine type detected in step 0, determine the configuration target to modify:

Detection result	Configuration method	Modification target
The script has `MM_TOOL` variable	tools.json	`MM_TOOL` JSON file pointed to
The YAML has `MM_TOOL_PATH` field	tools.json	`MM_TOOL_PATH` JSON file pointed to
In YAML there is `tools.profile` structure	YAML tools.profile	The YAML configuration file itself

Createnewscripts and configuration files (using timestamped filenames).

Modify the configuration according to the engine type:

Megatron engine (modify tools.json): Copy tools.json, add or override on top of the original content the profile field:

{
  "profile": {
    "enable": true,
    "profile_type": "static",
    "ranks": [0],
    "static_param": {
      "level": "level1",
      "with_stack": false,
      "with_memory": false,
      "record_shapes": false,
      "with_cpu": true,
      "save_path": "./npu_profiling",
      "start_step": 10,
      "end_step": 12,
      "data_simplification": false,
      "aic_metrics_type": "PipeUtilization",
      "analyse_flag": true
    }
  }
}

If tools.json originally contains configurations such as memory_profile , keep them unchanged.

Standalone FSDP2 engine (modify YAML): Copy the YAML configuration file, and add at the end the tools field:

tools:
  profile:
    enable: true
    profile_type: static
    ranks:
      - 0
    static_param:
      level: level1
      with_stack: false
      with_memory: false
      record_shapes: false
      with_cpu: true
      save_path: ./npu_profiling
      start_step: 10
      end_step: 12
      data_simplification: false
      aic_metrics_type: PipeUtilization
      analyse_flag: true

According to the user's Profiling configuration (the variable mapping in step 0), replace the values of the corresponding fields one by one.
Inscript copyfind the --train-iters parameter, and change its value to the number of training steps specified by the user. If the script does not contain --train-iters(such as when using train-iters), handle it the same way.
If a configuration file copy was created, also modify in syncscript copythe path variables (MM_TOOL / MM_TOOL_PATH) so that they point to the new copy.
Create the Profiling output directory:
```
mkdir -p <save_path>
```
Runscript copy. The original script and original configuration file must remain unchanged.
If training fails (non-zero exit code), report the error to the user:

Training execution failed (exit code: N). Profiling collection requires training to run normally.

Please fix the training configuration and try again.

Dynamic collection instructions

If the user chooses dynamic collection:

Set profile_type: dynamic。
In dynamic_param.config_path specify an empty directory (such as ./dynamic_profile_config）。
After training is run, profiler_config.json files will be automatically generated in that directory.
While training is running, modify the profiler_config.json in start_step / end_step, and Profiling will automatically start at the next step.

Offline parsing

If analyse_flag is false, or if re-parsing is needed, execute according to the engine type:

Megatron engine:

python mindspeed_mm/tools/profiler.py --mm-tool <tools.json 路径> --profiler-path <profiling 数据目录>

Standalone FSDP2 engine:

python mindspeed_mm/fsdp/tools/profiler.py --profiler-path <profiling 数据目录>

Step 3: Verify and report the results

After training is complete, verify the Profiling output:

Check whether the save_path directory exists, and whether it contains the *_ascend_pt/ subdirectory.
It should contain at least PROF_*/device_*/data/(NPU data) or the ASCEND_PROF_* directory.
Report to the user: output location, configuration used, directory structure.
It is recommended to use MindStudio Insight for visual analysis.

If the output is missing or empty, check:start_step >= 1, whether training reached the collection step, and whether the CANN environment has been loaded.

Troubleshooting

Symptom	Solution
`NPU out of memory`	Decrease `micro_batch_size` or enable `data_simplification`
Profiling directory is empty	`start_step` must be >= 1; training must reach that step
`Address already in use`	Modify the `MASTER_PORT`
Dynamic collection did not take effect	Check `config_path` directory's `profiler_config.json`
Only some ranks were collected	Check `ranks` whether the parameter configuration is correct
Cannot find `tools.json`	Confirm `MM_TOOL` or `MM_TOOL_PATH` the path is correct

Prohibitions

Do not modify the user's original training script (.sh); you must create a copy before making any modifications.
Do not modify the user's original configuration file (YAML/JSON); you must create a copy.
Do not install frameworks, download models, or convert weights.
Do not set start_step = 0 — must be >= 1.
Do not collect all steps in large-scale training — 2-3 steps are sufficient.
If there is no runnable training script, do not continue execution.
Unless the user explicitly requests it, do not collect all ranks (ranks: [-1]), to prevent excessive data volume.

Reference Documentation

See reference/mindspeed-mm-profiling-config.md for the complete configuration parameter table.