Gemma-3-270M-IT 昇腾 NPU 部署实战课程

┌──────────────────────────────────────────────────────┐
│              Flask API 服务层                          │
│  ┌────────────────────────────────────────────────┐    │
│  │  Flask (port 8100)                             │    │
│  │  ├─ /v1/models          → 模型列表             │    │
│  │  └─ /v1/chat/completions → 对话推理接口        │    │
│  └──────────────────┬───────────────────────────────┘    │
│                     │                                 │
│  ┌──────────────────▼───────────────────────────┐    │
│  │  Chat 函数（核心推理链路）                     │    │
│  │  ├─ tokenizer.apply_chat_template → 输入编码  │    │
│  │  ├─ model.generate → 自回归生成               │    │
│  │  └─ clean_response → 特殊 token 清洗          │    │
│  └──────────────────┬───────────────────────────┘    │
│                     │                                 │
├─────────────────────┼────────────────────────────────┤
│       torch_npu + CANN 8.5                           │
│  ┌──────────────┐  ┌─▼──────────────┐                │
│  │  GE 图编译   │  │ CUBE 矩阵乘    │                │
│  │  (算子融合)  │  │  (Attention)   │                │
│  └──────────────┘  └────────────────┘                │
├──────────────────────────────────────────────────────┤
│              昇腾 NPU 硬件层                          │
│  ┌──────────┐ ┌──────────┐ ┌──────────────┐          │
│  │ Scalar   │ │ Vector   │ │ CUBE 矩阵乘  │          │
│  │ 控制流   │ │ 激活函数 │ │ 自回归解码   │          │
│  └──────────┘ └──────────┘ └──────────────┘          │
└──────────────────────────────────────────────────────┘

# [工具安装] 安装 atomgit SDK
!pip install -U atomgit atomgit_hub -i https://mirrors.huaweicloud.com/repository/pypi/simple

# [模型下载] 从 AtomGit 平台下载 Gemma-3-270M-IT 模型
from atomgit_hub import snapshot_download

# 下载模型到本地
model_repo = "hf_mirrors/google/gemma-3-270m-it"
local_dir = "/opt/atomgit/gemma-3-270m-it"

snapshot_download(model_repo, local_dir=local_dir)
print(f"模型已下载到: {local_dir}")

# [模型加载] 加载 Tokenizer 和模型到 NPU
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

MODEL_PATH = "/opt/atomgit/gemma-3-270m-it"

# 加载 tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
print(f"Tokenizer 类: {type(tokenizer).__name__}")

# 检查 chat template
if tokenizer.chat_template:
    print(f"Chat template 已加载，长度: {len(tokenizer.chat_template)}")
else:
    import os
    template_path = os.path.join(MODEL_PATH, "chat_template.jinja")
    if os.path.exists(template_path):
        with open(template_path) as f:
            tokenizer.chat_template = f.read()
        print(f"已从文件加载 chat template，长度: {len(tokenizer.chat_template)}")

# 加载模型（不使用 device_map，完全手动控制）
print("\n正在加载模型到 CPU，请稍候...")

model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
)

print("正在将模型移到 NPU...")
model = model.to("npu:0")

print("✅ 模型加载完成")
print(f"模型所在设备: {next(model.parameters()).device}")

# [推理封装] 封装 Chat 函数和输出清洗
import re

def clean_response(text):
    """清理模型输出中的特殊 token 和乱码"""
    # 1. 移除已知的特殊 token
    text = re.sub(r'<(end_of_turn|start_of_turn|eos|pad|bos|\/s|s>)[^>]*>', '', text)
    text = text.replace("<end_of_turn>", "").replace("<start_of_turn>", "")
    text = text.replace("<eos>", "").replace("<pad>", "").replace("<bos>", "")

    # 2. 移除末尾的孤立非 ASCII 字符（tokenizer 未正确解码的残留）
    text = text.strip()
    while text and len(text) > 1:
        last_char = text[-1]
        if ord(last_char) > 127:
            if len(text) > 1 and ord(text[-2]) <= 127:
                text = text[:-1].strip()
            else:
                break
        else:
            break

    return text.strip()


def chat(messages, max_new_tokens=256, temperature=0.7, top_p=0.95):
    """核心推理函数"""
    inputs = tokenizer.apply_chat_template(
        messages,
        return_tensors="pt",
        add_generation_prompt=True
    )
    # 关键：apply_chat_template 返回 CPU 张量，必须手动迁移到 NPU
    inputs = inputs.to("npu:0")

    with torch.no_grad():
        outputs = model.generate(
            inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=temperature,
            top_p=top_p,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    # 截取生成部分（去除输入 prompt）
    response_tokens = outputs[0][inputs.shape[1]:]
    raw_response = tokenizer.decode(response_tokens, skip_special_tokens=True)
    response = clean_response(raw_response)

    return response

# [预热测试] 预热推理 + 性能测量
import time

messages = [{"role": "user", "content": "Tell me a short story about a robot"}]

# 预热（触发算子编译）
_ = chat(messages, max_new_tokens=50)

# 正式测试
start = time.time()
response = chat(messages, max_new_tokens=256)
elapsed = time.time() - start

tokens_generated = len(tokenizer.encode(response))
print(f"生成 {tokens_generated} tokens，耗时 {elapsed:.2f}s")
print(f"速度: {tokens_generated/elapsed:.1f} tokens/s")
print(f"\nResponse: {response[:500]}...")

# [API 服务] 启动 Flask API 服务
from flask import Flask, request, jsonify
import threading

app = Flask(__name__)

@app.route("/v1/chat/completions", methods=["POST"])
def chat_completions():
    data = request.json
    messages = data.get("messages", [])
    max_tokens = data.get("max_tokens", 256)
    temperature = data.get("temperature", 0.7)

    try:
        response_text = chat(messages, max_new_tokens=max_tokens, temperature=temperature)
        return jsonify({
            "id": "chatcmpl-local",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": "gemma-3-270m-it",
            "choices": [{
                "index": 0,
                "message": {"role": "assistant", "content": response_text},
                "finish_reason": "stop"
            }]
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route("/v1/models", methods=["GET"])
def list_models():
    return jsonify({
        "object": "list",
        "data": [{"id": "gemma-3-270m-it", "object": "model"}]
    })

# 在后台线程启动
def run_app():
    app.run(host="0.0.0.0", port=8100, debug=False)

thread = threading.Thread(target=run_app, daemon=True)
thread.start()
print("✅ API 服务已启动在 http://localhost:8100")

# [API 测试] 验证 API 服务
import requests

url = "http://localhost:8100/v1/chat/completions"
payload = {
    "model": "gemma-3-270m-it",
    "messages": [{"role": "user", "content": "hello, who are you"}],
    "max_tokens": 128,
    "temperature": 0.7
}

response = requests.post(url, json=payload, timeout=60)
print(response.json()["choices"][0]["message"]["content"])

# ❌ 错误：使用 device_map="auto"
model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="auto")

# ✅ 正确：手动控制设备迁移
model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=torch.bfloat16)
model = model.to("npu:0")

inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True)
inputs = inputs.to("npu:0")  # 关键：手动迁移到 NPU

GET /v1/models

{"object": "list", "data": [{"id": "gemma-3-270m-it", "object": "model"}]}

POST /v1/chat/completions
Content-Type: application/json