File size: 1,313 Bytes
5e21013 | 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 | """Download and test Qwen2.5-3B-Instruct on MPS."""
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_ID = "Qwen/Qwen2.5-3B-Instruct"
print(f"Downloading {MODEL_ID} (~6GB, one-time)...")
t0 = time.time()
tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID, trust_remote_code=True, dtype=torch.float16,
).to("mps")
model.eval()
n_params = sum(p.numel() for p in model.parameters()) / 1e6
print(f"Loaded: {n_params:.0f}M params on MPS (float16) in {time.time() - t0:.0f}s")
# Quick test
print("\nTesting generation...")
inputs = tok("What is quantum computing?", return_tensors="pt").to("mps")
with torch.no_grad():
t1 = time.time()
out = model.generate(
**inputs,
max_new_tokens=150,
temperature=0.7,
do_sample=True,
pad_token_id=tok.eos_token_id,
)
elapsed = time.time() - t1
gen_ids = out[0][inputs["input_ids"].shape[1]:]
gen_text = tok.decode(gen_ids, skip_special_tokens=True)
n_tokens = len(gen_ids)
tps = n_tokens / max(elapsed, 0.001)
print(f"Speed: {tps:.1f} tokens/sec ({n_tokens} tokens in {elapsed:.1f}s)")
print(f"Response:\n{gen_text[:500]}")
print(f"\nModel ready. M4 Max + 36GB + MPS = {MODEL_ID} runs perfectly.")
|