all-MiniLM-L6-v2 on Ascend NPU

1. 简介

本文档记录 sentence-transformers/all-MiniLM-L6-v2 句嵌入模型在昇腾 NPU（Ascend 910B3）上的迁移适配、精度评测与性能验证结果。

all-MiniLM-L6-v2 是一个基于 BERT 的轻量级句嵌入模型（6 层 Transformer，384 维隐藏层），通过 Mean Pooling + L2 归一化将文本映射到 384 维语义向量空间，广泛用于语义搜索、文本聚类和相似度计算等任务。

2. 验证环境

组件	版本
`torch`	`2.8.0`
`torch_npu`	`2.8.0.post4`
`transformers`	`5.8.1`
`sentence-transformers`	`5.5.0`
`CANN`	`8.5.1`

NPU：8 × Ascend 910B3
精度对比基准：CPU（x86, PyTorch 2.8.0）

3. 部署使用流程

3.1 环境准备

# 创建 conda 环境
conda create -n all-MiniLM-L6-v2 python=3.11 -y
conda activate all-MiniLM-L6-v2

# 安装依赖（使用华为源加速）
pip install torch==2.8.0 torch_npu==2.8.0.post4 \
    -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install transformers sentence-transformers numpy \
    -i https://pypi.tuna.tsinghua.edu.cn/simple

3.2 模型权重下载

# ModelScope 下载
modelscope download --model sentence-transformers/all-MiniLM-L6-v2 \
    --local_dir ./all-MiniLM-L6-v2

# 或 HuggingFace 镜像下载
HF_ENDPOINT=https://hf-mirror.com \
    huggingface-cli download sentence-transformers/all-MiniLM-L6-v2 \
    --local-dir ./all-MiniLM-L6-v2

3.3 推理脚本使用

提供 inference.py 作为独立的句嵌入推理脚本：

# 单条文本推理
python inference.py \
    --model_path ./all-MiniLM-L6-v2 \
    --text "Hello world" \
    --device npu

# 批量文本推理
python inference.py \
    --model_path ./all-MiniLM-L6-v2 \
    --batch_file texts.txt \
    --device npu \
    --batch_size 32

编程接口调用：

from inference import MiniLMEncoder

encoder = MiniLMEncoder(model_path="./all-MiniLM-L6-v2", device="npu")
embeddings = encoder.encode(["Hello world", "Another sentence"])
# embeddings.shape → (2, 384)

4. Smoke 验证

基础功能检查：

python inference.py \
    --model_path ./all-MiniLM-L6-v2 \
    --text "This is a test sentence." \
    --device npu

预期输出：

嵌入维度：384
嵌入向量为归一化浮点数
无运行时错误

验证 NPU 设备可用性：

import torch
import torch_npu
assert torch.npu.is_available(), "NPU 不可用"
print(f"NPU 数量: {torch.npu.device_count()}")
print(f"NPU 型号: {torch.npu.get_device_name(0)}")

5. 性能参考

测试条件：24 条多样化测试句子（含短句、长句、标点、多语言），batch_size=32，NPU 预热 3 轮后取 10 轮平均。

指标	数值
CPU 吞吐量	`42.32` sentences/s
NPU 吞吐量 (bs=32)	`1,123.75` sentences/s
NPU 延迟 (bs=1)	`9.83` ms
NPU 延迟 (bs=32)	`21.36` ms
CPU/NPU 加速比	`28.34` ×

不同 batch size 的 NPU 吞吐量：

Batch Size	延迟 (ms)	吞吐量 (sentences/s)
1	9.83	101.70
4	11.15	358.75
16	15.91	1,005.36
32	21.36	1,123.75
64	20.09	1,194.44

6. 精度评测

6.1 评测方法

采用 24 条多样化测试句子（英文、中文、法文、标点、数字、短句、长句），分别在 CPU 和 NPU 上执行推理，计算以下指标：

NPU-CPU 余弦相似度：每条句子 NPU 嵌入与 CPU 嵌入的余弦相似度，衡量逐句精度
元素级误差：NPU 嵌入与 CPU 嵌入间的最大绝对误差和均方误差（MSE）
相似度矩阵 Pearson 相关系数：NPU 与 CPU 各自计算的语义相似度矩阵之间的相关性，衡量语义结构一致性

6.2 评测结果

指标	数值
平均余弦相似度	`0.999999`
最小余弦相似度	`0.999998`
精度误差率	`0.0001%`
最大元素绝对误差	`0.000413`
MSE	`0.00000000`
相似度矩阵 Pearson r	`1.000000`

结论：精度误差率 0.0001%，远低于 1% 要求，评测通过。

7. 迁移适配说明

7.1 模型结构

all-MiniLM-L6-v2 是一个 sentence-transformers 格式的模型，由三个子模块组成：

Transformer（BERT-base，6 层）：文本编码
Pooling（Mean Pooling）：对所有 token 嵌入取加权平均
Normalize（L2 归一化）：将嵌入向量归一化到单位球面

7.2 适配要点

由于该模型为标准 BERT 架构，NPU 适配过程无需修改模型结构：

模型加载：使用 AutoModel.from_pretrained() 加载权重，model.to("npu:0") 迁移到 NPU
Pooling 实现：手动实现 Mean Pooling（sentence-transformers 库的 Pooling 模块依赖 torch.nn.Module，可正常在 NPU 上运行）
输入处理：AutoTokenizer 分词后在 CPU 执行，input_ids/attention_mask 通过 .to(device) 转移到 NPU
输出处理：嵌入向量通过 .cpu().numpy() 转回 CPU

7.3 关键代码片段

# NPU 推理核心流程
import torch
import torch_npu
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("all-MiniLM-L6-v2").to("npu:0")
model.eval()

tokenizer = AutoTokenizer.from_pretrained("all-MiniLM-L6-v2")
inputs = tokenizer("Hello world", return_tensors="pt")
inputs = {k: v.to("npu:0") for k, v in inputs.items()}

with torch.no_grad():
    outputs = model(**inputs)
    embeddings = mean_pooling(outputs, inputs["attention_mask"])

8. 注意事项

NPU 内存管理：首次推理时会触发算子编译，耗时较长。建议在正式推理前执行一次预热调用。
batch_size 选择：batch_size ≥ 16 可充分发挥 NPU 并行计算能力，吞吐量提升显著。
与 sentence-transformers 库兼容：NPU 上可直接使用 AutoModel + Mean Pooling 替代 SentenceTransformer.encode()，精度完全一致。
多 NPU 场景：当前推理使用单 NPU（npu:0），如需多卡并行可指定不同 device ID。
输入最大长度：默认 max_seq_length=256，与原始 sentence_bert_config.json 配置一致，超出部分将被截断。