Upload folder using huggingface_hub

src/mathcompose/eval/run.py

@@ -15,8 +15,12 @@ from __future__ import annotations
 
 import argparse
 import json
+import os
 from pathlib import Path
 
+# reduce fragmentation OOMs (set before torch/CUDA initializes)
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
 RESULTS_DIR = Path("results")
 
 

src/mathcompose/eval/runners.py

@@ -167,13 +167,15 @@ def auto_concurrent_seqs() -> int:
     if not torch.cuda.is_available():
         return 8
     vram = torch.cuda.get_device_properties(0).total_memory / 1e9
+    # conservative starting point (a 7B in bf16 leaves far less room than a 1.5B);
+    # verify_batch_predictions auto-halves on OOM as a backstop.
     if vram >= 70:
-        return 128
-    if vram >= 40:
         return 64
-    if vram >= 20:
+    if vram >= 40:
         return 32
-    return 12
+    if vram >= 20:
+        return 16
+    return 8
 
 
 def verify_batch_predictions(runner: HFRunner, items, n: int = 1, batch_size: int = 16,
@@ -189,12 +191,29 @@ def verify_batch_predictions(runner: HFRunner, items, n: int = 1, batch_size: in
         def tqdm(x, **k):
             return x
 
+    import torch
+
     preds: list[int] = []
-    chunks = list(range(0, len(items), batch_size))
-    for i in tqdm(chunks, desc=desc):
-        chunk = items[i:i + batch_size]
+    i, bs = 0, batch_size
+    pbar = tqdm(total=len(items), desc=desc)
+    while i < len(items):
+        chunk = items[i:i + bs]
         msgs = [verifier_messages(p, s) for p, s in chunk]
-        samples = runner.chat_many(msgs, n=n, temperature=temperature,
-                                   max_new_tokens=max_new_tokens)
+        try:
+            samples = runner.chat_many(msgs, n=n, temperature=temperature,
+                                       max_new_tokens=max_new_tokens)
+        except torch.cuda.OutOfMemoryError:
+            torch.cuda.empty_cache()
+            if bs == 1:
+                raise
+            bs = max(1, bs // 2)
+            try:
+                pbar.write(f"[oom] halving batch -> {bs}")
+            except Exception:
+                pass
+            continue
         preds.extend(majority_index(s) for s in samples)
+        i += len(chunk)
+        pbar.update(len(chunk))
+    pbar.close()
     return preds