MSProbe 调试指南#

在推理或训练过程中,我们经常会遇到准确性异常,例如输出偏离预期、数值行为不稳定(NaN/Inf)或预测结果不再与标签匹配。为了查明根本原因,我们必须监控和捕获模型执行过程中产生的中间数据——特征图、权重、激活和层输出。通过在特定阶段捕获关键张量、记录核心层的输入输出对,并保留上下文元数据(提示、张量数据类型、硬件配置等),我们可以系统地追踪准确性下降或数值错误开始的位置。本指南详细介绍了诊断 AI 模型(重点关注 vllm-ascend 服务)准确性问题的端到端工作流程:准备、数据捕获和分析与验证。

0. 背景概念#

msprobe 支持三个准确性级别

  • L0:在模块级别转储张量,并生成 construct.json,以便可视化工具可以重建网络结构。必须传入模型或子模块句柄。

  • L1:仅收集算子级别的统计信息,适用于轻量级故障排除。

  • 混合 (mix):捕获结构信息和算子统计信息,当您同时需要图重构和数值比较时很有用。

1. 先决条件#

1.1 安装 msprobe#

使用 pip 安装 msprobe

pip install mindstudio-probe==8.3.0

1.2 可选的可视化依赖项#

如果您需要可视化捕获的数据,请安装其他依赖项。

  1. 安装 tb_graph_ascend

    pip install tb_graph_ascend
    

2. 使用 msprobe 收集数据#

我们通常遵循一种粗到细的策略来捕获数据。首先识别问题出现的 token,然后决定围绕该 token 需要采样哪个范围。典型的流程描述如下。

2.1 准备转储配置文件#

创建一个可以被 PrecisionDebugger 解析的 config.json 文件,并将其放在一个可访问的路径。常用字段如下:

字段

描述

必需

task

转储任务的类型。常见的 PyTorch 值包括 "statistics""tensor"。统计任务收集张量统计信息(均值、方差、最大值、最小值等),而张量任务捕获任意张量。

dump_path

转储结果存储的目录。如果省略,msprobe 将使用默认路径。

排名

要采样的 rank。空列表表示采样所有 rank。对于单卡任务,您必须将此字段设置为 []

步骤

要采样的 token 迭代次数。空列表表示所有迭代。

级别

转储级别字符串("L0""L1""mix")。L0 针对 nn.ModuleL1 针对 torch.apimix 则两者都收集。

async_dump

是否启用异步转储(PyTorch statistics/tensor 任务支持)。默认为 false

scope

要采样的模块范围。空列表表示采样所有模块。

列表

要采样的算子范围。空列表表示采样所有算子。

要限制捕获的算子,请配置 list

  • scope (list[str]): 在 PyTorch pynative 场景中,此字段限制转储范围。提供两个遵循工具命名约定的模块或 API 名称以锁定一个范围;仅转储两个名称之间的数据。示例:

    "scope": ["Module.conv1.Conv2d.forward.0", "Module.fc2.Linear.forward.0"]
    "scope": ["Cell.conv1.Conv2d.forward.0", "Cell.fc2.Dense.backward.0"]
    "scope": ["Tensor.add.0.forward", "Functional.square.2.forward"]
    

    level 设置决定了可以提供什么——当 level=L0 时为模块,当 level=L1 时为 API,当 level=mix 时可以是模块或 API。

  • list (list[str]): 自定义算子列表。选项包括:

    • 在 PyTorch pynative 场景中,提供特定 API 的完整名称以仅转储这些 API。示例:"list": ["Tensor.permute.1.forward", "Tensor.transpose.2.forward", "Torch.relu.3.backward"]

    • level=mix 时,您可以提供模块名称,以便转储扩展到模块运行时产生的所有内容。示例:"list": ["Module.module.language_model.encoder.layers.0.mlp.ParallelMlp.forward.0"]

    • 提供一个子字符串,例如 "list": ["relu"],以转储名称包含该子字符串的所有 API。当 level=mix 时,名称包含该子字符串的模块也会被扩展。

配置示例

cat <<'JSON' > /data/msprobe_config.json
{
  "task": "statistics",
  "dump_path": "/home/data_dump",
  "rank": [],
  "step": [],
  "level": "L1",
  "async_dump": false,

  "statistics": {
    "scope": [],
    "list": [],
    "tensor_list": [],
    "data_mode": ["all"],
    "summary_mode": "statistics"
  }
}
JSON

2. 在 vllm-ascend 中启用 msprobe#

  1. 通过添加 --enforce-eager 来以 eager 模式启动 vLLM(静态图场景尚不支持),并通过 --additional-config 传递配置路径。

    vllm serve Qwen/Qwen2.5-0.5B-Instruct \
      --dtype float16 \
      --enforce-eager \
      --host 0.0.0.0 \
      --port 8000 \
      --additional-config '{"dump_config": "/data/msprobe_config.json"}' &
    

3. 发送请求并收集转储#

  1. 照常发送推理请求,例如:

    curl https://:8000/v1/completions \
      -H "Content-Type: application/json" \
      -d '{
            "model": "Qwen/Qwen2.5-0.5B-Instruct",
            "prompt": "Explain gravity in one sentence.",
            "max_tokens": 32,
            "temperature": 0
          }' | python -m json.tool
    
  2. 每个请求驱动 msprobe: start -> forward/backward -> stop -> step 的序列。运行器在每个代码路径上调用 step(),因此即使推理提前返回,您也能获得完整的数据集。

  3. 转储文件写入 dump_path。它们通常包含:

    • 按算子/模块分组的张量文件。

    • dump.json,其中记录了 dtype、shape、min/max 以及 requires_grad 等元数据。

    • construct.json,当 levelL0mix 时生成(可视化必需)。

    示例目录布局

    ├── dump_path
    │   ├── step0
    │   │   ├── rank0
    │   │   │   ├── dump_tensor_data
    │   │   │   │    ├── Tensor.permute.1.forward.pt
    │   │   │   │    ├── Functional.linear.5.backward.output.pt    # Format: {api_type}.{api_name}.{call_count}.{forward/backward}.{input/output}.{arg_index}.
    │   │   │   │    │                                              # arg_index is the nth input or output of the API. If an input is a list, keep numbering with decimals (e.g., 1.1 is the first element of the first argument).
    │   │   │   │    ├── Module.conv1.Conv2d.forward.0.input.0.pt          # Format: {Module}.{module_name}.{class_name}.{forward/backward}.{call_count}.{input/output}.{arg_index}.
    │   │   │   │    ├── Module.conv1.Conv2d.forward.0.parameters.bias.pt  # Module parameter data: {Module}.{module_name}.{class_name}.forward.{call_count}.parameters.{parameter_name}.
    │   │   │   │    └── Module.conv1.Conv2d.parameters_grad.weight.pt     # Module parameter gradients: {Module}.{module_name}.{class_name}.parameters_grad.{parameter_name}. Gradients do not include call_count because the same gradient updates all invocations.
    │   │   │   │                                                          # When the `model` argument passed to dump is a List[torch.nn.Module] or Tuple[torch.nn.Module], module-level data names also include the index inside the list ({Module}.{index}.*), e.g., Module.0.conv1.Conv2d.forward.0.input.0.pt.
    │   │   │   ├── dump.json
    │   │   │   ├── stack.json
    │   │   │   ├── dump_error_info.log
    │   │   │   └── construct.json
    │   │   ├── rank1
    │   │   │   ├── dump_tensor_data
    │   │   │   │   └── ...
    │   │   │   ├── dump.json
    │   │   │   ├── stack.json
    │   │   │   ├── dump_error_info.log
    │   │   │   └── construct.json
    │   │   ├── ...
    │   │   │
    │   │   └── rank7
    │   ├── step1
    │   │   ├── ...
    │   ├── step2
    
    • rank:设备 ID。每张卡将其数据写入相应的 rank{ID} 目录。在非分布式场景中,目录直接命名为 rank

    • dump_tensor_data:收集到的张量负载。

    • dump.json:每个 API 或模块的前向/后向数据的统计信息,包括名称、dtype、shape、max、min、mean、L2 范数(L2 方差的平方根),以及当 summary_mode="md5" 时,还有 CRC-32。有关详细信息,请参阅 dump.json 文件说明

    • dump_error_info.log:仅在转储工具遇到错误时出现,并记录失败日志。

    • stack.json:API/模块的调用堆栈。

    • construct.json:分层结构描述。当 level=L1 时为空。

4. 分析结果#

4.1 先决条件#

您通常需要两个转储数据集:一个来自“问题方”(暴露准确性或数值错误的情况)和一个来自“基准方”(良好的基线)。这些数据集不必相同——它们可以来自不同的分支、框架版本,甚至不同的实现(算子替换、不同的图优化开关等)。只要它们使用相同或相似的输入、硬件拓扑和采样点(步/token),msprobe 就可以比较它们并找到差异节点。如果您找不到完美的干净基准,可以先捕获问题方数据,手动创建一个最小的可复现案例,然后进行自我比较。下面我们假设问题转储是 problem_dump,基准转储是 bench_dump

4.2 可视化#

使用 msprobe graph_visualize 生成可以在 tb_graph_ascend 中打开的图像。

  1. 确保转储包含 construct.json(即 level = L0level = mix)。

  2. 准备一个比较文件,例如 compare.json。其格式和生成流程在 msprobe_visualization.md 的 3.1.3 节中描述。示例(最小可运行片段):

    {
      "npu_path": "./problem_dump",
      "bench_path": "./bench_dump",
      "is_print_compare_log": true
    }
    

    在调用 msprobe graph_visualize 之前,替换为您的转储目录的路径。**如果您只需要构建单个图**,请省略 bench_path 以可视化一个转储。
    也支持多 rank 场景(单 rank、多 rank 或多步多 rank)。npu_pathbench_path 必须包含名为 rank+数字 的文件夹,并且每个 rank 文件夹必须包含一个非空的 construct.json 以及 dump.jsonstack.json。如果任何 construct.json 为空,请验证转储级别是否包含 L0mix。在比较图时,npu_pathbench_path 都必须包含相同集合的 rank 文件夹,以便一对一匹配。

    ├── npu_path or bench_path
    |   ├── rank0
    |   |   ├── dump_tensor_data (only when the `tensor` option is enabled)
    |   |   |    ├── Tensor.permute.1.forward.pt
    |   |   |    ├── MyModule.0.forward.input.pt
    |   |   |    ...
    |   |   |    └── Function.linear.5.backward.output.pt
    |   |   ├── dump.json         # Tensor metadata
    |   |   ├── stack.json        # Operator call stack information
    |   |   └── construct.json    # Hierarchical structure; empty when `level=L1`
    |   ├── rank1
    |   |   ├── dump_tensor_data
    |   |   |   └── ...
    |   |   ├── dump.json
    |   |   ├── stack.json
    |   |   └── construct.json
    |   ├── ...
    |   |
    |   └── rankn
    
  3. 运行

    msprobe graph_visualize \
        --input_path ./compare.json \
        --output_path ./graph_output
    

    比较完成后,将在 graph_output 下创建一个 *.vis.db 文件。

    • 图构建:build_{timestamp}.vis.db

    • 图比较:compare_{timestamp}.vis.db

  4. 启动 tensorboard 并加载输出目录以检查结构差异、数值比较、溢出检测结果、跨设备通信节点以及过滤器/搜索。将包含 .vis.db 文件的目录传递给 --logdir

    tensorboard --logdir out_path --bind_all --port [optional_port]
    
  5. 检查可视化。UI 通常会显示具有算子、参数和张量输入输出的整体模型结构。单击任何节点以展开其子节点。

    • 差异可视化:比较结果用不同的颜色突出显示差异节点(差异越大,节点越红)。单击节点以查看其详细信息,包括张量输入/输出、参数和算子类型。分析数据差异和周围连接以精确定位具体差异。

    • 辅助功能:

      • 切换 rank/步:快速检查不同 rank 和步上的差异节点。

      • 搜索/过滤:使用搜索框按算子名称等过滤节点。

      • 手动映射:自动映射无法覆盖所有情况,因此该工具允许您在生成比较结果之前手动映射问题图和基准图中的节点。

5. 故障排除#

  • RuntimeError: Please enforce eager mode:重新启动 vLLM 并添加 --enforce-eager 标志。

  • 无转储文件:确认 JSON 路径正确且所有节点都有写入权限。在分布式场景中,设置 keep_all_ranks 以便每个 rank 写入自己的转储。

  • 转储文件过大:从 statistics 任务开始,以定位异常张量,然后使用 scope/list/tensor_listfilterstoken_range 等缩小范围。


附录#

dump.json 文件说明#

L0 级别#

L0 级别的 dump.json 包含模块的前向/后向输入输出,以及参数和参数梯度。以 PyTorch 的 Conv2d 为例,网络代码如下:

output = self.conv2(input)  # self.conv2 = torch.nn.Conv2d(64, 128, 5, padding=2, bias=True)

dump.json 包含以下条目:

  • Module.conv2.Conv2d.forward.0:模块的前向数据。input_args 表示位置输入,input_kwargs 表示关键字输入,output 存储前向输出,parameters 存储权重/偏置。

  • Module.conv2.Conv2d.parameters_grad:参数梯度(权重和偏置)。

  • Module.conv2.Conv2d.backward.0:模块的后向数据。input 表示流入模块的梯度(前向输出的梯度),output 表示流出的梯度(模块输入的梯度)。

注意:当传递给转储 API 的 model 参数是 List[torch.nn.Module]Tuple[torch.nn.Module] 时,模块级名称将包含列表中的索引({Module}.{index}.*)。例如:Module.0.conv1.Conv2d.forward.0

{
 "task": "tensor",
 "level": "L0",
 "framework": "pytorch",
 "dump_data_dir": "/dump/path",
 "data": {
  "Module.conv2.Conv2d.forward.0": {
   "input_args": [
    {
     "type": "torch.Tensor",
     "dtype": "torch.float32",
     "shape": [
      8,
      16,
      14,
      14
     ],
     "Max": 1.638758659362793,
     "Min": 0.0,
     "Mean": 0.2544615864753723,
     "Norm": 70.50277709960938,
     "requires_grad": true,
     "data_name": "Module.conv2.Conv2d.forward.0.input.0.pt"
    }
   ],
   "input_kwargs": {},
   "output": [
    {
     "type": "torch.Tensor",
     "dtype": "torch.float32",
     "shape": [
      8,
      32,
      10,
      10
     ],
     "Max": 1.6815717220306396,
     "Min": -1.5120246410369873,
     "Mean": -0.025344856083393097,
     "Norm": 149.65576171875,
     "requires_grad": true,
     "data_name": "Module.conv2.Conv2d.forward.0.output.0.pt"
    }
   ],
   "parameters": {
    "weight": {
     "type": "torch.Tensor",
     "dtype": "torch.float32",
     "shape": [
      32,
      16,
      5,
      5
     ],
     "Max": 0.05992485210299492,
     "Min": -0.05999220535159111,
     "Mean": -0.0006165213999338448,
     "Norm": 3.421217441558838,
     "requires_grad": true,
     "data_name": "Module.conv2.Conv2d.forward.0.parameters.weight.pt"
    },
    "bias": {
     "type": "torch.Tensor",
     "dtype": "torch.float32",
     "shape": [
      32
     ],
     "Max": 0.05744686722755432,
     "Min": -0.04894155263900757,
     "Mean": 0.006410328671336174,
     "Norm": 0.17263513803482056,
     "requires_grad": true,
     "data_name": "Module.conv2.Conv2d.forward.0.parameters.bias.pt"
    }
   }
  },
  "Module.conv2.Conv2d.parameters_grad": {
   "weight": [
    {
     "type": "torch.Tensor",
     "dtype": "torch.float32",
     "shape": [
      32,
      16,
      5,
      5
     ],
     "Max": 0.018550323322415352,
     "Min": -0.008627401664853096,
     "Mean": 0.0006675920449197292,
     "Norm": 0.26084786653518677,
     "requires_grad": false,
     "data_name": "Module.conv2.Conv2d.parameters_grad.weight.pt"
    }
   ],
   "bias": [
    {
     "type": "torch.Tensor",
     "dtype": "torch.float32",
     "shape": [
      32
     ],
     "Max": 0.014914230443537235,
     "Min": -0.006656786892563105,
     "Mean": 0.002657240955159068,
     "Norm": 0.029451673850417137,
     "requires_grad": false,
     "data_name": "Module.conv2.Conv2d.parameters_grad.bias.pt"
    }
   ]
  },
  "Module.conv2.Conv2d.backward.0": {
   "input": [
    {
     "type": "torch.Tensor",
     "dtype": "torch.float32",
     "shape": [
      8,
      32,
      10,
      10
     ],
     "Max": 0.0015069986693561077,
     "Min": -0.001139344065450132,
     "Mean": 3.3215508210560074e-06,
     "Norm": 0.020567523315548897,
     "requires_grad": false,
     "data_name": "Module.conv2.Conv2d.backward.0.input.0.pt"
    }
   ],
   "output": [
    {
     "type": "torch.Tensor",
     "dtype": "torch.float32",
     "shape": [
      8,
      16,
      14,
      14
     ],
     "Max": 0.0007466732058674097,
     "Min": -0.00044813455315306783,
     "Mean": 6.814070275140693e-06,
     "Norm": 0.01474067009985447,
     "requires_grad": false,
     "data_name": "Module.conv2.Conv2d.backward.0.output.0.pt"
    }
   ]
  }
 }
}

L1 级别#

L1 级别的 dump.json 记录 API 的前向/后向输入输出。以 PyTorch 的 relu 函数为例(output = torch.nn.functional.relu(input)),文件包含:

  • Functional.relu.0.forward:API 的前向数据。input_args 是位置输入,input_kwargs 是关键字输入,output 存储前向输出。

  • Functional.relu.0.backward:API 的后向数据。input 表示前向输出的梯度,output 表示流向前向输入的梯度。

{
 "task": "tensor",
 "level": "L1",
 "framework": "pytorch",
 "dump_data_dir":"/dump/path",
 "data": {
  "Functional.relu.0.forward": {
   "input_args": [
    {
     "type": "torch.Tensor",
     "dtype": "torch.float32",
     "shape": [
      32,
      16,
      28,
      28
     ],
     "Max": 1.3864083290100098,
     "Min": -1.3364859819412231,
     "Mean": 0.03711778670549393,
     "Norm": 236.20692443847656,
     "requires_grad": true,
     "data_name": "Functional.relu.0.forward.input.0.pt"
    }
   ],
   "input_kwargs": {},
   "output": [
    {
     "type": "torch.Tensor",
     "dtype": "torch.float32",
     "shape": [
      32,
      16,
      28,
      28
     ],
     "Max": 1.3864083290100098,
     "Min": 0.0,
     "Mean": 0.16849493980407715,
     "Norm": 175.23345947265625,
     "requires_grad": true,
     "data_name": "Functional.relu.0.forward.output.0.pt"
    }
   ]
  },
  "Functional.relu.0.backward": {
   "input": [
    {
     "type": "torch.Tensor",
     "dtype": "torch.float32",
     "shape": [
      32,
      16,
      28,
      28
     ],
     "Max": 0.0001815402356442064,
     "Min": -0.00013352684618439525,
     "Mean": 0.00011915402356442064,
     "Norm": 0.007598237134516239,
     "requires_grad": false,
     "data_name": "Functional.relu.0.backward.input.0.pt"
    }
   ],
   "output": [
    {
     "type": "torch.Tensor",
     "dtype": "torch.float32",
     "shape": [
      32,
      16,
      28,
      28
     ],
     "Max": 0.0001815402356442064,
     "Min": -0.00012117840378778055,
     "Mean": 2.0098118724831693e-08,
     "Norm": 0.006532244384288788,
     "requires_grad": false,
     "data_name": "Functional.relu.0.backward.output.0.pt"
    }
   ]
  }
 }
}  

混合级别#

mix 级别的 dump.json 包含 L0 和 L1 级别的数据;文件格式与上述示例相同。