File size: 6,698 Bytes
3b2d368 | 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 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | import os
import sys
import subprocess
import json
from pathlib import Path
from typing import Optional
import torch
import traceback
from transformers import GPT2LMHeadModel, GPT2Tokenizer
class Benchmark:
def __init__(self, benchmark_config, model, tokenizer, checkpointing, output_dir, device=None):
self.benchmark_config = benchmark_config
self.model = model
self.tokenizer = tokenizer
self.checkpointing = checkpointing
self.device = device if device is not None else torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.output_dir = Path(f"{benchmark_config.prefix if benchmark_config.prefix is not None else ''}{output_dir}")
self.output_dir.mkdir(parents=True, exist_ok=True)
self.lm_eval_dir = Path(__file__).resolve().parent / "lm-evaluation-harness"
self.results_dir = self.output_dir / "results"
self.summary_path = self.output_dir / "summary.json"
self.checkpointing.load_model_states("recent")
self._init_repo()
def _init_repo(self):
if not self.lm_eval_dir.exists():
print("Cloning lm-evaluation-harness...")
subprocess.run([
"git", "clone", "https://github.com/EleutherAI/lm-evaluation-harness", str(self.lm_eval_dir)
], check=True)
print("Installing lm-evaluation-harness...")
subprocess.run(
["pip", "install", "-e", "."],
cwd=str(self.lm_eval_dir),
check=True
)
else:
print("lm-evaluation-harness already installed.")
def run_benchmarks(self):
sys.path.insert(0, str(self.lm_eval_dir))
from lm_eval.evaluator import simple_evaluate
from .benchmark_wrapper import BenchmarkWrapper
summary = self._load_existing_summary()
status_report = {}
model = self.model
if self.benchmark_config.compile:
print(f"Compiling model with mode: {self.benchmark_config.compile_mode}")
model = torch.compile(model, mode=self.benchmark_config.compile_mode)
wrapped_model = BenchmarkWrapper(
model,
self.tokenizer,
device=self.device,
batch_size=self.benchmark_config.batch_size,
dtype=getattr(torch, self.benchmark_config.precision)
)
new_results = {}
for task in self.benchmark_config.tasks:
try:
if task in summary:
print(f"Skipping {task}, already benchmarked.")
status_report[task] = {"status": "cached", "result": summary[task]}
continue
print(f"Running benchmark for task: {task}")
result = simple_evaluate(
model=wrapped_model,
tasks=[task],
num_fewshot=self.benchmark_config.num_fewshot,
batch_size=self.benchmark_config.batch_size,
limit=self.benchmark_config.limit if hasattr(self.benchmark_config, "limit") else None,
bootstrap_iters=self.benchmark_config.bootstrap_iters,
device=self.device,
log_samples=False
)
self.results_dir.mkdir(parents=True, exist_ok=True)
def make_json_safe(obj):
if isinstance(obj, torch.device):
return str(obj)
elif isinstance(obj, dict):
return {k: make_json_safe(v) for k, v in obj.items()}
elif isinstance(obj, list):
return [make_json_safe(v) for v in obj]
elif isinstance(obj, tuple):
return tuple(make_json_safe(v) for v in obj)
else:
return obj
with open(self.results_dir / f"{task}.json", "w") as f:
json.dump(make_json_safe(result), f, indent=2)
task_result = result["results"].get(task, {})
summary[task] = task_result
new_results[task] = task_result
status_report[task] = {"status": "completed", "result": task_result}
except Exception as e:
print(f"⚠️ Warning: Benchmark for task '{task}' failed with error: {e}")
traceback.print_exc()
status_report[task] = {"status": "failed", "error": str(e)}
self._save_summary(summary)
self._print_status_report(status_report)
def _load_existing_summary(self) -> dict:
if self.summary_path.exists():
with open(self.summary_path, "r") as f:
return json.load(f)
return {}
def _save_summary(self, summary: dict):
with open(self.summary_path, "w") as f:
json.dump(summary, f, indent=2)
def _print_status_report(self, report):
def is_metric_key(k):
return any(key in k.lower() for key in ["acc", "f1", "loss", "perplexity", "exact", "score"])
def format_metric(k, v):
clean_key = k.replace(",none", "")
if "norm" in k:
return None
if isinstance(v, dict) and "value" in v:
val = v["value"]
stderr = v.get("stderr", None)
if stderr is not None:
return f"{clean_key}: {val:.5f} ± {stderr:.5f}"
else:
return f"{clean_key}: {val:.5f}"
elif isinstance(v, (float, int)):
return f"{clean_key}: {v:.5f}"
return None
print("\n=== Benchmark Summary Report ===")
for task, info in report.items():
status = info["status"]
if status in {"completed", "cached"}:
label = "✅" if status == "completed" else "📄"
metrics = info["result"]
filtered_metrics = {
k: v for k, v in metrics.items()
if is_metric_key(k)
}
formatted = [
format_metric(k, v)
for k, v in filtered_metrics.items()
if format_metric(k, v) is not None
]
metric_str = ", ".join(formatted)
print(f"{label} {task}: {status} | {metric_str}")
elif status == "failed":
print(f"❌ {task}: {status} | Error: {info['error']}") |