File size: 8,679 Bytes
6634b22 cc18230 09f9061 cc18230 09f9061 6634b22 5f83c26 6634b22 3ceb445 6634b22 3ceb445 6634b22 3ceb445 6634b22 cc18230 6634b22 3ceb445 6634b22 cc18230 6634b22 cc18230 6634b22 0cdf651 b70fd3c 5f83c26 fea27e9 5f83c26 6634b22 09f9061 15e60f5 09f9061 15e60f5 cc18230 6634b22 58e7ba8 6634b22 cc18230 6634b22 0cdf651 5f83c26 6634b22 5f83c26 b70fd3c 6634b22 | 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 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 | """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 a.get("stem"):
cmd += ["--stem-depth", str(a["stem"])]
if a.get("dense"):
cmd += ["--dense-supervision"]
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 stage5(args):
"""v2 aligned variants (user design): (A) R = k+4 margin, end-of-chain
supervision only - tests whether exact halting was the blocker;
(B) R = k+4 + untied stem block + dense per-iteration supervision with
hold-after-k targets f^{min(t,k)}(x) - the full fix package."""
lr = chosen_lr(args.hub_repo)
# resume support: a fresh container has no local checkpoints, so pull
# finished ones from the Hub; the skip-if-present check below then
# skips them and only unfinished variants train
if not glob.glob("out/rec1m4*.pt"):
from huggingface_hub import snapshot_download
snapshot_download(repo_id=args.hub_repo, repo_type="model",
allow_patterns=["ckpts/rec1m4*.pt"],
local_dir="out")
for f in glob.glob("out/ckpts/rec1m4*.pt"):
os.replace(f, os.path.join("out", os.path.basename(f)))
variants = [
("rec1m4", {"arch": "rec", "depth": 1, "iters": "aligned4"}),
("rec1m4dense", {"arch": "rec", "depth": 1, "iters": "aligned4",
"stem": 1, "dense": 1}),
]
for tmpl, base in variants:
for s in (0, 1, 2):
name = f"{tmpl}-seed{s}"
if not os.path.exists(f"out/{name}.pt"):
run(train_cmd({**base, "seed": s}, name, 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", "stage5", "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, "stage5": stage5,
"eval": stage_eval}[a.stage](a)
if __name__ == "__main__":
main() |