llm: first real-LLM prototype — recombining specialist LLMs (C2/C4)

First step from toy models toward real language models, on one 16 GB GPU.
New src/llm/ package: procedural task families + exact-match verifier
(tasks.py), batched eval (evaluate.py), LoRA specialisation (specialise.py,
manual answer-only SFT), weight-space merge via peft add_weighted_adapter
(merge.py: soup = averaged deltas, ties = sign-reconciled union), runner
(experiment.py, kind llm_merge). Base Qwen2.5-0.5B-Instruct (Apache-2.0);
three disjoint hard families (lists/strings/arith); one LoRA specialist each
(~90s total).

Result (seed 1), reported honestly:
- STRONG/robust: the merges are the ONLY models competent across ALL
  families -- worst-family ~0.25 vs <0.16 for every single specialist (the
  Fisher-Muller "generalist assembled from specialists" signature, in real
  LoRA weights).
- MARGINAL: "exceeds every parent overall" is only marginal at this scale
  (soup 0.64 vs best specialist 0.63; ties 0.61 below it).
- CAVEAT VISIBLE: averaging dilutes peaks (lists specialist 0.43 -> merge
  0.26) -- Layer-1's "merge, don't average" (E4) appearing in real weights.

The pipeline works end-to-end; the balance/retention half reproduces; the
strict overall-exceeds and soup-vs-ties distinction need scale (bigger base,
more/cleaner families, seeds, a dilution-resistant / offspring-selected
merge) -- the HPC step. Env: Python 3.14 + transformers 5.13 works;
note transformers-5.x apply_chat_template returns a dict. make env-llm /
make llm; adapters under gitignored models/llm/, base in the HF cache.
figures/plot_llm_merge.py, README, tests/test_llm.py (+3, 125 green).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

src/llm/evaluate.py

@@ -0,0 +1,61 @@
+"""Load models/adapters, batch-generate, and score against the task verifier.
+
+Kept deliberately small: a chat-formatted, left-padded batched generator (greedy by default) plus
+a per-family accuracy over a list of :class:`~llm.tasks.Task`. Accuracy on held-out tasks is the
+LLM prototype's capability readout — the analogue of the oracle-measured mode distribution.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .tasks import FAMILIES, Task, verify
+
+
+def load_model(name: str, device: str = "cuda", adapter_dir: str | None = None):
+    """Load a base causal-LM (bf16) and tokenizer; optionally attach a trained LoRA adapter."""
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(name)
+    if tok.pad_token_id is None:
+        tok.pad_token = tok.eos_token
+    tok.padding_side = "left"                              # left-pad for batched decoder generation
+    model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.bfloat16).to(device)
+    if adapter_dir is not None:
+        from peft import PeftModel
+        model = PeftModel.from_pretrained(model, adapter_dir)
+    model.eval()
+    return model, tok
+
+
+def generate(model, tok, prompts: list[str], max_new_tokens: int = 24,
+             batch_size: int = 32, device: str = "cuda") -> list[str]:
+    """Greedy chat-formatted batched generation; returns the decoded completions."""
+    import torch
+
+    outs: list[str] = []
+    for i in range(0, len(prompts), batch_size):
+        chunk = prompts[i:i + batch_size]
+        msgs = [[{"role": "user", "content": p}] for p in chunk]
+        enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt",
+                                      return_dict=True, padding=True).to(device)
+        with torch.no_grad():
+            gen = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False,
+                                 pad_token_id=tok.pad_token_id)
+        for j in range(len(chunk)):
+            outs.append(tok.decode(gen[j, enc["input_ids"].shape[1]:],
+                                   skip_special_tokens=True).strip())
+    return outs
+
+
+def evaluate(model, tok, tasks: list[Task], **gen_kw) -> dict:
+    """Per-family and overall exact-match accuracy over ``tasks``."""
+    outs = generate(model, tok, [t.prompt for t in tasks], **gen_kw)
+    correct = np.array([verify(o, t) for o, t in zip(outs, tasks)])
+    fam = np.array([t.family for t in tasks])
+    row = {"overall": float(correct.mean())}
+    for f in FAMILIES:
+        mask = fam == f
+        row[f] = float(correct[mask].mean()) if mask.any() else float("nan")
+    return row

src/llm/experiment.py

@@ -0,0 +1,109 @@
+"""The LLM merge experiment (blueprint C2/C4) — recombining specialist LLMs.
+
+Trains one LoRA specialist per task family on a small open-weight base, then evaluates the base,
+each specialist, and their weight-space **merges** (soup vs ties) on a held-out mixed test set.
+Produces a tidy long-form results frame and the standard artifact triple (via
+``knowledge.experiment.save_artifacts``), recording model/adapter provenance in the manifest.
+
+The claim under test (the real-LLM image of E8): a model recombined from decorrelated specialists is
+better than any single specialist overall, and — the sharper signature — competent across *all*
+families, which no single parent is (worst-family accuracy). ``python -m llm.experiment
+configs/llm/merge.yaml``.
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import yaml
+
+from knowledge.experiment import save_artifacts
+
+from .evaluate import evaluate, generate, load_model
+from .merge import load_specialists, make_merge
+from .specialise import train_specialist
+from .tasks import FAMILIES, make_tasks, verify
+
+
+def _rows(experiment: str, model: str, kind: str, acc: dict) -> list[dict]:
+    worst = min(acc[f] for f in FAMILIES)
+    out = [{"experiment": experiment, "model": model, "kind": kind, "metric": k, "accuracy": v}
+           for k, v in acc.items()]
+    out.append({"experiment": experiment, "model": model, "kind": kind,
+                "metric": "worst_family", "accuracy": worst})
+    return out
+
+
+def run_merge_experiment(cfg: dict) -> pd.DataFrame:
+    """Train specialists, evaluate base/specialists/merges, return long-form accuracies."""
+    import torch
+
+    name = cfg["experiment"]
+    base = cfg["base_model"]
+    fams = cfg.get("families", list(FAMILIES))
+    n_train, n_test = int(cfg.get("n_train", 700)), int(cfg.get("n_test", 100))
+    epochs = int(cfg.get("epochs", 3))
+    lora = cfg.get("lora", {})
+    merges = cfg.get("merges", ["soup", "ties"])
+    seed = int(cfg["seed"])
+    adapters_root = Path(cfg.get("adapters_dir", "models/llm"))
+
+    test = sum([make_tasks(f, n_test, seed=1000 + i) for i, f in enumerate(fams)], [])
+    rows: list[dict] = []
+
+    # base
+    m, tok = load_model(base)
+    rows += _rows(name, "base", "base", evaluate(m, tok, test))
+    del m; torch.cuda.empty_cache()
+
+    # one specialist per family
+    dirs = []
+    for i, f in enumerate(fams):
+        d = str(adapters_root / f"spec_{f}")
+        train_specialist(base, f, d, n_train=n_train, epochs=epochs, seed=seed + i,
+                         r=int(lora.get("r", 16)), alpha=int(lora.get("alpha", 32)))
+        dirs.append(d)
+        m, tok = load_model(base, adapter_dir=d)
+        rows += _rows(name, f"spec_{f}", "specialist", evaluate(m, tok, test))
+        del m; torch.cuda.empty_cache()
+
+    # weight-space merges (deployable single models)
+    model, tok = load_specialists(base, dirs)
+    prompts = [t.prompt for t in test]
+    fam = np.array([t.family for t in test])
+    for method in merges:
+        make_merge(model, len(dirs), method, method)
+        outs = generate(model, tok, prompts)
+        corr = np.array([verify(o, t) for o, t in zip(outs, test)])
+        acc = {"overall": float(corr.mean())}
+        acc.update({f: float(corr[fam == f].mean()) for f in fams})
+        rows += _rows(name, f"merge_{method}", "merge", acc)
+
+    return pd.DataFrame(rows)
+
+
+def run_and_save(config_path: str | Path) -> Path:
+    """Load an LLM experiment YAML, run it, and write the artifact triple."""
+    config_path = Path(config_path)
+    cfg = yaml.safe_load(config_path.read_text())
+    out_dir = Path(cfg.get("output", {}).get("dir", f"results/{cfg['experiment']}"))
+    cfg.setdefault("n_replicates", 1)
+    df = run_merge_experiment(cfg)
+    save_artifacts(cfg, df, out_dir, extra_libs=("torch", "transformers", "peft"),
+                   extra_manifest={"layer": "2", "tier": "llm", "base_model": cfg["base_model"]},
+                   grid=None)
+    return out_dir
+
+
+def main(argv: list[str] | None = None) -> None:
+    ap = argparse.ArgumentParser(description="Run an LLM merge experiment from a YAML config.")
+    ap.add_argument("config")
+    args = ap.parse_args(argv)
+    print(f"wrote artifacts to {run_and_save(args.config)}/")
+
+
+if __name__ == "__main__":
+    main()

src/llm/merge.py

@@ -0,0 +1,55 @@
+"""Combine specialist LoRA adapters into one deployable model — the recombination operator.
+
+The point of the prototype: several *specialist* parents (each a LoRA adapter) are combined into a
+single child model, and *how* you combine them decides whether the child inherits every specialty or
+a diluted average. Two deployable weight-space combinations, mirroring Layer 1's "merge, don't
+average":
+
+* **soup** — the averaging baseline (model soups): the child is ``base + mean_k(Δ_k)``. Each
+  specialist's contribution is diluted by ``1/K`` — the conservation law that cancels the benefit.
+* **merge** — union-preserving: ``ties`` (sign-reconciled, magnitude-pruned task arithmetic) or plain
+  ``task_arith`` (``base + Σ_k Δ_k``), which keeps each specialist's contribution rather than
+  averaging it away.
+
+Both produce a single model you can ship; the comparison is against that same single-model bar (the
+best individual specialist).
+"""
+
+from __future__ import annotations
+
+
+def load_specialists(base_name: str, adapter_dirs: list[str], device: str = "cuda"):
+    """Load the base model with every specialist adapter attached (named ``a0..a{K-1}``)."""
+    import torch
+    from peft import PeftModel
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(base_name)
+    if tok.pad_token_id is None:
+        tok.pad_token = tok.eos_token
+    tok.padding_side = "left"
+    base = AutoModelForCausalLM.from_pretrained(base_name, dtype=torch.bfloat16).to(device)
+    model = PeftModel.from_pretrained(base, adapter_dirs[0], adapter_name="a0")
+    for i, d in enumerate(adapter_dirs[1:], 1):
+        model.load_adapter(d, adapter_name=f"a{i}")
+    model.eval()
+    return model, tok
+
+
+def make_merge(model, k: int, method: str, name: str):
+    """Add a combined adapter ``name`` over ``a0..a{k-1}`` and make it active. Returns the model.
+
+    method ∈ {``soup`` (linear, 1/K weights), ``task_arith`` (linear, unit weights),
+    ``ties`` (sign-reconciled, pruned)}.
+    """
+    adapters = [f"a{i}" for i in range(k)]
+    if method == "soup":
+        model.add_weighted_adapter(adapters, [1.0 / k] * k, name, combination_type="linear")
+    elif method == "task_arith":
+        model.add_weighted_adapter(adapters, [1.0] * k, name, combination_type="linear")
+    elif method == "ties":
+        model.add_weighted_adapter(adapters, [1.0] * k, name, combination_type="ties", density=0.5)
+    else:
+        raise ValueError(f"unknown merge method {method!r} (expected soup|task_arith|ties)")
+    model.set_adapter(name)
+    return model

src/llm/specialise.py

@@ -0,0 +1,76 @@
+"""LoRA specialisation — train one adapter ("specialist parent") on one task family.
+
+A LoRA adapter is a small trainable patch on a frozen base model. Training one on a single family
+gives a *specialist*: better on its family, ~unchanged on the others (the decorrelated parent that
+recombination will later combine). Deliberately minimal: a manual supervised-fine-tuning loop with
+answer-only loss (the prompt tokens are masked), so it runs in minutes on a 16 GB GPU and has no
+dependence on the churning high-level trainer APIs.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+from .tasks import make_tasks
+
+
+def _encode(tok, task, device):
+    """Return (input_ids, labels) for one example with the prompt tokens masked out of the loss."""
+    text = tok.apply_chat_template([{"role": "user", "content": task.prompt}],
+                                   add_generation_prompt=True, tokenize=False)
+    prompt_ids = tok(text, add_special_tokens=False).input_ids
+    answer_ids = tok(task.answer + tok.eos_token, add_special_tokens=False).input_ids
+    ids = prompt_ids + answer_ids
+    labels = [-100] * len(prompt_ids) + answer_ids
+    return ids, labels
+
+
+def train_specialist(base_name: str, family: str, out_dir: str, *, n_train: int = 600,
+                     epochs: int = 3, lr: float = 2e-4, batch_size: int = 8, r: int = 16,
+                     alpha: int = 32, seed: int = 0, device: str = "cuda") -> str:
+    """Fine-tune a LoRA specialist on ``family`` and save the adapter to ``out_dir``.
+
+    Returns the adapter directory path.
+    """
+    import torch
+    from peft import LoraConfig, get_peft_model
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(base_name)
+    if tok.pad_token_id is None:
+        tok.pad_token = tok.eos_token
+    model = AutoModelForCausalLM.from_pretrained(base_name, dtype=torch.bfloat16).to(device)
+    lora = LoraConfig(r=r, lora_alpha=alpha, lora_dropout=0.0, bias="none",
+                      task_type="CAUSAL_LM",
+                      target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                                      "gate_proj", "up_proj", "down_proj"])
+    model = get_peft_model(model, lora)
+    model.train()
+
+    tasks = make_tasks(family, n_train, seed=seed)
+    encoded = [_encode(tok, t, device) for t in tasks]
+    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr)
+    rng = np.random.default_rng(seed)
+    pad = tok.pad_token_id
+
+    for _ in range(epochs):
+        order = rng.permutation(len(encoded))
+        for i in range(0, len(order), batch_size):
+            batch = [encoded[j] for j in order[i:i + batch_size]]
+            L = max(len(ids) for ids, _ in batch)
+            input_ids = torch.full((len(batch), L), pad, dtype=torch.long)
+            labels = torch.full((len(batch), L), -100, dtype=torch.long)
+            attn = torch.zeros((len(batch), L), dtype=torch.long)
+            for b, (ids, lab) in enumerate(batch):          # right-pad for training
+                input_ids[b, :len(ids)] = torch.tensor(ids)
+                labels[b, :len(lab)] = torch.tensor(lab)
+                attn[b, :len(ids)] = 1
+            out = model(input_ids=input_ids.to(device), attention_mask=attn.to(device),
+                        labels=labels.to(device))
+            opt.zero_grad(); out.loss.backward(); opt.step()
+
+    Path(out_dir).mkdir(parents=True, exist_ok=True)
+    model.save_pretrained(out_dir)
+    return out_dir

src/llm/tasks.py

@@ -0,0 +1,142 @@
+"""Procedural task families with an exact-match verifier — the LLM prototype's "reality".
+
+Three *disjoint* skill families so that a LoRA specialist trained on one is genuinely
+decorrelated from the others (the precondition for recombination to have anything to combine):
+
+* ``lists``  — operations on a short list of small integers (sort, reverse, sum, max, unique…);
+* ``strings`` — operations on a short lowercase word (reverse, uppercase, count a char, drop vowels…);
+* ``arith``  — small-integer arithmetic and sequences (a+b, a*b, next-in-sequence…).
+
+Each task is a natural-language instruction + input with a single canonical answer; the verifier
+normalises the model's output and checks exact equality. This is the LLM analogue of the synthetic
+sandbox's exact oracle — correctness is the "reality that can say no" (grounding / fitness).
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+import numpy as np
+
+FAMILIES = ("lists", "strings", "arith")
+_WORDS = ("apple", "table", "river", "cloud", "stone", "plant", "music", "green", "light", "brave",
+          "ocean", "tiger", "happy", "quiet", "smart", "brick", "chair", "dance", "eagle", "flame")
+
+
+@dataclass(frozen=True)
+class Task:
+    """One verifiable task: a prompt, its canonical answer, and its family."""
+
+    family: str
+    prompt: str
+    answer: str
+
+
+def _fmt_list(xs) -> str:
+    return "[" + ", ".join(str(int(x)) for x in xs) + "]"
+
+
+def _list_task(rng) -> Task:
+    n = int(rng.integers(6, 9))                            # longer lists (harder)
+    xs = rng.integers(0, 20, size=n).tolist()
+    op = rng.choice(["sort ascending", "sort descending", "reverse",
+                     "second largest value", "sum of the two largest values",
+                     "sorted then reversed", "unique values preserving order"])
+    if op == "sort ascending":
+        ans = _fmt_list(sorted(xs))
+    elif op == "sort descending":
+        ans = _fmt_list(sorted(xs, reverse=True))
+    elif op == "reverse":
+        ans = _fmt_list(list(reversed(xs)))
+    elif op == "second largest value":
+        ans = str(sorted(xs, reverse=True)[1])
+    elif op == "sum of the two largest values":
+        ans = str(sum(sorted(xs, reverse=True)[:2]))
+    elif op == "sorted then reversed":
+        ans = _fmt_list(sorted(xs, reverse=True))
+    else:
+        seen, out = set(), []
+        for x in xs:
+            if x not in seen:
+                seen.add(x); out.append(x)
+        ans = _fmt_list(out)
+    want = "the resulting list" if ans.startswith("[") else "the single number"
+    return Task("lists", f"Given the list {_fmt_list(xs)}, compute its {op}. "
+                f"Output only {want} and nothing else.", ans)
+
+
+def _string_task(rng) -> Task:
+    w = str(rng.choice(_WORDS))
+    op = rng.choice(["reverse", "uppercase", "count of the letter 'a'", "count of vowels",
+                     "remove all vowels", "first three characters", "last three characters"])
+    if op == "reverse":
+        ans = w[::-1]
+    elif op == "uppercase":
+        ans = w.upper()
+    elif op == "count of the letter 'a'":
+        ans = str(w.count("a"))
+    elif op == "count of vowels":
+        ans = str(sum(c in "aeiou" for c in w))
+    elif op == "remove all vowels":
+        ans = "".join(c for c in w if c not in "aeiou")
+    elif op == "first three characters":
+        ans = w[:3]
+    else:
+        ans = w[-3:]
+    return Task("strings", f"Given the word \"{w}\", compute its {op}. "
+                f"Output only the answer and nothing else.", ans)
+
+
+def _arith_task(rng) -> Task:
+    kind = rng.choice(["mul", "sub", "mul2", "seq"])       # multiplication-heavy (harder)
+    if kind == "mul":
+        a, b = rng.integers(11, 40, size=2).tolist(); ans = str(a * b)
+        q = f"What is {a} * {b}?"
+    elif kind == "sub":
+        a, b = sorted(rng.integers(100, 1000, size=2).tolist(), reverse=True); ans = str(a - b)
+        q = f"What is {a} - {b}?"
+    elif kind == "mul2":
+        a = int(rng.integers(11, 100)); b = int(rng.integers(3, 10)); ans = str(a * b)
+        q = f"What is {a} * {b}?"
+    else:
+        start = int(rng.integers(2, 12)); step = int(rng.integers(3, 12))
+        seq = [start + i * step for i in range(4)]; ans = str(seq[-1] + step)
+        q = f"What comes next in the sequence {', '.join(map(str, seq))}?"
+    return Task("arith", q + " Output only the number and nothing else.", ans)
+
+
+_GEN = {"lists": _list_task, "strings": _string_task, "arith": _arith_task}
+
+
+def make_tasks(family: str, n: int, seed: int) -> list[Task]:
+    """Generate ``n`` unique-ish tasks for a family (deterministic in ``seed``)."""
+    rng = np.random.default_rng(seed)
+    return [_GEN[family](rng) for _ in range(n)]
+
+
+def _normalise(s: str) -> str:
+    """Canonicalise an answer or a model output for exact comparison."""
+    s = s.strip()
+    m_list = re.findall(r"\[[^\[\]]*\]", s)           # prefer the last bracketed list
+    if m_list:
+        nums = re.findall(r"-?\d+", m_list[-1])
+        return "[" + ", ".join(nums) + "]"
+    m_num = re.findall(r"-?\d+", s)                    # else the last integer
+    if m_num and re.fullmatch(r"[^A-Za-z]*-?\d[\d\s,.\-]*", s):
+        return m_num[-1]
+    # else a word answer: last alphabetic token, lowercased
+    toks = re.findall(r"[A-Za-z]+", s)
+    return toks[-1].lower() if toks else s.lower()
+
+
+def verify(output: str, task: Task) -> bool:
+    """True iff the model output matches the task's canonical answer after normalisation."""
+    want = _normalise(task.answer)
+    got = _normalise(output)
+    if got == want:
+        return True
+    # numeric answers: also accept a bare-int match anywhere in the output
+    if re.fullmatch(r"-?\d+", want):
+        return want in re.findall(r"-?\d+", output)
+    return False