OLMoE-1B-7B-0125 NPU 适配验证
#+NPU
1. 简介
本仓库记录 OLMoE-1B-7B-0125 模型在 Ascend NPU 上的适配和验证结果。
2. 验证环境
| 组件 | 版本 |
|---|
| Python | 3.11.14 |
| PyTorch | 2.9.0+cpu |
| torch_npu | 2.9.0 |
| transformers | 4.57.6 |
| vllm-ascend | 0.18.0rc1 |
| CANN | 8.5.1 |
| NPU | Ascend 910B2 × 4 卡 |
| modelscope | 1.36.3 |
3. 模型下载
modelscope download --model allenai/OLMoE-1B-7B-0125
4. 模型架构
| 参数 | 值 |
|---|
| 架构 | OlmoeForCausalLM |
| 隐藏层大小 | 2048 |
| 层数 | 16 |
| 注意力头数 | 16 (键值头数: 16) |
| 中间层大小 | 1024 |
| 专家数量 | 64 (激活 8) |
| 词汇表大小 | 50304 |
| 总参数量 | ~7B (激活 ~1B) |
5. 基础推理验证
#!/usr/bin/env python3
import torch_npu
from vllm import LLM, SamplingParams
MODEL_PATH = "/home/openmind/volume/modelscope/hub/models/allenai/OLMoE-1B-7B-0125"
llm = LLM(
model=MODEL_PATH,
trust_remote_code=True,
dtype="float16",
tensor_parallel_size=4,
max_model_len=4096,
gpu_memory_utilization=0.8,
enforce_eager=True,
)
sampling = SamplingParams(max_tokens=64, temperature=0)
outputs = llm.generate(["The capital of France is"], sampling)
print(outputs[0].outputs[0].text)
冒烟测试结果 (4/4 通过)
| 测试 | 输入 | 输出 |
|---|
| 常识推理 | The capital of France is | Paris. |
| 代码生成 | def fibonacci(n): | a, b = 0, 1 |
| 翻译 | Translate to English: Bonjour le monde | , je suis une jeune fille de 16 ans, je suis une étudiante en médecine... |
| 数学运算 | 2 + 3 * 4 = | (代码块输出) |
6. 精度评测
ARC-Challenge 25轮测试
| 指标 | 硬件：NPU（vLLM-Ascend tp=4） | 基线（GPU） | 误差 |
|---|
| 准确率 | 26.11%（306/1172） | 26.50% | -0.39% |
| 评测时间 | 163秒 | - | - |
| 吞吐量 | 7.18 样本/秒 | - | - |
7. 性能基准
输入长度与输出吞吐量对比（输出长度=128，提示词数量=10）：
| 输入长度 | 吞吐量（tokens/秒） |
|---|
| 32 | 182.48 |
| 128 | 173.26 |
| 512 | 170.18 |
| 1024 | 165.69 |
| 2048 | 168.74 |
8. 注意事项
- 模型为MoE架构（64专家），需4卡TP加载（每卡约3.2GB）
enforce_eager=True 确保在NPU上使用eager模式- 吞吐量约165-182 tokens/秒
