BAAI/llm-embedder on Ascend 910B3

1. 简介

本文档记录 BAAI/llm-embedder 在昇腾 Ascend 910B3 NPU 上的迁移适配、推理部署与精度评测结果。

LLM-Embedder 是 BAAI 推出的统一文本嵌入模型，基于 BERT-base 架构（12 层 Transformer，768 隐藏维），参数量约 109.48M。该模型可将任意文本映射为低维稠密向量，用于检索、聚类、语义搜索等任务。

本次适配工作包括：

在 NPU（Ascend 910B3）上验证文本嵌入推理的正确性
对比 NPU 与 CPU 的输出精度，确保误差 < 1%（以余弦相似度衡量）
提供可直接使用的 NPU 推理脚本 inference.py
提供精度与性能评测脚本 eval.py

2. 验证环境

组件	版本
`Python`	`3.9.13`
`torch`	`2.8.0+cpu`
`torch_npu`	`2.8.0.post4`
`transformers`	`4.57.6`
`numpy`	`1.24.4`

NPU：Ascend 910B3 × 8 逻辑卡
驱动版本：25.5.2

3. 模型适配与部署

3.1 适配说明

LLM-Embedder 使用标准 BERT 架构，transformers 库原生支持。NPU 适配无需修改模型结构或权重。

已验证通过的适配流程：

import torch
import torch_npu
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("BAAI/llm-embedder")
model = model.npu()
model.eval()

tokenizer = AutoTokenizer.from_pretrained("BAAI/llm-embedder")
inputs = tokenizer(["Embed this text"], return_tensors="pt")
inputs = {k: v.npu() for k, v in inputs.items()}

with torch.no_grad():
    outputs = model(**inputs)
    embedding = outputs.last_hidden_state[:, 0, :].cpu()

3.2 环境准备

pip install torch torch_npu transformers -i https://repo.huaweicloud.com/repository/pypi/simple/
export HF_ENDPOINT=https://hf-mirror.com

3.3 推理脚本使用

# NPU 推理（CLS 池化）
python inference.py --text "Climate change is a global challenge."

# Mean 池化 + L2 归一化
python inference.py --text "Embed this text" --pooling mean --normalize

# 从文件批量推理
python inference.py --input texts.txt --output embeddings.json

# CPU 推理
python inference.py --text "Text" --device cpu

4. Smoke 验证

python inference.py --text "Climate change is a major global challenge."

预期输出：

Embedding Results (NPU)
  Pooling: cls, Normalize: False
  Embedding dim: 768
  Output shape: [1, 768]
  ...
  Total: 0.xxxs | Avg: xxx.xxms/text

5. 性能参考

测试条件：batch_size=20，seq_len=64，float32 精度，连续 10 次取平均。

指标	CPU	NPU (Ascend 910B3)
平均推理时间 (20 samples)	~1540 ms	~21 ms
单样本平均耗时	~77 ms	~1 ms
加速比	1x	~70x
参数量	109.48M	109.48M
模型大小	417.6 MB	417.6 MB

6. 精度评测

评测方法

在 CPU 上加载模型并推理得到参考嵌入向量
在 NPU 上加载同一权重并推理得到 NPU 嵌入向量
计算余弦相似度（嵌入模型的主要精度指标）

评测结果

使用 20 条科技文本进行评测：

指标	数值	要求	结果
MSE	2.02e-6	-	-
Cosine Similarity (mean)	0.99999666	> 0.999	✓
Cosine Similarity (min)	> 0.99999	> 0.999	✓
Max Absolute Error	6.03e-3	-	-

结论：NPU 输出与 CPU 输出的平均余弦相似度达 0.99999666，满足精度要求。

详细评测日志见 eval_log.txt。

7. 注意事项

权重文件：NPU 适配无需修改原始权重
池化方式：支持 CLS token 池化和 Mean 池化（--pooling cls|mean）
归一化：支持 L2 归一化（--normalize），用于余弦相似度检索
设备选择：脚本默认自动检测 NPU，若 NPU 不可用则回退到 CPU
torch_npu 版本：确保与 torch 版本匹配
单卡推理：当前使用单张 NPU 卡