本地模型推理：权重、设备、精度、量化与速度

python -m pip install "transformers" "torch" "accelerate"

import os
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


MODEL_ID = os.getenv("LOCAL_MODEL_ID", "Qwen/Qwen2.5-0.5B-Instruct")

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype="auto",
    device_map="auto",
)
model.eval()

messages = [
    {"role": "system", "content": "只整理输入事实，不补充建议。"},
    {"role": "user", "content": "节点 NX-8 网络超时，11:05 切换线路后恢复。"},
]

prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

started = time.perf_counter()
with torch.inference_mode():
    generated = model.generate(
        **inputs,
        max_new_tokens=80,
        do_sample=False,
        use_cache=True,
    )
elapsed = time.perf_counter() - started

new_tokens = generated[:, inputs["input_ids"].shape[1] :]
text = tokenizer.batch_decode(new_tokens, skip_special_tokens=True)[0]

print(text)
print({
    "input_tokens": int(inputs["input_ids"].shape[1]),
    "output_tokens": int(new_tokens.shape[1]),
    "elapsed_seconds": round(elapsed, 3),
    "tokens_per_second": round(new_tokens.shape[1] / elapsed, 2),
})

from threading import Thread
from transformers import TextIteratorStreamer


streamer = TextIteratorStreamer(
    tokenizer,
    skip_prompt=True,
    skip_special_tokens=True,
)
generation_args = {
    **inputs,
    "streamer": streamer,
    "max_new_tokens": 80,
    "do_sample": False,
    "use_cache": True,
}

if torch.cuda.is_available():
    torch.cuda.reset_peak_memory_stats()

started = time.perf_counter()
worker = Thread(target=model.generate, kwargs=generation_args)
worker.start()

pieces = []
first_token_at = None
for piece in streamer:
    if piece and first_token_at is None:
        first_token_at = time.perf_counter()
    pieces.append(piece)
worker.join()
finished = time.perf_counter()

result = "".join(pieces)
output_tokens = len(tokenizer.encode(result, add_special_tokens=False))
decode_seconds = max(finished - (first_token_at or finished), 1e-9)

metrics = {
    "first_token_seconds": None if first_token_at is None else round(first_token_at - started, 3),
    "total_seconds": round(finished - started, 3),
    "output_tokens": output_tokens,
    "decode_tokens_per_second": round(output_tokens / decode_seconds, 2),
}
if torch.cuda.is_available():
    metrics["peak_gpu_gib"] = round(torch.cuda.max_memory_allocated() / 1024 ** 3, 3)

print(result)
print(metrics)

from transformers import BitsAndBytesConfig

quantization = BitsAndBytesConfig(load_in_4bit=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="auto",
    quantization_config=quantization,
    dtype="auto",
)
print(model.get_memory_footprint())