ibm-granite/granite-embedding-small-english-r2 昇腾 NPU 适配部署文档

一、模型简介

• 模型名称：ibm-granite/granite-embedding-small-english-r2
• 类型：NLP 句嵌入（encoder-only）
• 架构：IBM Granite Transformer encoder
• 输入：文本句子
• 输出：句嵌入 [hidden=384]
• 参数量：~30M

二、环境依赖

• torch>=2.9.0
• torch_npu>=2.9.0
• transformers>=4.57.6
• numpy

三、分步推理操作流程

1. 环境检查

python inference.py --check-npu

2. 拉权重

pip install "huggingface_hub>=0.34.0,<1.0"
export HF_ENDPOINT=https://hf-mirror.com
export HF_HUB_DISABLE_XET=1
python -c "
import os
os.environ['HF_HUB_DISABLE_XET'] = '1'
from huggingface_hub import snapshot_download
snapshot_download(repo_id='ibm-granite/granite-embedding-small-english-r2', local_dir='./models/granite', allow_patterns=['config.json','tokenizer_config.json','vocab.txt','special_tokens_map.json','pytorch_model.bin'])
print('[OK] 权重拉完')
"
du -sh models/granite/pytorch_model.bin  # 预期 ~120MB

3. 跑推理

python inference.py --model_path ./models/granite

四、测试用例 + 输出结果

• 输入：2 条测试句子
• 输出：句嵌入 [2, 384] + 余弦相似度 cosine(s0, s1)

五、昇腾 NPU 适配关键点

• bf16 权重加载（torch_dtype=torch.bfloat16）
• input_ids 保留 int64 不转 bf16（否则 NPU embedding 报错）
• attn_implementation='eager' 避 torch_npu FlashAttention 不支持 mask 形状坑
• attention_mask 加权 mean pooling 屏蔽 PAD token（重要：不用 last_hidden.mean(dim=1) 会污染短句嵌入）
• 张量形状自检期望 input_ids [2, seq_len] / last_hidden [2, seq_len, 384] / 嵌入 [2, 384]

六、性能参考

• NPU: Ascend910_9362
• 推理时延：2 句 < 0.3s

七、精度验收

• bf16 推理输出与 fp32 数值差异 < 1e-2
• 嵌入范数合理

八、注意事项

• granite-embedding-small 是 encoder-only 无生成头，用 AutoModel 不是 AutoModelForCausalLM
• hidden_size 384（轻量模型）
• attention_mask 加权 mean pooling 屏蔽 PAD 是关键（直接 mean 会污染短句嵌入）

九、复现说明

git clone https://gitcode.com/qq_44936245/granite-embedding-small-english-r2.git
cd granite-embedding-small-english-r2
pip install -r requirements.txt
python inference.py --model_path ./models/granite