meta-llama/Llama-3.2-1B 昇腾 NPU 适配报告

# 1. 加载 CANN 环境（必须）
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 2. 配置华为镜像源
export PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple/
export HF_ENDPOINT=https://hf-mirror.com

# 3. 安装依赖
pip install transformers accelerate sentencepiece protobuf numpy==1.26.4

# 4. 验证 NPU 可用
python3 -c "import torch; import torch_npu; a = torch.randn(3,4).npu(); print(a + a)"

import torch_npu
from torch_npu.contrib import transfer_to_npu

# 显式初始化 NPU（比 torch.cuda.is_available() 更可靠）
try:
    torch.npu.init()
    npu_available = True
except Exception:
    npu_available = False

device = torch.device("cuda" if npu_available else "cpu")

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "unsloth/Llama-3.2-1B"  # 或 meta-llama/Llama-3.2-1B（需授权）

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    dtype=torch.bfloat16,  # Ascend910 原生支持 bf16
    trust_remote_code=True,
)
model = model.to(device)
model.eval()

torch.manual_seed(42)
inputs = tokenizer(prompt, return_tensors="pt", padding=True)
inputs = {k: v.to(device) for k, v in inputs.items()}

with torch.no_grad():
    # device_type 传 "cuda"，transfer_to_npu 内部映射为 npu
    with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
        outputs = model.generate(
            **inputs,
            max_new_tokens=50,
            do_sample=True,
            top_k=50,
            top_p=0.95,
            temperature=0.8,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

python3 llama32_npu_greedy_verify.py

[Prompt] The future of artificial intelligence is
[Output] The future of artificial intelligence is here: The new artificial intelligence system uses the power of music to predict how you will respond to information...

[Prompt] Once upon a time in a distant galaxy
[Output] Once upon a time in a distant galaxy, there was a great warrior. He was known for his courage, strength, and skill...

[Prompt] The key to success is
[Output] The key to success is to choose the right kind of marketing campaign that will effectively promote your products...

python3 llama32_npu_perf_test.py

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export HF_ENDPOINT=https://hf-mirror.com
python3 llama32_npu_inference.py

python3 llama32_npu_greedy_verify.py

python3 llama32_npu_perf_test.py