bert-tiny-chinese on Ascend NPU

1. 简介

本文档记录 HFL（哈工大讯飞联合实验室）发布的 bert-tiny-chinese 中文轻量 BERT 模型在昇腾 NPU（Ascend 910B3）上的迁移适配、精度评测与性能验证结果。

该模型是 BERT-tiny 架构的中文版本（L=2, H=128，仅 3.2M 参数），是目前最小的中文 BERT 变体之一。通过 Mean Pooling 可提取 128 维句嵌入，适合资源极度受限的中文 NLP 场景（如移动端、边缘设备、IoT）。

相关获取地址：

• 权重下载地址（ModelScope）：hfl/bert-tiny-chinese

2. 验证环境

• NPU：8 × Ascend 910B3
• 精度对比基准：CPU（x86, PyTorch 2.8.0）
• 说明：仅含 vocab.txt，使用 slow tokenizer（BertTokenizer）

3. 部署使用流程

3.1 环境准备

conda create -n bert-tiny-chinese python=3.11 -y
conda activate bert-tiny-chinese

pip install torch==2.8.0 torch_npu==2.8.0.post4 \
    -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install transformers numpy \
    -i https://pypi.tuna.tsinghua.edu.cn/simple

3.2 推理

python inference.py --text "今天天气很好。" --device npu

from inference import BGEEncoder
encoder = BGEEncoder(model_path="./bert-tiny-chinese", device="npu")
embeddings = encoder.encode(["今天天气很好。"])
# embeddings.shape → (1, 128)

4. Smoke 验证

python inference.py --text "今天天气很好。" --device npu

预期输出：128 维嵌入向量，无运行时错误。

5. 性能参考

测试条件：23 条多样化测试句子，batch_size=32，NPU 预热 1 轮。

BERT-tiny 仅 2 层 128 维，推理速度极快（2037 sent/s），是已适配模型中最快的之一。

6. 精度评测

6.1 评测方法

分别在 CPU 和 NPU 上对 23 条测试句子推理，比较 128 维嵌入向量的余弦相似度、MSE、语义相似度矩阵 Pearson 相关系数。

6.2 评测结果

结论：精度误差率 0.0000%，NPU 与 CPU 输出完全等价，评测通过。

7. 迁移适配说明

7.1 模型结构

• Backbone：BertModel（L=2, H=128，BERT-tiny 极轻量架构）
• Pooling：Mean Pooling（mask 加权平均）
• Normalize：L2 归一化
• Tokenizer：BERT WordPiece（vocab.txt），中文按字分词
• 参数量：3.2M（极轻量，BERT-base 的 1/35）

7.2 适配要点

1. AutoModel.from_pretrained() 加载，model.to("npu:0") 迁移
2. 仅含 vocab.txt，需 use_fast=False 使用 slow tokenizer
3. Mean Pooling 手动实现
4. 2 层 128 维使算子编译和推理都极快

7.3 关键代码

from transformers import AutoModel, AutoTokenizer
model = AutoModel.from_pretrained("bert-tiny-chinese").to("npu:0")
tokenizer = AutoTokenizer.from_pretrained("bert-tiny-chinese", use_fast=False)
inputs = tokenizer("今天天气很好。", return_tensors="pt")
inputs = {k:v.to("npu:0") for k,v in inputs.items()}
with torch.no_grad():
    emb = model(**inputs).last_hidden_state
    mask = inputs["attention_mask"].unsqueeze(-1).float()
    embedding = torch.nn.functional.normalize((emb*mask).sum(1)/mask.sum(1), p=2, dim=1)

8. 注意事项

1. 极轻量模型：仅 3.2M 参数（bert-base 的 1/35），推理速度 2000+ sent/s，适合移动端和边缘设备
2. slow tokenizer：仅含 vocab.txt，需使用 BertTokenizer（slow），不支持 BertTokenizerFast
3. 中文按字分词：中文文本按字符切分，128 维嵌入表达能力有限，适合简单语义匹配
4. 首次 NPU 推理：2 层 Transformer 算子编译约 1-2 秒，极快
5. HFL 系列：哈工大讯飞联合实验室发布的中文轻量 BERT 系列，含 tiny/mini/small/medium 等变体