ModernBERT-base on Ascend NPU

1. 简介

本文档记录 answerdotai/ModernBERT-base 模型在昇腾 NPU（Ascend 910B3）上的迁移适配、精度评测与性能验证结果。

ModernBERT 是 Answer.AI 推出的现代化 BERT 变体，采用 22 层 Transformer、768 维隐藏层、8192 最大位置编码，支持长文本编码。本适配使用 Mean Pooling 提取句嵌入，适用于语义搜索、文本聚类等场景。

相关获取地址：

• 权重下载地址（ModelScope）：answerdotai/ModernBERT-base
• 权重下载地址（HuggingFace）：https://huggingface.co/answerdotai/ModernBERT-base

2. 验证环境

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

3. 部署使用流程

3.1 环境准备

conda create -n ModernBERT-base python=3.11 -y
conda activate ModernBERT-base
pip install torch==2.8.0 torch_npu==2.8.0.post4 tiktoken \
    -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install transformers numpy \
    -i https://pypi.tuna.tsinghua.edu.cn/simple

ModernBERT 使用 tiktoken 分词器，需要额外安装 tiktoken。

3.2 推理脚本使用

python inference.py --text "Hello world" --device npu
python inference.py --batch_file texts.txt --device npu --batch_size 16

编程接口：

from inference import ModernBERTEncoder
encoder = ModernBERTEncoder(model_path="./ModernBERT-base", device="npu")
embeddings = encoder.encode(["Hello world"])
# embeddings.shape -> (1, 768)

4. Smoke 验证

python inference.py --text "This is a test." --device npu

预期输出：嵌入维度 768，无运行时错误。

5. 性能参考

测试条件：23 条多样化测试句子，batch_size=16。

6. 精度评测

6.1 评测方法

分别在 CPU 和 NPU 上对 23 条多样化测试句子推理（含中英文、短句、长句、特殊字符），计算 NPU 与 CPU 嵌入的余弦相似度。

6.2 评测结果

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

7. 迁移适配说明

7.1 模型结构

ModernBERT-base 由以下核心组件构成：

• Backbone：ModernBertModel（22 层 Transformer，768 维隐藏层，8192 最大位置编码）
• 参数量：149M（Base 规模，接近 BERT-large 的参数量但推理效率更高）
• Head：MLM Head（head.dense, head.norm, decoder.bias），仅预训练使用，推理时通过 AutoModel 加载忽略
• Tokenizer：tiktoken BPE（不同于 BERT 的 WordPiece），需要安装 tiktoken 包

7.2 适配要点

1. 使用 AutoModel.from_pretrained() 加载，MLM head 权重标记为 UNEXPECTED（安全忽略，不影响推理）
2. model.to("npu:0") 迁移到 NPU
3. 使用 Mean Pooling（而非 CLS token）提取句嵌入，提供更丰富的语义表示
4. tiktoken tokenizer 在 CPU 端执行分词，tensor 转移至 NPU
5. 长文本（8192 tokens）场景需调整 batch_size 以避免 NPU 显存溢出

7.3 关键代码

import torch, torch_npu
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("ModernBERT-base").to("npu:0")
tokenizer = AutoTokenizer.from_pretrained("ModernBERT-base")

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 = outputs.last_hidden_state  # shape: (1, seq_len, 768)
    # Mean pooling over token dimension
    mask = inputs["attention_mask"].unsqueeze(-1).float()
    sentence_embedding = (embeddings * mask).sum(1) / mask.sum(1)

8. 注意事项

1. tiktoken 依赖：ModernBERT 使用 tiktoken 分词器，环境需预先安装 pip install tiktoken。缺少 tiktoken 会导致 tokenizer 初始化失败。
2. UNEXPECTED 权重警告：使用 AutoModel 加载时 MLM head 权重（head.dense, decoder.bias, head.norm）会标记为 UNEXPECTED，这是预期行为，不影响编码器推理。
3. NPU 预热：首次推理会触发 NPU 算子编译（22 层 Transformer 编译约 5-8 秒），建议生产环境预热后使用。
4. batch_size 选择：由于模型较大（149M 参数，22 层），建议 batch_size ≤ 16。大 batch 可能导致 NPU HBM 内存溢出。
5. 最大序列长度：ModernBERT 支持 8192 tokens，远超 BERT 的 512。长文本推理时注意 NPU 内存占用随长度平方增长。
6. 与 BERT 对比：ModernBERT 在推理速度上优于同规模 BERT，但 NPU 上 22 层的编译/预热开销高于 12 层的 BERT-base。