Giorgio Gilestro
e433e48860
Adds the "directed sex" operator (E10) the moe regime-flip pointed to: don't
commit to one a-priori blend — breed a population of recombinant offspring
(specialists merged at Dirichlet-sampled weights), score each on a held-out
validation split with the verifier, and keep the fittest, reported on a fresh
test split. Two breeding objectives: best-overall and best-worst-family.
src/llm/directed.py + kind llm_directed, reusing the cached specialists.
Result — refinements pay off in proportion to how far the uniform soup is from
optimal:
- 0.5B (soup dilutes): directed selection beats soup on the bred objective —
directed_overall 0.69 > soup 0.64; directed_balanced worst-family 0.37 > 0.26.
Riders: single-objective selection trades off the other axis (overall-breed
tanks lists to 0.17); a global blend still trails per-input routing (0.74).
- 7B (Imperial CX3, soup already composes to ceiling on near-saturated families,
strings/arith 1.00): directed ~= soup (0.868 ~ 0.873, marginally below via a
val/test overfit gap) — no fitter offspring to breed.
Through-line across all four LLM runs: "merge, don't average" and its refinements
(routing, directed selection) are weak-base / suboptimal-default phenomena — they
help at 0.5B and are inert at 7B. Honest limitation kept in the writeup: the 7B
families are near-saturated, which caps the headroom; a harder unsaturated
benchmark is the fair next test.
Also folds in the two llm_moe local manifest/config files missed in 8da0dac.
+3 directed unit tests (130 green). Results in results/llm_directed{,_hpc}/
(parquet gitignored).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
94 lines
4.6 KiB
Python
94 lines
4.6 KiB
Python
"""LLM-prototype tests — the pure, always-runnable parts (task generation + verifier).
|
|
|
|
The model/LoRA/merge path is heavy (downloads a base model, trains on a GPU) and is validated by the
|
|
experiment run itself, not in CI. What *is* unit-testable — and worth locking, since it is the
|
|
prototype's "reality that says no" — is that tasks are well-formed and the exact-match verifier
|
|
accepts correct answers (including verbose model phrasings) and rejects wrong ones.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from llm.directed import sample_merge_weights, select_winners
|
|
from llm.moe import learned_routes
|
|
from llm.tasks import FAMILIES, make_tasks, verify
|
|
|
|
|
|
def test_make_tasks_wellformed_and_deterministic():
|
|
for fam in FAMILIES:
|
|
tasks = make_tasks(fam, 20, seed=0)
|
|
assert len(tasks) == 20 and all(t.family == fam for t in tasks)
|
|
assert all(t.prompt and t.answer for t in tasks)
|
|
a = make_tasks("arith", 10, seed=3)
|
|
b = make_tasks("arith", 10, seed=3)
|
|
assert [t.answer for t in a] == [t.answer for t in b] # deterministic in the seed
|
|
|
|
|
|
def test_verifier_accepts_correct_including_verbose():
|
|
tasks = make_tasks("lists", 40, seed=1) + make_tasks("arith", 40, seed=2)
|
|
assert all(verify(t.answer, t) for t in tasks) # the canonical answer verifies
|
|
# a verbose but correct model phrasing still verifies (the verifier extracts the answer)
|
|
num_task = next(t for t in tasks if t.family == "arith")
|
|
assert verify(f"The answer is {num_task.answer}.", num_task)
|
|
list_task = next(t for t in tasks if t.family == "lists" and t.answer.startswith("["))
|
|
assert verify(f"Here you go: {list_task.answer}", list_task)
|
|
|
|
|
|
def test_verifier_rejects_wrong():
|
|
t = make_tasks("arith", 1, seed=5)[0]
|
|
wrong = str(int(t.answer) + 1) if t.answer.lstrip("-").isdigit() else "zzz"
|
|
assert not verify(wrong, t)
|
|
lt = next(x for x in make_tasks("lists", 30, seed=6) if x.answer.startswith("["))
|
|
assert not verify("[9, 9, 9]", lt) or lt.answer == "[9, 9, 9]"
|
|
|
|
|
|
def test_learned_router_assigns_nearest_centroid():
|
|
# Three well-separated families in a 4-D "embedding" space; the nearest-centroid router
|
|
# (the MoE expert-selection gene) must route each test prompt to its own family's specialist.
|
|
rng = np.random.default_rng(0)
|
|
fams = ["lists", "strings", "arith"]
|
|
anchors = {"lists": [5, 0, 0, 0], "strings": [0, 5, 0, 0], "arith": [0, 0, 5, 0]}
|
|
train_emb = np.array([anchors[f] for f in fams for _ in range(8)], dtype=float)
|
|
train_emb += rng.normal(scale=0.1, size=train_emb.shape)
|
|
train_fam = np.array([f for f in fams for _ in range(8)])
|
|
test_fam = np.array(["arith", "lists", "strings", "arith"])
|
|
test_emb = np.array([anchors[f] for f in test_fam], dtype=float) + rng.normal(scale=0.1, size=(4, 4))
|
|
routes = learned_routes(train_emb, train_fam, test_emb, fams)
|
|
assert [fams[r] for r in routes] == list(test_fam) # each routed to its own family
|
|
|
|
|
|
def test_learned_router_is_cosine_scale_invariant():
|
|
# Cosine routing must ignore prompt-embedding magnitude (long vs short prompts): a test point on a
|
|
# family's ray routes there regardless of its norm.
|
|
fams = ["a", "b"]
|
|
train_emb = np.array([[1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]])
|
|
train_fam = np.array(["a", "a", "b", "b"])
|
|
test_emb = np.array([[10.0, 0.0], [0.0, 0.01]]) # very different magnitudes
|
|
routes = learned_routes(train_emb, train_fam, test_emb, fams)
|
|
assert [fams[r] for r in routes] == ["a", "b"]
|
|
|
|
|
|
def test_merge_weights_population_pins_baselines_and_diversifies():
|
|
# The offspring population must contain the two canonical baselines (uniform soup, unit task-arith)
|
|
# and be diverse + reproducible for the rest.
|
|
rng = np.random.default_rng(0)
|
|
w = sample_merge_weights(3, 16, rng)
|
|
assert w.shape == (16, 3)
|
|
assert np.allclose(w[0], 1 / 3) # candidate 0 = uniform soup
|
|
assert np.allclose(w[1], 1.0) # candidate 1 = task arithmetic
|
|
assert np.unique(w[2:].round(3), axis=0).shape[0] > 5 # the random offspring are diverse
|
|
assert np.allclose(sample_merge_weights(3, 16, np.random.default_rng(0)), w) # deterministic
|
|
|
|
|
|
def test_select_winners_picks_argmax_per_objective():
|
|
val_overall = np.array([0.5, 0.9, 0.7])
|
|
val_worst = np.array([0.4, 0.1, 0.6]) # a different candidate is most balanced
|
|
w = select_winners(val_overall, val_worst)
|
|
assert w == {"overall": 1, "balanced": 2}
|
|
|
|
|
|
def test_merge_weights_requires_two_candidates():
|
|
import pytest
|
|
with pytest.raises(ValueError):
|
|
sample_merge_weights(3, 1, np.random.default_rng(0))
|