diff --git a/.gitignore b/.gitignore index 6b4c5c1..aa4505a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,6 @@ +# machine-local Claude Code artifacts +CLAUDE.local.md + # caches / build __pycache__/ *.pyc @@ -25,3 +28,4 @@ results/*.out *.log conductor_*_out/ conductor_toolscale_*/ +.tokensave diff --git a/HANDOFF.md b/HANDOFF.md new file mode 100644 index 0000000..943b9dc --- /dev/null +++ b/HANDOFF.md @@ -0,0 +1,18 @@ +# HANDOFF: TRINITY materialized-weights loader + +Repo: /Users/russ/Projects/snowy2/OpenFugu (venv .venv with torch/transformers; `pip install safetensors huggingface_hub` into .venv if missing). Do NOT git commit. + +Problem: openfugu/mini.py loads a 19,456-dim CMA-ES vector (env FUGU_VECTOR, model_iter_60.npy) and reconstructs 9 SVF-adapted matrices + linear router head. That .npy no longer exists upstream. HF dataset `nshkrdotcom/trinity-coordinator-adapted-qwen3-0.6b` now ships MATERIALIZED weights: +- checkpoints/0001_embedder.token_embedding.kernel.safetensors ... 0009_language_modeling_head.output.kernel.safetensors (embedder; decoder.blocks.26 self_attention query/key/value/output; ffn gate/intermediate/output; lm head) +- router_head.safetensors (tensor `trinity.router_head.linear.weight`, [10,1024]) +- manifest.json: 7 agent_labels + 3 role_labels = 10 outputs, hidden 1024, opt layer [26] +CAUTION: exported from Elixir/Bumblebee ("kernel" naming) — tensors may be TRANSPOSED vs PyTorch nn.Linear convention; verify by shape against the Qwen3-0.6B state dict and transpose where needed. + +Steps: +1. Read openfugu/mini.py fully (FUGU_VECTOR load path, weight installation, head application) + docs/HOW_FUGU_IS_IMPLEMENTED.md (SVF + head sections). +2. Download the 11 files via huggingface_hub.hf_hub_download(repo_type="dataset") into artifacts/materialized/ (~650MB). +3. New file openfugu/materialized.py: `load_materialized(checkpoint_dir) -> (weights: dict[qwen_module_weight_name, torch.Tensor], head: torch.Tensor)`; map Bumblebee names → HF Qwen3 module paths as used by mini.py; handle dtype + transposition. +4. Minimal hook in mini.py: env FUGU_CHECKPOINT_DIR set → install materialized weights + head instead of SVF reconstruction. Smallest possible diff. +5. Verify: FUGU_MODEL=/Users/russ/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/c1899de289a04d12100db370d81485cdf75e47ca FUGU_CHECKPOINT_DIR=$PWD/artifacts/materialized FUGU_FIXTURE=$PWD/artifacts/qwen_router_prompt_eval_cases.json .venv/bin/python -m openfugu.mini --self-test (check exact flag in mini.py). Target ≈95% agent / 100% role on the 37-case fixture; chance-level accuracy ⇒ suspect transposition/name mapping, iterate. + +Report: self-test numbers + mapping/transposition decisions. diff --git a/openfugu/materialized.py b/openfugu/materialized.py new file mode 100644 index 0000000..2abfccd --- /dev/null +++ b/openfugu/materialized.py @@ -0,0 +1,94 @@ +import os +import json +import torch +from safetensors.torch import load_file + +def load_materialized(checkpoint_dir: str) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + """ + Loads materialized weights and router head from the specified checkpoint directory. + Maps Bumblebee names -> HF Qwen3 PyTorch module paths. + Handles transposition for weights whose shapes do not match PyTorch convention. + """ + manifest_path = os.path.join(checkpoint_dir, "manifest.json") + with open(manifest_path, "r") as f: + manifest = json.load(f) + + # 1. Load the 9 selected tensors + weights = {} + selected_tensors = manifest.get("selected_tensors", []) + + expected_shapes = { + "model.embed_tokens.weight": (151936, 1024), + "model.layers.26.self_attn.q_proj.weight": (2048, 1024), + "model.layers.26.self_attn.k_proj.weight": (1024, 1024), + "model.layers.26.self_attn.v_proj.weight": (1024, 1024), + "model.layers.26.self_attn.o_proj.weight": (1024, 2048), + "model.layers.26.mlp.gate_proj.weight": (3072, 1024), + "model.layers.26.mlp.up_proj.weight": (3072, 1024), + "model.layers.26.mlp.down_proj.weight": (1024, 3072), + "lm_head.weight": (151936, 1024), + } + + for entry in selected_tensors: + chk_path = entry["checkpoint_path"] + abs_chk_path = os.path.join(checkpoint_dir, chk_path) + source_name = entry["source_name"] + path_key = entry["path"] + + if not os.path.exists(abs_chk_path): + raise FileNotFoundError(f"Checkpoint file not found: {abs_chk_path}") + + tensors_dict = load_file(abs_chk_path) + if path_key not in tensors_dict: + # fallback to checking if any key in tensors_dict contains the name + matching_keys = [k for k in tensors_dict if path_key in k] + if matching_keys: + tensor = tensors_dict[matching_keys[0]] + else: + raise KeyError( + f"Tensor key {path_key!r} for {source_name!r} not found in " + f"{abs_chk_path} (available keys: {sorted(tensors_dict)})" + ) + else: + tensor = tensors_dict[path_key] + + if source_name not in expected_shapes: + raise ValueError(f"No expected shape registered for tensor {source_name!r}") + expected_shape = expected_shapes[source_name] + if source_name in [ + "model.layers.26.self_attn.k_proj.weight", + "model.layers.26.self_attn.v_proj.weight", + ]: + tensor = tensor.T.contiguous() + elif tensor.shape != expected_shape: + # Check if transposing makes shapes match + if tuple(tensor.shape[::-1]) == expected_shape: + tensor = tensor.T.contiguous() + else: + raise ValueError( + f"Tensor {source_name} shape {tensor.shape} does not match expected shape {expected_shape}" + ) + + weights[source_name] = tensor + + # 2. Load the router head + router_head_path = os.path.join(checkpoint_dir, "router_head.safetensors") + if not os.path.exists(router_head_path): + raise FileNotFoundError(f"Router head file not found: {router_head_path}") + + router_head_dict = load_file(router_head_path) + + # Try common keys + if "trinity_router_head" in router_head_dict: + head = router_head_dict["trinity_router_head"] + else: + head = list(router_head_dict.values())[0] + + # Expected shape is (10, 1024) + if head.shape != (10, 1024): + if tuple(head.shape[::-1]) == (10, 1024): + head = head.T.contiguous() + else: + raise ValueError(f"Router head shape {head.shape} is not (10, 1024)") + + return weights, head diff --git a/openfugu/mini.py b/openfugu/mini.py index 9f1d37a..0b044a0 100644 --- a/openfugu/mini.py +++ b/openfugu/mini.py @@ -105,17 +105,13 @@ class FuguRouter: which is what makes a routing decision ~one forward pass. [EXEC] """ - def __init__(self, model_dir: str, vector_path: str, dtype="float32", + def __init__(self, model_dir: str, vector_path: str | None = None, dtype="float32", device: str | None = None, seed: int | None = None): import torch from transformers import AutoModelForCausalLM, AutoTokenizer self.torch = torch self.rng = np.random.default_rng(seed) - vec = np.load(vector_path).astype(np.float64) - if vec.shape != (VEC_LEN,): - raise ValueError(f"router vector must be {VEC_LEN} floats, got {vec.shape}") - self.tok = AutoTokenizer.from_pretrained(model_dir) # transformers >=5 uses dtype=, <5 uses torch_dtype= — support both td = getattr(torch, dtype) @@ -127,9 +123,28 @@ def __init__(self, model_dir: str, vector_path: str, dtype="float32", self.model.to(device) self.device = next(self.model.parameters()).device - self._apply_svf(vec[:SVF_LEN]) - # head: last 10240 -> (10, 1024) [EXEC] - self.head = torch.from_numpy(vec[SVF_LEN:].copy()).float().reshape(HEAD_ROWS, HIDDEN).to(self.device) + checkpoint_dir = os.environ.get("FUGU_CHECKPOINT_DIR") + if checkpoint_dir: + try: + from .materialized import load_materialized + except ImportError: + from materialized import load_materialized + weights, head = load_materialized(checkpoint_dir) + with torch.no_grad(): + sd = self.model.state_dict() + for k, w in weights.items(): + sd[k].copy_(w.to(device=self.device, dtype=td)) + self.svf_keys = list(weights.keys()) + self.head = head.to(device=self.device, dtype=torch.float32) + else: + if not vector_path: + raise ValueError("vector_path must be provided if FUGU_CHECKPOINT_DIR is not set") + vec = np.load(vector_path).astype(np.float64) + if vec.shape != (VEC_LEN,): + raise ValueError(f"router vector must be {VEC_LEN} floats, got {vec.shape}") + self._apply_svf(vec[:SVF_LEN]) + # head: last 10240 -> (10, 1024) [EXEC] + self.head = torch.from_numpy(vec[SVF_LEN:].copy()).float().reshape(HEAD_ROWS, HIDDEN).to(self.device) # SVF: scale only singular values, freeze U/V, energy-preserving renorm. [CODE] # Matrices consumed in state_dict order: embed_tokens, layer-26 {q,k,v,o, diff --git a/train/train_trinity_coding.py b/train/train_trinity_coding.py new file mode 100644 index 0000000..6c5d855 --- /dev/null +++ b/train/train_trinity_coding.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +# OpenFugu — Apache-2.0. Part of an independent, open reimplementation of +# the Fugu orchestrator. NOT affiliated with Sakana AI. See NOTICE. +# Reference: TRINITY (arXiv:2512.04695). Self-trains the coordinator via +# sep-CMA-ES on REAL coding data (HumanEval) with a real worker pool — the +# coding-benchmark variant of train_trinity_real.py. Original code. +""" +train_trinity_coding.py — TRINITY self-training on HumanEval coding benchmark. + +Same sep-CMA-ES loop as train_trinity_real.py, but for coding: + - tasks : HumanEval problems (openai/openai_humaneval), split="test" + - features: a REAL Qwen3-0.6B penultimate hidden state of the prompt + - workers : a real pool via litellm (Novita), differing models + - reward : real code execution (verifiable signal via Python subprocess execution) + +The coordinator (bias-free linear head over the hidden state) learns which +worker to send each coding problem to, to maximize solved rate. +""" +from __future__ import annotations +import argparse, os, re, subprocess, sys, tempfile, time +import numpy as np + +HIDDEN = 1024 +HIDDEN_POS = -2 + + +def extract_completion(text: str) -> str: + """Extract candidate completion from fenced block or use raw response as-is.""" + m = re.search(r"```python\n(.*?)\n?```", text, re.DOTALL) + if m: + return m.group(1) + m = re.search(r"```\n(.*?)\n?```", text, re.DOTALL) + if m: + return m.group(1) + return text + + +class Backbone: + """Real Qwen3-0.6B -> penultimate hidden state of a question (the router feature).""" + def __init__(self, model_dir, device=None): + import torch + from transformers import AutoModelForCausalLM, AutoTokenizer + self.torch = torch + self.tok = AutoTokenizer.from_pretrained(model_dir) + try: + self.model = AutoModelForCausalLM.from_pretrained(model_dir, dtype=torch.float32).eval() + except TypeError: + self.model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=torch.float32).eval() + if device: + self.model.to(device) + self.device = next(self.model.parameters()).device + self._cache = {} + + def feature(self, question: str) -> np.ndarray: + if question in self._cache: + return self._cache[question] + torch = self.torch + ids = self.tok(f"user: {question}", return_tensors="pt").to(self.device) + with torch.no_grad(): + h = self.model.model(**ids).last_hidden_state[0, HIDDEN_POS, :] + v = h.float().cpu().numpy() + self._cache[question] = v + return v + + +def route(head_vec, feat, n_workers): + """Bias-free linear head -> worker id (argmax), faithful to mini.py.""" + W = head_vec.reshape(n_workers, HIDDEN) + return int(np.argmax(W @ feat)) + + +def main(): + ap = argparse.ArgumentParser(description="TRINITY self-train on HumanEval (minimal).") + ap.add_argument("--model", default=os.environ.get("FUGU_MODEL", "Qwen/Qwen3-0.6B")) + ap.add_argument("--slot-models", required=True, help="csv of litellm worker ids (the pool)") + ap.add_argument("--n-train", type=int, default=40, help="HumanEval questions (kept small/cheap)") + ap.add_argument("--iters", type=int, default=25) + ap.add_argument("--sigma0", type=float, default=0.5) + ap.add_argument("--max-tokens", type=int, default=768) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--out", default="trinity_humaneval.npy") + args = ap.parse_args() + + import cma, litellm + from datasets import load_dataset + + workers = args.slot_models.split(",") + n_workers = len(workers) + api_key = os.environ.get("FUGU_API_KEY") or os.environ.get("NOVITA_API_KEY") or os.environ.get("OPENAI_API_KEY") + api_base = os.environ.get("FUGU_BASE_URL") or os.environ.get("OPENAI_BASE_URL") + + ds = load_dataset("openai/openai_humaneval", split="test") + tasks = list(ds)[:args.n_train] + print(f"[real-train] {len(tasks)} HumanEval tasks, {n_workers} workers: {workers}", flush=True) + + bb = Backbone(args.model) + feats = [bb.feature(t["prompt"]) for t in tasks] # cache real hidden states once + print(f"[real-train] cached {len(feats)} real Qwen3-0.6B features (dim {feats[0].shape[0]})", flush=True) + + # worker call cache: (worker_id, task_id) -> solved? so CMA candidates reuse answers + solve_cache: dict = {} + + def worker_solves(wid, task): + task_id = task["task_id"] + key = (wid, task_id) + if key in solve_cache: + return solve_cache[key] + + prompt = task["prompt"] + test = task["test"] + entry_point = task["entry_point"] + + worker_prompt = prompt + "\nComplete this Python function. Output ONLY the indented function body (do NOT repeat the `def` line or docstring), in a single ```python fenced code block, no explanation." + + kw = dict(model=workers[wid], + messages=[{"role": "user", + "content": worker_prompt}], + max_tokens=args.max_tokens, temperature=0.0, timeout=45) + if api_key: kw["api_key"] = api_key + if api_base: kw["api_base"] = api_base + ok = 0.0 + call_succeeded = False + # Candidate code is untrusted model output, executed via subprocess (the whole + # point of HumanEval-style grading, not a shell/injection risk: no shell=True, + # args are a list). Give it an allowlisted env, not the parent's full env with + # secrets denylisted after the fact. + run_env = {k: os.environ[k] for k in ("PATH", "PYTHONPATH") if k in os.environ} + for attempt in range(5): + try: + out = litellm.completion(**kw).choices[0].message.content or "" + candidate = extract_completion(out) + call_succeeded = True + + # Grading (the reward) + full_code = prompt + "\n" + candidate + "\n" + test + f"\ncheck({entry_point})\n" + tmp_file = tempfile.NamedTemporaryFile(suffix=".py", mode="w", delete=False, encoding="utf-8") + tmp_path = tmp_file.name + try: + tmp_file.write(full_code) + tmp_file.close() + res = subprocess.run([sys.executable, tmp_path], capture_output=True, timeout=10, env=run_env) + ok = 1.0 if res.returncode == 0 else 0.0 + except subprocess.TimeoutExpired: + ok = 0.0 + except Exception as e: + print(f" [warn] execution error: {str(e)[:60]}", flush=True) + ok = 0.0 + finally: + try: + if os.path.exists(tmp_path): + os.unlink(tmp_path) + except Exception: + pass + break + except litellm.RateLimitError: + wait = 2 ** attempt + print(f" [rate-limit] worker {wid} 429, retry {attempt+1}/5 in {wait}s", flush=True) + time.sleep(wait) + except Exception as e: + print(f" [warn] worker {wid} call failed: {str(e)[:60]}", flush=True) + break + if not call_succeeded: + # Transient/exhausted call failure, not a graded answer — don't poison the + # cache with a fake "unsolved" that would stick for the rest of the run. + print(f" [call] worker {wid} task SKIPPED (call never succeeded)", flush=True) + return 0.0 + solve_cache[key] = ok + print(f" [call] worker {wid} task done -> {'OK' if ok else 'miss'} (cache={len(solve_cache)})", flush=True) + return ok + + def fitness(head_vec): + tot = 0.0 + for task, feat in zip(tasks, feats): + wid = route(head_vec, feat, n_workers) + tot += worker_solves(wid, task) + return tot / len(tasks) + + # baseline: each worker alone + random (uses the same cache -> cheap) + rng = np.random.default_rng(args.seed) + per_worker = [] + for w in range(n_workers): + per_worker.append(np.mean([worker_solves(w, t) for t in tasks])) + best_single = max(per_worker) + print("[baseline] per-worker solved rate: " + + ", ".join(f"{workers[w].split('/')[-1]}={per_worker[w]:.2f}" for w in range(n_workers)), flush=True) + + # sep-CMA-ES over the head + dim = n_workers * HIDDEN + es = cma.CMAEvolutionStrategy(np.zeros(dim), args.sigma0, + {"seed": args.seed, "verbose": -9, "CMA_diagonal": True}) + best_vec, best_fit = None, -1.0 + for it in range(args.iters): + cands = es.ask() + fits = [fitness(c) for c in cands] + es.tell(cands, [-f for f in fits]) + i = int(np.argmax(fits)) + if fits[i] > best_fit: + best_fit, best_vec = fits[i], cands[i].copy() + print(f"[iter {it}] best_solved={best_fit:.3f} " + f"(best single worker {best_single:.3f}) cache={len(solve_cache)}", flush=True) + + np.save(args.out, best_vec) + print(f"\n[result] coordinator solved {best_fit:.3f} vs best single worker {best_single:.3f}") + print(f"[result] learned routing per task:") + for task, feat in zip(tasks, feats): + w = route(best_vec, feat, n_workers) + print(f" -> {workers[w].split('/')[-1]:24s} | {task['prompt'][:50]}") + if best_fit >= best_single: + print("PASS — sep-CMA-ES self-trained TRINITY on REAL HumanEval, " + "coordinator >= best single worker") + + +if __name__ == "__main__": + main() diff --git a/train/train_trinity_real.py b/train/train_trinity_real.py index cf3d228..a047096 100644 --- a/train/train_trinity_real.py +++ b/train/train_trinity_real.py @@ -21,7 +21,7 @@ reward is just a number compare. """ from __future__ import annotations -import argparse, os, re, sys +import argparse, os, re, sys, time import numpy as np HIDDEN = 1024 @@ -106,19 +106,34 @@ def worker_solves(wid, q, gold): key = (wid, q) if key in solve_cache: return solve_cache[key] - try: - kw = dict(model=workers[wid], - messages=[{"role": "user", - "content": q + "\nGive the final numeric answer at the end."}], - max_tokens=args.max_tokens, temperature=0.0) - if api_key: kw["api_key"] = api_key - if api_base: kw["api_base"] = api_base - out = litellm.completion(**kw).choices[0].message.content or "" - ok = 1.0 if numeric_answer(out) == gold else 0.0 - except Exception as e: - print(f" [warn] worker {wid} call failed: {str(e)[:60]}", flush=True) - ok = 0.0 + kw = dict(model=workers[wid], + messages=[{"role": "user", + "content": q + "\nGive the final numeric answer at the end."}], + max_tokens=args.max_tokens, temperature=0.0, timeout=45) + if api_key: kw["api_key"] = api_key + if api_base: kw["api_base"] = api_base + ok = 0.0 + call_succeeded = False + for attempt in range(5): + try: + out = litellm.completion(**kw).choices[0].message.content or "" + ok = 1.0 if numeric_answer(out) == gold else 0.0 + call_succeeded = True + break + except litellm.RateLimitError: + wait = 2 ** attempt + print(f" [rate-limit] worker {wid} 429, retry {attempt+1}/5 in {wait}s", flush=True) + time.sleep(wait) + except Exception as e: + print(f" [warn] worker {wid} call failed: {str(e)[:60]}", flush=True) + break + if not call_succeeded: + # Transient/exhausted call failure, not a graded answer — don't poison the + # cache with a fake "unsolved" that would stick for the rest of the run. + print(f" [call] worker {wid} task SKIPPED (call never succeeded)", flush=True) + return 0.0 solve_cache[key] = ok + print(f" [call] worker {wid} task done -> {'OK' if ok else 'miss'} (cache={len(solve_cache)})", flush=True) return ok def fitness(head_vec):