Upload Evaluation/run_benchmarks.py with huggingface_hub

Evaluation/run_benchmarks.py

@@ -17,6 +17,7 @@ import json
 import argparse
 import time
 import re
+import gc
 from pathlib import Path
 
 import torch
@@ -472,10 +473,93 @@ ALL_TASKS = [
 ]
 
 
-def run_benchmarks(model_path, tokenizer_path, tasks, batch_size, device, output_file):
+def configure_runtime_threads(reserve_cores=2, max_thread_fraction=0.85):
+    """Use most CPU threads while keeping headroom for system stability."""
+    total = os.cpu_count() or 4
+    usable = max(1, total - max(0, reserve_cores))
+    target = max(1, int(usable * max_thread_fraction))
+
+    # Keep BLAS/OpenMP stacks aligned to avoid oversubscription.
+    for key in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS"):
+        os.environ[key] = str(target)
+
+    try:
+        torch.set_num_threads(target)
+    except Exception:
+        pass
+    try:
+        torch.set_num_interop_threads(max(1, min(8, target // 2)))
+    except Exception:
+        pass
+
+    print(f"  CPU threads: {target}/{total} (reserve={reserve_cores}, fraction={max_thread_fraction:.2f})")
+    return target
+
+
+def auto_tune_batch_size(model, seq_len, vocab_size, device, headroom=0.80, max_search=1024):
+    """Probe a safe evaluation batch size with GPU memory headroom.
+
+    Uses forward-only passes because lm-eval does not backprop.
+    """
+    if str(device) != "cuda" or not torch.cuda.is_available():
+        return 8
+
+    low, high = 1, 1
+    best = 1
+    p = next(model.parameters())
+    probe_dtype = p.dtype if p.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16
+
+    def can_run(bs):
+        x = None
+        try:
+            torch.cuda.empty_cache()
+            gc.collect()
+            x = torch.randint(0, vocab_size, (bs, seq_len), device=device, dtype=torch.long)
+            with torch.no_grad(), autocast(device_type="cuda", dtype=probe_dtype, enabled=True):
+                logits, _ = model(x)
+                # Touch output to avoid lazy behavior in some backends.
+                _ = logits[:, -1, :].float().mean().item()
+            del logits
+            del x
+            torch.cuda.empty_cache()
+            return True
+        except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
+            if "out of memory" in str(e).lower() or "cuda" in str(e).lower():
+                if x is not None:
+                    del x
+                torch.cuda.empty_cache()
+                return False
+            raise
+
+    while high < max_search and can_run(high):
+        best = high
+        high *= 2
+    low = best
+
+    lo, hi = low, max(low, min(high - 1, max_search))
+    while lo <= hi:
+        mid = (lo + hi) // 2
+        if can_run(mid):
+            best = mid
+            lo = mid + 1
+        else:
+            hi = mid - 1
+
+    safe = max(1, int(best * headroom))
+    print(f"  Batch probe: max={best}, safe={safe} (headroom={headroom:.2f})")
+    return safe
+
+
+def run_benchmarks(model_path, tokenizer_path, tasks, batch_size, device, output_file,
+                   auto_resources=True, reserve_cores=2, cpu_thread_fraction=0.85,
+                   batch_headroom=0.80):
     """Run all specified benchmarks and print results."""
     import lm_eval
 
+    if auto_resources:
+        configure_runtime_threads(reserve_cores=reserve_cores,
+                                  max_thread_fraction=cpu_thread_fraction)
+
     # Load model
     model, mcfg = load_luna_model(model_path, device=device)
     total_params = sum(p.numel() for p in model.parameters())
@@ -485,6 +569,16 @@ def run_benchmarks(model_path, tokenizer_path, tasks, batch_size, device, output
     from transformers import AutoTokenizer
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
 
+    # Batch auto-tuning: 0 or negative means auto.
+    if batch_size <= 0:
+        batch_size = auto_tune_batch_size(
+            model,
+            seq_len=mcfg.get("seq_len", 1024),
+            vocab_size=mcfg.get("vocab_size", 50304),
+            device=device,
+            headroom=batch_headroom,
+        )
+
     # Create eval wrapper
     lm = get_lm_eval_model(model, tokenizer, device,
                             max_length=mcfg.get("seq_len", 1024),
@@ -574,9 +668,18 @@ def parse_args():
                    help="Path to tokenizer directory")
     p.add_argument("--tasks", type=str, default=",".join(ALL_TASKS),
                    help=f"Comma-separated tasks (default: all)")
-    p.add_argument("--batch_size", type=int, default=16,
-                   help="Batch size for evaluation")
+    p.add_argument("--batch_size", type=int, default=0,
+                   help="Batch size for evaluation (<=0 means auto-tune)")
     p.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--auto_resources", type=lambda x: x.lower() in ("1", "true", "yes"),
+                   default=True,
+                   help="Auto-tune CPU threads and batch size with headroom")
+    p.add_argument("--reserve_cores", type=int, default=2,
+                   help="CPU cores to keep free for system headroom")
+    p.add_argument("--cpu_thread_fraction", type=float, default=0.85,
+                   help="Fraction of usable CPU threads to use")
+    p.add_argument("--batch_headroom", type=float, default=0.80,
+                   help="Safety factor applied to probed max batch")
     p.add_argument("--output", type=str, default=None,
                    help="Save results to JSON file (default: auto-named)")
     return p.parse_args()
@@ -599,5 +702,9 @@ if __name__ == "__main__":
         tasks=tasks,
         batch_size=args.batch_size,
         device=args.device,
+        auto_resources=args.auto_resources,
+        reserve_cores=args.reserve_cores,
+        cpu_thread_fraction=args.cpu_thread_fraction,
+        batch_headroom=args.batch_headroom,
         output_file=output_file,
     )