145 lines
4.8 KiB
Python
Executable File
145 lines
4.8 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""本地 GPU 推理服务（ONNX Runtime + CUDA）
|
||
RTX 3050 4GB 最优方案：Qwen2.5-0.5B ONNX + CUDAExecutionProvider
|
||
用法:
|
||
python3 gpu_local_infer.py "你的问题" [max_tokens]
|
||
python3 gpu_local_infer.py --bench # 性能基准
|
||
"""
|
||
import os
|
||
import sys
|
||
import time
|
||
import subprocess
|
||
import numpy as np
|
||
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
|
||
os.environ['LD_LIBRARY_PATH'] = '/home/muc/.local/lib:/usr/local/cuda/lib64:' + os.environ.get('LD_LIBRARY_PATH', '')
|
||
|
||
import onnxruntime as ort
|
||
from transformers import AutoTokenizer
|
||
|
||
MODEL_DIR = os.path.expanduser("~/models/onnx-community-Qwen2.5-0.5B-Instruct")
|
||
ONNX_FILE = os.path.join(MODEL_DIR, "onnx", "model.onnx")
|
||
|
||
NUM_LAYERS = 24
|
||
NUM_KV_HEADS = 2
|
||
HEAD_DIM = 64
|
||
|
||
class LocalGPUInfer:
|
||
def __init__(self, verbose=True):
|
||
self.tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
|
||
self.session = ort.InferenceSession(
|
||
ONNX_FILE,
|
||
providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
|
||
)
|
||
self.providers = self.session.get_providers()
|
||
if verbose:
|
||
print(f"✅ ONNX Runtime: {ort.__version__}")
|
||
print(f"✅ Provider: {self.providers}")
|
||
if 'CUDAExecutionProvider' in self.providers:
|
||
print("🎉 CUDA 加速已启用！")
|
||
|
||
def generate(self, prompt, max_new_tokens=100, temperature=0.7):
|
||
inputs = self.tokenizer(prompt, return_tensors="np")
|
||
input_ids = inputs['input_ids'].astype(np.int64)
|
||
seq_len = input_ids.shape[1]
|
||
batch = 1
|
||
|
||
# 初始 past_key_values 为空
|
||
past_key_values = {}
|
||
for i in range(NUM_LAYERS):
|
||
for kv in ['key', 'value']:
|
||
past_key_values[f'past_key_values.{i}.{kv}'] = np.zeros(
|
||
(batch, NUM_KV_HEADS, 0, HEAD_DIM), dtype=np.float32
|
||
)
|
||
|
||
generated = []
|
||
current_input_ids = input_ids
|
||
past_len = 0
|
||
|
||
for step in range(max_new_tokens):
|
||
seq = current_input_ids.shape[1]
|
||
ort_inputs = {
|
||
'input_ids': current_input_ids,
|
||
'attention_mask': np.ones((batch, past_len + seq), dtype=np.int64),
|
||
'position_ids': np.arange(past_len, past_len + seq).reshape(1, seq).astype(np.int64),
|
||
}
|
||
for name, val in past_key_values.items():
|
||
ort_inputs[name] = val
|
||
|
||
outputs = self.session.run(None, ort_inputs)
|
||
logits = outputs[0]
|
||
|
||
if temperature <= 0:
|
||
next_token = np.argmax(logits[:, -1, :], axis=-1)
|
||
else:
|
||
# 温度采样
|
||
probs = logits[:, -1, :] / temperature
|
||
probs = np.exp(probs - probs.max())
|
||
probs = probs / probs.sum()
|
||
next_token = np.array([np.random.choice(len(probs[0]), p=probs[0])], dtype=np.int64)
|
||
|
||
next_token = next_token.astype(np.int64)
|
||
generated.append(next_token[0])
|
||
|
||
for i in range(NUM_LAYERS):
|
||
past_key_values[f'past_key_values.{i}.key'] = outputs[1 + i*2]
|
||
past_key_values[f'past_key_values.{i}.value'] = outputs[2 + i*2]
|
||
|
||
current_input_ids = next_token.reshape(1, 1)
|
||
past_len += seq
|
||
|
||
if next_token[0] == self.tokenizer.eos_token_id:
|
||
break
|
||
|
||
return self.tokenizer.decode(generated, skip_special_tokens=True)
|
||
|
||
def benchmark(self):
|
||
"""性能基准"""
|
||
prompts = [
|
||
"用一句话解释什么是过拟合",
|
||
"你好",
|
||
"中国的首都是哪里？",
|
||
]
|
||
print("\n=== 性能基准 ===")
|
||
for prompt in prompts:
|
||
# 预热
|
||
self.generate(prompt, max_new_tokens=10)
|
||
# 计时
|
||
start = time.time()
|
||
result = self.generate(prompt, max_new_tokens=50)
|
||
elapsed = time.time() - start
|
||
print(f"\n输入: {prompt[:20]}...")
|
||
print(f"输出: {result[:60]}...")
|
||
print(f"耗时: {elapsed:.3f}s")
|
||
|
||
print("\n=== GPU 状态 ===")
|
||
r = subprocess.run(
|
||
["nvidia-smi", "--query-gpu=memory.used,memory.total,utilization.gpu", "--format=csv,noheader"],
|
||
capture_output=True, text=True
|
||
)
|
||
print(f" {r.stdout.strip()}")
|
||
|
||
def main():
|
||
if len(sys.argv) < 2:
|
||
print(__doc__)
|
||
return
|
||
|
||
if sys.argv[1] == "--bench":
|
||
infer = LocalGPUInfer()
|
||
infer.benchmark()
|
||
return
|
||
|
||
prompt = sys.argv[1]
|
||
max_tokens = int(sys.argv[2]) if len(sys.argv) > 2 else 100
|
||
|
||
infer = LocalGPUInfer()
|
||
start = time.time()
|
||
result = infer.generate(prompt, max_new_tokens=max_tokens)
|
||
elapsed = time.time() - start
|
||
print(f"\n输入: {prompt}")
|
||
print(f"输出: {result}")
|
||
print(f"\n耗时: {elapsed:.3f}s")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|