aclgraph_dump 使用指南

简介

PyTorch ACLGraph 图模式下执行精度对齐时,整体策略如下:先整网筛查,后单点深挖。通过整网采集快速收敛异常范围,针对疑似问题算子进行单点采集并保存其 Tensor 数据,以便于细粒度排查。aclgraph_dump 提供如下采集能力:

  • 单点采集:acl_save

  • 整网采集:AclGraphDumper

使用前准备

环境准备

  1. 安装并正确配置 Ascend Extension for PyTorch(torch_npu)。

  2. 安装 msProbe 工具,详见《msProbe 安装指南》。

[!NOTE]

源码编译安装时需包含 aclgraph_dump 模块,通过如下命令安装:

python3 setup.py bdist_wheel --include-mod=aclgraph_dump --no-check

约束

  • 仅支持 PyTorch 框架。

  • 构建 aclgraph_dump 需要 torch_npu 参与编译;若未包含该模块,将无法正常使用该功能。

整网采集

快速入门

  1. 在使用整网采集功能前,需要配置文件(config.json):

    {
    "task": "statistics",
    "dump_path": "./L0_dump",
    "rank": [],
    "level": "L0",
    "statistics": {
        "list": ["linear", "attention"]
    }
    }
    

    参考说明

    整网 aclgraph dump 当前支持的配置项如下:

    配置项

    可选/必选

    说明

    task

    可选

    采集任务类型,str类型。默认值为statistics,整网 aclgraph dump 当前仅支持 statistics

    dump_path

    必选

    dump 结果输出目录,str类型。工具会检查并创建该目录。

    rank

    可选

    指定采集的 rank,list[int |str]类型。默认值为空,表示采集所有 rank;字符串仅支持 "start-end" 范围格式。非目标 rank 不开启整网采集。

    level

    可选

    根级采集级别,str类型。支持 L0L1mix,默认值为L0
    L0 采集 module 输入/输出统计值;L1 采集 API 输入/输出统计值;mix 同时采集 module 和 API 统计值。

    list

    可选

    模块名关键词过滤列表,list[str]类型。默认值为空,表示采集所有模块。

  2. 完成文件(config.json)配置后,下面示例展示如何使用整网采集功能:

      import torch
      import torch_npu
    + from msprobe.pytorch import AclGraphDumper
    
      N,D_in, H, D_out = 640, 4096, 2048, 1024
      # 模型初始化
      model = torch.nn.Sequential(
        torch.nn.Linear(D_in, H),
        torch.nn.ReLU(),
        torch.nn.Linear(H, D_out)
      ).npu()
    + # 初始化配置
    + dumper = AclGraphDumper('./acl_config.json')
    + # 在编图前配置采集任务
    + dumper.start(model)
      static_input = torch.randn(N, D_in).npu()
      static_target = torch.randn(N, D_out).npu()
    
      g = torch.npu.NPUGraph()
      # 编图
      with torch.npu.graph(g):
        static_target = model(static_input)
    
      real_inputs = [torch.rand_like(static_input) for _ in range(10)]
      real_targets = [torch.rand_like(static_target) for _ in range(10)]
    
      for data, target in zip(real_inputs, real_targets):
        static_input.copy_(data)
        static_target.copy_(target)
        # 图replay
        g.replay()
    +   # 数据落盘
    +   dumper.step()
    

整网采集功能介绍

功能说明

AclGraphDumper 用于采集整网中间数据,当前支持 module 级别、API 级别以及 module+API 混合级别的统计值采集,结果包括张量形状、数据类型、统计值等信息。
AclGraphDumper 的初始化与 start 调用需在模型编图( 如torch.npu.graphtorch.compile )之前完成。

接口说明

函数原型

AclGraphDumper(config_path: str | None = None)

参数说明

参数名

可选/必选

说明

config_path

可选

配置文件路径,str类型。若不传,默认读取 msprobe 包内置 config.jsondump_pathtaskranklevellist 从该配置文件中读取。

函数原型

AclGraphDumper.start(model: torch.nn.Module) -> None

参数说明

参数名

可选/必选

说明

model

必选

待采集模型,torch.nn.Module类型。

函数原型

AclGraphDumper.step(dump: bool = True) -> None

参数说明

参数名

类型

说明

是否必选

dump

bool

是否将当前统计结果落盘到 dump.jsonTrue:清理统计并落盘,step_id 增加 1;False:仅清理统计不落盘,step_id 不增加(可用于 dummy_run 预热阶段)。

若未启动采集则直接返回。

输出说明

单卡场景

AclGraphDumper 单卡场景输出路径为:dump_path/step{step_id}/pid{pid}/dump.json

生成目录示例:

L0_dump
├── step0
│   └── pid9527
│       └── dump.json
├── step1
│   └── pid9527
│       └── dump.json
├── step2
|   └── pid9527
|       └── dump.json

多卡场景

AclGraphDumper 输出路径为:dump_path/step{step_id}/rank{rank_id}/dump.json

生成目录示例:

L0_dump
├── step0
│   └── rank0
│   |    └── dump.json
│   └── rank1
│   |    └── dump.json
│   └── rank2
│         └── dump.json

比对说明

可直接通过 msprobe compare 对整网采集结果进行比对。
比对完成后会生成 csv 报告文件,例如:compare_result_{rank_id}_{timestamp}.csv

在分布式多进程场景中,通常会按 rank 生成对应的 compare 结果文件,请结合 rank 维度查看结果。

单点采集

快速入门

下面示例展示如何在前向过程中保存某个张量:

  import torch
  import torch_npu
 
+ from msprobe.pytorch import acl_save
 
 
  class ToyModel(torch.nn.Module):
      def __init__(self):
          super().__init__()
          self.linear = torch.nn.Linear(8, 4)
 
      def forward(self, x):
          y = self.linear(x)
+         # 保存中间张量
+         acl_save(y, "./dump/linear_out.pt")
          return y
 
 
  if __name__ == "__main__":
      model = ToyModel().to("npu:0")
      x = torch.randn(2, 8, device="npu:0")
      out = model(x)

单点采集功能介绍

功能说明

acl_save 用于保存张量数据,调用后会生成 .pt 文件。

接口说明

函数原型

acl_save(x: torch.Tensor, path: str) -> torch.Tensor

参数说明

参数名

可选/必选

说明

x

必选

待保存张量,torch.Tensor类型。

path

必选

保存路径(支持相对/绝对路径),str类型。实际落盘文件名会在该路径文件名基础上追加序号,格式为 {base}_{seq}.pt。例如传入 ./dump/act.pt,实际落盘为 ./dump/act_0.pt./dump/act_1.pt

返回值

返回一个与输入形状一致的张量,仅用于触发保存操作。

使用示例

  1. 推理过程中的单点保存

    from msprobe.pytorch import acl_save
    
    logits = model(x)
    acl_save(logits, "./dump/logits.pt")
    
  2. 多卡场景下的单点采集

    # 多卡场景需要区分rank,使用参考如下。
    # 需要保证“./dump/rank{torch.distributed.get_rank()}”目录已创建,否则会出现目录不存在问题。
    acl_save(tensor,f'./dump/rank{torch.distributed.get_rank()}/tensor.pt')
    

输出说明

调用 acl_save 后,会在 path 指定目录下生成 .pt 文件(文件名自动追加序号),例如生成:./dump/act_0.pt、./dump/act_1.pt、./dump/act_2.pt。

数据解析

.pt 文件为 PyTorch 序列化格式,可通过 torch.load 读取:

import torch

tensor = torch.load("./dump/act_0.pt")

附录

常见问题

1. 导入报错:Failed to import msprobe.lib.aclgraph_dump_ext

请确认:

  • 编译安装时已包含 --include-mod=aclgraph_dump

  • 已安装 torch_npu 且环境变量配置正确;

  • 当前系统为 Linux。

2. Allocate SQ failed 问题

CANN 8.5 以下(不含 8.5)可能出现 Allocate SQ failed,这是老版本 SQ 不复用导致。可将 ccsrc/aclgraph_dump/aclgraph_dump.cppCurrentNPUStream 改为 DefaultNPUStream 规避,或升级至 CANN 8.5.0及以上版本。