ai-for-science-ai4s-profiling

Ascend NPU Profiling Collection and Performance Analysis Skill

This Skill provides a standardized process for collecting performance Profiling data on Huawei Ascend NPUs, supports three collection levels: L0 (minimal overhead), L1 (operator level), and L2 (full call stack), covers both training and inference scenarios, as well as integration methods for various training frameworks.

Important Default Behaviors

1. Default collection level: When the user only says "collect profiling" without explicitly specifying the collection level, by default collect the NPU L0 (minimal overhead) level, i.e., the template in Section 3.1. Only when the user explicitly requests deeper requirements such as operator analysis, call stacks, or memory analysis, upgrade to L1 or L2.

2. Do not modify CUDA code in training/inference scripts: When reviewing the user's training or inference scripts, even if the code contains cuda、.cuda()、torch.device("cuda") and similar terms, alsodo nothelp the user change them to npu -related syntax. Because in the Ascend environment, through import torch_npu combined with automatic migration (transfer_to_npu), these CUDA calls will be automatically converted to NPU calls at runtime, with no need to manually modify the source code. The Profiling code itself using torch_npu.profiler is necessary, but the training/inference business code can remain as is.

Prerequisites

Before executing Profiling, confirm that the following environment is ready:

Item	Requirement
Hardware	Ascend910 series (at least 1 card)
CANN	≥ 8.0 (8.2+ recommended)
Python	3.8 – 3.10
PyTorch	Matches the CANN version
torch_npu	Consistent with the PyTorch version
Disk space	≥ 10GB available recommended (L2 data volume is relatively large)

Process Overview

0. 环境准备与校验
→ 1. 确定采集场景（训练 / 推理）
→ 2. 选择采集级别（L0 / L1 / L2）
→ 3. 植入 Profiling 代码
→ 4. 执行采集
→ 5. 数据解析与可视化
→ 6. 性能分析与调优建议
→ 7. GPU 对比采集（可选）

---

0. Environment Preparation and Verification

0.1 CANN Environment Initialization

source /usr/local/Ascend/ascend-toolkit/set_env.sh
cat /usr/local/Ascend/ascend-toolkit/latest/version.cfg
npu-smi info

0.2 torch_npu Availability Verification

python3 -c "import torch; import torch_npu; print(torch.npu.is_available()); a = torch.randn(3,4).npu(); print(a.device)"

Output True and npu:0 indicates the environment is normal.

0.3 Performance Optimization Environment Variables (must be set before collection)

Before executing the Profiling collection command, it is recommended to set the following environment variables to obtain more accurate performance data:

# CPU 绑核：将训练进程绑定到固定 CPU 核心，减少跨核调度带来的性能抖动
export CPU_AFFINITY_CONF=1

# 流水优化：开启 Host 侧任务下发与 Device 侧执行的异步流水，
# 减少 Host 等待 Device 的空闲时间，提升整体吞吐
# 0=关闭, 1=开启, 2=开启增强模式（推荐）
export TASK_QUEUE_ENABLE=2

Note:

• CPU_AFFINITY_CONF=1 Avoid the operating system frequently scheduling training threads across different CPU cores, reduce cache misses and NUMA cross-node access, making the collected performance data more stable and reproducible.
• TASK_QUEUE_ENABLE=2 Enable enhanced-mode Host-Device asynchronous pipeline scheduling, which can effectively hide operator dispatch latency. Setting this item during Profiling allows observation of real performance closer to the production environment.
• These two environment variables should be setbeforestarting the training/inference script, and executed together with source set_env.sh .

---

1. Determine the Collection Scenario

1.1 Training Scenario

In the training scenario, use schedule to control which iteration steps to collect, avoiding data bloat caused by full collection:

• skip_first: Skip the initial iterations (including non-steady-state overhead such as compilation and data loading)
• wait: The number of steps to wait after skipping
• warmup: warm-up steps (collected but not recorded)
• active: actual recorded steps
• repeat: number of repeated collection rounds

Typical configuration:

schedule = torch_npu.profiler.schedule(
    wait=1, warmup=1, active=1, repeat=1, skip_first=20
)

This configuration means: skip the first 20 steps → wait 1 step → warm up 1 step → collect 1 step, repeated for 1 round.

Note: After collection ends, the data will be parsed automatically, and after parsing the model will continue running the remaining iterations. If you do not need to continue training, you can stop promptly after collection is complete.

1.2 Inference Scenarios

In inference scenarios, you also need to call prof.step() after inference ends to trigger trace data export:

with torch_npu.profiler.profile(
    activities=[torch_npu.profiler.ProfilerActivity.NPU],
    on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./npu-profiling-inference")
) as prof:
    model(input_data)
    prof.step()

---

2. Collection Level Description

Level	Description	Data Volume	Typical Use
L0 (default)	Only collect NPU activity, minimal overhead	Small	Quickly locate hotspots and overall time distribution
L1	Collect CPU + NPU + operator details	Medium	Analyze Cube/Vector/MatMul/Conv operator time consumption
L2	L1 + call stack + memory	Large	In-depth analysis of call chains and memory bottlenecks

---

3. Inserting Profiling Code

3.1 L0 — Minimal Overhead Collection (NPU Only)

Only collect NPU-side activity, with minimal data overhead, suitable for preliminary analysis:

import torch
import torch_npu

with torch_npu.profiler.profile(
    activities=[torch_npu.profiler.ProfilerActivity.NPU],
    schedule=torch_npu.profiler.schedule(
        wait=1, warmup=1, active=1, repeat=1, skip_first=20
    ),
    on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(
        "./npu-profiling-L0"
    )
) as prof:
    for step, batch in enumerate(train_dataloader):
        train_one_step(batch)
        prof.step()

3.2 L1 — Operator-Level Collection (CPU + NPU + Operator Details)

Collect both CPU and NPU activity, enable ProfilerLevel.Level1, to analyze the detailed time consumption of operator structures such as Cube, Vector, MatMul, and Conv:

import torch
import torch_npu

with torch_npu.profiler.profile(
    activities=[
        torch_npu.profiler.ProfilerActivity.CPU,
        torch_npu.profiler.ProfilerActivity.NPU,
    ],
    with_stack=False,
    record_shapes=True,
    profile_memory=False,
    schedule=torch.profiler.schedule(
        wait=1, warmup=1, active=1, repeat=1, skip_first=20
    ),
    experimental_config=torch_npu.profiler._ExperimentalConfig(
        profiler_level=torch_npu.profiler.ProfilerLevel.Level1
    ),
    on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(
        "./npu-profiling-L1"
    )
) as prof:
    for step, batch in enumerate(train_dataloader):
        train_one_step(batch)
        prof.step()

Parameter description:

• record_shapes=True: collect the input shape and input type of torch ops
• with_stack=False: do not collect function call stacks (to reduce space usage)
• profile_memory=False: do not collect memory-related data
• experimental_config: NPU-specific parameter, sets the profiler level

3.3 L2 — Full Call Stack + Memory Collection

Enable call stack and memory collection on top of L1, suitable for in-depth performance analysis:

import torch
import torch_npu

with torch_npu.profiler.profile(
    activities=[
        torch_npu.profiler.ProfilerActivity.CPU,
        torch_npu.profiler.ProfilerActivity.NPU,
    ],
    with_stack=True,
    record_shapes=True,
    profile_memory=True,
    schedule=torch.profiler.schedule(
        wait=1, warmup=1, active=1, repeat=1, skip_first=20
    ),
    experimental_config=torch_npu.profiler._ExperimentalConfig(
        profiler_level=torch_npu.profiler.ProfilerLevel.Level1
    ),
    on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(
        "./npu-profiling-L2"
    )
) as prof:
    for step, batch in enumerate(train_dataloader):
        train_one_step(batch)
        prof.step()

Note: with_stack=True and profile_memory=True will significantly increase the amount of output data, only enable them when you need to analyze call stacks or memory bottlenecks.

---

4. Framework Adaptation

4.1 PyTorch Lightning Training Collection

For projects using PyTorch Lightning, you cannot directly use a with statement to wrap the training loop; you need to integrate the Profiler through the Callback mechanism:

import torch
import torch_npu
import pytorch_lightning as pl


class NPUProfilingCallback(pl.Callback):
    """昇腾 NPU Profiling 回调，支持 L0/L1/L2 级别切换。"""

    def __init__(self, output_dir="./npu-profiling", level="L0",
                 skip_first=20, wait=1, warmup=1, active=1, repeat=1):
        super().__init__()
        self.output_dir = output_dir
        self.level = level
        self.skip_first = skip_first
        self.wait = wait
        self.warmup = warmup
        self.active = active
        self.repeat = repeat
        self.prof = None

    def _build_profile_kwargs(self):
        kwargs = {
            "activities": [torch_npu.profiler.ProfilerActivity.NPU],
            "schedule": torch_npu.profiler.schedule(
                wait=self.wait, warmup=self.warmup,
                active=self.active, repeat=self.repeat,
                skip_first=self.skip_first,
            ),
            "on_trace_ready": torch_npu.profiler.tensorboard_trace_handler(
                self.output_dir
            ),
        }
        if self.level in ("L1", "L2"):
            kwargs["activities"].insert(
                0, torch_npu.profiler.ProfilerActivity.CPU
            )
            kwargs["record_shapes"] = True
            kwargs["experimental_config"] = (
                torch_npu.profiler._ExperimentalConfig(
                    profiler_level=torch_npu.profiler.ProfilerLevel.Level1
                )
            )
        if self.level == "L2":
            kwargs["with_stack"] = True
            kwargs["profile_memory"] = True
        return kwargs

    def on_train_start(self, trainer, pl_module):
        self.prof = torch_npu.profiler.profile(**self._build_profile_kwargs())
        self.prof.start()

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        if self.prof is not None:
            torch.npu.synchronize()
            self.prof.step()

    def on_train_end(self, trainer, pl_module):
        if self.prof is not None:
            self.prof.stop()
            self.prof = None

Usage: Add the Callback to the Trainer's callbacks list:

trainer = pl.Trainer(
    callbacks=[
        ...,
        NPUProfilingCallback(
            output_dir="./npu-profiling-L1",
            level="L1",
            skip_first=20,
        ),
    ],
    ...
)

4.2 HuggingFace Transformers Trainer

For projects using HuggingFace Trainer , you can integrate through a custom TrainerCallback :

import torch
import torch_npu
from transformers import TrainerCallback


class NPUProfilingTrainerCallback(TrainerCallback):

    def __init__(self, output_dir="./npu-profiling", skip_first=20):
        self.output_dir = output_dir
        self.skip_first = skip_first
        self.prof = None

    def on_train_begin(self, args, state, control, **kwargs):
        self.prof = torch_npu.profiler.profile(
            activities=[torch_npu.profiler.ProfilerActivity.NPU],
            schedule=torch_npu.profiler.schedule(
                wait=1, warmup=1, active=1, repeat=1,
                skip_first=self.skip_first,
            ),
            on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(
                self.output_dir
            ),
        )
        self.prof.start()

    def on_step_end(self, args, state, control, **kwargs):
        if self.prof is not None:
            torch.npu.synchronize()
            self.prof.step()

    def on_train_end(self, args, state, control, **kwargs):
        if self.prof is not None:
            self.prof.stop()
            self.prof = None

4.3 DeepSpeed Scenarios

When using DeepSpeed, the Profiling code is still wrapped around the outer layer of the training loop, but note:

• Ensure that prof.step() is called after model_engine.step() call after
• In multi-card scenarios, it is recommended to collect only on rank 0 to reduce data volume

import torch
import torch_npu

if local_rank == 0:
    prof = torch_npu.profiler.profile(
        activities=[torch_npu.profiler.ProfilerActivity.NPU],
        schedule=torch_npu.profiler.schedule(
            wait=1, warmup=1, active=1, repeat=1, skip_first=20
        ),
        on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(
            "./npu-profiling-deepspeed"
        ),
    )
    prof.start()

for step, batch in enumerate(train_dataloader):
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()
    if local_rank == 0:
        prof.step()

if local_rank == 0:
    prof.stop()

---

5. GPU Comparison Collection (Optional)

When GPU vs NPU performance comparison is needed, use the same schedule configuration on the GPU environment to collect data. GPU collection uses native torch.profiler, differences from NPU:

1. Use torch.profiler.profile instead of torch_npu.profiler.profile
2. Use torch.profiler.ProfilerActivity.CUDA instead of NPU
3. Do not use experimental_config parameter (GPU does not have this parameter)

5.1 GPU minimal expansion collection

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(
        wait=1, warmup=1, active=1, repeat=1, skip_first=20
    ),
    on_trace_ready=torch.profiler.tensorboard_trace_handler(
        "./gpu-profiling-L0"
    )
) as prof:
    for step, batch in enumerate(train_dataloader):
        train_one_step(batch)
        prof.step()

5.2 GPU collection with CPU data

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    schedule=torch.profiler.schedule(
        wait=1, warmup=1, active=1, repeat=1, skip_first=20
    ),
    on_trace_ready=torch.profiler.tensorboard_trace_handler(
        "./gpu-profiling-with-cpu"
    )
) as prof:
    for step, batch in enumerate(train_dataloader):
        train_one_step(batch)
        prof.step()

---

6. Quick reference: collection code template selection

Select the appropriate template based on the actual scenario:

Scenario	Level	Template location
Training + quick overview	L0	Section 3.1
Training + operator analysis	L1	Section 3.2
Training + in-depth analysis	L2	Section 3.3
Inference	L0	Section 1.2
PyTorch Lightning	L0/L1/L2	Section 4.1
HuggingFace Trainer	L0	Section 4.2
DeepSpeed multi-GPU	L0	Section 4.3
GPU comparison	-	Section 5

Companion script

• Profiling pre-check:python scripts/validate_profiling_env.py --device npu:0 --output-dir ./profiling_output

Reference materials

• Profiling collection and conclusion template:references/analysis-checklist.md