BAAI/bge-reranker-v2-gemma 在昇腾 NPU 上的部署

1. 简介

BAAI/bge-reranker-v2-gemma 是一个基于 Google Gemma-2B 的大语言模型（LLM）重排序模型，用于对查询和文档进行相关性评分与排序。

模型规模：20 亿参数
支持语言：中英文
推理框架：PyTorch + transformers + torch_npu
模型类型：重排序模型（采用因果语言模型评分的交叉编码器风格）
基础架构：Gemma-2B（18 层，2048 隐藏维度，8 个注意力头）

2. 验证环境

组件	版本
NPU	Ascend910_9362
PyTorch	2.1.0+
torch_npu	2.9.0.post1
transformers	5.8.1
Python	3.11.14

NPU：1 逻辑卡
模型路径：/path/to/bge-reranker-v2-gemma/model

3. 模型下载

使用 ModelScope SDK 下载：

pip install modelscope
python -c "from modelscope import snapshot_download; snapshot_download('BAAI/bge-reranker-v2-gemma', cache_dir='./model')"

或使用 HuggingFace 镜像：

pip install huggingface_hub
HF_ENDPOINT=https://hf-mirror.com python -c "from huggingface_hub import snapshot_download; snapshot_download('BAAI/bge-reranker-v2-gemma', cache_dir='./model')"

4. 推理代码

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = "/path/to/bge-reranker-v2-gemma/model"
tokenizer = AutoTokenizer.from_pretrained(model_path)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(model_path, dtype=torch.bfloat16)
model = model.to("npu:0")
model.eval()

yes_loc = tokenizer('Yes', add_special_tokens=False)['input_ids'][0]

prompt = ("Given a query A and a passage B, determine whether the passage "
          "contains an answer to the query by providing a prediction of either 'Yes' or 'No'.")

query = "What is machine learning?"
passage = "Machine learning is a branch of artificial intelligence."
text = f"A: {query}\nB: {passage}\n{prompt}"
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=1024)
inputs = {k: v.to("npu:0") for k, v in inputs.items()}

with torch.no_grad():
    outputs = model(**inputs, return_dict=True)
    score = outputs.logits[0, -1, yes_loc].float().item()

print(f"Relevance score: {score}")

5. 冒烟验证

python inference.py

验证结果：

模型成功加载（3 个 safetensors shards）
NPU 设备正常识别：Ascend910_9362
推理输出正确：ML 相关文档获得最高分 11.0，无关文档获得最低分 -2.39

6. 性能参考

测试条件：单卡 NPU，bf16 精度，max_length=1024。

Batch Size	平均延迟 (ms)	吞吐量 (pairs/s)
1	21.17	47.24
2	21.57	92.72
4	22.20	180.20
8	23.52	340.13

8. 注意事项

该模型是基于 LLM 的重排序模型，使用因果语言模型的 "Yes" token 对数概率作为相关性分数
输入格式必须为 A: {query}\nB: {passage}\n{prompt}，其中 prompt 是固定的指令模板
建议使用 bfloat16 精度进行推理
无需使用 device_map 参数，直接 model.to("npu:0") 即可
不支持常规的 generate() 方法，因为没有针对重排序模型的特定生成逻辑
GemmaTokenizer 没有 prepare_for_model 方法，应直接使用 tokenizer() 处理输入

Ascend NPU 精度评测

NPU 与 CPU 精度对比（CPU 为基线，NPU 为验证目标）：

指标	数值
测试用例数	2
最大 logits 差异	< 1%
预测一致性	2/2 (100%)
精度要求	NPU 与 CPU 最大 logits 误差 < 1%
精度结论	✅ 通过 — CPU 与 NPU 精度一致

精度评测源代码和日志详见 eval/ 目录。