Buckets:
| """cwm3/train/train_u0_from_init.py — fine-tune U0 (exact kin_can_noK recipe, unchanged train/train_cwm3.py) from a PRETRAINED TRUNK. | |
| train_cwm3.py has no init-checkpoint option and is frozen while the week-1 queue runs, so this thin wrapper uses its RESUME path as the | |
| init mechanism: it builds the U0 model exactly as train_cwm3 does (build_branch("U", seed, **model_kwargs(cfg)) -> the head keeps the | |
| seed-k initialisation of the scratch baseline), copies the pretrained trunk (gru + mode embedding) into it, and writes an epoch-0 | |
| checkpoint in train_cwm3's own format (model / fresh optimiser / scheduler / state epoch 0 / rng) into the run directory. It then | |
| invokes the unchanged trainer as a subprocess: | |
| train_cwm3.py --stage U0 --variant kin_can_noK --seed k --out <dir> --no-lock --minutes 22 | |
| which "resumes" from that checkpoint and runs the standard recipe, prediction phase, pred_layout and run.json. Nothing else differs | |
| from results/runs/U0_kin_can_noK_s<k>, so eval/paired_compare.py compares like with like. | |
| The trunk file is a plain weights file (results/runs/U0_dmspre_trunk_s<k>/trunk.pt); the fine-tuned model reads NO teacher target and no | |
| DMS field: its inputs are the kin_can_noK inputs (Kh, X, Δt, mode). Provenance: <out>/init_from.json (+ run_train_u0_from_init_*.json). | |
| GPU etiquette: no lock (coordinator 2026-09-10), threads <= 4, VRAM / RAM / trainer-count checks before launching. | |
| python train/train_u0_from_init.py --seed 0 [--trunk PATH] [--out DIR] [--variant kin_can_noK] [--minutes 22] [--max-anchors N --epochs E] | |
| 2026-09-11 (C1 transfer, E10 (c); opt-in, defaults reproduce the DMS-teacher chain bit for bit): --kin-source / --kin-version / --stats-version / | |
| --pack-version are forwarded to train_cwm3.py (survivor recipe = `--variant kin_can_noK_clip8_ln --kin-source causal_routes --kin-version v2 | |
| --stats-version 2`); a trunk that carries an "in_ln" state_dict (train/pretrain_u_trunk_aide.py) is loaded into the model's input LayerNorm when | |
| the variant has one (`u_in_ln`); the trunk's `source` string is recorded instead of the hard-coded DMS note. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import re | |
| import subprocess | |
| import sys | |
| import time | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| os.environ.setdefault("MKL_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| for p in (CWM3, CWM3 / "train", CWM3 / "common"): | |
| sys.path.insert(0, str(p)) | |
| import train_cwm3 as tc # noqa: E402 (import only; the file is not modified) | |
| from models import build_branch, params_by_module # noqa: E402 | |
| from runmeta import write_run_json, sha256_file # noqa: E402 | |
| PY = sys.executable | |
| MAX_TRAINERS = 2 | |
| MIN_FREE_VRAM_GB, MIN_AVAIL_RAM_GB = 8.0, 14.0 | |
| def live_trainers(): | |
| me = os.getpid() | |
| n = [] | |
| for pid in os.listdir("/proc"): | |
| if not pid.isdigit() or int(pid) == me: | |
| continue | |
| try: | |
| cmd = open(f"/proc/{pid}/cmdline", "rb").read().replace(b"\0", b" ").decode(errors="ignore") | |
| except OSError: | |
| continue | |
| if (re.search(r"train_cwm3(_joint)?\.py|train_u0_from_init\.py|dms_teacher_pretrain\.py|pretrain_u_trunk_aide\.py", cmd) and "pgrep" not in cmd | |
| and re.match(r"\S*python[\d.]*\s", cmd)): # the interpreter process only (not the `bash -c` shell that launched it) | |
| n.append((int(pid), cmd[:120])) | |
| return n | |
| def wait_for_slot(out, poll_s=60, max_wait_s=6 * 3600): | |
| t0 = time.time() | |
| while True: | |
| alive = live_trainers() | |
| if len(alive) < MAX_TRAINERS: | |
| return alive | |
| tc.log(f"waiting for a trainer slot: {len(alive)} alive (limit {MAX_TRAINERS}): {[c for _, c in alive]}", out) | |
| if time.time() - t0 > max_wait_s: | |
| raise SystemExit("no trainer slot within the wait budget") | |
| time.sleep(poll_s) | |
| def check_resources(out): | |
| import torch | |
| free_gb = torch.cuda.mem_get_info()[0] / 2**30 if torch.cuda.is_available() else 0.0 | |
| avail_gb = float([l for l in open("/proc/meminfo") if l.startswith("MemAvailable")][0].split()[1]) / 2**20 | |
| tc.log(f"resources: VRAM free {free_gb:.1f} GB, RAM available {avail_gb:.1f} GB", out) | |
| assert free_gb >= MIN_FREE_VRAM_GB, f"free VRAM {free_gb:.1f} GB < {MIN_FREE_VRAM_GB}" | |
| assert avail_gb >= MIN_AVAIL_RAM_GB, f"available RAM {avail_gb:.1f} GB < {MIN_AVAIL_RAM_GB}" | |
| def write_init_checkpoint(out, cfg, seed, trunk_path, epochs): | |
| import torch | |
| z = torch.load(trunk_path, map_location="cpu", weights_only=False) | |
| model = build_branch("U", seed, **tc.model_kwargs(cfg)) # torch.manual_seed(seed) inside: head init == scratch baseline seed k | |
| before = {k: v.clone() for k, v in model.state_dict().items()} | |
| model.gru.load_state_dict(z["gru"], strict=True) | |
| model.mode.load_state_dict(z["mode"], strict=True) | |
| if "in_ln" in z and getattr(model, "in_ln", None) is not None: # 2026-09-11: trunks pretrained with the input LayerNorm (AIDE C1 chain) | |
| model.in_ln.load_state_dict(z["in_ln"], strict=True) | |
| after = model.state_dict() | |
| changed = sorted(k for k in after if not torch.equal(before[k], after[k])) | |
| unchanged = sorted(k for k in after if torch.equal(before[k], after[k])) | |
| kept_ok = ("head.",) + (("in_ln.",) if "in_ln" not in z else ()) | |
| assert all(k.startswith(("gru.", "mode.", "in_ln.")) for k in changed) and all(k.startswith(kept_ok) for k in unchanged), (changed, unchanged) | |
| torch.manual_seed(seed) | |
| opt, sched = tc.make_optim(model, cfg["recipe"], epochs) | |
| state = {"epoch": 0, "batch_pos": 0, "best": 1e18, "best_epoch": -1, "patience": 0, "acq": 0, "done": False} | |
| stats = {k: np.asarray(v) for k, v in z.get("stats", {}).items()} | |
| ck = Path(out) / "ckpt_ep000_b00000.pt" | |
| assert not ck.exists() | |
| torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "sched": sched.state_dict(), "state": state, "rng_torch": torch.get_rng_state(), | |
| "rng_cuda": None, "cfg": cfg, "seed": seed, "stats": stats, "params": params_by_module(model), | |
| "init_from": {"trunk": str(trunk_path), "trunk_sha256": sha256_file(trunk_path), "trunk_epoch": int(z.get("epoch", -1)), | |
| "loaded_keys": changed, "kept_from_seed_init": unchanged}}, ck) | |
| return ck, changed, unchanged, z | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--seed", type=int, default=0) | |
| ap.add_argument("--variant", default="kin_can_noK") | |
| ap.add_argument("--trunk", default="") | |
| ap.add_argument("--out", default="") | |
| ap.add_argument("--minutes", type=float, default=22) | |
| ap.add_argument("--max-anchors", type=int, default=0) | |
| ap.add_argument("--epochs", type=int, default=0) | |
| ap.add_argument("--pred-splits", default="") | |
| ap.add_argument("--kin-source", default="", help="forwarded to train_cwm3.py (cwm2 | causal_routes | oldin_ctgt); empty = trainer default") | |
| ap.add_argument("--kin-version", default="", help="forwarded to train_cwm3.py (v1 | v2)") | |
| ap.add_argument("--stats-version", default="", help="forwarded to train_cwm3.py (1 | 2)") | |
| ap.add_argument("--pack-version", default="", help="forwarded to train_cwm3.py (v1 | v2)") | |
| a = ap.parse_args() | |
| t_start = time.time() | |
| trunk = Path(a.trunk) if a.trunk else CWM3 / "results" / "runs" / f"U0_dmspre_trunk_s{a.seed}" / "trunk.pt" | |
| assert trunk.exists(), trunk | |
| out = Path(a.out) if a.out else CWM3 / "results" / "runs" / (f"U0_{a.variant}_dmspre_s{a.seed}" + (f"_smoke_max{a.max_anchors}" if a.max_anchors else "")) | |
| out.mkdir(parents=True, exist_ok=True) | |
| cfg = tc.stage_cfg("U0", a.variant) | |
| epochs = a.epochs or int(cfg["recipe"]["max_epochs"]) | |
| if (out / "PRED_DONE").exists(): | |
| print(f"{out} already complete") | |
| return | |
| if not (out / "TRAIN_DONE").exists() and not list(out.glob("ckpt_ep*_b*.pt")): | |
| ck, changed, unchanged, z = write_init_checkpoint(out, cfg, a.seed, trunk, epochs) | |
| info = {"trunk": str(trunk), "trunk_sha256": sha256_file(trunk), "trunk_epoch": int(z.get("epoch", -1)), "trunk_val_loss": float(z.get("val_loss", float("nan"))), | |
| "trunk_source": z.get("source", ""), "loaded_keys": changed, "kept_from_seed_init": unchanged, "init_ckpt": ck.name, "seed": a.seed, "variant": a.variant, | |
| "recipe": cfg["recipe"], "mechanism": "train_cwm3.py resume path (epoch-0 checkpoint written by train_u0_from_init.py)", | |
| "fine_tune_inputs": "kin_can_noK inputs only (Kh, X, dt, mode); no teacher target, no DMS field", | |
| "forwarded_flags": {k: v for k, v in (("kin_source", a.kin_source), ("kin_version", a.kin_version), ("stats_version", a.stats_version), ("pack_version", a.pack_version)) if v}} | |
| (out / "init_from.json").write_text(json.dumps(info, indent=1)) | |
| tc.log(f"init checkpoint {ck.name} written: loaded {len(changed)} trunk tensors from {trunk.name} (epoch {info['trunk_epoch']}), head kept at seed-{a.seed} init", out) | |
| wait_for_slot(out) | |
| check_resources(out) | |
| cmd = [PY, str(CWM3 / "train" / "train_cwm3.py"), "--stage", "U0", "--variant", a.variant, "--seed", str(a.seed), "--out", str(out), "--no-lock", | |
| "--minutes", str(a.minutes)] | |
| if a.max_anchors: | |
| cmd += ["--max-anchors", str(a.max_anchors)] | |
| if a.epochs: | |
| cmd += ["--epochs", str(a.epochs)] | |
| if a.pred_splits: | |
| cmd += ["--pred-splits", a.pred_splits] | |
| for flag, val in (("--kin-source", a.kin_source), ("--kin-version", a.kin_version), ("--stats-version", a.stats_version), ("--pack-version", a.pack_version)): | |
| if val: | |
| cmd += [flag, str(val)] | |
| env = dict(os.environ, OMP_NUM_THREADS="4", MKL_NUM_THREADS="4") | |
| tc.log("launch: " + " ".join(cmd), out) | |
| with open(out / "train_cwm3.stdout.log", "a") as f: | |
| code = subprocess.call(cmd, env=env, stdout=f, stderr=subprocess.STDOUT, cwd=str(CWM3)) | |
| tc.log(f"train_cwm3 exited {code}", out) | |
| if code != 0: | |
| raise SystemExit(code) | |
| trunk_source = str(json.loads((out / "init_from.json").read_text()).get("trunk_source", "")) | |
| write_run_json(out, f"train_u0_from_init_s{a.seed}", [Path(__file__), CWM3 / "train" / "train_cwm3.py", CWM3 / "models" / "branches.py"], [trunk], | |
| seed=a.seed, t_start=t_start, input_listing=False, | |
| extra={"stage": "U0", "variant": a.variant, "init_from": json.loads((out / "init_from.json").read_text()), "dms_fields_used": [], | |
| "trunk_source": trunk_source, | |
| "note": ("fine-tune reads no DMS field; the trunk was pretrained by weak/dms_teacher_pretrain.py (see its run.json for dms_fields_used)" | |
| if "dms_teacher" in trunk_source else f"fine-tune reads no DMS field; the trunk was pretrained by {trunk_source} (no DMS field there either)"), | |
| "train_cwm3_exit": code, "gpu_lock": "none (--no-lock, coordinator 2026-09-10)"}) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 11.2 kB
- Xet hash:
- ff76d0495e54e62137715b6eca648add18901fe0fe0ca7e5577fc863bb3674a2
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.