all-MiniLM-L6-v2 on Ascend NPU

1. 简介

本文档记录 all-MiniLM-L6-v2 在昇腾NPU上的适配与验证结果。

all-MiniLM-L6-v2 是 sentence-transformers 库中的一款轻量级句子嵌入模型，将句子映射到384维稠密向量空间，可用于聚类、语义搜索等任务。

2. 验证环境

组件	版本
`python`	`3.10`
`torch`	`2.5.1`
`torch-npu`	`2.5.1`
`transformers`	`5.8.1`
`sentence-transformers`	`5.5.0`

NPU：8 卡（型号 910B3）
操作系统：Linux aarch64

3. 环境准备

3.1 创建Conda环境

conda create -n all-MiniLM-L6-v2 python=3.10 -y
conda activate all-MiniLM-L6-v2

3.2 安装依赖

安装PyTorch和torch-npu（使用华为源加速）：

pip install torch==2.5.1 torchvision torchaudio -f https://download.pytorch.org/whl/torch_npu.html
pip install torch-npu==2.5.1 -f https://download.pytorch.org/whl/torch_npu.html

安装NLP相关库：

pip install transformers sentence-transformers -i https://repo.huaweicloud.com/repository/pypi/simple

安装系统依赖（部分昇腾组件需要）：

pip install decorator psutil attrs cloudpickle tornado ml-dtypes scipy

4. 模型目录结构

原始模型目录包含以下文件：

all-MiniLM-L6-v2/
├── 1_Pooling/                 # 池化层配置
├── config.json                # 模型配置
├── config_sentence_transformers.json
├── configuration.json
├── data_config.json
├── model.safetensors          # 模型权重（safetensors格式）
├── modules.json
├── pytorch_model.bin          # 模型权重（PyTorch格式）
├── sentence_bert_config.json
├── special_tokens_map.json
├── tokenizer.json
├── tokenizer_config.json
└── vocab.txt

5. 推理脚本使用

5.1 基本用法

# 默认使用NPU推理（如可用）
python inference.py

# 使用CPU推理
python inference.py --device cpu

# 自定义输入文本
python inference.py --texts "Hello world" "NPU is fast"

# 指定模型路径
python inference.py --model-path /path/to/model

5.2 脚本参数说明

参数	说明	默认值
`--model-path`, `-m`	模型路径	`/data/xxy/all-MiniLM-L6-v2`
`--device`, `-d`	运行设备：`cpu`、`npu`、`auto`	`auto`
`--texts`, `-t`	输入文本列表	内置测试句
`--batch-size`, `-b`	批处理大小	`32`
`--show-similarity`, `-s`	显示相似度矩阵	`False`

5.3 推理示例输出

======================================================================
all-MiniLM-L6-v2 推理脚本 (昇腾NPU适配版)
======================================================================
[INFO] 检测到 8 个NPU设备可用
[INFO] 正在加载模型: /data/xxy/all-MiniLM-L6-v2
[INFO] 使用设备: npu
[INFO] 模型加载完成
[INFO] 开始编码 6 条文本...
[INFO] 编码完成，耗时: 0.023秒
[INFO] 输出向量维度: (6, 384)

[RESULT] 向量统计:
  - 形状: (6, 384)
  - 数据类型: float32
  - 均值范围: [-0.0249, 0.0312]
  - 标准差范围: [0.9982, 1.0000]

[INFO] 推理完成!

6. 精度与性能评测

6.1 运行评测

python benchmark.py --log-file log.txt

6.2 评测结果

评测日期：2026-05-17

指标	数值
测试样本数	20
批处理大小	32
平均相对误差	0.1172%
最大相对误差	0.1853%
平均余弦相似度	0.999999
精度检查	PASS ✓

6.3 性能对比

设备	平均延迟	吞吐量
CPU	717.76 ms	202.02 tokens/s
NPU	22.88 ms	6336.07 tokens/s
加速比	31.36x

6.4 精度评估标准

精度要求：NPU与CPU输出误差 < 1%

评测方法：

使用相同输入文本分别在CPU和NPU上推理
计算NPU输出与CPU输出的相对误差
计算余弦相似度

评测结果满足精度要求。

7. 文件结构

适配后的文件组织如下：

all-MiniLM-L6-v2-ascend/
├── inference.py              # NPU适配推理脚本
├── benchmark.py             # 精度与性能评测脚本
├── log.txt                  # 评测运行日志
└── README.md                # 本文档

8. 注意事项

环境隔离：建议为每个模型创建独立的Conda环境，避免依赖冲突。
NPU初始化：首次在NPU上运行模型时需要进行图编译，耗时较长，后续运行会显著加速。
精度波动：NPU与CPU的少量数值误差（<1%）是正常现象，主要由浮点运算精度差异导致。
性能调优：如需进一步提升NPU性能，可调整以下环境变量：
```
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
```

9. 常见问题

Q1: NPU不可用

检查torch-npu是否正确安装：

python -c "import torch; print(torch.npu.is_available())"

Q2: 精度误差超标

如果精度误差超过1%，请检查：

模型权重是否完整
输入数据预处理是否一致
是否使用了相同的归一化设置

Q3: 内存不足

减小批处理大小：

python inference.py --batch-size 16

10. 参考资料