File size: 4,012 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 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """Run a short WikiText perplexity smoke benchmark for base/quantized checkpoints."""
from __future__ import annotations
import argparse
import gc
import json
import sys
import time
from pathlib import Path
import torch
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from scripts.eval_quantized import apply_quantized_weights
def eval_limited_ppl(
model,
tokenizer,
text: str,
device: str,
tokens: int,
max_length: int,
stride: int,
) -> dict:
encodings = tokenizer(text, return_tensors="pt")
input_ids = encodings["input_ids"][:, :tokens].to(device)
seq_len = input_ids.shape[1]
nlls = []
prev_end_loc = 0
for begin_loc in range(0, seq_len, stride):
end_loc = min(begin_loc + max_length, seq_len)
trg_len = end_loc - prev_end_loc
batch = input_ids[:, begin_loc:end_loc]
target = batch.clone()
target[:, :-trg_len] = -100
with torch.no_grad():
outputs = model(batch, labels=target)
nlls.append(outputs.loss.detach() * trg_len)
prev_end_loc = end_loc
if end_loc >= seq_len:
break
ppl = torch.exp(torch.stack(nlls).sum() / seq_len).item()
return {"ppl": ppl, "seq_len": seq_len, "chunks": len(nlls)}
def run(args: argparse.Namespace) -> dict:
from transformers import AutoModelForCausalLM, AutoTokenizer
model_dir = Path(args.model_dir)
device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16 if device == "cuda" else torch.float32
tokenizer = AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
str(model_dir),
dtype=dtype,
device_map=device,
trust_remote_code=True,
)
model.eval()
apply_stats = {"replaced": 0, "skipped": []}
checkpoint_stats = None
if args.quantized_pt:
q_data = torch.load(args.quantized_pt, map_location="cpu", weights_only=True)
checkpoint_stats = q_data.get("stats")
apply_stats = apply_quantized_weights(
model,
q_data["quantized"],
device=device,
model_dir=model_dir,
checkpoint_weight_keys=q_data.get("weight_keys"),
strict=False,
)
del q_data
gc.collect()
text = Path(args.wikitext).read_text(encoding="utf-8")
metrics = eval_limited_ppl(
model,
tokenizer,
text,
device,
tokens=args.tokens,
max_length=args.max_length,
stride=args.stride,
)
metrics.update(
{
"label": args.label,
"mode": "quantized" if args.quantized_pt else "base",
"quantized_pt": args.quantized_pt,
"apply_stats": apply_stats,
"checkpoint_stats": checkpoint_stats,
"device": device,
}
)
return metrics
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--label", required=True)
parser.add_argument("--model-dir", default="models/gemma-4-E2B")
parser.add_argument("--wikitext", default="data/wiki.test.txt")
parser.add_argument("--quantized-pt", default=None)
parser.add_argument("--tokens", type=int, default=4096)
parser.add_argument("--max-length", type=int, default=512)
parser.add_argument("--stride", type=int, default=512)
parser.add_argument("--device", default=None)
parser.add_argument("--output", required=True)
return parser.parse_args()
def main() -> None:
args = parse_args()
start = time.time()
result = run(args)
result["elapsed_s"] = round(time.time() - start, 1)
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
print("RESULT=" + json.dumps(result, indent=2), flush=True)
if __name__ == "__main__":
main()
|