Flan-T5 在昇腾 NPU 上的部署与应用

# 查询 NPU 状态
npu-smi info

# 预期输出示例
+------------------------------------------------------------------------------------------------+
| npu-smi 25.5.2                   Version: 25.5.2                                               |
+---------------------------+---------------+----------------------------------------------------+
| NPU   Name                | Health        | Power(W)    Temp(C)           Hugepages-Usage(page)|
| Chip  Phy-ID              | Bus-Id        | AICore(%)   Memory-Usage(MB)  HBM-Usage(MB)        |
+===========================+===============+====================================================+
| 3     Ascend910           | OK            | 163.2       47                0    / 0             |
| 0     6                   | 0000:0A:00.0  | 0           0    / 0          3107 / 65536         |
+------------------------------------------------------------------------------------------------+

# 根据实际安装路径调整
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 关键环境变量
export ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest
export LD_LIBRARY_PATH=${ASCEND_HOME_PATH}/lib64:${LD_LIBRARY_PATH}
export PATH=${ASCEND_HOME_PATH}/bin:${PATH}

# 创建虚拟环境（推荐）
python -m venv flan-t5-npu-env
source flan-t5-npu-env/bin/activate

# 安装 PyTorch 与 torch_npu（已预装则跳过）
pip install torch==2.9.0
pip install torch_npu==2.9.0.post1

# 安装 Transformers 与配套库
pip install transformers>=4.23.1
pip install sentencepiece accelerate

# 安装国内模型下载工具
pip install modelscope

python -c "
import torch
import torch_npu
print(f'PyTorch: {torch.__version__}')
print(f'torch_npu: {torch_npu.__version__}')
print(f'NPU available: {torch.npu.is_available()}')
print(f'Device count: {torch.npu.device_count()}')
"

python -c "
from modelscope import snapshot_download
model_dir = snapshot_download('google/flan-t5-base', cache_dir='./models')
print(f'Downloaded to: {model_dir}')
"

# AtomGit 镜像（如可用）
git clone https://atomgit.com/google/flan-t5.git

python -c "
from transformers import T5ForConditionalGeneration
model = T5ForConditionalGeneration.from_pretrained('google/flan-t5-base')
model.save_pretrained('./models/flan-t5-base')
"

flan-t5-base/
├── config.json              # 模型配置
├── pytorch_model.bin        # PyTorch 权重 (~990MB)
├── model.safetensors        # SafeTensors 格式权重
├── tokenizer_config.json    # 分词器配置
├── spiece.model             # SentencePiece 模型
├── tokenizer.json           # 快速分词器
└── generation_config.json   # 生成配置

import torch
import torch_npu
from transformers import T5ForConditionalGeneration, T5Tokenizer

# 配置
MODEL_PATH = "./models/google/flan-t5-base"
DEVICE = "npu:0"

# 加载模型与分词器
tokenizer = T5Tokenizer.from_pretrained(MODEL_PATH, local_files_only=True)
model = T5ForConditionalGeneration.from_pretrained(
    MODEL_PATH,
    local_files_only=True,
    torch_dtype=torch.float16  # 推荐 FP16 加速
).to(DEVICE)
model.eval()

# 推理示例
text = "translate English to German: The house is wonderful."
inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=50,
        num_beams=4,
        early_stopping=True
    )

result = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(result)  # Das Haus ist wunderbar.

batch_texts = [
    "translate English to German: Hello world.",
    "summarize: The quick brown fox jumps over the lazy dog.",
    "question: What is the capital of France? context: France is in Europe.",
]

inputs = tokenizer(batch_texts, return_tensors="pt", padding=True, truncation=True, max_length=512)
inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=50)

results = tokenizer.batch_decode(outputs, skip_special_tokens=True)
for r in results:
    print(r)

# scripts/inference_server.py
from fastapi import FastAPI
from pydantic import BaseModel
import torch
import torch_npu
from transformers import T5ForConditionalGeneration, T5Tokenizer

app = FastAPI()

MODEL_PATH = "./models/google/flan-t5-base"
DEVICE = "npu:0"

tokenizer = T5Tokenizer.from_pretrained(MODEL_PATH, local_files_only=True)
model = T5ForConditionalGeneration.from_pretrained(
    MODEL_PATH, local_files_only=True, torch_dtype=torch.float16
).to(DEVICE).eval()

class InferRequest(BaseModel):
    text: str
    max_new_tokens: int = 50
    num_beams: int = 4

@app.post("/generate")
def generate(req: InferRequest):
    inputs = tokenizer(req.text, return_tensors="pt", max_length=512, truncation=True)
    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=req.max_new_tokens,
            num_beams=req.num_beams,
            early_stopping=True
        )
    return {"result": tokenizer.decode(outputs[0], skip_special_tokens=True)}

# 启动: uvicorn inference_server:app --host 0.0.0.0 --port 8000

pip install datasets

import torch
import torch_npu
from torch.utils.data import DataLoader, Dataset
from transformers import T5ForConditionalGeneration, T5Tokenizer
from torch.optim import AdamW

DEVICE = "npu:0"
MODEL_PATH = "./models/google/flan-t5-base"

# 加载
tokenizer = T5Tokenizer.from_pretrained(MODEL_PATH, local_files_only=True)
model = T5ForConditionalGeneration.from_pretrained(MODEL_PATH, local_files_only=True)
model = model.to(DEVICE)

optimizer = AdamW(model.parameters(), lr=5e-5)

# 自定义数据集
class MyDataset(Dataset):
    def __init__(self, data, tokenizer):
        self.data = data
        self.tokenizer = tokenizer

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        src, tgt = self.data[idx]
        inputs = self.tokenizer(src, max_length=128, truncation=True, padding="max_length", return_tensors="pt")
        labels = self.tokenizer(tgt, max_length=64, truncation=True, padding="max_length", return_tensors="pt")
        return {
            "input_ids": inputs["input_ids"].squeeze(),
            "attention_mask": inputs["attention_mask"].squeeze(),
            "labels": labels["input_ids"].squeeze(),
        }

# 训练循环
model.train()
for epoch in range(3):
    for batch in dataloader:
        batch = {k: v.to(DEVICE) for k, v in batch.items()}
        labels = batch.pop("labels")
        labels[labels == tokenizer.pad_token_id] = -100

        optimizer.zero_grad()
        outputs = model(**batch, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()

        print(f"Epoch {epoch+1} | Loss: {loss.item():.4f}")

# 保存
model.save_pretrained("./flan-t5-finetuned-npu")
tokenizer.save_pretrained("./flan-t5-finetuned-npu")

from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments, DataCollatorForSeq2Seq

args = Seq2SeqTrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    learning_rate=5e-5,
    logging_steps=10,
    save_steps=500,
    fp16=True,  # 启用 FP16
    report_to="none",
)

trainer = Seq2SeqTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    data_collator=DataCollatorForSeq2Seq(tokenizer, model=model),
)

trainer.train()

============================================================
Flan-T5 NPU Inference Test
============================================================

[1] NPU Available: True
[1] NPU Count: 2
[1] Using device: npu:0

[2] Loading tokenizer...
[2] Tokenizer vocab size: 32100

[3] Loading model to NPU...
[3] Model loaded in 1.56s
[3] Model device: npu:0
[3] Total parameters: 247.6M

[4] Running inference tests...

  Test 1: translate English to German: The house is wonderful....
    Output: Das Haus ist wunderbar.
    Time: 0.541s

  Test 2: summarize: The quick brown fox jumps over the lazy dog. This...
    Output: The quick brown fox jumps over the lazy dog.
    Time: 0.281s

  Test 3: question: What is the capital of France? context: France is ...
    Output: Paris
    Time: 0.099s

[5] Warmup + throughput test (10 iterations)...
[5] Avg latency: 87.88 ms
[5] Throughput: 11.38 infer/s

============================================================
Inference test PASSED
============================================================

============================================================
Flan-T5 NPU Training Test
============================================================

[1] Device: npu:0

[2] Loading model and tokenizer...
[2] Model on npu:0

[3] Preparing dataset...
[3] Dataset size: 32, Batches: 8

[4] Running 3 training epochs...
    Epoch 1 Batch 1: loss=38.7053
  Epoch 1: avg_loss=29.6944, time=1.86s
    Epoch 2 Batch 1: loss=22.1720
  Epoch 2: avg_loss=18.9105, time=0.97s
    Epoch 3 Batch 1: loss=13.8752
  Epoch 3: avg_loss=10.6407, time=0.97s

[5] Training summary:
    Epoch 1: loss=29.6944, time=1.86s
    Epoch 2: loss=18.9105, time=0.97s
    Epoch 3: loss=10.6407, time=0.97s

[5] Average epoch time: 1.27s
[5] Samples/sec: 25.24

[6] Checkpoint saved: /opt/atomgit/flan-t5-npu-checkpoint.pt

============================================================
Training test PASSED
============================================================

UserWarning: Cannot create tensor with interal format while allow_internel_format=False,
tensor will be created with base format.

torchrun --nproc_per_node=2 --master_port=29500 train.py

import torch.distributed as dist
dist.init_process_group(backend="hccl")  # 昇腾使用 HCCL 通信库
torch_npu.npu.set_device(local_rank)

from modelscope import snapshot_download
snapshot_download('google/flan-t5-base', cache_dir='./models')

# 查看 NPU 状态
npu-smi info

# 实时监控
watch -n 1 npu-smi info

Flan-T5/
├── README.md                    # 本文件
├── scripts/
│   ├── inference_npu.py         # 推理验证脚本
│   ├── training_npu.py          # 训练验证脚本
│   └── benchmark_npu.py         # 性能基准测试
├── docs/
│   └── (预留文档目录)
└── models/                      # 模型权重存放目录（下载后）
    └── google/
        └── flan-t5-base/
            ├── config.json
            ├── pytorch_model.bin
            └── ...

Copyright 2023 Google LLC
Licensed under the Apache License, Version 2.0

Copyright 2026 Ascend-SACT Contributors
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0