Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208#!/usr/bin/env python3"""Validation gates for the colony harness — run before trusting any colony result.
1. K=1, identity standpoint, mean aggregation must match TRM *exactly* (same weights, same inputs -> same logits). If this fails, the colony is not a generalization of TRM and no comparison against the 84.5% anchor is valid.2. K>1 must run forward and backward in every blackboard/aggregate/halt mode.3. Clonal colonies must not grow parameters with K (only role embeddings may).4. Standpoint transforms must round-trip: rho^-1(rho(x)) == x.
Usage: DSEM_DEVICE=cpu python scripts/check_colony.py"""
import osimport sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))os.environ.setdefault("DSEM_DEVICE", "cpu")
import torch # noqa: E402
from dsem.models.recursive_reasoning.colony import ( # noqa: E402 ColonyModel_ACTV1, build_standpoints, invert_perm,)from dsem.models.recursive_reasoning.trm import TinyRecursiveReasoningModel_ACTV1 # noqa: E402
BATCH, SEQ, VOCAB = 4, 81, 11
def base_cfg(**over): cfg = dict( batch_size=BATCH, seq_len=SEQ, vocab_size=VOCAB, num_puzzle_identifiers=1, puzzle_emb_ndim=64, puzzle_emb_len=16, H_cycles=2, L_cycles=2, H_layers=0, L_layers=2, hidden_size=64, expansion=4, num_heads=1, pos_encodings="none", halt_max_steps=2, halt_exploration_prob=0.1, forward_dtype="float32", mlp_t=True, no_ACT_continue=True, causal=False, ) cfg.update(over) return cfg
def make_batch(): g = torch.Generator().manual_seed(0) return { "inputs": torch.randint(1, VOCAB, (BATCH, SEQ), generator=g), "labels": torch.randint(1, VOCAB, (BATCH, SEQ), generator=g), "puzzle_identifiers": torch.zeros(BATCH, dtype=torch.int32), }
def gate_1_equivalence(): batch = make_batch() torch.manual_seed(0) trm = TinyRecursiveReasoningModel_ACTV1(base_cfg()) torch.manual_seed(0) col = ColonyModel_ACTV1(base_cfg(num_agents=1, blackboard="latent", standpoints="none", aggregate="mean", halt_rule="boss"))
# transplant weights: colony's `level` corresponds to TRM's `L_level` mapping = {k.replace("inner.L_level", "inner.level"): v for k, v in trm.state_dict().items()} mapping = {k.replace("inner.H_init", "inner.y_init").replace("inner.L_init", "inner.z_init"): v for k, v in mapping.items()} missing, unexpected = col.load_state_dict(mapping, strict=False) # colony-only buffers (views/masks) have no TRM counterpart, so they stay at init colony_only = ("inner.standpoint", "inner.standpoint_inv", "inner.view_mask") unexpected = [k for k in unexpected if not k.startswith(colony_only)] assert not unexpected, f"unexpected keys: {unexpected}" assert all(k.startswith(colony_only) for k in missing), f"missing: {missing}"
trm.eval(); col.eval() with torch.no_grad(): c1 = trm.initial_carry(batch); c1.halted[:] = True c2 = col.initial_carry(batch); c2.halted[:] = True _, o1 = trm(carry=c1, batch=batch) _, o2 = col(carry=c2, batch=batch) diff = (o1["logits"] - o2["logits"]).abs().max().item() assert diff < 1e-4, f"GATE 1 FAILED: K=1 colony differs from TRM by {diff}" print(f"gate 1 OK — K=1 colony reproduces TRM (max logit diff {diff:.2e})")
def gate_2_modes(): batch = make_batch() for blackboard in ("latent", "tokens"): for aggregate in ("mean", "vote", "confidence"): for halt_rule in ("quorum", "any", "boss"): model = ColonyModel_ACTV1(base_cfg( num_agents=3, blackboard=blackboard, standpoints="sudoku", aggregate=aggregate, halt_rule=halt_rule, role_embeddings=True)) model.train() carry = model.initial_carry(batch) carry.halted[:] = True carry, out = model(carry=carry, batch=batch) assert out["logits"].shape == (BATCH, SEQ, VOCAB), out["logits"].shape assert out["q_halt_logits"].shape == (BATCH,), out["q_halt_logits"].shape # touch every head so a missing gradient means a real disconnection loss = (out["logits"].square().mean() + out["q_halt_logits"].square().mean() + out["q_continue_logits"].square().mean()) loss.backward() bad = [n for n, p in model.named_parameters() if p.requires_grad and (p.grad is None or not torch.isfinite(p.grad).all())] assert not bad, f"bad grads in {blackboard}/{aggregate}/{halt_rule}: {bad}" print("gate 2 OK — all blackboard x aggregate x halt modes run forward+backward")
def gate_3_clonal_params(): counts = [] for k in (1, 2, 4, 8): m = ColonyModel_ACTV1(base_cfg(num_agents=k, standpoints="sudoku")) counts.append(sum(p.numel() for p in m.parameters())) assert len(set(counts)) == 1, f"GATE 3 FAILED: params vary with K: {counts}" print(f"gate 3 OK — clonal colony keeps {counts[0]:,} params for K=1,2,4,8")
def gate_6_distributed_evidence(): """view_fraction<1 must give units genuinely different questions, cover every cell, and leave the K=1/full-evidence path untouched.""" from dsem.models.recursive_reasoning.colony import build_view_masks m = build_view_masks(4, SEQ, fraction=0.5, seed=0) assert m.shape == (4, SEQ) assert m.any(0).all(), "some cell is invisible to every unit — evidence lost, not distributed" frac = m.float().mean().item() assert 0.4 < frac < 0.65, f"unexpected visible fraction {frac}" pairwise = (m[0] == m[1]).float().mean().item() assert pairwise < 0.95, "units see near-identical evidence"
full = build_view_masks(4, SEQ, fraction=1.0, seed=0) assert full.all(), "full evidence must mean everyone sees everything"
batch = make_batch() model = ColonyModel_ACTV1(base_cfg(num_agents=4, view_fraction=0.5, standpoints="none")) model.eval() with torch.no_grad(): carry = model.initial_carry(batch); carry.halted[:] = True inner = model.inner.reset_carry(carry.halted, carry.inner_carry) views = model.inner._question_views(batch) spread = (views[:, 0] - views[:, 1]).abs().max().item() assert spread > 0, "units received identical question embeddings despite masking" _ = inner print(f"gate 6 OK — distributed evidence: {frac:.0%} visible per unit, all cells covered, " f"units see different questions (max emb diff {spread:.2f})")
def gate_7_instrumentation_matches(): """forward_instrumented must compute exactly what forward computes.
It once didn't: it bypassed the view masks, gave every unit the full question, and so reported perfect consensus and redundant members that were artefacts of the measurement rather than properties of the model. """ batch = make_batch() for frac in (1.0, 0.5): torch.manual_seed(0) model = ColonyModel_ACTV1(base_cfg(num_agents=3, view_fraction=frac, standpoints="sudoku")) model.eval() with torch.no_grad(): carry = model.initial_carry(batch); carry.halted[:] = True ic = model.inner.reset_carry(carry.halted, carry.inner_carry) _, logits, _ = model.inner(ic, batch)
ic2 = model.inner.reset_carry(carry.halted, carry.inner_carry) out = model.inner.forward_instrumented(ic2, batch) agg = out["rounds"][-1]["aggregate"] diff = (logits.argmax(-1) != agg).float().mean().item() assert diff < 1e-6, ( f"GATE 7 FAILED at view_fraction={frac}: instrumented path disagrees with " f"forward on {diff:.1%} of cells") print("gate 7 OK — instrumented path reproduces forward() at full and partial evidence")
def gate_5_device_context(): """Buffers must land on the ambient device — pretrain.py builds models inside `with torch.device(...)`, and a CPU-pinned buffer only fails there, not here.""" from dsem.device import get_device dev = get_device() if dev.type == "cpu": print("gate 5 skipped (no accelerator on this host)") return with torch.device(dev): m = ColonyModel_ACTV1(base_cfg(num_agents=2, standpoints="sudoku")) bad = [n for n, b in m.named_buffers() if b.device.type != dev.type] assert not bad, f"GATE 5 FAILED: buffers not on {dev}: {bad}" print(f"gate 5 OK — all buffers land on {dev} when built in a device context")
def gate_4_standpoints(): perms = build_standpoints("sudoku", 8, SEQ, seed=0) inv = invert_perm(perms) x = torch.arange(SEQ).view(1, 1, SEQ).expand(2, 8, SEQ) fwd = torch.gather(x, 2, perms.unsqueeze(0).expand(2, -1, -1)) back = torch.gather(fwd, 2, inv.unsqueeze(0).expand(2, -1, -1)) assert torch.equal(back, x), "GATE 4 FAILED: standpoint round-trip broken" assert torch.equal(perms[0], torch.arange(SEQ)), "agent 0 must be the identity view" uniq = len({tuple(p.tolist()) for p in perms}) print(f"gate 4 OK — standpoints round-trip; {uniq}/8 distinct views; agent 0 = identity")
if __name__ == "__main__": gate_4_standpoints() gate_3_clonal_params() gate_1_equivalence() gate_2_modes() gate_5_device_context() gate_6_distributed_evidence() gate_7_instrumentation_matches() print("\nall gates passed")