| """Per-node GR00T relay: per-layout batch-size probe (largest that fits), then 20-epoch production per task. |
| |
| Run on every node (nohup), node rank from the DLC env RANK: python groot_chain.py [task ...] |
| Default tasks: b6 a3 a7 a18 (basket postponed). A phase succeeds only if all 4 node launchers exit 0. |
| |
| Batch (user rule: as large as memory allows, 2026-11-10): descending CANDIDATES, 4 steps each on all 32 GPUs; a |
| candidate fits when all launchers exit 0 AND the peak reserved memory of every rank (written by hb_launch.py to |
| logs/peak-<phase>-rank<r>.json) is <= MAX_FRAC of the GPU. Node 0 alone judges each probe and writes a verdict file |
| the other nodes wait for, so all nodes take the same decision. One decision per robot layout (g1 / h2). |
| Outputs: $G/groot-runs-bmax/ (the earlier per-GPU-64 runs stay in $G/groot-runs/). |
| """ |
| import fcntl |
| import json |
| import math |
| import os |
| import shutil |
| import subprocess |
| import sys |
| import time |
| from pathlib import Path |
|
|
| H = Path(__file__).resolve().parent |
| G = H.parent |
| RUNS = G / os.environ.get("HB_GROOT_RUNS_NAME", "groot-runs-bmax") |
| NODE = int(os.environ["RANK"]) |
| TASKS = sys.argv[1:] or ["b6", "a3", "a7", "a18"] |
| CANDIDATES = [320, 288, 256, 224, 192, 160, 128, 96, 64, 48, 32] |
| MAX_FRAC = 0.92 |
| EPOCHS = 20 |
| PORT_BASE = int(os.environ.get("HB_GROOT_PORT_BASE", "32500")) |
| RUNS.mkdir(parents=True, exist_ok=True) |
| os.environ["HB_GROOT_RUNS"] = str(RUNS) |
| lock = (RUNS / f"chain-node-{NODE}.lock").open("a") |
| fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) |
|
|
|
|
| def log(msg): |
| line = f"{time.strftime('%Y-%m-%d %H:%M:%S')} node{NODE} {msg}" |
| print(line, flush=True) |
| with (RUNS / f"chain-node-{NODE}.log").open("a") as f: |
| f.write(line + "\n") |
|
|
|
|
| def run_phase(task, phase, bs, steps, save, port): |
| logs = RUNS / task / "logs" |
| exits = [logs / f"{phase}-node-{r}.exit" for r in range(4)] |
| if all(p.is_file() for p in exits): |
| return all(p.read_text().strip() == "0" for p in exits) |
| subprocess.run(["bash", str(H / "groot_node.sh"), task, phase, str(bs), str(steps), str(save), str(port)]) |
| deadline = time.monotonic() + 600 |
| while not all(p.is_file() for p in exits): |
| if time.monotonic() > deadline: |
| log(f"{task}/{phase}: timed out waiting for peer exit files") |
| return False |
| time.sleep(5) |
| ok = all(p.read_text().strip() == "0" for p in exits) |
| log(f"{task}/{phase}: bs={bs} steps={steps} -> {'OK' if ok else 'FAIL'}") |
| return ok |
|
|
|
|
| def peak_frac(task, phase): |
| """Max over ranks of peak reserved / total; None when no rank reported.""" |
| fracs = [] |
| for f in (RUNS / task / "logs").glob(f"peak-{phase}-rank*.json"): |
| try: |
| d = json.loads(f.read_text()) |
| fracs.append(d["peak_reserved_gb"] / d["total_gb"]) |
| except (ValueError, KeyError, OSError): |
| continue |
| return (max(fracs), len(fracs)) if fracs else (None, 0) |
|
|
|
|
| def probe_verdict(task, cand, port): |
| phase = f"probe-bs{cand}" |
| verdict = RUNS / task / "logs" / f"{phase}.verdict.json" |
| ok = run_phase(task, phase, cand, 4, 1000000, port) |
| if NODE == 0 and not verdict.is_file(): |
| time.sleep(10) |
| frac, n = peak_frac(task, phase) if ok else (None, 0) |
| fits = bool(ok and frac is not None and frac <= MAX_FRAC) |
| verdict.write_text(json.dumps({"per_gpu_batch": cand, "ran_ok": ok, "peak_frac": frac, "ranks_reported": n, |
| "max_frac": MAX_FRAC, "fits": fits}) + "\n") |
| deadline = time.monotonic() + 600 |
| while not verdict.is_file(): |
| if time.monotonic() > deadline: |
| log(f"{task}/{phase}: no verdict from node 0") |
| return False |
| time.sleep(5) |
| v = json.loads(verdict.read_text()) |
| log(f"{task}/{phase}: verdict {v}") |
| return v["fits"] |
|
|
|
|
| def prune_resume_state(out): |
| """Keep weights of every checkpoint but resume state (DeepSpeed optimizer shards, scheduler, RNG) only in the latest.""" |
| ckpts = sorted((p for p in out.glob("checkpoint-*") if p.name.split("-")[-1].isdigit()), |
| key=lambda p: int(p.name.split("-")[-1])) |
| for c in ckpts[:-1]: |
| for p in list(c.glob("global_step*")) + list(c.glob("optimizer*.pt")) + list(c.glob("scheduler.pt")) \ |
| + list(c.glob("rng_state*.pth")) + list(c.glob("latest")) + list(c.glob("zero_to_fp32.py")): |
| if p.is_dir(): |
| shutil.rmtree(p, ignore_errors=True) |
| else: |
| p.unlink(missing_ok=True) |
| log(f"{c.name}: resume state dropped (weights kept)") |
|
|
|
|
| def batch_for(task): |
| layout = "h2" if task in ("a17", "a18") else "g1" |
| decision = RUNS / f"batch-decision-{layout}.json" |
| if decision.is_file(): |
| return json.loads(decision.read_text())["per_gpu_batch"] |
| for i, cand in enumerate(CANDIDATES): |
| if probe_verdict(task, cand, PORT_BASE + i + (20 if layout == "h2" else 0)): |
| if NODE == 0: |
| decision.write_text(json.dumps({"per_gpu_batch": cand, "global_batch": cand * 32, "probe_task": task, |
| "layout": layout, "candidates": CANDIDATES, |
| "max_frac": MAX_FRAC}) + "\n") |
| return cand |
| time.sleep(20) |
| log(f"{layout}: no batch size fit; stopping") |
| sys.exit(1) |
|
|
|
|
| for k, task in enumerate(TASKS): |
| bs = batch_for(task) |
| log(f"{task}: per-GPU batch {bs} (global {bs * 32})") |
| receipt = json.loads((G / "groot-data" / task / "hb-dataset-receipt.json").read_text()) |
| steps = math.ceil(EPOCHS * receipt["frames"] / (bs * 32)) |
| save = math.ceil(steps / 2) |
| if not run_phase(task, "production", bs, steps, save, PORT_BASE + 50 + k): |
| log(f"{task}: production failed; stopping chain") |
| sys.exit(1) |
| if NODE == 0: |
| prune_resume_state(RUNS / task / "production") |
| (RUNS / task / "production-complete.json").write_text(json.dumps( |
| {"task": task, "frames": receipt["frames"], "epochs": EPOCHS, "per_gpu_batch": bs, |
| "global_batch": bs * 32, "max_steps": steps, "save_steps": save}) + "\n") |
| log("all tasks done") |
|
|
