persona: fix what the persona/ move broke, and add the gates that would have caught it
- phase3_lora_and_transplant: slice_qwen4exp_gguf stayed at top-level scripts/ - trainer REPO vs ROOT: vendor/ and models/ are not under persona/, and ensure_gguf_bridge mkdirs lib_path.parent, so a directory probe on persona/vendor self-validated a wrong root - --validate-only / --export-only were never added to main()'s dispatch, so both exited 0 having printed --help: success with no work done - --export-only: re-export a checkpoint at another quantization without a 28 min val sweep - evaluate_validation_splits: guard each window against OOM and log its CE, because an unguarded 28 min silent sweep reads exactly like a crashed process - transplant_persona_writers: CLI bundle/out paths measured: epoch-2 heldout win rate 1.0 (49/49), ood_ce 0.8774. q8_0 writer export flattens the voice on 3/3 probes against in_place at pinned MoCE routing; in_place (3.90 GB) matches base in voice and is 440 MB smaller, so it is the only correct default export type for this model.
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980#!/usr/bin/env python3"""Run code inside a Colab runtime through its Jupyter kernel, the way `colab exec` does.
Colab's ssh proxy serves one WebSocket client per runtime and wedges server-side when a client dieswithout closing cleanly, which leaves no way in short of destroying the VM. The kernel path has nosuch limit: `jupyter_kernel_client` speaks the framing (it is the library google-colab-cli uses), sothis only has to hand it the runtime URL and the proxy token.
Usage: colab_kernel_exec.py --host HOST --token TOKEN --code "python""""import argparseimport jsonimport urllib.request
import jupyter_kernel_client as jkc
def make_client(host, token, kernel_id=None): kwargs = { "subprotocol": jkc.JupyterSubprotocol.DEFAULT, "extra_params": {"colab-runtime-proxy-token": token}, } server_url = f"https://{host}" headers = {"X-Colab-Client-Agent": "albedo", "X-Colab-Runtime-Proxy-Token": token} if hasattr(jkc, "ColabKernelClient"): cls, tok = jkc.ColabKernelClient, "proxy_token" elif hasattr(jkc, "KernelClient"): cls, tok = jkc.KernelClient, "token" else: cls, tok = jkc.JupyterKernelClient, "token" client = cls(server_url=server_url, **{tok: token}, kernel_id=kernel_id, client_kwargs=kwargs, headers=headers) # Never let the library reap a kernel it did not create. client._own_kernel = False return client
def run(host, token, code, timeout, kernel_id=None): if kernel_id: client = make_client(host, token, kernel_id) client.start() return client, kernel_id req = urllib.request.Request( f"https://{host}/api/kernels?colab-runtime-proxy-token={token}", data=json.dumps({"name": "python3"}).encode(), method="POST", headers={"Content-Type": "application/json"}) with urllib.request.urlopen(req, timeout=30) as resp: kernel_id = json.loads(resp.read())["id"] client = make_client(host, token, kernel_id) client.start() return client, kernel_id
def main(): ap = argparse.ArgumentParser() ap.add_argument("--host", required=True) ap.add_argument("--token", required=True) ap.add_argument("--code", default="") ap.add_argument("--file", default="") ap.add_argument("--timeout", type=float, default=120.0) ap.add_argument("--kernel-id", default="") a = ap.parse_args()
code = Path(a.file).read_text() if a.file else a.code client, kernel_id = run(a.host, a.token, code, a.timeout, a.kernel_id or None) print(f"KERNEL_ID {kernel_id}", flush=True) result = client.execute(code, timeout=a.timeout) for out in result.get("outputs", []): kind = out.get("output_type") if kind == "stream": print(out.get("text", ""), end="", flush=True) elif kind in ("execute_result", "display_data"): print(out.get("data", {}).get("text/plain", ""), flush=True) elif kind == "error": print("\n".join(out.get("traceback", [])), flush=True) print(f"\nEXEC_STATUS {result.get('status')}", flush=True)
if __name__ == "__main__": main()42 unmodified lines434445464748495051525350 unmodified lines42 unmodified lines
def pull(url, local, headers, resume=True, quiet=False): os.makedirs(os.path.dirname(os.path.abspath(local)) or ".", exist_ok=True) # A checkpoint written every optimizer step keeps identical byte sizes (same tensor shapes), so # "same size" is NOT "same content": resume-by-size would silently never refresh latest.pt. # Pass resume=False for anything that is rewritten in place. have = os.path.getsize(local) if (resume and os.path.exists(local)) else 0 total, status = size_and_status(url, headers) if status == 404: raise SystemExit(f"{local}: remote 404")50 unmodified lines35 unmodified lines363738394041424344454647612 unmodified linesExpand all35 unmodified linesimport tempfilefrom pathlib import Pathfrom typing import Dict, Iterable, List, Optional, Set, Tuple, Union
# `slice_qwen4exp_gguf` stayed at the repo's top-level scripts/ when the persona scripts moved into# persona/scripts/, so this folder alone no longer resolves the import.sys.path.insert(0, str(Path(__file__).resolve().parent))sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "scripts"))from slice_qwen4exp_gguf import ( GGML_TYPE_TRAITS, GGUF_MAGIC, GGUF_TYPE_STRING,612 unmodified linesExpand all35 unmodified lines3637383940414243446 unmodified lines515253545556575859606162636433 unmodified lines98991001019596102103104105106100107108109110111126 unmodified linesExpand all23823924024123524224323723824424524624724824924324425025125225325425568 unmodified lines32432532632732132232832933033133233362 unmodified lines3963973983993934004014024034041640 unmodified linesExpand all2045204620472048204220432044204920502051205220532054205520562057205820592060206120622063206420652066206720472048204920502051205220532068206920702071207220732074207520762077207820792080208120822083267 unmodified linesExpand all2351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375125 unmodified linesExpand all2501250225032504246625052506250725082509140 unmodified linesExpand all265026512652265326542655265626572658265926602661266226172663266426652666266714 unmodified lines268226832684268526402686268726882689269026915 unmodified lines35 unmodified linesfrom __future__ import annotations
import argparseimport ctypesimport gcimport gzipimport jsonimport mathimport mmap6 unmodified linesimport urllib.requestfrom pathlib import Path
ROOT = Path(__file__).resolve().parent.parent# `vendor/` and `models/` stayed at the repo root when these scripts moved into persona/scripts/,# so ROOT (the persona dir, where data/ and checkpoints/ live) is the wrong base for them.# Probe a file, not the directory: the bridge build mkdirs lib_path.parent, which would# create persona/vendor/ and make a directory probe succeed inside the wrong tree.REPO = next((p for p in (ROOT, ROOT.parent) if (p / "vendor" / "strata" / "metal" / "src" / "gguf_bridge.cpp").is_file()), ROOT)DATA_DIR = ROOT / "data" / "agent_persona_v1"PACK_PATH = DATA_DIR / "colab_a100_tokenized_pack.json.gz"
33 unmodified lines return limit, (a[limit] if len(a) > limit else -1), (b[limit] if len(b) > limit else -1)
def build_tokenized_pack() -> int: sys.path.insert(0, str(ROOT / "vendor" / "strata" / "serve")) sys.path.insert(0, str(ROOT / "vendor" / "strata" / "tools")) sys.path.insert(0, str(REPO / "vendor" / "strata" / "serve")) sys.path.insert(0, str(REPO / "vendor" / "strata" / "tools")) from frontend import ChatTemplate, literal_tags, mark_think_literals, unmark_think_literals # noqa: E402 import strata_tokenizer as ST # noqa: E402
tok_dir = ROOT / "models" / "K512_iq2_xs" / "tokenizer" tok_dir = REPO / "models" / "K512_iq2_xs" / "tokenizer" vocab = json.loads((tok_dir / "vocab.json").read_text(encoding="utf-8")) toks = [None] * len(vocab) for t_s, i_s in vocab.items(): toks[i_s] = t_s126 unmodified linesExpand all
def ensure_gguf_bridge() -> ctypes.CDLL: """Compiles and loads `libgguf_bridge` (.dylib on macOS, .so on Linux/Colab).""" ext = ".dylib" if sys.platform == "darwin" else ".so" lib_path = ROOT / "vendor" / "strata" / "build-metal" / f"libgguf_bridge{ext}" lib_path = REPO / "vendor" / "strata" / "build-metal" / f"libgguf_bridge{ext}" lib_path.parent.mkdir(parents=True, exist_ok=True) src_bridge = ROOT / "vendor" / "strata" / "metal" / "src" / "gguf_bridge.cpp" src_ref = ROOT / "vendor" / "strata" / "metal" / "src" / "ref_quant.cpp" src_bridge = REPO / "vendor" / "strata" / "metal" / "src" / "gguf_bridge.cpp" src_ref = REPO / "vendor" / "strata" / "metal" / "src" / "ref_quant.cpp" if not lib_path.exists() or src_bridge.stat().st_mtime > lib_path.stat().st_mtime: cxx = "/usr/bin/clang++" if sys.platform == "darwin" and Path("/usr/bin/clang++").exists() else "g++" cmd = [ cxx, "-O3", "-std=c++17", "-shared", "-fPIC", f"-I{ROOT / 'vendor' / 'llama.cpp' / 'ggml' / 'src'}", f"-I{ROOT / 'vendor' / 'llama.cpp' / 'ggml' / 'include'}", f"-I{REPO / 'vendor' / 'llama.cpp' / 'ggml' / 'src'}", f"-I{REPO / 'vendor' / 'llama.cpp' / 'ggml' / 'include'}", str(src_bridge), str(src_ref), "-o", str(lib_path), ] if sys.platform == "darwin": cmd.extend(["-framework", "Accelerate"])68 unmodified lines import shutil nvcc = shutil.which("nvcc") or ("/usr/local/cuda/bin/nvcc" if Path("/usr/local/cuda/bin/nvcc").exists() else None) if not nvcc: return None so_path = ROOT / "vendor" / "strata" / "build-metal" / "libgguf_cuda.so" cu_src = ROOT / "vendor" / "strata" / "metal" / "src" / "gguf_gdn_cuda.cu" so_path = REPO / "vendor" / "strata" / "build-metal" / "libgguf_cuda.so" cu_src = REPO / "vendor" / "strata" / "metal" / "src" / "gguf_gdn_cuda.cu" if not cu_src.exists(): return None so_path.parent.mkdir(parents=True, exist_ok=True) if not so_path.exists() or cu_src.stat().st_mtime > so_path.stat().st_mtime:62 unmodified lines ple_path = Path(args.ple_gguf).resolve() if args.ple_gguf else None if ple_path is None or not ple_path.exists(): cand_ple = [ gguf_path.parent / "ple-iq4_nl.gguf", ROOT / "models" / "K256_reap_l2_sum" / "Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", REPO / "models" / "K256_reap_l2_sum" / "Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf", gguf_path, ] for cp in cand_ple: if cp.exists():1640 unmodified linesExpand all sft_id_slice = pack["sft_val_id"][:max_items] if max_items > 0 else pack["sft_val_id"] sft_ood_slice = pack["sft_val_ood_repo"][:max_items] if max_items > 0 else pack["sft_val_ood_repo"] ftpo_val_slice = pack["ftpo_val_id"][:max_items] if max_items > 0 else pack["ftpo_val_id"]
with torch.no_grad(): for w in sft_id_slice: l_val = compute_sft_window_loss(model, w) def _sweep(losses, windows, label): """Score each window, skipping (and logging) one the GPU cannot hold.
Training guards every window against OutOfMemoryError (see the epoch loop) but this sweep did not, so one long validation window ended the process after the final training step and before any metrics were written. Log per window too: a silent death is then locatable from the log. """ for wi, w in enumerate(windows): try: l_val = compute_sft_window_loss(model, w) except torch.OutOfMemoryError: print(f" [{label} {wi + 1}/{len(windows)}] OOM, skipping", flush=True) model.clear_seq_cache() gc.collect() torch.cuda.empty_cache() continue model.clear_seq_cache() ce = float(l_val.item()) if l_val is not None else float("nan") if l_val is not None: val_id_losses.append(float(l_val.item())) for w in sft_ood_slice: l_val = compute_sft_window_loss(model, w) model.clear_seq_cache() if l_val is not None: val_ood_losses.append(float(l_val.item())) for p in ftpo_val_slice: losses.append(ce) print(f" [{label} {wi + 1}/{len(windows)}] ce={ce:.4f}", flush=True)
with torch.no_grad(): _sweep(val_id_losses, sft_id_slice, "sft_val_id") _sweep(val_ood_losses, sft_ood_slice, "sft_val_ood") for pi, p in enumerate(ftpo_val_slice): _, margin = compute_ftpo_dpo_pair_loss(model, p, beta=args.ftpo_beta, do_backward=False) ftpo_margins.append(margin) if margin > 0.0: ftpo_wins += 1 print(f" [ftpo_val {pi + 1}/{len(ftpo_val_slice)}] margin={margin:+.4f}", flush=True)
model.train() ood_skipped = max(0, (max_items if max_items > 0 else total_ood_raw) - len(val_ood_losses)) if not args.smoke_test:267 unmodified linesExpand all tmp_p = path.with_suffix(".pt.tmp") torch.save(payload, tmp_p) tmp_p.replace(path)
if args.validate_only: if args.export_only: # Re-exporting a checkpoint at a different quantization is a merge over the 96 writer # tensors and needs no validation pass, which is 28 minutes of forward passes. if args.out_writers_bin: exp_res = export_merged_writers_bin(model, Path(args.out_writers_bin), export_type=args.export_type) print(f"Exported merged GGUF writer replacement bundle: {json.dumps(exp_res)}", flush=True) return 0 val_metrics = evaluate_validation_splits(model, pack) _log(f"=== Epoch {start_epoch} Validation (validate-only) === {json.dumps(val_metrics)}") _save_atomic_ckpt(start_epoch, resume_idx, global_opt_step, extra={"val_metrics": val_metrics}, path=ckpt_dir / f"epoch_{start_epoch}.pt") if args.out_writers_bin: exp_res = export_merged_writers_bin(model, Path(args.out_writers_bin), export_type=args.export_type) print(f"Exported merged GGUF writer replacement bundle: {json.dumps(exp_res)}", flush=True) return val_metrics
stop_early = False for epoch in range(start_epoch, args.epochs): if stop_early: break125 unmodified linesExpand all ) tp_path = Path(args.transplant_out_gguf).resolve() print(f"Transplanted 96 merged writers into {args.transplant_out_gguf}: {json.dumps({'transplanted_tensors_count': tp_res['transplanted_tensors_count'], 'file_size': tp_res['file_size']})}", flush=True) # If strata_window_cli is present on macOS and .experts sidecar exists, verify transplanted GGUF logits parity vs PyTorch lora_enabled=True! strata_cli = ROOT / "vendor" / "strata" / "build-metal" / "strata_window_cli" strata_cli = REPO / "vendor" / "strata" / "build-metal" / "strata_window_cli" if strata_cli.exists() and ple_path is not None: tp_exp_link = Path(str(tp_path) + ".experts") if exp_sidecar_path.exists() and not tp_exp_link.exists(): try:140 unmodified linesExpand all ap.add_argument("--check-gradients", action="store_true", help="Run 48-layer directional finite-difference gradient check at T=2255 (>2051) on L43_attn_output & L44_ssm_out and exit") ap.add_argument("--benchmark-window", type=int, default=0, help="Run full 48-layer forward+backward on longest train window <= N tokens and log peak memory + timing") ap.add_argument("--train", action="store_true", help="Run GGUF-backed 96-writer LoRA SFT + FTPO/DPO training") ap.add_argument("--smoke-test", action="store_true", help="Run local 4-window + 4-pair smoke test asserting 96/96 grads & CE << ln V") ap.add_argument("--export-only", action="store_true", help="With --validate-only, skip the heldout sweep and only re-export " "--out-writers-bin at --export-type. Merging writers needs no forward pass.") ap.add_argument("--validate-only", action="store_true", help="Score the heldout splits against --resume and exit, without training. The " "post-loop validation only runs inside an epoch that actually trained, so a " "run that died after its last step needs this to produce its metrics.") ap.add_argument("--eval-live", type=str, default="", help="Evaluate val_live_prompts against live Strata URL") ap.add_argument("--eval-temp", type=float, default=0.0, help="Sampling temperature for --eval-live (default: 0.0 greedy)") ap.add_argument("--gguf", type=str, default=str(ROOT / "models" / "K512_iq2_xs" / "K512-iq2_xs-metal.gguf")) ap.add_argument("--gguf", type=str, default=str(REPO / "models" / "K512_iq2_xs" / "K512-iq2_xs-metal.gguf")) ap.add_argument("--ple-gguf", type=str, default="") ap.add_argument("--checkpoint-dir", type=str, default=str(ROOT / "checkpoints" / "agent_persona_lora_v1")) ap.add_argument("--out-writers-bin", type=str, default=str(ROOT / "checkpoints" / "agent_persona_lora_v1" / "agent_persona_writers.bin")) ap.add_argument("--export-type", type=str, choices=["bf16", "q8_0", "in_place"], default="q8_0", help="Quantization type for exported 96 writer tensors")14 unmodified lines args = ap.parse_args()
if args.pack: return build_tokenized_pack() if args.verify_parity or args.check_gradients or args.benchmark_window > 0 or args.train or args.smoke_test: if (args.verify_parity or args.check_gradients or args.benchmark_window > 0 or args.train or args.smoke_test or args.validate_only or args.export_only): return run_gguf_lora_training(args) if args.eval_live: return evaluate_live_prompts(args.eval_live, temperature=args.eval_temp) ap.print_help()5 unmodified lines1234567891011121314151617181920212223242526#!/usr/bin/env python3"""Transplant a trained writer bundle into a copy of the K512 GGUF.
`transplant_gguf_tensors` is whitelist-only and byte-level, so this is the whole job: load the bundle,hand it over with the writers scope, and report what changed."""import sysfrom pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))from phase3_lora_and_transplant import load_writers_bin, transplant_gguf_tensors # noqa: E402
BASE = Path("models/K512_iq2_xs/K512-iq2_xs-metal.gguf")BUNDLE = Path("persona/checkpoints/agent_persona_lora_v1/agent_persona_writers.bin")OUT = Path("models/K512_iq2_xs/K512-iq2_xs-metal-persona-q8_0.gguf")
if __name__ == "__main__": import sys as _sys bundle = Path(_sys.argv[1]) if len(_sys.argv) > 1 else BUNDLE out = Path(_sys.argv[2]) if len(_sys.argv) > 2 else OUT replacements = load_writers_bin(str(bundle)) print(f"bundle: {len(replacements)} writers from {bundle}", flush=True) result = transplant_gguf_tensors(str(BASE), str(out), replacements, scope="swift_writers", allow_type_change=True) print(f"transplant: {result}", flush=True) print(f"wrote {out} ({out.stat().st_size:,} bytes)", flush=True)