SuperAPI AI 工具集

pytorch-profiling-collection

编程与开发

查看项目

更新于 10-15 17:19

技能介绍

用途:本 Skill 面向通用场景(明确排除 MindSpeed-LLM 与 MindSpeed-MM),通过 torch_npu.profiler 采集 Ascend NPU 的 Profiling 性能数据,供 profiling-analysis 系列 Skill 做后续分析。

主要能力:支持 level0/level1(推荐)/level2 采集级别;支持训练循环、单次推理、指定代码段三种采集方式;可配置 active 步数(默认 3)、CPU 采集(默认开启)、内存采集(默认关闭)、Tensor Shape 记录(默认开启)、with_stack 堆栈采集(默认关闭)、输出目录(默认 ./profiling_result)等;兼容 .py 脚本与 deepspeed/torchrun 的 .sh 启动脚本。

使用方式:先收集并确认配置表;再检查环境,执行 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 或 torch_npu 不可用则停止。随后基于用户脚本创建带时间戳的新脚本并注入 profiling 代码,不修改原始脚本;.py 走分支 A 直接注入后 python 运行,.sh 走分支 B 先提取被引用的主入口 Python 脚本再注入,并修改新 shell 脚本中的引用路径后 bash 运行。最后用 ls {OUTPUT_DIR}/.pt.json 与 ls {OUTPUT_DIR}/worker_0_/trace_*.json 验证产物并报告。

限制:必须提供可运行的 .py 或 .sh 脚本,否则停止;仅支持 .py 与 .sh 两种类型;采集级别、步数等未提供时使用默认值;需要可用的 Ascend NPU 与 torch_npu 环境;产物格式与具体字段以实际输出为准;其他环境要求、版本与性能收益资料暂未提供。