File size: 3,591 Bytes
9c41926 | 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | """
Perplexity evaluation for quantized .pt checkpoints.
This runs a real transformer forward pass after applying quantized weights.
For GGUF/GEMV latency checks, use run_gemv_chain.py instead.
"""
import argparse
import gc
from pathlib import Path
import torch
from scripts.eval_quantized import apply_quantized_weights, eval_perplexity
def evaluate_quantized_perplexity(
model_name="models/gemma-4-E2B",
quantized_path="quantized/gemma-4-E2B-sub1bit.pt",
wikitext_path="data/wiki.test.txt",
device=None,
max_length=512,
stride=512,
):
"""Evaluate WikiText perplexity after applying quantized weights."""
if device is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
if not Path(wikitext_path).exists():
raise FileNotFoundError(f"WikiText not found: {wikitext_path}")
print("=" * 60)
print("QUANTIZED MODEL PERPLEXITY EVALUATION")
print("=" * 60)
print(f"Device: {device}")
print(f"Model: {model_name}")
print(f"Quantized: {quantized_path}")
print(f"WikiText: {wikitext_path}")
print("\n[1] Loading quantized checkpoint...")
q_data = torch.load(quantized_path, map_location="cpu", weights_only=True)
quantized = q_data["quantized"]
print(f" {len(quantized)} quantized entries")
print("\n[2] Loading base model...")
from transformers import AutoModelForCausalLM, AutoTokenizer
torch_dtype = torch.float16 if device == "cuda" else torch.float32
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map=device,
torch_dtype=torch_dtype,
trust_remote_code=True,
)
model.eval()
print("\n[3] Applying quantized weights...")
apply_stats = apply_quantized_weights(
model,
quantized,
device=device,
model_dir=model_name if Path(model_name).is_dir() else None,
checkpoint_weight_keys=q_data.get("weight_keys"),
)
print(f" Replaced {apply_stats['replaced']}/{len(quantized)} weights")
if apply_stats["skipped"]:
print(f" Skipped {len(apply_stats['skipped'])} shared-KV checkpoint entries")
print("\n[4] Evaluating perplexity...")
ppl, stats = eval_perplexity(
model,
tokenizer,
wikitext_path,
device,
max_length=max_length,
stride=stride,
)
print()
print("=" * 60)
print("RESULTS")
print("=" * 60)
print(f" Perplexity: {ppl:.4f}")
print(f" Chunks: {stats['n_chunks']}")
print(f" Target: <= 10.5")
print(f" Status: {'PASS' if ppl <= 10.5 else 'FAIL'}")
print("=" * 60)
del model
gc.collect()
if device == "cuda":
torch.cuda.empty_cache()
return ppl, stats
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Evaluate quantized model perplexity")
parser.add_argument("--model", default="models/gemma-4-E2B")
parser.add_argument("--quantized", default="quantized/gemma-4-E2B-sub1bit.pt")
parser.add_argument("--wikitext", default="data/wiki.test.txt")
parser.add_argument("--device", default=None)
parser.add_argument("--max-length", type=int, default=512)
parser.add_argument("--stride", type=int, default=512)
args = parser.parse_args()
evaluate_quantized_perplexity(
model_name=args.model,
quantized_path=args.quantized,
wikitext_path=args.wikitext,
device=args.device,
max_length=args.max_length,
stride=args.stride,
)
|