| |
| """Compare FP32 FLAN-T5 Small with a quantized checkpoint on fixed prompts.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import difflib |
| import gc |
| import json |
| import platform |
| import statistics |
| import time |
| from datetime import datetime, timezone |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| import transformers |
| from transformers import AutoModelForSeq2SeqLM, AutoTokenizer |
|
|
|
|
| DEFAULT_PROMPTS = [ |
| "translate English to German: How old are you?", |
| "Answer this question: What is the capital of France?", |
| "Classify the sentiment as positive or negative: I loved the thoughtful story and acting.", |
| "summarize: The James Webb Space Telescope observes the universe in infrared light, allowing it to see through dust and study very distant galaxies.", |
| "Premise: All roses are flowers. Some flowers fade quickly. Question: Does it follow that some roses fade quickly? Explain briefly.", |
| ] |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--base-model", default="google/flan-t5-small") |
| parser.add_argument("--base-revision", default="main") |
| parser.add_argument("--quantized-model", default="ShinpacheShimura/t5-smaller") |
| parser.add_argument("--quantized-revision", default="main") |
| parser.add_argument("--quantized-subfolder", default=None) |
| parser.add_argument( |
| "--prompts-file", |
| type=Path, |
| help="Optional UTF-8 text file with one non-empty prompt per line.", |
| ) |
| parser.add_argument("--output-dir", type=Path, default=Path("benchmark-results")) |
| parser.add_argument("--warmup-runs", type=int, default=2) |
| parser.add_argument("--runs", type=int, default=10) |
| parser.add_argument("--max-new-tokens", type=int, default=64) |
| parser.add_argument("--num-beams", type=int, default=1) |
| parser.add_argument("--seed", type=int, default=42) |
| return parser.parse_args() |
|
|
|
|
| def load_prompts(path: Path | None) -> list[str]: |
| if path is None: |
| return DEFAULT_PROMPTS |
| prompts = [line.strip() for line in path.read_text(encoding="utf-8").splitlines()] |
| prompts = [prompt for prompt in prompts if prompt] |
| if not prompts: |
| raise SystemExit("The prompts file contains no non-empty prompts.") |
| return prompts |
|
|
|
|
| def synchronize() -> None: |
| if torch.cuda.is_available(): |
| torch.cuda.synchronize() |
|
|
|
|
| def reset_peak_memory() -> None: |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
| torch.cuda.reset_peak_memory_stats() |
|
|
|
|
| def input_device(model: torch.nn.Module) -> torch.device: |
| try: |
| return model.device |
| except AttributeError: |
| return next(model.parameters()).device |
|
|
|
|
| def generate_once( |
| model: torch.nn.Module, |
| tokenizer: Any, |
| prompt: str, |
| max_new_tokens: int, |
| num_beams: int, |
| ) -> tuple[str, int, float]: |
| encoded = tokenizer(prompt, return_tensors="pt").to(input_device(model)) |
| synchronize() |
| started = time.perf_counter() |
| with torch.inference_mode(): |
| generated = model.generate( |
| **encoded, |
| max_new_tokens=max_new_tokens, |
| num_beams=num_beams, |
| do_sample=False, |
| ) |
| synchronize() |
| elapsed = time.perf_counter() - started |
| text = tokenizer.decode(generated[0], skip_special_tokens=True) |
| generated_tokens = int(generated.shape[-1]) |
| return text, generated_tokens, elapsed |
|
|
|
|
| def benchmark_model( |
| label: str, |
| model_id: str, |
| revision: str, |
| prompts: list[str], |
| warmup_runs: int, |
| runs: int, |
| max_new_tokens: int, |
| num_beams: int, |
| quantized: bool, |
| subfolder: str | None = None, |
| ) -> dict[str, Any]: |
| tokenizer_kwargs: dict[str, Any] = {"revision": revision} |
| model_kwargs: dict[str, Any] = {"revision": revision} |
| if subfolder: |
| tokenizer_kwargs["subfolder"] = subfolder |
| model_kwargs["subfolder"] = subfolder |
|
|
| tokenizer = AutoTokenizer.from_pretrained(model_id, **tokenizer_kwargs) |
| if quantized: |
| model_kwargs["device_map"] = "auto" |
| else: |
| model_kwargs["torch_dtype"] = torch.float32 |
| model_kwargs["device_map"] = "auto" |
|
|
| reset_peak_memory() |
| load_started = time.perf_counter() |
| model = AutoModelForSeq2SeqLM.from_pretrained(model_id, **model_kwargs) |
| model.eval() |
| synchronize() |
| load_seconds = time.perf_counter() - load_started |
|
|
| for index in range(warmup_runs): |
| generate_once( |
| model, |
| tokenizer, |
| prompts[index % len(prompts)], |
| max_new_tokens, |
| num_beams, |
| ) |
|
|
| timings: list[float] = [] |
| total_output_tokens = 0 |
| outputs: list[dict[str, str]] = [] |
| for prompt in prompts: |
| output, generated_tokens, _ = generate_once( |
| model, tokenizer, prompt, max_new_tokens, num_beams |
| ) |
| outputs.append({"prompt": prompt, "output": output}) |
|
|
| for index in range(runs): |
| _, generated_tokens, elapsed = generate_once( |
| model, |
| tokenizer, |
| prompts[index % len(prompts)], |
| max_new_tokens, |
| num_beams, |
| ) |
| timings.append(elapsed) |
| total_output_tokens += generated_tokens |
|
|
| footprint = int(model.get_memory_footprint()) |
| peak_cuda = int(torch.cuda.max_memory_allocated()) if torch.cuda.is_available() else None |
| result = { |
| "label": label, |
| "model_id": model_id, |
| "revision": revision, |
| "subfolder": subfolder, |
| "load_seconds": load_seconds, |
| "model_memory_footprint_bytes": footprint, |
| "peak_cuda_allocated_bytes": peak_cuda, |
| "latency_seconds": { |
| "mean": statistics.mean(timings), |
| "median": statistics.median(timings), |
| "min": min(timings), |
| "max": max(timings), |
| }, |
| "examples_per_second": runs / sum(timings), |
| "output_tokens_per_second": total_output_tokens / sum(timings), |
| "outputs": outputs, |
| } |
|
|
| del model, tokenizer |
| gc.collect() |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
| return result |
|
|
|
|
| def output_agreement(base: dict[str, Any], quantized: dict[str, Any]) -> dict[str, Any]: |
| pairs = [] |
| exact_count = 0 |
| similarities: list[float] = [] |
| for base_item, quantized_item in zip(base["outputs"], quantized["outputs"], strict=True): |
| base_text = base_item["output"].strip() |
| quantized_text = quantized_item["output"].strip() |
| exact = base_text == quantized_text |
| similarity = difflib.SequenceMatcher(None, base_text, quantized_text).ratio() |
| exact_count += int(exact) |
| similarities.append(similarity) |
| pairs.append( |
| { |
| "prompt": base_item["prompt"], |
| "fp32_output": base_text, |
| "quantized_output": quantized_text, |
| "exact_match": exact, |
| "text_similarity": similarity, |
| } |
| ) |
| return { |
| "exact_match_rate": exact_count / len(pairs), |
| "mean_text_similarity": statistics.mean(similarities), |
| "note": "These are regression diagnostics, not ground-truth quality metrics.", |
| "pairs": pairs, |
| } |
|
|
|
|
| def mib(value: int | None) -> str: |
| return "N/A" if value is None else f"{value / (1024**2):.2f} MiB" |
|
|
|
|
| def render_markdown(results: dict[str, Any]) -> str: |
| base = results["models"]["fp32"] |
| quantized = results["models"]["quantized"] |
| agreement = results["agreement"] |
| lines = [ |
| "# Benchmark results", |
| "", |
| f"Generated: `{results['created_at_utc']}`", |
| "", |
| "| Measurement | FP32 base | NF4 checkpoint |", |
| "|---|---:|---:|", |
| f"| Model memory footprint | {mib(base['model_memory_footprint_bytes'])} | {mib(quantized['model_memory_footprint_bytes'])} |", |
| f"| Peak CUDA allocation | {mib(base['peak_cuda_allocated_bytes'])} | {mib(quantized['peak_cuda_allocated_bytes'])} |", |
| f"| Load time | {base['load_seconds']:.4f} s | {quantized['load_seconds']:.4f} s |", |
| f"| Median generation latency | {base['latency_seconds']['median']:.4f} s | {quantized['latency_seconds']['median']:.4f} s |", |
| f"| Examples/second | {base['examples_per_second']:.3f} | {quantized['examples_per_second']:.3f} |", |
| f"| Output tokens/second | {base['output_tokens_per_second']:.3f} | {quantized['output_tokens_per_second']:.3f} |", |
| "", |
| f"Exact output agreement: **{agreement['exact_match_rate']:.1%}** ", |
| f"Mean text similarity: **{agreement['mean_text_similarity']:.3f}**", |
| "", |
| "> Agreement and text similarity compare outputs with FP32. They do not measure correctness against labels.", |
| "", |
| "## Environment", |
| "", |
| "```json", |
| json.dumps(results["environment"], indent=2), |
| "```", |
| "", |
| "## Outputs", |
| "", |
| ] |
| for index, pair in enumerate(agreement["pairs"], start=1): |
| lines.extend( |
| [ |
| f"### Prompt {index}", |
| "", |
| f"**Input:** {pair['prompt']}", |
| "", |
| f"**FP32:** {pair['fp32_output']}", |
| "", |
| f"**NF4:** {pair['quantized_output']}", |
| "", |
| ] |
| ) |
| return "\n".join(lines) |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| if args.runs < 1 or args.warmup_runs < 0: |
| raise SystemExit("--runs must be at least 1 and --warmup-runs cannot be negative.") |
|
|
| torch.manual_seed(args.seed) |
| prompts = load_prompts(args.prompts_file) |
| args.output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| environment: dict[str, Any] = { |
| "python": platform.python_version(), |
| "platform": platform.platform(), |
| "torch": torch.__version__, |
| "transformers": transformers.__version__, |
| "cuda_available": torch.cuda.is_available(), |
| "torch_cuda_version": torch.version.cuda, |
| "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, |
| "seed": args.seed, |
| "warmup_runs": args.warmup_runs, |
| "timed_runs": args.runs, |
| "max_new_tokens": args.max_new_tokens, |
| "num_beams": args.num_beams, |
| "prompt_count": len(prompts), |
| } |
|
|
| print("Benchmarking FP32 base model...") |
| base = benchmark_model( |
| "fp32", |
| args.base_model, |
| args.base_revision, |
| prompts, |
| args.warmup_runs, |
| args.runs, |
| args.max_new_tokens, |
| args.num_beams, |
| quantized=False, |
| ) |
| print("Benchmarking NF4 checkpoint...") |
| quantized = benchmark_model( |
| "nf4", |
| args.quantized_model, |
| args.quantized_revision, |
| prompts, |
| args.warmup_runs, |
| args.runs, |
| args.max_new_tokens, |
| args.num_beams, |
| quantized=True, |
| subfolder=args.quantized_subfolder, |
| ) |
|
|
| results = { |
| "created_at_utc": datetime.now(timezone.utc).isoformat(), |
| "environment": environment, |
| "models": {"fp32": base, "quantized": quantized}, |
| "agreement": output_agreement(base, quantized), |
| } |
| json_path = args.output_dir / "results.json" |
| markdown_path = args.output_dir / "results.md" |
| json_path.write_text(json.dumps(results, indent=2) + "\n", encoding="utf-8") |
| markdown_path.write_text(render_markdown(results), encoding="utf-8") |
| print(f"Wrote {json_path} and {markdown_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
