pai-bert-tiny-zh on Ascend NPU

1. 简介

本文档记录 PAI-BERT-tiny 中文预训练语言模型在昇腾 NPU (Ascend 910B3) 上的迁移适配与精度验证结果。

PAI-BERT-tiny 是阿里 PAI 团队基于 BERT-tiny 架构训练的中文预训练模型。模型仅 4 层 Transformer，128 维隐层，约 4.4M 参数，是极轻量级的中文语言模型，适用于资源受限场景的文本表示和特征提取。

相关获取地址：

• 权重下载地址（ModelScope）：https://modelscope.cn/models/damo/pai-bert-tiny-zh
• 权重下载地址（HuggingFace Mirror）：https://hf-mirror.com/damo/pai-bert-tiny-zh

参考文档：

• https://modelscope.cn/models/damo/pai-bert-tiny-zh

2. 验证环境

• NPU：Ascend 910B3，单卡推理
• 模型大小：约 4.4M 参数（极轻量）
• 输入：中文文本（最大 512 tokens）
• 输出：[1, seq_len, 128] 隐层表示

3. 推理启动

3.1 环境准备

conda create -n pai-bert-tiny-zh python=3.10 -y
conda activate pai-bert-tiny-zh

pip install torch==2.8.0 torch_npu==2.8.0.post4 \
    -i https://mirrors.huaweicloud.com/repository/pypi/simple \
    --trusted-host mirrors.huaweicloud.com

pip install transformers modelscope numpy

3.2 下载模型权重

modelscope download --model damo/pai-bert-tiny-zh \
    --local_dir ./pai-bert-tiny-zh

3.3 命令行推理

python inference.py --text "测试文本" --device npu

3.4 Python API

from inference import run_inference
result = run_inference("测试文本", device="npu")
print(f"向量维度: {result['embedding'].shape}, 耗时: {result['time_ms']:.1f}ms")

4. Smoke 验证

python evaluate.py

验证结果：

• CPU 推理正常，输出 [1, seq_len, 128]
• NPU 推理正常，输出同形状
• 精度评测通过

5. 性能参考

测试条件：输入中文短句。

6. 精度评测

6.1 评测方法

对比 CPU (PyTorch) 与 NPU (torch_npu) 在相同输入下 BERT 模型的 last_hidden_state 输出。

6.2 评测结果

7. 注意事项

1. 极轻量模型：仅 4.4M 参数，128 维隐层，CPU 推理已很快（51ms），NPU 适合批量推理
2. 标准 HuggingFace 格式：使用 AutoModel.from_pretrained 直接加载
3. 应用场景：适合资源受限的文本分类、特征提取等下游任务
4. 首次推理：NPU 首次推理有算子编译开销，后续稳定在 ~5ms