Granite 4.1-3B on Ascend NPU

1. 简介

本文档记录 ibm-granite/granite-4.1-3b 在华为昇腾 NPU（Ascend910B4）上的适配与验证结果。

Granite 4.1-3B 是 IBM 开源的稠密 Transformer 语言模型，架构为标准的 GraniteForCausalLM（transformers 原生支持），无自定义 CUDA 算子。适配策略采用轻量级脚本化迁移：通过 torch_npu + transfer_to_npu 自动完成 cuda → npu 映射，结合 accelerate 的 device_map 将模型加载至 NPU 显存。

模型关键配置：

• 架构：GraniteForCausalLM
• 参数量：3B（约 6GB @ bfloat16）
• 层数：40
• 隐藏维度：2560
• KV Heads：8（GQA）
• 上下文长度：131072
• 原生精度：bfloat16

相关获取地址：

• 权重下载地址（HuggingFace）：https://huggingface.co/ibm-granite/granite-4.1-3b
• 权重下载地址（HF Mirror）：https://hf-mirror.com/ibm-granite/granite-4.1-3b

参考文档：

• https://huggingface.co/ibm-granite/granite-4.1-3b
• https://www.ibm.com/granite/docs/

2. 验证环境

• NPU：1 逻辑卡（Ascend910B4）
• NPU 显存：32GB HBM
• 模型路径：ibm-granite/granite-4.1-3b（HuggingFace 缓存）
• 操作系统：Linux 5.10.0-60.18.0.50.aarch64

3. 推理运行

运行前确保 CANN 环境已加载：

source /usr/local/Ascend/ascend-toolkit/set_env.sh

已验证通过的运行命令：

export HF_ENDPOINT=https://hf-mirror.com
export ASCEND_RT_VISIBLE_DEVICES=0

cd /opt/atomgit/granite-4.1-3b-npu
python3 inference_npu.py \
  --model ibm-granite/granite-4.1-3b \
  --prompt "def fibonacci(n):" \
  --max-new-tokens 64 \
  --device npu

启动脚本关键特性：

• 顶部自动注入 torch_npu + transfer_to_npu，完成 CUDA→NPU 自动映射
• device_map="npu" 将模型权重自动分配至 NPU
• dtype=torch.bfloat16，Ascend910B4 原生支持 bf16，无需精度转换
• trust_remote_code=False，标准 transformers 架构，无需自定义代码

4. Smoke 验证

基础检查（NPU 可用性与设备信息）：

python3 -c "import torch; import torch_npu; print(torch.npu.is_available()); print(torch.npu.get_device_properties(0))"

验证结果：

• torch.npu.is_available() 返回 True
• 模型加载成功，设备为 npu:0
• model.generate() 返回 200（正常输出）
• 生成内容符合预期（代码续写通顺）

验证现象示例：

NPU available: True, device count: 1
  NPU:0 name=Ascend910B4, mem=32.0GB
Loading model from: ibm-granite/granite-4.1-3b
Using dtype: torch.bfloat16
Model loaded in 271.0s
Model device: npu:0
Prompt: 'def fibonacci(n):'
Generated 64 tokens in 5.76s (11.1 tok/s)

5. 性能参考

测试条件：单卡 Ascend910B4，bfloat16，batch_size=1，greedy / sample 生成。

6. 精度说明

Granite 4.1-3B 原生权重为 bfloat16，Ascend910B 系列 NPU 对 bf16 有原生硬件支持，因此：

• 无需进行 fp16/bf16 精度转换
• 无需 PTQ / 量化即可直接部署
• 生成结果与 CUDA 基线语义一致（英文 / 代码续写场景）

注：当前验证聚焦于推理通路与基础功能，未执行大规模评测集（如 HumanEval、MMLU）的精度对标。如需严格精度评测，建议使用 lm-evaluation-harness 或 EvalScope 在 NPU 与 CUDA 环境下进行对比。

7. 注意事项

1. 设备可见性

   • 若环境中仅有单卡 NPU（设备 0），请勿设置 ASCEND_RT_VISIBLE_DEVICES=1，否则会导致 NPU is not available 错误。

2. 依赖安装

   • device_map 功能依赖 accelerate 库，若未安装会报错：ValueError: Using a device_map requires accelerate。
   • 建议通过华为镜像安装：pip install accelerate -i https://repo.huaweicloud.com/repository/pypi/simple/

3. 网络与下载

   • 模型权重约 6GB，首次下载建议配置 HuggingFace 镜像：export HF_ENDPOINT=https://hf-mirror.com
   • 若下载超时，huggingface-cli 支持断点续传，可重复执行。

4. 自动迁移限制

   • transfer_to_npu 可覆盖绝大多数 torch.cuda.* 自动映射，但无法处理 .major / .get_device_capability() 等 CUDA 专有属性访问。
   • Granite 4.1-3B 为标准 transformers 模型，内部无此类调用，因此无需额外手动修复。

5. 性能优化建议

   • 当前数据为单 batch、单序列的 baseline。如需提升吞吐，可考虑：
     • 使用 vLLM-Ascend 进行 continuous batching 服务化部署
     • 开启图编译缓存（ASCEND_AICPU_PATH）减少二次启动耗时
     • 调整 max_position_embeddings 与实际输入长度匹配，避免过量内存预留