File size: 3,920 Bytes
ffdcfe7 | 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 | """LoRA readout on the 64 final MAE checkpoints, 3 pinned seeds, recalibrated recipe.
Recipe is the 15M recalibration, not the scale_base defaults:
rank 16, targets qkv/proj/fc1/fc2, LayerNorm trainable, lr 3e-3,
**4 epochs** (the knee at 15M; 3 is the knee at scale_base and costs ~0.007 here).
Seeds are pinned 0/1/2 across every arm so the common-mode component of sigma_ft
(0.00076 of 0.00167 at this scale) cancels in the arm contrast, exactly as the
base campaign did. Same CPU-thread capping as the probe stage: this is a
BLAS-bound workload and uncapped OpenMP costs ~10x.
"""
import argparse
import json
import os
import queue
import subprocess
import threading
from pathlib import Path
REPO = "/workspace/code/eat-map-regmix"
EPOCHS = 4
LR = 3e-3
SEEDS = (0, 1, 2)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("root", default="/workspace/runs/mae-64", nargs="?")
ap.add_argument("--gpus", default="0,1,2,3")
ap.add_argument("--per-gpu", type=int, default=2)
ap.add_argument("--threads", type=int, default=12)
args = ap.parse_args()
jobs: queue.Queue = queue.Queue()
for run in sorted(Path(args.root).iterdir()):
exports = sorted((run / "exports").glob("step_*")) if run.is_dir() else []
if not exports:
continue
final = exports[-1]
for seed in SEEDS:
if not (final / f"lora_s{seed}.json").exists():
jobs.put((run.name, final, seed))
total = jobs.qsize()
print(f"{total} LoRA fine-tunes ({EPOCHS} epochs, lr {LR:g}, seeds {SEEDS})")
lock = threading.Lock()
done = [0]
def worker(gpu):
while True:
try:
name, exp, seed = jobs.get_nowait()
except queue.Empty:
return
t = str(args.threads)
proc = subprocess.run(
["python", "-m", "eatmap.cli.lora_finetune", "--export", str(exp),
"--tag", f"s{seed}", "--epochs", str(EPOCHS), "--lr", str(LR),
"--seed", str(seed)],
cwd=REPO, capture_output=True,
env={**os.environ, "CUDA_VISIBLE_DEVICES": str(gpu), "PYTHONPATH": REPO,
"OMP_NUM_THREADS": t, "OPENBLAS_NUM_THREADS": t,
"MKL_NUM_THREADS": t, "NUMEXPR_NUM_THREADS": t},
)
with lock:
done[0] += 1
if proc.returncode:
print(f" FAIL {name} s{seed}: {proc.stderr.decode()[-200:]}")
elif done[0] % 20 == 0:
print(f" {done[0]}/{total}")
jobs.task_done()
threads = [threading.Thread(target=worker, args=(int(g),))
for g in args.gpus.split(",") for _ in range(args.per_gpu)]
for t in threads:
t.start()
for t in threads:
t.join()
# ---- collect, and screen for the collapse mode ranking-transfer.md found --
rows = {}
for run in sorted(Path(args.root).iterdir()):
exports = sorted((run / "exports").glob("step_*")) if run.is_dir() else []
if not exports:
continue
vals = []
for seed in SEEDS:
p = exports[-1] / f"lora_s{seed}.json"
if p.exists():
vals.append(json.loads(p.read_text())["lora/map"])
if len(vals) == len(SEEDS):
import statistics as st
rows[run.name] = {"vals": vals, "mean": st.mean(vals), "sd": st.stdev(vals)}
bad = [k for k, v in rows.items() if v["sd"] > 0.02]
print(f"\n{len(rows)} arms scored")
print(f"collapse screen (lora sd > 0.02): {bad if bad else 'none'}")
print(" ranking-transfer.md saw ~1% incidence at lr 3e-3, non-deterministic; re-run these")
Path(args.root, "lora_summary.json").write_text(json.dumps(rows, indent=1))
print(f"wrote {args.root}/lora_summary.json")
if __name__ == "__main__":
main()
|