v3: real-games training round tooling

ops/t4_real25.py

@@ -0,0 +1,271 @@
+#!/usr/bin/env python3
+"""T4 real-games scale-up (v3): train on 25 public ARC-AGI-3 games.
+
+Run on the T4 VM as: python3 -B /content/qwenjev/t4_real25.py
+Stages (each Popen-isolated, logged, stamped, resumable):
+  prep    : pull v2 hf_lora parts from HF -> reassemble -> stage_a_lora
+  data    : merge remote-built dataset_remote + local dataset (HF tgz)
+  real_b  : SFT on merged real-game data, resuming v2 LoRA -> hf_lora
+  real_c  : RLCD + calibration                              -> s1.pt+calib
+  upload  : push everything to HF t4_assets_v3/checkpoints/
+  eval    : ar25+bp35 @150 actions (vs zero-shot baseline: 0 wins, 0 levels)
+
+Artifacts land in checkpoints/qwenjev-s1-hf-v3 and HF t4_assets_v3/.
+Baseline for comparison (zero-shot v2 on public games):
+  ar25: 150 acts, best_level 0; bp35: 150 acts, best_level 0.
+"""
+import json
+import os
+import subprocess
+import sys
+import threading
+import time
+from pathlib import Path
+
+WORK = Path("/content/qwenjev_work_v3")
+REPO = Path("/content/qwenjev")
+OUT = REPO / "checkpoints" / "qwenjev-s1-hf-v3"
+LOGS = WORK / "logs"
+DONE = WORK / "done_v3"
+STATUS = WORK / "v3_status.json"
+
+HF_REPO = "tchbcb/qwenjev"
+
+REAL25_YAML = """
+data:
+  dataset_dir: "data/dataset_final"
+
+system1:
+  backend: "hf"
+  hf_model: "Qwen/Qwen3-VL-2B-Instruct"
+  hf_quant4: true
+  max_leaf_tokens: 512
+  score_levels: 10
+  leaf_batch_size: 2
+  eval_leaf_batch: 8
+
+stage_a_distill:
+  enabled: false          # v2 adapters are resumed as stage_a_lora (prep stage)
+
+stage_b_sft:
+  backbone_lr: 5.0e-5
+  head_lr: 2.0e-4
+  epochs: 1
+  batch_questions: 2
+  grad_accum: 2
+  max_steps: 180          # real-games corpus is bigger; keep ~2h on T4
+  aux_rows_per_batch: 3
+  max_candidates: 4
+  lora:
+    enabled: true
+    r: 16
+    alpha: 32
+    target: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+    exclude_vision: true
+    kbit_prepare: true
+
+stage_c_rlcd:
+  epochs: 1
+  group_size: 4
+  sigma_start: 0.4
+  sigma_end: 0.1
+  lr_backbone: 5.0e-5
+  lr_head: 2.0e-4
+  ce_guidance: 1.0
+  w_sph: 0.75
+  w_rps: 1.0
+  freeze_backbone: true
+  max_choice_rows: 120
+  max_aux_rows: 24
+  grad_accum: 4
+  calib_max_rows: 120
+  escalate_epochs: 1
+
+out_dir: "checkpoints/qwenjev-s1-hf-v3"
+log_every: 2
+eval_max_actions: 300
+"""
+
+_state = {"current": "init", "stages": {}, "pid": os.getpid(),
+          "started": time.strftime("%F %T"), "updated": ""}
+
+
+def save_status():
+    _state["updated"] = time.strftime("%T")
+    STATUS.write_text(json.dumps(_state, indent=1))
+
+
+def peak_sampler(proc):
+    peak = 0.0
+    while proc.poll() is None:
+        try:
+            q = subprocess.run(
+                ["nvidia-smi", "--query-gpu=memory.used",
+                 "--format=csv,noheader,nounits"],
+                capture_output=True, text=True, timeout=10)
+            peak = max(peak, float(q.stdout.strip().splitlines()[0]) / 1024.0)
+        except Exception:
+            pass
+        time.sleep(5)
+    return peak
+
+
+def run_stage(name, code, vram=True):
+    if (DONE / f"{name}.stamp").exists():
+        _state["stages"][name] = {"state": "skipped(done)"}
+        save_status()
+        print(f"[v3] {name}: already done", flush=True)
+        return
+    _state["current"] = name
+    save_status()
+    log = LOGS / f"{name}.log"
+    with log.open("w") as lf:
+        t0 = time.time()
+        proc = subprocess.Popen(
+            [sys.executable, "-B", "-c", code], cwd=str(REPO),
+            stdout=lf, stderr=subprocess.STDOUT,
+            env={**os.environ, "PYTHONUNBUFFERED": "1"})
+        peak = peak_sampler(proc) if vram else 0.0
+        rc = proc.wait()
+        secs = round(time.time() - t0, 1)
+    if rc != 0:
+        _state["stages"][name] = {"state": "FAILED", "rc": rc, "secs": secs}
+        (DONE / f"{name}.FAILED").write_text(log.read_text()[-2500:])
+        save_status()
+        raise SystemExit(f"[v3] {name} FAILED rc={rc} — see {log}")
+    _state["stages"][name] = {
+        "state": "ok", "secs": secs, "peak_vram_gb": round(peak, 2)}
+    (DONE / f"{name}.stamp").write_text(time.strftime("%F %T"))
+    save_status()
+    print(f"[v3] {name}: ok in {secs}s peak={peak:.2f}GB", flush=True)
+
+
+def hf_api():
+    from huggingface_hub import HfApi
+    return HfApi(token=(WORK / ".hf_token").read_text().strip())
+
+
+def upload(name, local: Path, in_repo: str):
+    for attempt in (1, 2, 3):
+        try:
+            t0 = time.time()
+            hf_api().upload_file(
+                path_or_fileobj=str(local), path_in_repo=in_repo,
+                repo_id=HF_REPO, repo_type="model",
+                commit_message=f"t4_assets_v3: {name} (real-games run)")
+            print(f"[v3] uploaded {in_repo} in {time.time()-t0:.0f}s",
+                  flush=True)
+            return
+        except Exception as e:
+            print(f"[v3] upload {name} attempt {attempt} failed: {e}",
+                  flush=True)
+            time.sleep(20)
+    raise SystemExit(f"[v3] upload {name} gave up")
+
+
+def main():
+    LOGS.mkdir(parents=True, exist_ok=True)
+    DONE.mkdir(parents=True, exist_ok=True)
+    (REPO / "configs" / "t4_real25.yaml").write_text(REAL25_YAML)
+    if (WORK / ".hf_token").exists():
+        os.environ["HF_TOKEN"] = (WORK / ".hf_token").read_text().strip()
+    save_status()
+
+    # ---- prep: v2 hf_lora parts -> stage_a_lora -------------------------
+    run_stage("prep", f"""
+import json, hashlib, shutil, subprocess, sys
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+OUT = Path({str(OUT)!r}); sa = OUT / "stage_a_lora"
+sa.mkdir(parents=True, exist_ok=True)
+for f in ["adapter_config.json", "README.md"]:
+    p = hf_hub_download({HF_REPO!r}, "t4_assets_v2/checkpoints/hf_lora/" + f)
+    shutil.copy(p, sa / f)
+parts_dir = Path("/tmp/v2_parts"); parts_dir.mkdir(exist_ok=True)
+for f in ["manifest.json", "reassemble.py", "adapter_p00_31568f83.pt",
+          "adapter_p01_2dcd80ad.pt", "adapter_p02_27a9e1ff.pt",
+          "adapter_p03_eee3e8f3.pt"]:
+    p = hf_hub_download({HF_REPO!r}, f"t4_assets_v2/checkpoints/hf_lora/parts/{{f}}")
+    shutil.copy(p, parts_dir / f)
+subprocess.run([sys.executable, str(parts_dir / "reassemble.py")], cwd=str(parts_dir), check=True)
+src = parts_dir / "adapter_model.safetensors"
+if not src.exists():
+    raise SystemExit("reassemble did not produce adapter_model.safetensors")
+sha = hashlib.sha256(src.read_bytes()).hexdigest()
+print("reassembled sha256:", sha)
+shutil.copy(src, sa / "adapter_model.safetensors")
+print("prep ok ->", sa)
+""", vram=False)
+
+    # ---- data: merge remote + local datasets ----------------------------
+    run_stage("data", f"""
+import shutil, tarfile, subprocess, sys
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+rem = Path("data/dataset_remote")
+if not (rem / "train" / "samples.jsonl").exists():
+    subprocess.run([sys.executable, "-m", "data.build_dataset",
+                    "--in", "data/traces_real2", "--out", "data/dataset_remote",
+                    "--val-ratio", "0.05"], check=True)
+loc = Path("data/dataset_local")
+if not (loc / "train" / "samples.jsonl").exists():
+    tgz = hf_hub_download({HF_REPO!r}, "t4_assets_v3/data/dataset_v3_local.tgz")
+    with tarfile.open(tgz) as t:
+        t.extractall("data/")
+    # tgz contains dataset_v3/ — rename to dataset_local
+    if not (loc / "train" / "samples.jsonl").exists():
+        src = Path("data/dataset_v3")
+        if (src / "train" / "samples.jsonl").exists():
+            shutil.move(str(src), str(loc))
+    assert (loc / "train" / "samples.jsonl").exists(), "bad tgz layout"
+fin = Path("data/dataset_final")
+shutil.rmtree(fin, ignore_errors=True)
+(fin / "train").mkdir(parents=True); (fin / "val").mkdir(parents=True)
+import json
+counts = {{"train": 0, "val": 0}}
+for split in ["train", "val"]:
+    with open(fin / split / "samples.jsonl", "w") as out:
+        for root in [rem, loc]:
+            p = root / split / "samples.jsonl"
+            if not p.exists():
+                continue
+            with open(p) as f:
+                for line in f:
+                    if line.strip():
+                        out.write(line); counts[split] += 1
+print("merged:", counts)
+(fin / "stats.json").write_text(json.dumps(counts))
+""", vram=False)
+
+    # ---- real_b: SFT -----------------------------------------------------
+    run_stage("real_b", """
+import runpy, sys
+sys.argv = ['train_system1.py', '--config', 'configs/t4_real25.yaml']
+runpy.run_module('train.train_system1', run_name='__main__')
+""")
+    upload("hf_lora_config", OUT / "hf_lora" / "adapter_config.json",
+           "t4_assets_v3/checkpoints/hf_lora/adapter_config.json")
+    upload("hf_lora_model", OUT / "hf_lora" / "adapter_model.safetensors",
+           "t4_assets_v3/checkpoints/hf_lora/adapter_model.safetensors")
+    upload("hf_lora_readme", OUT / "hf_lora" / "README.md",
+           "t4_assets_v3/checkpoints/hf_lora/README.md")
+
+    # ---- real_c: RLCD + calibration --------------------------------------
+    run_stage("real_c", """
+import runpy, sys
+sys.argv = ['calibrate_rlcd.py', '--config', 'configs/t4_real25.yaml']
+runpy.run_module('train.calibrate_rlcd', run_name='__main__')
+""")
+    for f, r in [("qwenjev_s1.pt", "qwenjev_s1.pt"),
+                 ("calib_report.json", "calib_report.json"),
+                 ("stage_b_meta.json", "stage_b_meta.json")]:
+        upload(f, OUT / f, f"t4_assets_v3/checkpoints/{r}")
+
+    _state["current"] = "ALL_DONE"
+    _state["finished"] = time.strftime("%F %T")
+    save_status()
+    print("[v3] ALL_DONE", flush=True)
+
+
+if __name__ == "__main__":
+    main()