code/gpu/job2_main.py

"""HF Job 2 driver — main scaled experiments (paper Wfe1iJocjF / FlashOptim).

Stage A (Claims 2, 3): memory-scaling sweep — one variant per process —
  sizes gpt-30m/82m/124m/355m x {ref, flash, flash_gr, split_only, quant_only}
  + resnet50 x {ref, flash} (tests the one deviating Table 6 row).
Stage B (Claims 4, 6): GPT-2 124M on real FineWeb10B tokens:
  ref_s0 / flash_s0 (matched init + data), ref_s1 (seed-noise yardstick),
  linear_s0 (companding ablation, expect instability per paper Fig. 5).

Time-budgeted; writes everything under /out/job2/.
"""

import json
import os
import subprocess
import sys
import time
import urllib.request
import zipfile
import io

SHA = "43548ceebbbdb869c4ea9db7b21bba1cf9ae7004"
OUT = os.environ.get("OUT_DIR", "/out/job2")
CODE = os.path.dirname(os.path.abspath(__file__))
T0 = time.time()
RESULTS = {"stages": {}, "sha": SHA}

# wall-clock budgets (s)
BUDGET_ARMS = {"ref_s0": 2400, "flash_s0": 2400, "ref_s1": 1200, "linear_s0": 900}


def dump():
    os.makedirs(OUT, exist_ok=True)
    with open(f"{OUT}/results.json", "w") as f:
        json.dump(RESULTS, f, indent=2, default=str)


def install():
    url = f"https://github.com/databricks/flashoptim/archive/{SHA}.zip"
    buf = urllib.request.urlopen(url, timeout=120).read()
    zipfile.ZipFile(io.BytesIO(buf)).extractall("/work")
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", "--no-deps",
                    f"/work/flashoptim-{SHA}"], check=True)
    import torch, triton, flashoptim  # noqa
    RESULTS["env"] = dict(gpu=torch.cuda.get_device_name(0),
                          torch=torch.__version__, triton=triton.__version__,
                          cuda=torch.version.cuda, flashoptim=flashoptim.__version__)
    dump()


def stage_a():
    runs = [(s, v) for s in ["gpt-30m", "gpt-82m", "gpt-124m", "gpt-355m"]
            for v in ["ref", "flash", "flash_gr", "split_only", "quant_only"]]
    runs += [("resnet50", "ref"), ("resnet50", "flash")]
    out = []
    for model, variant in runs:
        p = subprocess.run([sys.executable, f"{CODE}/mem_probe.py",
                            "--variant", variant, "--model", model],
                           capture_output=True, text=True, timeout=900)
        line = p.stdout.strip().splitlines()[-1] if p.stdout.strip() else ""
        rec = (json.loads(line) if p.returncode == 0 and line.startswith("{")
               else {"variant": variant, "model": model, "exit": p.returncode,
                     "stderr": p.stderr[-500:]})
        out.append(rec)
        print("MEM", json.dumps(rec), flush=True)
    RESULTS["mem_sweep"] = out
    with open(f"{OUT}/mem_sweep.json", "w") as f:
        json.dump(out, f, indent=2)
    dump()


def stage_b():
    from huggingface_hub import hf_hub_download
    os.makedirs("/work/data", exist_ok=True)
    train = hf_hub_download("kjj0/fineweb10B-gpt2", "fineweb_train_000001.bin",
                            repo_type="dataset", local_dir="/work/data")
    val = hf_hub_download("kjj0/fineweb10B-gpt2", "fineweb_val_000000.bin",
                          repo_type="dataset", local_dir="/work/data")
    arms = [("ref_s0", "ref", 0), ("flash_s0", "flash", 0),
            ("ref_s1", "ref", 1), ("linear_s0", "linear", 0)]
    outs = {}
    for name, arm, seed in arms:
        remaining = 10200 - (time.time() - T0) - 300  # keep 5 min margin, 170m job
        budget = min(BUDGET_ARMS[name], max(0, remaining))
        if budget < 300:
            outs[name] = {"skipped": "insufficient time remaining"}
            continue
        cmd = [sys.executable, f"{CODE}/gpt2_train.py", "--arm", arm,
               "--size", "gpt-124m", "--init-seed", str(seed),
               "--batch", "8", "--ctx", "1024", "--lr", "6e-4",
               "--warmup", "300", "--compile", "1",
               "--time-budget-s", str(int(budget)),
               "--train-shard", train, "--val-shard", val,
               "--out", f"{OUT}/{name}"]
        print("LAUNCH", name, "budget", int(budget), flush=True)
        p = subprocess.run(cmd, capture_output=True, text=True,
                           timeout=int(budget) + 900)
        with open(f"{OUT}/{name}.log", "w") as f:
            f.write(p.stdout[-100000:] + "\n--stderr--\n" + p.stderr[-20000:])
        tail = [ln for ln in p.stdout.splitlines() if ln.startswith("SUMMARY")]
        outs[name] = (json.loads(tail[-1][8:]) if tail else
                      {"exit": p.returncode, "stderr": p.stderr[-800:]})
        RESULTS["gpt_arms"] = outs
        dump()
    RESULTS["gpt_arms"] = outs
    dump()


if __name__ == "__main__":
    for name, fn in [("install", install), ("stage_a", stage_a), ("stage_b", stage_b)]:
        t = time.time()
        try:
            fn()
            RESULTS["stages"][name] = {"status": "ok", "secs": round(time.time() - t, 1)}
        except Exception as e:
            import traceback
            RESULTS["stages"][name] = {"status": "FAIL", "error": str(e),
                                       "trace": traceback.format_exc()[-2000:]}
            print(f"STAGE {name} FAILED: {e}", flush=True)
        dump()
    print("JOB2_RESULTS " + json.dumps(RESULTS, default=str))
    sys.exit(0 if RESULTS["stages"].get("install", {}).get("status") == "ok" else 1)