Pix2Struct-Base NPU 适配与性能调优报告

本报告记录 google/pix2struct-base 在华为昇腾 NPU (Ascend 910B) 上的推理适配、精度验证与性能调优结果。

模型信息

硬件与软件环境

NPU 推理跑通验证

验证方法

使用随机权重初始化模型，构造模拟预处理后的图像 patch 输入，在 CPU 和 NPU 上分别执行完整的前向传播与 .generate() 自回归解码，验证模型在 NPU 上可正常完成推理。

验证脚本

python run_pix2struct_test.py

验证结果

前向传播验证

完整生成验证 (.generate())

结论：Pix2Struct-Base 在昇腾 NPU 上推理链路完全跑通，encoder 与 decoder 均正常工作。

精度验证

验证策略

• 对比基准：CPU FP32（同权重、同输入、固定 seed）
• 验证维度：前向传播 logits、完整 generate 输出、AMP FP16 混合精度
• 阈值声明：Token 一致率 >= 99%（误差 < 1%）

前向传播精度

完整生成精度 (.generate(), max_new_tokens=32)

精度异常说明

• 直接 .half() 转 FP16 会导致数值崩溃（输出全为 0），一致率仅 3.03%。
• 解决方案：使用 torch.autocast(device_type="npu", dtype=torch.float16) 混合精度，可在保持 100% Token 一致率 的同时启用 FP16 计算。

精度结论

所有有效配置下，NPU 与 CPU 的 Token 一致率均为 100.00%，满足误差 < 1% 的要求。

精度深入分析

在基础精度验证通过后，进一步对逐层误差累积、长序列生成漂移、不同 batch 规模以及多数据类型进行了边界测试。

逐层误差累积 (FP32)

使用 output_hidden_states=True 提取 encoder 和 decoder 每层输出，对比 CPU 与 NPU 的逐层数值差异。

发现：误差随网络深度缓慢累积（Encoder 从 3e-6 增至 1e-5，Decoder 从 0 增至 2e-5），但绝对值仍远小于 1% 阈值。

长序列生成漂移 (64 steps)

手动执行 64 步自回归解码，逐 step 对比 logits 相对误差，观察误差是否随时间累积。

关键发现：贪心解码下 64 步无 Token 分歧，但 logits 数值差异呈上升趋势。对于 temperature > 0 的采样生成，长时间生成可能产生不同结果。

Batch 精度一致性

结论：并行计算未引入额外精度损失，batch 规模对精度无显著影响。

数据类型精度边界

重要结论：

• BF16 在当前模型与硬件组合下精度完全不可用
• Direct FP16 有 3.12% Token 不一致风险
• AMP FP16 是 FP16 的唯一安全选择，但仅限于贪心解码场景

性能基准与调优

性能基线定义

本报告在昇腾 NPU (Ascend 910B) 上定义并验证了以下可复现的性能基线指标。所有指标均在固定测试口径下测得：FP32 精度、seq_len=1024、seed=42、预热 3 轮、测试 10 轮取平均、torch.npu.synchronize() 同步计时。

基线达成结论：全部 7 项定义基线均已达成，其中批处理吞吐（26.68 img/s）远超设定目标（> 6 img/s）。

与公开社区数据对比

以下为 AtomGit AI 社区在同等硬件（Ascend 910B）上的公开数据，用于外部对标。

口径对齐说明：社区精细基线 172.8 ms 经本报告验证，实际对应 16tok 生成长度（本报告 16tok 为 166.84 ms，已优于该基线），而非 32tok。此前直接将 32tok 结果（343.75 ms）与 172.8 ms 对比属于口径不一致，现已修正。

本报告结果与社区基线对比（同口径）

测试方法

• 固定 seed=42，贪心解码 (do_sample=False, num_beams=1)
• 预热 3 轮，测试 10 轮取平均
• 同步方式：torch.npu.synchronize()

优化前后对比

性能提升：延迟从 723.97ms 降至 343.75ms，提升约 52.5%。

不同生成长度的延迟对比 (seq=1024)

注：decoder 自回归生成是主要瓶颈，延迟与生成 token 数近似线性相关。

不同序列长度的延迟对比 (max_new_tokens=32)

注：encoder attention 为 O(n²) 复杂度，缩短 seq_len 可显著降低延迟。

NPU 通用性能调优（流水优化）

使用 /ai4s-perf-tuning Skill 进行标准化性能调优，测试条件：seq=1024, max_new_tokens=16, Batch=1, FP32。

结论：TASK_QUEUE_ENABLE=2 带来显著正向收益（延迟 -26.6%，吞吐 +36.3%）；CPU_AFFINITY_CONF=1 在本场景下为负优化，已回退；tcmalloc 因环境限制未能完成测试。

优化手段总结

已知问题与限制

1. FP16 数值稳定性：直接 model.half() 会导致输出全 0，必须使用 torch.autocast 混合精度。
2. torch.compile 收益有限：当前 torch_npu 对 torch.compile 的优化收益较小（<2%），可能与模型规模（282M）较小有关。
3. Batch 推理未加速：在 seq_len=1024/4096 配置下，batch=4 的延迟增长超过了并行收益，建议根据实际硬件内存和输入尺寸重新评估 batch 策略。
4. Internal Format Warning：运行中出现 Cannot create tensor with internal format 警告，当前不影响正确性，但可能影响极致性能。
5. CPU 绑核负优化：CPU_AFFINITY_CONF=1 在本模型场景下导致延迟增加约 1.8%，已回退。
6. tcmalloc 安装受限：当前环境无 root 权限且外网访问受限，未能完成 tcmalloc 安装与测试。

快速开始

环境准备

pip install torch==2.9.0 transformers==4.57.6
# 安装对应 CANN 版本的 torch_npu

推理示例

import torch
import torch_npu
from transformers import Pix2StructForConditionalGeneration, Pix2StructConfig

config = Pix2StructConfig()
config.vision_config.seq_len = 1024  # 推荐使用优化后的序列长度

model = Pix2StructForConditionalGeneration(config).eval().to("npu:0")

# 构造输入 (实际使用时应由 Pix2StructProcessor 预处理)
flattened_patches = torch.zeros(1, 1024, 770)
flattened_patches[:, :, 0] = torch.randint(0, 1024, (1, 1024)).float()
flattened_patches[:, :, 1] = torch.randint(0, 1024, (1, 1024)).float()
flattened_patches[:, :, 2:] = torch.randn(1, 1024, 768)
attention_mask = torch.ones(1, 1024)

with torch.no_grad():
    outputs = model.generate(
        flattened_patches=flattened_patches.to("npu:0"),
        attention_mask=attention_mask.to("npu:0"),
        max_new_tokens=32,
        do_sample=False,
        num_beams=1,
    )
print(outputs)

附录：测试脚本清单

报告信息