昇腾社区首页
中文
注册
开发者
下载
支持
积分兑换 NEW

模型数据转储

工具对外提供“DumpConfig”及“register_hook”接口用于数据Dump，使用前需要将接口插入到模型脚本中，接口说明请参见API参考。

使用样例

在模型Python文件中文件开头导入对应函数，在模型初始化后添加Dump代码，使用样例参考如下。

from ait_llm import DumpConfig, register_hook
dump_config = DumpConfig(dump_path="./ait_dump")
register_hook(model, dump_config) # model是要dump中间tensor的模型示例

转储落盘位置

Dump默认落盘路径{DUMP_DIR}为当前目录下，如果在“DumpConfig”中指定“dump_path”，落盘路径则为指定的“{DUMP_PATH}”

  • Tensor信息会生成在默认落盘路径的“ait_dump”目录下，具体路径是 {DUMP_DIR}/ait_dump/torch_tensors/{device_id}_{PID}/{TID}目录下。
  • Model信息会生成在默认落盘路径的“ait_dump”目录下，具体路径是 {DUMP_DIR}/ait_dump/torch_tensors/{device_id}_{PID}/model_tree.json

{device_id}为设备号，{PID}为进程号，{TID}为token_id。