"""Job driver: run experiment stages sequentially inside one job. Stages: lrsweep - 6 short runs (3k steps): {fixed-4, rec-1x4} x lr {1e-4,3e-4,1e-3} stage1 - matched-FLOP family at the chosen lr: fixed-4, rec-1x4, rec-2x2, rec-1xdyn (each 3 seeds, 25k steps) + their evaluations stage2 - iso-param pair: rec-4x4 (3 seeds), fixed-4-long (100k steps, 2 seeds) + evaluations eval - evaluate every local checkpoint not yet evaluated """ import argparse import glob import json import os import subprocess import sys STEPS = 25000 LR_FILE = "out/chosen_lr.json" def run(cmd): print("+", " ".join(cmd), flush=True) subprocess.run(cmd, check=True) def run_eval(args, ckpt, out_dir="results"): """Run evaluate.py on one checkpoint, then upload the JSON to the Hub. Without the upload the eval JSONs die with the job's ephemeral disk.""" name = os.path.splitext(os.path.basename(ckpt))[0] out = f"{out_dir}/{name}_eval.json" if os.path.exists(out): return run([sys.executable, "evaluate.py", "--ckpt", ckpt, "--out-dir", out_dir]) if args.hub_repo: run([sys.executable, "-c", f"from huggingface_hub import HfApi; HfApi().upload_file(" f"path_or_fileobj={out!r}, " f"path_in_repo='results/{name}_eval.json', " f"repo_id='{args.hub_repo}', repo_type='model')"]) def train_cmd(a, run_name, steps, lr, hub_repo): cmd = [sys.executable, "train.py", "--arch", a["arch"], "--depth", str(a["depth"]), "--iters", a["iters"], "--seed", str(a["seed"]), "--steps", str(steps), "--lr", str(lr), "--run-name", run_name, "--out-dir", "out"] if hub_repo: cmd += ["--hub-repo", hub_repo, "--space-id", os.environ.get("TRACKIO_SPACE_ID", "")] return cmd def stage_lrsweep(args): for lr in (1e-4, 3e-4, 1e-3): for a in ({"arch": "fixed", "depth": 4, "iters": "1", "seed": 0}, {"arch": "rec", "depth": 1, "iters": "4", "seed": 0}): name = f"lr{lr:g}-{a['arch']}{a['depth']}x{a['iters']}" run(train_cmd(a, name, 3000, lr, args.hub_repo)) # pick lr by mean id accuracy over the last eval of each pair scores = {} for lr in (1e-4, 3e-4, 1e-3): tot = 0.0 for a in ("fixed4x1", "rec1x4"): h = json.load(open(f"out/lr{lr:g}-{a}_history.json")) tot += h["id_acc"][-1]["mean"] scores[str(lr)] = tot / 2 best = max(scores, key=scores.get) json.dump({"scores": scores, "chosen": best}, open(LR_FILE, "w")) if args.hub_repo: run([sys.executable, "-c", "import sys; from huggingface_hub import " "HfApi; HfApi().upload_file(path_or_fileobj='out/chosen_lr.json', " "path_in_repo='results/chosen_lr.json', " f"repo_id='{args.hub_repo}', repo_type='model')"]) print("[lrsweep] scores", scores, "-> chosen", best) def chosen_lr(hub_repo=None): if not os.path.exists(LR_FILE): # stage1/stage2 run in fresh containers: fetch the lr choice made by # the lrsweep job from the Hub repo from huggingface_hub import hf_hub_download p = hf_hub_download(repo_id=hub_repo, repo_type="model", filename="results/chosen_lr.json", local_dir="out") os.replace(os.path.join("out", "results", "chosen_lr.json"), LR_FILE) return json.load(open(LR_FILE))["chosen"] def stage1(args): lr = chosen_lr(args.hub_repo) seeds = (0, 1, 2) archs = [("fixed", 4, "1"), ("rec", 1, "4"), ("rec", 2, "2"), ("rec", 1, "dynamic")] for arch, depth, iters in archs: for s in seeds: name = f"{arch}{depth}x{iters}-seed{s}" ck = f"out/{name}.pt" if os.path.exists(ck): continue run(train_cmd({"arch": arch, "depth": depth, "iters": iters, "seed": s}, name, STEPS, lr, args.hub_repo)) run_eval(args, ck) def stage2(args): lr = chosen_lr(args.hub_repo) for s in (0, 1, 2): name = f"rec4x4-seed{s}" if not os.path.exists(f"out/{name}.pt"): run(train_cmd({"arch": "rec", "depth": 4, "iters": "4", "seed": s}, name, STEPS, lr, args.hub_repo)) run_eval(args, f"out/{name}.pt") for s in (0, 1): name = f"fixed4-long-seed{s}" if not os.path.exists(f"out/{name}.pt"): run(train_cmd({"arch": "fixed", "depth": 4, "iters": "1", "seed": s}, name, 4 * STEPS, lr, args.hub_repo)) run_eval(args, f"out/{name}.pt") def stage3(args): """Aligned-unroll experiment: train with R = k (one unroll per hop), testing whether the models' implicit R~k tuning can be made explicit and extended beyond the training range.""" lr = chosen_lr(args.hub_repo) for s in (0, 1, 2): name = f"rec1aligned-seed{s}" if not os.path.exists(f"out/{name}.pt"): run(train_cmd({"arch": "rec", "depth": 1, "iters": "aligned", "seed": s}, name, STEPS, lr, args.hub_repo)) run_eval(args, f"out/{name}.pt") def stage4(args): """Aligned objective at 4x the step budget (1 seed): distinguishes 'harder optimization, needs more steps' from 'fundamentally harder' after the 25k-step aligned runs plateaued near 0.10 id_acc.""" lr = chosen_lr(args.hub_repo) name = "rec1aligned-long-seed0" if not os.path.exists(f"out/{name}.pt"): run(train_cmd({"arch": "rec", "depth": 1, "iters": "aligned", "seed": 0}, name, 4 * STEPS, lr, args.hub_repo)) run_eval(args, f"out/{name}.pt") def stage_eval(args): if not glob.glob("out/*.pt"): # recovery mode (fresh container): pull every checkpoint from the Hub, # plus existing eval JSONs so the skip-if-present logic works from huggingface_hub import snapshot_download snapshot_download(repo_id=args.hub_repo, repo_type="model", allow_patterns=["ckpts/*.pt", "results/*_eval.json"], local_dir=".") cks = sorted(glob.glob("out/*.pt") + glob.glob("ckpts/*.pt")) for ck in cks: name = os.path.splitext(os.path.basename(ck))[0] if name.startswith(("lr", "smoke")): continue # not part of the main comparison grid if os.path.exists(f"results/{name}_eval.json"): continue run_eval(args, ck) def main(): p = argparse.ArgumentParser() p.add_argument("stage", choices=["lrsweep", "stage1", "stage2", "stage3", "stage4", "eval"]) p.add_argument("--hub-repo", default=None) a = p.parse_args() if a.hub_repo: os.environ.setdefault("TRACKIO_HUB_REPO", a.hub_repo) {"lrsweep": stage_lrsweep, "stage1": stage1, "stage2": stage2, "stage3": stage3, "stage4": stage4, "eval": stage_eval}[a.stage](a) if __name__ == "__main__": main()