本文以某自定义 NPU 后端为例，方法论适用于任何基于 PrivateUse1 机制实现的自定义 PyTorch 后端的精度调试。

在自定义 AI 加速器后端上调试模型精度：以 Qwen3.5-0.8B GatedDeltaNet 为例

适用范围：本文以某自定义 NPU 后端为例，方法论适用于任何基于 PrivateUse1 机制实现的自定义 PyTorch 后端的精度调试。

背景

Qwen3.5 是一个混合架构模型：每4层一个 GQA 注意力，其余层是 GatedDeltaNet（线性注意力）。
后者在 prefill 阶段调用 chunk_gated_delta_rule 核函数，其中包含一段逐行修正的 for 循环。
当我们把模型从 CPU 移到 NPU 上运行时，最终 logit 的 cosine similarity 从 1.0 降到了约 0.999，
看起来不显眼，但对于精确复现语言模型输出来说是不可接受的。

本文记录从”发现精度下降”到”锁定根因、写出回归测试并修复”的完整过程，
重点介绍可复用的 tensor dump 注入方法。

第一步：建立可重复的对比基线

在深入任何细节之前，先确保两件事：

  1. 结果可复现：固定随机种子，使用 model.eval() + torch.no_grad()，两次运行结果完全一致。
  2. 量化指标：用 cosine similarity 作为精度度量——它归一化了量级差异，只看方向偏差。
1
2
3
4
def _cosine(a, b):
    a = a.float().flatten()
    b = b.float().flatten()
    return torch.nn.functional.cosine_similarity(a, b, dim=0).item()

最终 logit 相似度是一个全局指标，发现”有问题”，但无法告诉你问题在哪。
接下来需要一个能在模型内部任意位置采样的工具。

第二步：设计轻量 Tensor Dump 基础设施

核心思想

在模型源码中插入”探针”，当 dump 开关打开时将张量快照到一个全局字典。
同一套代码分别在 CPU 和 NPU 上跑，然后逐键对比。

实现（约 20 行代码）

在模型文件的模块级别（import 之后，类定义之前）加入：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
# ===== Precision Dump Support =====
_dump_enabled: bool = False
_dump_store: dict = {}
_gdr_layer_idx: int = -1   # 辅助变量：供嵌套函数知道当前层号

def enable_tensor_dump() -> None:
    global _dump_enabled
    _dump_enabled = True
    _dump_store.clear()          # 每次 enable 时清空，防止串数据

def disable_tensor_dump() -> None:
    global _dump_enabled
    _dump_enabled = False

def get_tensor_dump() -> dict:
    return dict(_dump_store)     # 返回副本，防止外部修改

def _record(key: str, tensor: torch.Tensor) -> None:
    if _dump_enabled:
        # detach：不影响计算图
        # cpu()：统一在 CPU 上对比
        # float()：消除 bf16/fp16 精度格式差异，只看数值
        # clone()：防止后续 in-place 操作修改快照
        _dump_store[key] = tensor.detach().cpu().float().clone()
# ===== End Precision Dump Support =====

设计要点：

决策 原因
detach() 不污染 autograd 图
.cpu().float() 统一设备和精度，让 CPU/NPU 结果可直接相减
.clone() 快照时刻的值，不受后续 in-place 操作影响
_dump_store.clear() 在 enable 时 防止多次 forward 的数据叠加
全局开关 _dump_enabled 零开销关闭，不影响生产路径

在模型中插入检查点

找到关键的 forward 函数，在每个中间结果之后调用 _record。
键名采用 层号 + 子位置 的层次化命名方案：

1
2
3
4
5
6
# Qwen3_5DecoderLayer.forward() 中
L = self.layer_idx
_record(f"L{L:02d}.0_input",         hidden_states)
_record(f"L{L:02d}.1.0_input_norm",  residual)
# ...
_record(f"L{L:02d}.A_output",        hidden_states)
1
2
3
4
5
# torch_chunk_gated_delta_rule() 中（子函数，需要通过全局变量传递层号）
_pfx = f"L{_gdr_layer_idx:02d}.1.11"
_record(f"{_pfx}.1_after_l2norm_q",  query)
_record(f"{_pfx}.7_intra_attn_raw",  attn)
# ...

命名规范：

  • 格式 L{层:02d}.{阶段}.{子步骤}_{描述} 确保字典键可排序
  • 数字补零（02d）让字典序 = 执行序
  • 子函数通过 _gdr_layer_idx 全局变量感知层号，在调用子函数前设置它：
1
2
3
4
# Qwen3_5GatedDeltaNet.forward() 中，调用前一行
global _gdr_layer_idx
_gdr_layer_idx = self.layer_idx
output = self.chunk_gated_delta_rule(...)

第三步：编写对比脚本

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def run_step_comparison(model, inputs):
    # --- CPU pass ---
    enable_tensor_dump()
    cpu_logits = single_forward(model, inputs, "cpu")
    cpu_dumps = get_tensor_dump()
    disable_tensor_dump()

    # --- NPU pass ---
    enable_tensor_dump()
    npu_logits = single_forward(model, inputs, "npu")
    npu_dumps = get_tensor_dump()
    disable_tensor_dump()

    # --- 逐检查点对比 ---
    common_keys = sorted(set(cpu_dumps) & set(npu_dumps), key=_natural_key)
    for key in common_keys:
        cos = _cosine(cpu_dumps[key], npu_dumps[key])
        max_diff = (cpu_dumps[key] - npu_dumps[key]).abs().max().item()
        flag = " <<<" if cos < 0.999 else ""
        print(f"{key:<45} {cos:>10.6f} {max_diff:>14.6e}{flag}")

陷阱：字典键的排序问题

Python 默认的字符串排序是字典序（lexicographic），导致两个经典 bug：

Bug 1：L00.1.10 < L00.1.9（因为 "10"[1] = "0" < "9"）

Bug 2："_" 的 ASCII 码（95）> "." 的 ASCII 码（46），
所以 L00.1_input_norm 会排在 L00.1.9_xxx 之后！
（解决方案：将子步骤也统一用点分隔，L00.1.0_input_norm）

修复：使用自然排序，把连续数字段解析为整数再比较：

1
2
3
4
5
6
import re

def _natural_key(s: str):
    return [int(c) if c.isdigit() else c for c in re.split(r'(\d+)', s)]

common_keys = sorted(..., key=_natural_key)

第四步：分层缩小——从层到算子

粗粒度：层间对比

首先只在每个 Decoder Layer 的输入/输出插入检查点，观察哪一层首次偏离：

1
2
3
4
5
L00.0_input          1.000002   0.000e+00
L00.1.0_input_norm   1.000003   0.000e+00
...
L00.1.11.8_intra_attn_corrected  0.999778  MaxAbsDiff=0.1587  <<<
L00.1.11.9_value_corrected       0.997791                     <<<

结论：精度损失发生在第 0 层的 GatedDeltaNet 的 intra_attn_corrected 步骤。

细粒度：子函数内部

在 torch_chunk_gated_delta_rule 内部插入每个中间步骤的检查点：

1
2
3
4
L00.1.11.7_intra_attn_raw      0.999993   ← 计算 attn 初值，正常
L00.1.11.7L01_attn_row_written 0.999990   ← i=1，写入后读回，正常？
L00.1.11.7L04_sub              0.993488   ← i=4，读出 sub 时已经错了！
L00.1.11.8_intra_attn_corrected 0.999778  ← for 循环后

关键发现：sub = attn[..., :i, :i].clone() 在 i=4 时读出了错误值。
这说明 i=2、i=3 的写入操作 attn[..., i, :i] = ... 并没有真正生效。

精确复现：最小化验证脚本

一旦锁定了可疑操作，立刻写一个 10 行的独立脚本来验证假设：

1
2
3
4
5
6
7
8
a = torch.zeros(1, 2, 1, 8, 8, dtype=torch.float32, device='npu')
a_cpu = a.cpu().clone()

a_cpu[..., 4, :4] = torch.ones(1, 2, 1, 4)
a[..., 4, :4]     = torch.ones(1, 2, 1, 4, device='npu')

diff = (a.cpu() - a_cpu).abs().max().item()
print(f"5D write diff: {diff}")  # → 0.822  BROKEN!

第五步：深挖根因——理解 aten::select 的 fallback 行为

现象

4D 及以上张量的 tensor[..., i, :i] = value 写入失败，原张量未被修改。

调用链分析

Python 的 a[..., i, :i] = val 在 PyTorch 内部会分解为：

1
2
aten::select.int(a, dim=-2, index=i)  → 得到 row_view
aten::copy_(row_view[:i], val)         → 写入 row_view 的前 i 列

select 本应返回原张量的一个 view（别名），使得 copy_ 直接写入原始内存。

NPU 后端的实现

查看后端 memory.cpp 中的 select_int：

1
2
3
4
5
6
7
8
9
10
at::Tensor select_int(const at::Tensor &self, int64_t dim, c10::SymInt index) {
    ensure_contiguous(self);
    // ... 计算 sizes, strides ...

    if (!view_result_is_contiguous(sizes, strides)) {
        // 结果 view 不连续 → 触发 CPU fallback
        return cpu_fallback_tensor(op_handler, self, dim, index);
    }
    return make_view_tensor(self, sizes, strides, storage_offset);
}

view_result_is_contiguous 检查结果张量的 strides 是否满足 C-contiguous 规则：

1
2
3
4
shape  = (1, H, C, S, S)
select dim=-2 (第 i 行) 后:
  sizes   = (1, H, C, S)
  strides = (H·C·S², C·S², S², 1)   ← stride[2] = S² ≠ S, 不连续！

因此 cpu_fallback_tensor 被触发，它把 self 拷到 CPU 执行 select，
再把结果转回 NPU——返回的是一个独立的新张量，不是原张量的 view。

后续的 copy_ 写入这个临时拷贝，原始 attn 张量毫发无损，
63 次 for 循环迭代全部静默失效。

2D 为何正常？select(dim=0) 后 strides = (1,)，是连续的，走 make_view_tensor 路径。

第六步：修复

Python 层 Workaround（立即可用）

在 for 循环前将张量临时移到 CPU 执行，循环后移回原设备：

1
2
3
4
5
6
7
8
9
_attn_device = attn.device
attn = attn.cpu()          # in-place write 在 CPU 上原生正确

for i in range(1, chunk_size):
    row = attn[..., i, :i].clone()
    sub = attn[..., :i, :i].clone()
    attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)

attn = attn.to(_attn_device)   # 移回 NPU

代价：一次 D2H（chunk_size × chunk_size 的 float32 矩阵，约 16 KB）
和一次 H2D，可接受。

C++ 层根本修复（长期方案）

在 select_int 中，当结果 view 不连续时，不再 fallback 到 CPU，
而是允许后端持有非连续张量并正确处理 copy_ 的写回语义。
这需要后端整体支持非连续 stride 布局，是更大的工程任务。

第七步：写回归测试

发现了 bug 就要固化为测试，防止将来的修复引入退化：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# tests/integration/models/qwen/test_select_inplace_write.py

@pytest.mark.parametrize("shape,index_fn,val_fn,dtype,label", [
    ((4, 4),          ...,  torch.float32, "2D-fp32"),          # baseline
    ((1,2,1,64,64),   ...,  torch.float32, "5D-fp32-qwen"),     # 复现 bug
    ((1,2,1,64,64),   ...,  torch.bfloat16,"5D-bf16-qwen"),
    ...
])
def test_select_inplace_write_propagates(shape, index_fn, val_fn, dtype, label):
    cpu_ref = torch.rand(shape, dtype=dtype)
    npu_t = cpu_ref.to("npu")
    index_fn(cpu_ref, val_fn("cpu"))
    index_fn(npu_t, val_fn("npu"))
    diff = (npu_t.cpu() - cpu_ref).abs().max().item()
    assert diff == 0.0, f"[{label}] write did not propagate (diff={diff})"

运行结果（修复前）：

1
2
3
4
5
PASSED  test_select_inplace_write_propagates[2D-fp32]
FAILED  test_select_inplace_write_propagates[4D-fp32]
FAILED  test_select_inplace_write_propagates[5D-fp32-qwen-shape]
FAILED  test_select_inplace_write_propagates[5D-bf16-qwen-shape]
FAILED  test_intra_chunk_causal_correction_loop

精确复现了边界条件：2D 正常，4D+ 全部 broken。

方法论总结

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
全局 cosine 下降
     │
     ▼
插入层间检查点 → 定位出问题的 Layer
     │
     ▼
插入子函数检查点 → 定位出问题的步骤
     │
     ▼
插入循环迭代检查点 → 定位出问题的迭代
     │
     ▼
写 10 行最小复现脚本 → 确认操作级 bug
     │
     ▼
阅读 C++ 后端实现 → 理解根因
     │
     ▼
写 pytest → 固化复现
     │
     ▼
修复 + 验证精度恢复

可复用的最佳实践

  1. 先建全局指标，再往下钻。不要一开始就看 C++ 代码。
  2. 检查点键名要可排序：使用数字补零 + 统一分隔符，配合自然排序函数。
  3. 每个检查点快照四个属性：.detach().cpu().float().clone()，缺一不可。
  4. 发现可疑操作后立刻最小化：把怀疑的操作提取为 10-20 行独立脚本验证。
  5. 对比 schema 和 C++ 实现：
    • 用 torch.ops.aten.xxx._schema 查 alias 注解
    • 在后端 memory.cpp / fallback.cpp 里确认实际行为
  6. 所有发现都写成 pytest：一个 bug，一个 test case，终身受益。

常见陷阱速查

现象 可能原因
in-place 操作后原张量不变 操作触发了 fallback，返回的是临时拷贝而非 view
仅高维张量出错，低维正常 C++ 的连续性检查：高维 select/slice 结果 stride 不连续
Float32 比 BF16 差异更大 后端可能不支持 Float32，走 CPU fallback，精度取决于 tensor copy 路径
某些迭代正确，某些迭代错误 读-改-写循环中”读”走了正确路径，”写”没有真正生效
多次 enable_tensor_dump 后数据混乱 忘记在 enable 时 clear()，或 disable 前没有保存副本
License

本文采用 署名-非商业性使用-相同方式共享 4.0 国际 许可协议，转载请注明出处。

Share

快来参与讨论吧~