Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118#!/usr/bin/env python3"""Validation gates for heterogeneous colonies.
The point of mixed sizes is that members cannot collapse into being the sameunit. These gates check the machinery actually delivers that, and that asame-width colony is a usable control.
Usage: DSEM_DEVICE=cpu python scripts/check_hetero.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.losses import ACTLossHead # noqa: E402from dsem.models.recursive_reasoning.hetero_colony import HeteroColonyModel_ACTV1 # noqa: E402
BATCH, SEQ, VOCAB = 4, 81, 11
def cfg(widths, **over): c = dict( batch_size=BATCH, seq_len=SEQ, vocab_size=VOCAB, H_cycles=2, L_cycles=2, L_layers=2, expansion=4, halt_max_steps=2, halt_exploration_prob=0.1, forward_dtype="float32", mlp_t=True, widths=list(widths), ) c.update(over) return c
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_a_mixed_widths_run(): batch = make_batch() for widths in ([64, 64, 64, 64], [256, 64, 64, 64], [128, 64], [256, 128, 64]): model = ACTLossHead(HeteroColonyModel_ACTV1(cfg(widths, view_fraction=0.5)), loss_type="stablemax_cross_entropy") model.train() carry = model.initial_carry(batch) carry, loss, _m, _p, _f = model(carry=carry, batch=batch, return_keys=[]) 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"widths={widths}: bad grads {bad}" print("gate A OK — mixed-width colonies run forward+backward, every member gets gradient")
def gate_b_members_are_distinct(): """Members must be genuinely different functions — no weight sharing anywhere.""" model = HeteroColonyModel_ACTV1(cfg([256, 64, 64, 64])) ptrs = {} for i, m in enumerate(model.inner.members): for p in m.parameters(): key = p.data_ptr() assert key not in ptrs or ptrs[key] == i, "members share a parameter tensor" ptrs[key] = i sizes = [sum(p.numel() for p in m.parameters()) for m in model.inner.members] assert sizes[0] > 3 * sizes[1], f"the elder should dwarf a juvenile: {sizes}" print(f"gate B OK — no shared tensors; member params {[f'{s:,}' for s in sizes]}")
def gate_c_untied_control(): """All-equal widths gives an untied homogeneous colony: same member size, K independent weight sets. This is the control for the mixed run.""" model = HeteroColonyModel_ACTV1(cfg([64, 64, 64, 64])) sizes = [sum(p.numel() for p in m.parameters()) for m in model.inner.members] assert len(set(sizes)) == 1, sizes total = sum(sizes) assert total == 4 * sizes[0], "untied colony should cost K x one member" print(f"gate C OK — untied control: 4 x {sizes[0]:,} = {total:,} params")
def gate_d_blackboard_is_shared(): """Different-width members must all read and write the same token blackboard.""" batch = make_batch() model = HeteroColonyModel_ACTV1(cfg([256, 64], view_fraction=0.5)) model.eval() with torch.no_grad(): carry = model.initial_carry(batch) ic = model.inner.reset_carry(torch.ones(BATCH, dtype=torch.bool), carry.inner_carry) new_ic, logits, _q = model.inner(ic, batch) assert logits.shape == (BATCH, SEQ, VOCAB), logits.shape assert new_ic.y.shape == (BATCH, SEQ, VOCAB), new_ic.y.shape assert [z.shape[-1] for z in new_ic.zs] == [256, 64], "members keep their own latent widths" print("gate D OK — one shared token blackboard, per-member latents at their own widths")
def gate_e_distributed_evidence(): """With partial views the members must receive different questions.""" batch = make_batch() model = HeteroColonyModel_ACTV1(cfg([256, 64, 64, 64], view_fraction=0.5)) qx = model.inner._question_tokens(batch["inputs"]) same = (qx[:, 0] == qx[:, 1]).float().mean().item() assert same < 0.95, f"members see near-identical questions ({same:.0%} agreement)" assert model.inner.view_mask.any(0).all(), "some cell invisible to every member" print(f"gate E OK — members see different questions ({same:.0%} token overlap), all cells covered")
if __name__ == "__main__": gate_b_members_are_distinct() gate_c_untied_control() gate_d_blackboard_is_shared() gate_e_distributed_evidence() gate_a_mixed_widths_run() print("\nall hetero gates passed")