File size: 2,840 Bytes
504da3d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | # -*- coding: utf-8 -*-
"""inference.py — samai-pnet-dmoe-2b 拉取即跑推理示例
用法:
python inference.py # 默认 HuggingFace 仓
python inference.py --repo /path/to/local # 本地目录
python inference.py --lru 4 # 启用 SSD-LRU 专家缓存模拟 (每层容量 4)
依赖: torch>=2.4, transformers>=5.0 (trust_remote_code 加载仓内 modeling_samai_pnet.py)
"""
import argparse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--repo", default="tchbcb/samai-pnet-dmoe-2b")
ap.add_argument("--lru", type=int, default=0,
help=">0 启用 SSD-LRU 专家缓存 (每层容量; 需仓内 modeling 模块)")
ap.add_argument("--max-new-tokens", type=int, default=256)
args = ap.parse_args()
tok = AutoTokenizer.from_pretrained(args.repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(args.repo, trust_remote_code=True,
dtype="auto").eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
if args.lru:
# SSDLRUExpertManager: 专家驻留 CPU(模拟 SSD), 路由激活时搬入, LRU 驱逐
from modeling_samai_pnet import SSDLRUExpertManager # remote_code 模块
mgr = SSDLRUExpertManager.install(model, capacity_per_layer=args.lru)
print(f"[ssd-lru] capacity/layer={args.lru} (storage=cpu)")
messages = [{"role": "user", "content": "一个长方形长 12 米宽 8 米，面积是多少？请简要推理。"}]
enc = tok.apply_chat_template(messages, add_generation_prompt=True,
tokenize=True, return_tensors="pt",
return_dict=True).to(device)
with torch.no_grad():
out = model.generate(**enc, max_new_tokens=args.max_new_tokens,
do_sample=False, pad_token_id=tok.pad_token_id or 1)
text = tok.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True)
print("=== 回答 ===")
print(text.strip())
# Ponder 诊断: 每次 forward 追加一条 {mode, executed, steps_mean, ...}
decode_entries = [e for e in model._ponder_log if e.get("mode") == "decode"]
if decode_entries:
steps = sum(e["steps_mean"] for e in decode_entries) / len(decode_entries)
print(f"\n=== Ponder 诊断 ===\n本次生成平均思考步数: {steps:.2f} "
f"(max_ponder_steps={model.config.max_ponder_steps})")
if args.lru:
s = mgr.summary()
print(f"=== SSD-LRU ===\nhit={s['hit']} miss={s['miss']} "
f"evict={s['evict']} hit_rate={s['hit_rate']:.3f}")
mgr.remove()
if __name__ == "__main__":
main()
|