Upload folder using huggingface_hub

eval/prune_weights.py

@@ -0,0 +1,481 @@
+"""
+BATCHED WEIGHT PRUNING (GPU-optimized)
+======================================
+Phase 1: Batch eval all candidates in parallel
+Phase 2: Apply all successes at once, binary search if conflicts
+"""
+
+import torch
+import time
+import argparse
+from safetensors.torch import save_file
+from iron_eval import BatchedFitnessEvaluator, create_population, load_model
+
+torch.manual_seed(0)
+
+
+def format_time(seconds):
+    if seconds < 60:
+        return f"{seconds:.1f}s"
+    elif seconds < 3600:
+        return f"{seconds/60:.1f}m"
+    else:
+        return f"{seconds/3600:.1f}h"
+
+
+def format_eta(elapsed, done, total):
+    if done == 0:
+        return "calculating..."
+    rate = done / elapsed
+    remaining = (total - done) / rate
+    return format_time(remaining)
+
+
+def apply_reductions(model, reductions):
+    """Apply a list of (name, flat_idx, shape, old_val) reductions."""
+    for name, flat_idx, shape, old_val in reductions:
+        new_val = old_val - 1 if old_val > 0 else old_val + 1
+        flat = model[name].flatten()
+        if flat[flat_idx].item() == old_val:
+            flat[flat_idx] = new_val
+            model[name] = flat.view(shape)
+
+
+def revert_reductions(model, reductions):
+    """Revert a list of reductions."""
+    for name, flat_idx, shape, old_val in reductions:
+        flat = model[name].flatten()
+        new_val = old_val - 1 if old_val > 0 else old_val + 1
+        if flat[flat_idx].item() == new_val:
+            flat[flat_idx] = old_val
+            model[name] = flat.view(shape)
+
+
+def check_fitness(model, evaluator, device):
+    """Check model fitness."""
+    torch.manual_seed(0)
+    pop = create_population(model, 1, device)
+    return evaluator.evaluate(pop, debug=False)[0].item()
+
+
+def sequential_conflict_resolution(model, evaluator, device, candidates, base_magnitude):
+    """
+    Sequential fallback - tests and applies reductions one at a time.
+    Slower but guarantees no interaction bugs.
+    """
+    accepted = []
+    for i, (name, flat_idx, shape, old_val) in enumerate(candidates):
+        apply_reductions(model, [(name, flat_idx, shape, old_val)])
+        fitness = check_fitness(model, evaluator, device)
+        if fitness >= 0.9999:
+            accepted.append((name, flat_idx, shape, old_val))
+            if (i + 1) % 50 == 0:
+                current_mag = sum(t.abs().sum().item() for t in model.values())
+                reduction_pct = 100 * (1 - current_mag / base_magnitude)
+                print(f"          Sequential: {len(accepted)}/{i+1} accepted | mag={current_mag:.0f} (-{reduction_pct:.2f}%)")
+        else:
+            revert_reductions(model, [(name, flat_idx, shape, old_val)])
+    return accepted
+
+
+def batched_conflict_resolution(model, evaluator, device, candidates, base_magnitude):
+    """
+    Batched binary search - evaluates multiple branches in parallel.
+    Uses BFS instead of DFS to maximize batching opportunities.
+    Verifies cumulative effect after each batch to prevent interaction bugs.
+    """
+    if len(candidates) == 0:
+        return []
+
+    # First try all at once
+    print(f"        Trying {len(candidates)} reductions at once...")
+    apply_reductions(model, candidates)
+    fitness = check_fitness(model, evaluator, device)
+
+    if fitness >= 0.9999:
+        current_mag = sum(t.abs().sum().item() for t in model.values())
+        reduction_pct = 100 * (1 - current_mag / base_magnitude)
+        print(f"        ALL {len(candidates)} OK | fitness={fitness:.6f} | "
+              f"mag={current_mag:.0f} (-{reduction_pct:.2f}%)")
+        return candidates
+
+    # Conflict - revert and use batched BFS
+    revert_reductions(model, candidates)
+    print(f"        CONFLICT (fitness={fitness:.6f}), starting batched resolution...")
+
+    accepted = []
+    # Queue of (candidate_list, depth) to process
+    pending = [(candidates, 0)]
+
+    while pending:
+        # Collect all pending groups for batch evaluation
+        to_eval = []
+        for group, depth in pending:
+            if len(group) == 0:
+                continue
+            elif len(group) == 1:
+                to_eval.append((group, depth, 'single'))
+            else:
+                to_eval.append((group, depth, 'group'))
+
+        pending = []
+
+        if not to_eval:
+            break
+
+        # Build batch: create model variants for each group
+        batch_size = len(to_eval)
+        print(f"          Batch evaluating {batch_size} groups...")
+
+        # Create population for batch eval
+        pop = {}
+        for name, tensor in model.items():
+            pop[name] = tensor.unsqueeze(0).expand(batch_size, *tensor.shape).clone().to(device)
+
+        # Apply each group's reductions to its population slot
+        for idx, (group, depth, gtype) in enumerate(to_eval):
+            for name, flat_idx, shape, old_val in group:
+                new_val = old_val - 1 if old_val > 0 else old_val + 1
+                flat_view = pop[name][idx].flatten()
+                # Check if not already modified in base model
+                base_val = model[name].flatten()[flat_idx].item()
+                if base_val == old_val:
+                    flat_view[flat_idx] = new_val
+
+        # Batch evaluate
+        torch.manual_seed(0)
+        fitnesses = evaluator.evaluate(pop, debug=False)
+
+        # Process results - collect accepted groups first, then verify
+        batch_accepted = []
+        ok_count = 0
+        conflict_count = 0
+        fail_count = 0
+
+        for idx, (group, depth, gtype) in enumerate(to_eval):
+            fit = fitnesses[idx].item()
+            indent = "          " + "  " * depth
+
+            if fit >= 0.9999:
+                batch_accepted.append((group, depth, indent))
+                ok_count += len(group)
+            else:
+                if len(group) == 1:
+                    name, flat_idx, shape, old_val = group[0]
+                    print(f"{indent}[1/1] FAIL {name}[{flat_idx}] | fitness={fit:.6f}")
+                    fail_count += 1
+                else:
+                    mid = len(group) // 2
+                    left = group[:mid]
+                    right = group[mid:]
+                    print(f"{indent}CONFLICT ({len(group)}) fitness={fit:.6f} -> split {len(left)}+{len(right)}")
+                    pending.append((left, depth + 1))
+                    pending.append((right, depth + 1))
+                    conflict_count += 1
+
+        # Apply all batch-accepted reductions
+        all_batch_reductions = []
+        for group, depth, indent in batch_accepted:
+            apply_reductions(model, group)
+            all_batch_reductions.extend(group)
+
+        # Verify cumulative effect
+        if all_batch_reductions:
+            verify_fitness = check_fitness(model, evaluator, device)
+            if verify_fitness >= 0.9999:
+                # All good - commit these reductions
+                for group, depth, indent in batch_accepted:
+                    current_mag = sum(t.abs().sum().item() for t in model.values())
+                    reduction_pct = 100 * (1 - current_mag / base_magnitude)
+                    if len(group) == 1:
+                        name, flat_idx, shape, old_val = group[0]
+                        print(f"{indent}[1/1] OK {name}[{flat_idx}] | mag={current_mag:.0f} (-{reduction_pct:.2f}%)")
+                    else:
+                        print(f"{indent}ALL {len(group)} OK | mag={current_mag:.0f} (-{reduction_pct:.2f}%)")
+                accepted.extend(all_batch_reductions)
+                print(f"          Batch result: {ok_count} accepted, {conflict_count} split, {fail_count} failed")
+            else:
+                # Interaction bug detected - revert and use sequential fallback
+                print(f"          INTERACTION BUG detected (batch fitness={verify_fitness:.6f})")
+                print(f"          Reverting {len(all_batch_reductions)} reductions, falling back to sequential...")
+                revert_reductions(model, all_batch_reductions)
+
+                # Process each group sequentially
+                seq_accepted = sequential_conflict_resolution(
+                    model, evaluator, device, all_batch_reductions, base_magnitude
+                )
+                accepted.extend(seq_accepted)
+                print(f"          Sequential fallback: {len(seq_accepted)}/{len(all_batch_reductions)} accepted")
+        else:
+            print(f"          Batch result: {ok_count} accepted, {conflict_count} split, {fail_count} failed")
+
+    return accepted
+
+
+def prune_weights(
+    passes: int = 10,
+    batch_size: int = 5000,
+    device: str = 'cuda',
+    checkpoint_path: str = "D:/8bit-threshold-computer/pruned.safetensors"
+):
+    print("=" * 80)
+    print(" BATCHED WEIGHT PRUNING (GPU-optimized)")
+    print("=" * 80)
+    print(f" Device: {device}")
+    print(f" Batch size: {batch_size}")
+    print(f" Max passes: {passes}")
+    print("=" * 80)
+
+    # Load model
+    print("\n[1/4] LOADING MODEL...")
+    load_start = time.perf_counter()
+    model = load_model()
+    load_time = time.perf_counter() - load_start
+
+    n_params = sum(t.numel() for t in model.values())
+    n_tensors = len(model)
+    base_magnitude = sum(t.abs().sum().item() for t in model.values())
+    base_max = max(t.abs().max().item() for t in model.values())
+    nonzero_params = sum((t != 0).sum().item() for t in model.values())
+
+    print(f"       Loaded in {load_time:.2f}s")
+    print(f"       Tensors: {n_tensors}")
+    print(f"       Parameters: {n_params}")
+    print(f"       Non-zero parameters: {nonzero_params}")
+    print(f"       Total magnitude: {base_magnitude:.0f}")
+    print(f"       Max weight: {base_max:.0f}")
+
+    # Initialize evaluator
+    print("\n[2/4] INITIALIZING EVALUATOR...")
+    eval_start = time.perf_counter()
+    evaluator = BatchedFitnessEvaluator(device=device)
+    eval_time = time.perf_counter() - eval_start
+    print(f"       Initialized in {eval_time:.2f}s")
+
+    # Verify initial fitness
+    print("\n[3/4] VERIFYING BASE MODEL...")
+    initial_fitness = check_fitness(model, evaluator, device)
+    print(f"       Fitness: {initial_fitness:.6f}")
+
+    if initial_fitness < 0.9999:
+        print(f"       ERROR: Base model fitness {initial_fitness:.6f} < 0.9999")
+        return None
+
+    print(f"       STATUS: PASS")
+
+    # Build parameter list
+    print("\n[4/4] BUILDING PARAMETER INDEX...")
+    param_list = []
+    for name, tensor in model.items():
+        flat = tensor.flatten()
+        for i in range(len(flat)):
+            param_list.append((name, i, tensor.shape))
+    print(f"       Indexed {len(param_list)} parameters")
+
+    # Main pruning loop
+    print("\n" + "=" * 80)
+    print(" PRUNING STARTED")
+    print("=" * 80)
+
+    total_reductions = 0
+    pruning_start = time.perf_counter()
+
+    for pass_num in range(passes):
+        torch.manual_seed(0)
+        pass_start = time.perf_counter()
+
+        print(f"\n{'='*80}")
+        print(f" PASS {pass_num + 1}/{passes}")
+        print(f"{'='*80}")
+
+        # Count candidates
+        candidates = []
+        for name, idx, shape in param_list:
+            flat = model[name].flatten()
+            val = flat[idx].item()
+            if val != 0:
+                candidates.append((name, idx, shape, val))
+
+        n_candidates = len(candidates)
+        print(f"\n    Candidates: {n_candidates} non-zero weights")
+
+        if n_candidates == 0:
+            print(f"    No candidates remaining. Stopping.")
+            break
+
+        # Phase 1: Batch evaluation
+        print(f"\n    PHASE 1: Batch evaluation (testing each reduction independently)")
+        print(f"    " + "-" * 60)
+        phase1_start = time.perf_counter()
+        successful_candidates = []
+        n_batches = (n_candidates + batch_size - 1) // batch_size
+
+        for batch_idx, batch_start_idx in enumerate(range(0, n_candidates, batch_size)):
+            batch = candidates[batch_start_idx:batch_start_idx + batch_size]
+            batch_len = len(batch)
+            batch_start_time = time.perf_counter()
+
+            # Build population
+            pop = {}
+            for name, tensor in model.items():
+                pop[name] = tensor.unsqueeze(0).expand(batch_len, *tensor.shape).clone().to(device)
+
+            # Apply reductions
+            for pop_idx, (name, flat_idx, shape, old_val) in enumerate(batch):
+                new_val = old_val - 1 if old_val > 0 else old_val + 1
+                flat_view = pop[name][pop_idx].flatten()
+                flat_view[flat_idx] = new_val
+
+            # Evaluate
+            torch.manual_seed(0)
+            if device == 'cuda':
+                torch.cuda.synchronize()
+            fitness = evaluator.evaluate(pop, debug=False)
+            if device == 'cuda':
+                torch.cuda.synchronize()
+
+            # Collect successes
+            batch_successes = 0
+            for pop_idx, (name, flat_idx, shape, old_val) in enumerate(batch):
+                if fitness[pop_idx].item() >= 0.9999:
+                    successful_candidates.append((name, flat_idx, shape, old_val))
+                    batch_successes += 1
+
+            batch_time = time.perf_counter() - batch_start_time
+            elapsed = time.perf_counter() - phase1_start
+            done = batch_start_idx + batch_len
+            eta = format_eta(elapsed, done, n_candidates)
+            throughput = batch_len / batch_time
+
+            print(f"      Batch {batch_idx + 1}/{n_batches}: "
+                  f"{batch_successes}/{batch_len} passed ({100*batch_successes/batch_len:.1f}%) | "
+                  f"Total OK: {len(successful_candidates)} | "
+                  f"Progress: {done}/{n_candidates} ({100*done/n_candidates:.1f}%) | "
+                  f"Speed: {throughput:.0f}/s | "
+                  f"ETA: {eta}")
+
+        phase1_time = time.perf_counter() - phase1_start
+        print(f"\n    Phase 1 complete: {len(successful_candidates)}/{n_candidates} candidates "
+              f"({100*len(successful_candidates)/n_candidates:.1f}%) in {format_time(phase1_time)}")
+
+        # Phase 2: Apply with conflict resolution
+        if len(successful_candidates) == 0:
+            print(f"\n    No reductions possible. Stopping.")
+            break
+
+        print(f"\n    PHASE 2: Apply reductions with conflict resolution")
+        print(f"    " + "-" * 60)
+        phase2_start = time.perf_counter()
+
+        accepted = batched_conflict_resolution(model, evaluator, device, successful_candidates, base_magnitude)
+        pass_reductions = len(accepted)
+
+        phase2_time = time.perf_counter() - phase2_start
+        print(f"\n    Phase 2 complete: {pass_reductions} reductions applied in {format_time(phase2_time)}")
+
+        # Pass summary
+        total_reductions += pass_reductions
+        current_magnitude = sum(t.abs().sum().item() for t in model.values())
+        current_nonzero = sum((t != 0).sum().item() for t in model.values())
+        pass_time = time.perf_counter() - pass_start
+        reduction_pct = 100 * (1 - current_magnitude / base_magnitude)
+
+        print(f"\n    PASS {pass_num + 1} SUMMARY:")
+        print(f"      Reductions this pass: {pass_reductions}")
+        print(f"      Total reductions: {total_reductions}")
+        print(f"      Current magnitude: {current_magnitude:.0f} (-{reduction_pct:.2f}%)")
+        print(f"      Current non-zero: {current_nonzero}")
+        print(f"      Pass time: {format_time(pass_time)}")
+
+        # Verify after pass
+        print(f"\n    Verifying model integrity...")
+        fitness = check_fitness(model, evaluator, device)
+        print(f"      Fitness: {fitness:.6f} {'PASS' if fitness >= 0.9999 else 'FAIL'}")
+
+        # Save checkpoint after each pass
+        checkpoint_name = checkpoint_path.replace('.safetensors', f'_pass{pass_num + 1}.safetensors')
+        print(f"\n    Saving checkpoint: {checkpoint_name}")
+        save_file(model, checkpoint_name)
+        print(f"      Saved. Magnitude: {current_magnitude:.0f} (-{reduction_pct:.2f}%)")
+
+        # Also save as "latest" for easy access
+        latest_path = checkpoint_path.replace('.safetensors', '_latest.safetensors')
+        save_file(model, latest_path)
+        print(f"      Also saved as: {latest_path}")
+
+        if pass_reductions == 0:
+            print(f"\n    No reductions achieved. Stopping early.")
+            break
+
+    # Final summary
+    pruning_time = time.perf_counter() - pruning_start
+    final_magnitude = sum(t.abs().sum().item() for t in model.values())
+    final_max = max(t.abs().max().item() for t in model.values())
+    final_nonzero = sum((t != 0).sum().item() for t in model.values())
+    reduction_pct = 100 * (1 - final_magnitude / base_magnitude)
+
+    print("\n" + "=" * 80)
+    print(" PRUNING COMPLETE")
+    print("=" * 80)
+    print(f"\n RESULTS:")
+    print(f"   Original magnitude:  {base_magnitude:.0f}")
+    print(f"   Final magnitude:     {final_magnitude:.0f}")
+    print(f"   Reduction:           {reduction_pct:.2f}%")
+    print(f"   Total reductions:    {total_reductions}")
+    print(f"   Original non-zero:   {nonzero_params}")
+    print(f"   Final non-zero:      {final_nonzero}")
+    print(f"   Zeros created:       {nonzero_params - final_nonzero}")
+    print(f"   Max weight:          {final_max:.0f}")
+    print(f"   Total time:          {format_time(pruning_time)}")
+
+    # Save
+    print(f"\n SAVING to {checkpoint_path}...")
+    save_file(model, checkpoint_path)
+    print(f"   Saved.")
+
+    # Final verification
+    print(f"\n FINAL VERIFICATION...")
+    from safetensors import safe_open
+    f = safe_open(checkpoint_path, framework='numpy')
+    verify_model = {name: torch.tensor(f.get_tensor(name)).float() for name in f.keys()}
+    verify_fitness = check_fitness(verify_model, evaluator, device)
+    print(f"   Fitness: {verify_fitness:.6f}")
+
+    if verify_fitness >= 0.9999:
+        print(f"   STATUS: PASS")
+    else:
+        print(f"   STATUS: FAIL - Model corrupted!")
+
+    print("\n" + "=" * 80)
+    return model
+
+
+MAX_BATCH_SIZE = 80000
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Batched Weight Pruning')
+    parser.add_argument('--passes', type=int, default=10,
+                        help='Maximum pruning passes (default: 10)')
+    parser.add_argument('--batch_size', type=int, default=80000,
+                        help=f'Batch size for parallel evaluation (default: 80000, max: {MAX_BATCH_SIZE})')
+    parser.add_argument('--device', type=str, default='cuda',
+                        help='Device: cuda or cpu (default: cuda)')
+    parser.add_argument('--output', type=str,
+                        default='D:/8bit-threshold-computer/pruned.safetensors',
+                        help='Output path')
+    args = parser.parse_args()
+
+    if args.batch_size > MAX_BATCH_SIZE:
+        print(f"WARNING: batch_size {args.batch_size} exceeds maximum {MAX_BATCH_SIZE}. Clamping.")
+        args.batch_size = MAX_BATCH_SIZE
+
+    print(f"\nStarting at {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
+
+    prune_weights(
+        passes=args.passes,
+        batch_size=args.batch_size,
+        device=args.device,
+        checkpoint_path=args.output
+    )
+
+    print(f"\nFinished at {time.strftime('%Y-%m-%d %H:%M:%S')}")

llm/circuit_llm.py

@@ -0,0 +1,606 @@
+"""
+Circuit-Augmented LLM: Embedding threshold logic circuits into SmolLM2
+======================================================================
+
+Replaces/augments MLP layers with frozen threshold circuits for exact arithmetic.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, Optional, Tuple
+from safetensors.torch import load_file
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import warnings
+warnings.filterwarnings('ignore')
+
+
+# =============================================================================
+# HEAVISIDE WITH STRAIGHT-THROUGH ESTIMATOR
+# =============================================================================
+
+class HeavisideSTE(torch.autograd.Function):
+    """Heaviside step function with straight-through estimator for backprop."""
+
+    @staticmethod
+    def forward(ctx, x):
+        return (x >= 0).float()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        # STE: pass gradient through unchanged
+        return grad_output
+
+
+def heaviside(x: torch.Tensor) -> torch.Tensor:
+    """Heaviside step: 1 if x >= 0, else 0. Uses STE for training."""
+    return HeavisideSTE.apply(x)
+
+
+# =============================================================================
+# CIRCUIT EXECUTOR - Runs the frozen threshold circuits
+# =============================================================================
+
+class CircuitExecutor(nn.Module):
+    """
+    Executes threshold logic circuits from the safetensors file.
+    All circuit weights are frozen - only interface layers train.
+    """
+
+    def __init__(self, circuit_path: str, device: str = 'cpu'):
+        super().__init__()
+        self.device = device
+
+        # Load all circuit tensors
+        raw_circuits = load_file(circuit_path)
+
+        # Store as frozen parameters (use underscores for valid param names)
+        self.circuits = {}
+        for k, v in raw_circuits.items():
+            safe_name = k.replace('.', '__')
+            self.register_buffer(safe_name, v.float().to(device))
+            self.circuits[k] = safe_name
+
+    def _get(self, name: str) -> torch.Tensor:
+        """Get circuit tensor by original dotted name."""
+        return getattr(self, self.circuits[name])
+
+    # -------------------------------------------------------------------------
+    # Boolean Gates
+    # -------------------------------------------------------------------------
+
+    def eval_and(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        """AND gate: output 1 iff both inputs are 1."""
+        inp = torch.stack([a, b], dim=-1)
+        w = self._get('boolean.and.weight')
+        bias = self._get('boolean.and.bias')
+        return heaviside(inp @ w + bias)
+
+    def eval_or(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        """OR gate: output 1 if either input is 1."""
+        inp = torch.stack([a, b], dim=-1)
+        w = self._get('boolean.or.weight')
+        bias = self._get('boolean.or.bias')
+        return heaviside(inp @ w + bias)
+
+    def eval_xor(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        """XOR gate: two-layer network (not linearly separable)."""
+        inp = torch.stack([a, b], dim=-1)
+
+        # Layer 1: OR and NAND neurons
+        w1_n1 = self._get('boolean.xor.layer1.neuron1.weight')
+        b1_n1 = self._get('boolean.xor.layer1.neuron1.bias')
+        w1_n2 = self._get('boolean.xor.layer1.neuron2.weight')
+        b1_n2 = self._get('boolean.xor.layer1.neuron2.bias')
+
+        h1 = heaviside(inp @ w1_n1 + b1_n1)
+        h2 = heaviside(inp @ w1_n2 + b1_n2)
+        hidden = torch.stack([h1, h2], dim=-1)
+
+        # Layer 2: AND of hidden
+        w2 = self._get('boolean.xor.layer2.weight')
+        b2 = self._get('boolean.xor.layer2.bias')
+
+        return heaviside(hidden @ w2 + b2)
+
+    # -------------------------------------------------------------------------
+    # Arithmetic: Full Adder
+    # -------------------------------------------------------------------------
+
+    def eval_full_adder(self, a: torch.Tensor, b: torch.Tensor,
+                        cin: torch.Tensor, prefix: str) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Full adder: sum = a XOR b XOR cin, cout = (a AND b) OR (cin AND (a XOR b))
+        Returns (sum_bit, carry_out)
+        """
+        inp_ab = torch.stack([a, b], dim=-1)
+
+        # HA1: a XOR b
+        w1_or = self._get(f'{prefix}.ha1.sum.layer1.or.weight')
+        b1_or = self._get(f'{prefix}.ha1.sum.layer1.or.bias')
+        w1_nand = self._get(f'{prefix}.ha1.sum.layer1.nand.weight')
+        b1_nand = self._get(f'{prefix}.ha1.sum.layer1.nand.bias')
+        w2 = self._get(f'{prefix}.ha1.sum.layer2.weight')
+        b2 = self._get(f'{prefix}.ha1.sum.layer2.bias')
+
+        h_or = heaviside(inp_ab @ w1_or + b1_or)
+        h_nand = heaviside(inp_ab @ w1_nand + b1_nand)
+        hidden = torch.stack([h_or, h_nand], dim=-1)
+        ha1_sum = heaviside(hidden @ w2 + b2)
+
+        # HA1 carry
+        w_c1 = self._get(f'{prefix}.ha1.carry.weight')
+        b_c1 = self._get(f'{prefix}.ha1.carry.bias')
+        ha1_carry = heaviside(inp_ab @ w_c1 + b_c1)
+
+        # HA2: ha1_sum XOR cin
+        inp_ha2 = torch.stack([ha1_sum, cin], dim=-1)
+        w1_or = self._get(f'{prefix}.ha2.sum.layer1.or.weight')
+        b1_or = self._get(f'{prefix}.ha2.sum.layer1.or.bias')
+        w1_nand = self._get(f'{prefix}.ha2.sum.layer1.nand.weight')
+        b1_nand = self._get(f'{prefix}.ha2.sum.layer1.nand.bias')
+        w2 = self._get(f'{prefix}.ha2.sum.layer2.weight')
+        b2 = self._get(f'{prefix}.ha2.sum.layer2.bias')
+
+        h_or = heaviside(inp_ha2 @ w1_or + b1_or)
+        h_nand = heaviside(inp_ha2 @ w1_nand + b1_nand)
+        hidden = torch.stack([h_or, h_nand], dim=-1)
+        ha2_sum = heaviside(hidden @ w2 + b2)
+
+        # HA2 carry
+        w_c2 = self._get(f'{prefix}.ha2.carry.weight')
+        b_c2 = self._get(f'{prefix}.ha2.carry.bias')
+        ha2_carry = heaviside(inp_ha2 @ w_c2 + b_c2)
+
+        # Carry out = ha1_carry OR ha2_carry
+        inp_cout = torch.stack([ha1_carry, ha2_carry], dim=-1)
+        w_or = self._get(f'{prefix}.carry_or.weight')
+        b_or = self._get(f'{prefix}.carry_or.bias')
+        cout = heaviside(inp_cout @ w_or + b_or)
+
+        return ha2_sum, cout
+
+    # -------------------------------------------------------------------------
+    # Arithmetic: 8-bit Ripple Carry Adder
+    # -------------------------------------------------------------------------
+
+    def add_8bit(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        8-bit ripple carry addition.
+        a_bits, b_bits: [..., 8] tensors of bits (LSB first)
+        Returns: (result_bits [..., 8], carry_out [...])
+        """
+        batch_shape = a_bits.shape[:-1]
+        carry = torch.zeros(batch_shape, device=a_bits.device)
+        result_bits = []
+
+        for i in range(8):
+            a_i = a_bits[..., i]
+            b_i = b_bits[..., i]
+            sum_bit, carry = self.eval_full_adder(
+                a_i, b_i, carry,
+                f'arithmetic.ripplecarry8bit.fa{i}'
+            )
+            result_bits.append(sum_bit)
+
+        return torch.stack(result_bits, dim=-1), carry
+
+    # -------------------------------------------------------------------------
+    # Arithmetic: 8-bit Comparators
+    # -------------------------------------------------------------------------
+
+    def greater_than_8bit(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> torch.Tensor:
+        """Returns 1 if a > b, else 0. Bits are MSB first."""
+        diff = a_bits - b_bits  # [..., 8]
+        w = self._get('arithmetic.greaterthan8bit.comparator')
+        score = (diff * w).sum(dim=-1)
+        return (score > 0).float()
+
+    def less_than_8bit(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> torch.Tensor:
+        """Returns 1 if a < b, else 0. Bits are MSB first."""
+        diff = b_bits - a_bits  # [..., 8]
+        w = self._get('arithmetic.lessthan8bit.comparator')
+        score = (diff * w).sum(dim=-1)
+        return (score > 0).float()
+
+    def equal_8bit(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> torch.Tensor:
+        """Returns 1 if a == b, else 0."""
+        gt = self.greater_than_8bit(a_bits, b_bits)
+        lt = self.less_than_8bit(a_bits, b_bits)
+        return (1 - gt) * (1 - lt)
+
+
+# =============================================================================
+# BIT EXTRACTION / INJECTION INTERFACES
+# =============================================================================
+
+class BitExtractor(nn.Module):
+    """
+    Learns to extract 8-bit operands from token embeddings.
+    Maps embedding -> 16 bits (two 8-bit operands).
+    """
+
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.d_model = d_model
+
+        # Project to logits, then binarize
+        self.proj = nn.Linear(d_model, 16)
+
+        # Learnable temperature for sigmoid approximation during training
+        self.temperature = nn.Parameter(torch.tensor(1.0))
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        x: [..., d_model]
+        Returns: a_bits [..., 8], b_bits [..., 8] (LSB first for arithmetic)
+        """
+        logits = self.proj(x)  # [..., 16]
+
+        # Binarize with STE
+        bits = heaviside(logits)
+
+        # Split into two operands
+        a_bits = bits[..., :8]
+        b_bits = bits[..., 8:]
+
+        return a_bits, b_bits
+
+
+class BitInjector(nn.Module):
+    """
+    Learns to inject circuit results back into embedding space.
+    Maps 16 bits (result + flags) -> embedding delta.
+    """
+
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.d_model = d_model
+
+        # Project bits to embedding
+        self.proj = nn.Linear(16, d_model)
+
+        # Learnable scale
+        self.scale = nn.Parameter(torch.tensor(0.1))
+
+    def forward(self, result_bits: torch.Tensor, flags: torch.Tensor) -> torch.Tensor:
+        """
+        result_bits: [..., 8]
+        flags: [..., 8] (carry, overflow, zero, negative, etc.)
+        Returns: [..., d_model]
+        """
+        combined = torch.cat([result_bits, flags], dim=-1)  # [..., 16]
+        return self.proj(combined) * self.scale
+
+
+# =============================================================================
+# CIRCUIT-AUGMENTED MLP BLOCK
+# =============================================================================
+
+class CircuitAugmentedMLP(nn.Module):
+    """
+    MLP block augmented with frozen threshold circuits.
+
+    The original MLP path runs in parallel with the circuit path.
+    A learned router decides how much to use each.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        intermediate_size: int,
+        circuit_path: str,
+        device: str = 'cpu'
+    ):
+        super().__init__()
+        self.d_model = d_model
+
+        # Original MLP components (will be loaded from pretrained)
+        self.gate_proj = nn.Linear(d_model, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(d_model, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, d_model, bias=False)
+        self.act_fn = nn.SiLU()
+
+        # Circuit components
+        self.circuits = CircuitExecutor(circuit_path, device)
+        self.bit_extractor = BitExtractor(d_model)
+        self.bit_injector = BitInjector(d_model)
+
+        # Router: decides circuit vs MLP contribution
+        self.router = nn.Sequential(
+            nn.Linear(d_model, 64),
+            nn.ReLU(),
+            nn.Linear(64, 2),
+            nn.Softmax(dim=-1)
+        )
+
+        # Operation selector (which arithmetic op to perform)
+        self.op_selector = nn.Sequential(
+            nn.Linear(d_model, 32),
+            nn.ReLU(),
+            nn.Linear(32, 4),  # add, sub, compare, passthrough
+            nn.Softmax(dim=-1)
+        )
+
+    def _compute_flags(self, result_bits: torch.Tensor, carry: torch.Tensor) -> torch.Tensor:
+        """Compute status flags from result."""
+        batch_shape = result_bits.shape[:-1]
+
+        # Zero flag: all bits are 0
+        zero = (result_bits.sum(dim=-1) == 0).float()
+
+        # Negative flag: MSB is 1 (two's complement)
+        negative = result_bits[..., 7]
+
+        # Carry flag
+        carry_flag = carry
+
+        # Pad to 8 flags
+        flags = torch.zeros(*batch_shape, 8, device=result_bits.device)
+        flags[..., 0] = zero
+        flags[..., 1] = negative
+        flags[..., 2] = carry_flag
+
+        return flags
+
+    def _circuit_forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Run input through threshold circuits."""
+        # Extract operands
+        a_bits, b_bits = self.bit_extractor(x)
+
+        # Get operation weights
+        op_weights = self.op_selector(x)  # [..., 4]
+
+        # Compute addition
+        add_result, add_carry = self.circuits.add_8bit(a_bits, b_bits)
+        add_flags = self._compute_flags(add_result, add_carry)
+
+        # Compute subtraction (a + (~b) + 1, simplified: just use add for now)
+        # For MVP, we'll focus on addition
+
+        # Inject result back
+        circuit_delta = self.bit_injector(add_result, add_flags)
+
+        return circuit_delta
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: [batch, seq_len, d_model]
+        Returns: [batch, seq_len, d_model]
+        """
+        # Original MLP path
+        mlp_out = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+        # Circuit path
+        circuit_out = self._circuit_forward(x)
+
+        # Route between paths
+        route_weights = self.router(x)  # [..., 2]
+        mlp_weight = route_weights[..., 0:1]
+        circuit_weight = route_weights[..., 1:2]
+
+        # Combine: MLP output + weighted circuit contribution
+        output = mlp_out + circuit_weight * circuit_out
+
+        return output
+
+
+# =============================================================================
+# MODEL SURGERY: Insert circuits into SmolLM2
+# =============================================================================
+
+def augment_smollm2_with_circuits(
+    model: AutoModelForCausalLM,
+    circuit_path: str,
+    layer_indices: list = None,
+    device: str = 'cpu'
+) -> AutoModelForCausalLM:
+    """
+    Surgically insert circuit blocks into SmolLM2's MLP layers.
+
+    Args:
+        model: Pretrained SmolLM2 model
+        circuit_path: Path to neural_computer.safetensors
+        layer_indices: Which layers to augment (default: middle layers)
+        device: Device for circuit tensors
+
+    Returns:
+        Modified model with circuit-augmented MLPs
+    """
+    config = model.config
+    num_layers = config.num_hidden_layers
+
+    # Default: augment middle third of layers
+    if layer_indices is None:
+        start = num_layers // 3
+        end = 2 * num_layers // 3
+        layer_indices = list(range(start, end))
+
+    print(f"Augmenting layers {layer_indices} with threshold circuits...")
+
+    for idx in layer_indices:
+        layer = model.model.layers[idx]
+        old_mlp = layer.mlp
+
+        # Create augmented MLP
+        new_mlp = CircuitAugmentedMLP(
+            d_model=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            circuit_path=circuit_path,
+            device=device
+        )
+
+        # Copy pretrained weights
+        new_mlp.gate_proj.weight.data = old_mlp.gate_proj.weight.data.clone()
+        new_mlp.up_proj.weight.data = old_mlp.up_proj.weight.data.clone()
+        new_mlp.down_proj.weight.data = old_mlp.down_proj.weight.data.clone()
+
+        # Replace
+        layer.mlp = new_mlp
+
+    # Freeze circuit weights, keep interfaces trainable
+    for name, param in model.named_parameters():
+        if 'circuits' in name:
+            param.requires_grad = False
+
+    print(f"Done. Circuit weights frozen, interfaces trainable.")
+
+    return model
+
+
+# =============================================================================
+# TRAINING UTILITIES
+# =============================================================================
+
+def generate_arithmetic_batch(batch_size: int, max_val: int = 255) -> Tuple[list, list]:
+    """Generate batch of arithmetic problems and solutions."""
+    prompts = []
+    targets = []
+
+    for _ in range(batch_size):
+        a = torch.randint(0, max_val + 1, (1,)).item()
+        b = torch.randint(0, max_val + 1, (1,)).item()
+        result = (a + b) % 256
+
+        prompts.append(f"{a} + {b} =")
+        targets.append(f" {result}")
+
+    return prompts, targets
+
+
+def evaluate_arithmetic(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    n_problems: int = 100,
+    device: str = 'cpu'
+) -> dict:
+    """Evaluate model on random arithmetic problems."""
+    correct = 0
+    total = 0
+    errors = []
+
+    model.eval()
+
+    for _ in range(n_problems):
+        a = torch.randint(0, 256, (1,)).item()
+        b = torch.randint(0, 256, (1,)).item()
+        expected = (a + b) % 256
+
+        prompt = f"{a} + {b} ="
+        inputs = tokenizer(prompt, return_tensors='pt').to(device)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=10,
+                do_sample=False,
+                pad_token_id=tokenizer.eos_token_id
+            )
+
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        # Extract number from response
+        try:
+            # Find the part after "="
+            answer_part = response.split('=')[-1].strip()
+            # Extract first number
+            predicted = int(''.join(c for c in answer_part.split()[0] if c.isdigit()))
+
+            if predicted == expected:
+                correct += 1
+            else:
+                errors.append((a, b, expected, predicted))
+        except:
+            errors.append((a, b, expected, "parse_error"))
+
+        total += 1
+
+    return {
+        'accuracy': correct / total,
+        'correct': correct,
+        'total': total,
+        'errors': errors[:10]  # First 10 errors
+    }
+
+
+# =============================================================================
+# MAIN: Demo
+# =============================================================================
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description='Circuit-Augmented LLM Demo')
+    parser.add_argument('--circuit-path', type=str,
+                        default='./neural_computer.safetensors',
+                        help='Path to circuit weights')
+    parser.add_argument('--device', type=str, default='cpu',
+                        help='Device (cpu or cuda)')
+    parser.add_argument('--eval-only', action='store_true',
+                        help='Only evaluate, do not augment')
+    args = parser.parse_args()
+
+    print("=" * 70)
+    print(" CIRCUIT-AUGMENTED LLM")
+    print("=" * 70)
+
+    # Load tokenizer and model
+    print("\n[1] Loading SmolLM2-360M...")
+    model_id = "HuggingFaceTB/SmolLM2-360M"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
+
+    print(f"    Parameters: {sum(p.numel() for p in model.parameters()):,}")
+
+    # Baseline evaluation
+    print("\n[2] Baseline arithmetic evaluation...")
+    baseline = evaluate_arithmetic(model, tokenizer, n_problems=50, device=args.device)
+    print(f"    Accuracy: {baseline['accuracy']*100:.1f}% ({baseline['correct']}/{baseline['total']})")
+    if baseline['errors']:
+        print(f"    Sample errors:")
+        for a, b, exp, got in baseline['errors'][:5]:
+            print(f"      {a} + {b} = {exp}, model said {got}")
+
+    if args.eval_only:
+        print("\nDone (eval only mode).")
+        exit(0)
+
+    # Augment with circuits
+    print(f"\n[3] Augmenting with threshold circuits...")
+    print(f"    Circuit path: {args.circuit_path}")
+    model = augment_smollm2_with_circuits(
+        model,
+        args.circuit_path,
+        device=args.device
+    )
+
+    new_params = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"    Total parameters: {new_params:,}")
+    print(f"    Trainable parameters: {trainable:,}")
+
+    # Test circuit execution directly
+    print("\n[4] Testing circuit execution...")
+    circuit_exec = CircuitExecutor(args.circuit_path, args.device)
+
+    test_cases = [(127, 128), (255, 1), (0, 0), (100, 55)]
+    for a, b in test_cases:
+        # Convert to bits (LSB first)
+        a_bits = torch.tensor([(a >> i) & 1 for i in range(8)], dtype=torch.float32)
+        b_bits = torch.tensor([(b >> i) & 1 for i in range(8)], dtype=torch.float32)
+
+        result_bits, carry = circuit_exec.add_8bit(
+            a_bits.unsqueeze(0),
+            b_bits.unsqueeze(0)
+        )
+
+        # Convert result bits back to int
+        result = sum(int(result_bits[0, i].item()) * (2**i) for i in range(8))
+        expected = (a + b) % 256
+
+        status = "OK" if result == expected else "FAIL"
+        print(f"    {a} + {b} = {result} (expected {expected}) [{status}]")
+
+    print("\n[5] Model ready for fine-tuning.")
+    print("    Next: Train interface layers on arithmetic examples.")
+    print("=" * 70)

llm/guide.md

@@ -0,0 +1,615 @@
+# Embedding Threshold Logic Circuits into Transformer MLPs
+
+## Technical Implementation Guide
+
+---
+
+## 1. Core Thesis
+
+Standard LLMs fail at arithmetic because they're interpolators—they approximate functions over training distributions rather than compute exact results. A 360M parameter model trained on internet text has seen "127 + 128 = 255" zero or few times, so it guesses "140" based on pattern matching.
+
+We solve this by embedding **frozen, proven-correct arithmetic circuits** directly into the transformer's MLP layers. The circuits use threshold logic (weighted sums + step activation), which is structurally compatible with neural network layers. We train only the **interface layers** that learn to:
+
+1. Extract operands from token embeddings
+2. Route computation through the circuits
+3. Inject results back into the residual stream
+
+The model learns **call dispatch**, not arithmetic. The arithmetic is already solved.
+
+---
+
+## 2. Threshold Logic Fundamentals
+
+### 2.1 Single Threshold Gate
+
+A threshold gate computes:
+
+```
+output = 1  if  (Σ wᵢxᵢ + b) ≥ 0
+         0  otherwise
+```
+
+This is a neuron with Heaviside step activation. With integer weights `w` and bias `b`, it computes a Boolean function of binary inputs.
+
+**Example: AND gate**
+```
+w = [1, 1], b = -2
+AND(0,0) = H(0 + 0 - 2) = H(-2) = 0
+AND(0,1) = H(0 + 1 - 2) = H(-1) = 0
+AND(1,0) = H(1 + 0 - 2) = H(-1) = 0
+AND(1,1) = H(1 + 1 - 2) = H(0)  = 1
+```
+
+**Example: OR gate**
+```
+w = [1, 1], b = -1
+OR(0,0) = H(0 + 0 - 1) = H(-1) = 0
+OR(0,1) = H(0 + 1 - 1) = H(0)  = 1
+OR(1,0) = H(1 + 0 - 1) = H(0)  = 1
+OR(1,1) = H(1 + 1 - 1) = H(1)  = 1
+```
+
+### 2.2 Multi-Layer Circuits
+
+XOR is not linearly separable—it requires two layers:
+
+```
+Layer 1:
+  neuron1 (OR):   w=[1,1], b=-1   → fires if a OR b
+  neuron2 (NAND): w=[-1,-1], b=1  → fires if NOT(a AND b)
+
+Layer 2:
+  neuron3 (AND): w=[1,1], b=-2   → fires if both layer1 outputs are 1
+
+XOR(a,b) = AND(OR(a,b), NAND(a,b))
+```
+
+### 2.3 Full Adder
+
+A full adder computes `sum` and `carry_out` from inputs `a`, `b`, `carry_in`:
+
+```
+sum = a XOR b XOR cin
+cout = (a AND b) OR (cin AND (a XOR b))
+```
+
+Implementation uses two half-adders chained:
+
+```
+HA1: (a, b) → (sum1 = a XOR b, carry1 = a AND b)
+HA2: (sum1, cin) → (sum2 = sum1 XOR cin, carry2 = sum1 AND cin)
+cout = carry1 OR carry2
+final_sum = sum2
+```
+
+Each XOR is 2 layers, each AND/OR is 1 layer. Total depth: ~4 layers per full adder.
+
+### 2.4 8-bit Ripple Carry Adder
+
+Chain 8 full adders, propagating carry:
+
+```
+FA0: (a[0], b[0], 0)      → (sum[0], c0)
+FA1: (a[1], b[1], c0)     → (sum[1], c1)
+FA2: (a[2], b[2], c1)     → (sum[2], c2)
+...
+FA7: (a[7], b[7], c6)     → (sum[7], c7)
+```
+
+Total circuit depth: ~32 threshold layers (8 FAs × 4 layers each).
+
+---
+
+## 3. Circuit Inventory
+
+The `neural_computer.safetensors` contains 3,122 tensors / 5,648 parameters implementing:
+
+| Category | Circuits | Tensors |
+|----------|----------|---------|
+| Boolean | AND, OR, NOT, NAND, NOR, XOR, XNOR, IMPLIES, BIIMPLIES | ~30 |
+| Arithmetic | Half adder, Full adder, Ripple carry 2/4/8-bit, 8×8 multiplier | ~800 |
+| Comparators | GT, LT, GEQ, LEQ, EQ (8-bit) | ~50 |
+| ALU | 16-operation ALU, opcode decoder, flag computation | ~400 |
+| Control | JMP, JZ, JNZ, JC, JNC, JN, JP, CALL, RET, PUSH, POP | ~200 |
+| Modular | Divisibility by 2-12 | ~600 |
+| Error Detection | Parity, CRC, Hamming, checksum | ~200 |
+| Pattern | Popcount, leading zeros, symmetry | ~150 |
+| Threshold | k-of-n gates, majority, minority | ~100 |
+
+All weights are integers. All activations are Heaviside. Verified with 6,590 exhaustive tests.
+
+---
+
+## 4. Transformer Integration Architecture
+
+### 4.1 Target: SmolLM2-360M
+
+```
+Architecture: LlamaForCausalLM
+Hidden dim:   960
+Layers:       32
+Heads:        15
+MLP expansion: 4x (intermediate = 3840)
+Vocab:        49152
+Parameters:   361,821,120
+```
+
+Standard MLP block:
+```python
+def forward(x):  # x: [batch, seq, 960]
+    gate = self.gate_proj(x)      # [batch, seq, 3840]
+    up = self.up_proj(x)          # [batch, seq, 3840]
+    hidden = silu(gate) * up      # SwiGLU activation
+    return self.down_proj(hidden) # [batch, seq, 960]
+```
+
+### 4.2 Augmented MLP Block
+
+```python
+def forward(x):  # x: [batch, seq, 960]
+    # Original MLP path (unchanged)
+    mlp_out = self.down_proj(silu(self.gate_proj(x)) * self.up_proj(x))
+
+    # Circuit path (new)
+    a_bits, b_bits = self.bit_extractor(x)       # [batch, seq, 8] each
+    result_bits, carry = self.circuits.add_8bit(a_bits, b_bits)
+    flags = self.compute_flags(result_bits, carry)
+    circuit_delta = self.bit_injector(result_bits, flags)
+
+    # Routing
+    route_weights = self.router(x)  # [batch, seq, 2] softmax
+
+    # Combine
+    return mlp_out + route_weights[..., 1:2] * circuit_delta
+```
+
+### 4.3 Layer Selection
+
+We augment the **middle third** of layers (10-20 of 32):
+
+- Early layers (0-9): Token/position encoding, not arithmetic-relevant
+- Middle layers (10-20): Abstract reasoning, computation
+- Late layers (21-31): Output formatting, vocabulary projection
+
+Rationale: Arithmetic computation happens in middle layers where the model processes relationships between tokens. Early layers haven't built sufficient representations; late layers are committed to output tokens.
+
+---
+
+## 5. Interface Layers (Trainable)
+
+### 5.1 BitExtractor
+
+Maps token embedding → two 8-bit operands.
+
+```python
+class BitExtractor(nn.Module):
+    def __init__(self, d_model=960):
+        self.proj = nn.Linear(d_model, 16)  # 960 → 16
+
+    def forward(self, x):
+        logits = self.proj(x)           # [batch, seq, 16]
+        bits = heaviside(logits)        # binarize with STE
+        a_bits = bits[..., :8]          # first operand
+        b_bits = bits[..., 8:]          # second operand
+        return a_bits, b_bits           # both [batch, seq, 8], LSB first
+```
+
+**What it learns**: Which embedding dimensions encode numeric magnitude. For token "127", it must learn that certain activation patterns correspond to bits `[1,1,1,1,1,1,1,0]`.
+
+**Parameters**: 960 × 16 + 16 = 15,376
+
+### 5.2 BitInjector
+
+Maps circuit outputs → embedding delta.
+
+```python
+class BitInjector(nn.Module):
+    def __init__(self, d_model=960):
+        self.proj = nn.Linear(16, d_model)  # 16 → 960
+        self.scale = nn.Parameter(torch.tensor(0.1))
+
+    def forward(self, result_bits, flags):
+        combined = torch.cat([result_bits, flags], dim=-1)  # [batch, seq, 16]
+        return self.proj(combined) * self.scale              # [batch, seq, 960]
+```
+
+**What it learns**: How to inject the result bits back into embedding space such that subsequent layers (and the final vocabulary projection) produce the correct output tokens.
+
+**Parameters**: 16 × 960 + 960 + 1 = 16,321
+
+### 5.3 Router
+
+Decides when to use circuit path.
+
+```python
+class Router(nn.Module):
+    def __init__(self, d_model=960):
+        self.net = nn.Sequential(
+            nn.Linear(d_model, 64),
+            nn.ReLU(),
+            nn.Linear(64, 2),
+            nn.Softmax(dim=-1)
+        )
+
+    def forward(self, x):
+        return self.net(x)  # [batch, seq, 2]: [mlp_weight, circuit_weight]
+```
+
+**What it learns**: "This position contains arithmetic" → route through circuits. "This is prose" → use normal MLP.
+
+**Parameters**: 960 × 64 + 64 + 64 × 2 + 2 = 61,698
+
+### 5.4 Total Trainable Parameters
+
+Per augmented layer:
+```
+BitExtractor:  15,376
+BitInjector:   16,321
+Router:        61,698
+OpSelector:    ~31,000
+───────────────────────
+Total:         ~124,395 per layer
+```
+
+For 11 augmented layers: **~1.37M trainable parameters**
+
+This is 0.38% of the model. The other 99.62% (including all circuit weights) is frozen.
+
+---
+
+## 6. Gradient Flow Through Heaviside
+
+### 6.1 The Problem
+
+Heaviside has zero gradient almost everywhere:
+
+```
+H(x) = 1 if x ≥ 0 else 0
+dH/dx = 0 for x ≠ 0, undefined at x = 0
+```
+
+Standard backprop would give zero gradients to BitExtractor.
+
+### 6.2 Straight-Through Estimator (STE)
+
+We use STE: forward pass uses true Heaviside, backward pass pretends it's identity.
+
+```python
+class HeavisideSTE(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x):
+        return (x >= 0).float()  # true step function
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output  # pass gradient through unchanged
+```
+
+**Intuition**: "If making the input larger would have helped the output, increase the input." The gradient tells us the direction even though the function is flat.
+
+### 6.3 Alternative: Sigmoid Annealing
+
+During training, use sigmoid with increasing temperature:
+
+```python
+def soft_heaviside(x, temperature):
+    return torch.sigmoid(x * temperature)
+
+# temperature: 1 → 10 → 100 over training
+# At high temperature, sigmoid ≈ step function
+```
+
+This provides smoother gradients early in training, then sharpens to true binary at inference.
+
+---
+
+## 7. Training Strategy
+
+### 7.1 Data Generation
+
+Generate arithmetic problems exhaustively:
+
+```python
+def generate_batch(batch_size):
+    a = torch.randint(0, 256, (batch_size,))
+    b = torch.randint(0, 256, (batch_size,))
+    result = (a + b) % 256
+
+    prompts = [f"{a[i]} + {b[i]} =" for i in range(batch_size)]
+    targets = [f" {result[i]}" for i in range(batch_size)]
+
+    return prompts, targets
+```
+
+For 8-bit addition, there are 256 × 256 = 65,536 unique problems. We can cover the entire space.
+
+### 7.2 Loss Function
+
+Standard cross-entropy on next-token prediction:
+
+```python
+outputs = model(input_ids, attention_mask=mask, labels=labels)
+loss = outputs.loss  # CE loss, only on target tokens
+```
+
+Labels are masked for prompt tokens (`-100`), so loss only backprops through the answer.
+
+### 7.3 Optimizer Configuration
+
+```python
+# Only train interface layers
+interface_params = [p for n, p in model.named_parameters()
+                    if any(x in n for x in ['bit_extractor', 'bit_injector', 'router'])]
+
+optimizer = AdamW(interface_params, lr=1e-4, weight_decay=0.01)
+scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs)
+```
+
+### 7.4 Curriculum Learning
+
+Start simple, increase difficulty:
+
+```
+Phase 1 (epochs 1-2):   Single-digit addition (0-9 + 0-9)
+Phase 2 (epochs 3-4):   Two-digit addition (0-99 + 0-99)
+Phase 3 (epochs 5-7):   Full 8-bit addition (0-255 + 0-255)
+Phase 4 (epochs 8-10):  Adversarial cases (carry chains: 127+128, 255+1)
+```
+
+This helps the interface layers learn the basic extraction pattern before tackling hard cases.
+
+### 7.5 Training Hyperparameters
+
+```
+Model:          SmolLM2-360M
+Augmented:      Layers 10-20 (11 layers)
+Trainable:      1.37M parameters
+Frozen:         362M parameters (including 5.6K circuit params)
+
+Batch size:     32
+Learning rate:  1e-4
+Epochs:         10
+Samples:        10,000 per epoch
+Warmup:         500 steps
+Device:         RTX 6000 Ada (48GB)
+
+Expected time:  ~30 minutes total
+```
+
+---
+
+## 8. Forward Pass Walkthrough
+
+Input: `"127 + 128 ="`
+
+### 8.1 Tokenization
+
+```
+Tokens: ["127", " +", " 128", " ="]
+IDs:    [12700, 489, 13824, 284]  # hypothetical
+```
+
+### 8.2 Embedding
+
+```
+embeddings = embed(input_ids)  # [1, 4, 960]
+```
+
+### 8.3 Layers 0-9 (Unchanged)
+
+Standard attention + MLP, building representations.
+
+### 8.4 Layer 10 (Augmented)
+
+```python
+# After attention
+x = layer_norm(attn_output + residual)  # [1, 4, 960]
+
+# MLP path
+mlp_out = down_proj(silu(gate_proj(x)) * up_proj(x))
+
+# Circuit path
+a_bits, b_bits = bit_extractor(x)
+# Position 0 ("127"): a_bits ≈ [1,1,1,1,1,1,1,0] if well-trained
+# Position 2 ("128"): b_bits ≈ [0,0,0,0,0,0,0,1]
+# (In practice, extraction happens per-position; aggregation is learned)
+
+result_bits, carry = circuits.add_8bit(a_bits, b_bits)
+# result_bits = [1,1,1,1,1,1,1,1] = 255
+
+flags = compute_flags(result_bits, carry)
+# zero=0, negative=1, carry=1
+
+circuit_delta = bit_injector(result_bits, flags)  # [1, 4, 960]
+
+# Routing
+route = router(x)  # [1, 4, 2]
+# Position 3 ("="): route ≈ [0.1, 0.9] → use circuits
+# Position 1 ("+"): route ≈ [0.8, 0.2] → mostly MLP
+
+# Combine
+output = mlp_out + route[..., 1:2] * circuit_delta
+```
+
+### 8.5 Layers 11-31
+
+Continue processing, eventually projecting to vocabulary.
+
+### 8.6 Output
+
+```
+logits = lm_head(final_hidden)  # [1, 4, 49152]
+next_token = argmax(logits[0, 3, :])  # token after "="
+# Should decode to "255" (possibly as " 255" or "255")
+```
+
+---
+
+## 9. Inference Characteristics
+
+### 9.1 Exactness
+
+At inference, Heaviside is true step function—no approximation. If BitExtractor correctly maps "127" → bits and "128" → bits, the circuit **will** output 255. The only failure mode is incorrect extraction.
+
+### 9.2 Latency
+
+Circuit computation adds ~5-10% overhead:
+- BitExtractor: 1 linear layer (960→16)
+- Circuits: ~32 threshold layers, but sparse and tiny
+- BitInjector: 1 linear layer (16→960)
+- Router: 2 linear layers
+
+The circuits have only 5,648 parameters total—negligible versus the 361M in the base model.
+
+### 9.3 Generalization
+
+Once the interface learns the mapping, it generalizes to **all** 65,536 8-bit additions. There's no memorization—the circuits compute.
+
+---
+
+## 10. Evaluation Metrics
+
+### 10.1 Arithmetic Accuracy
+
+```python
+def eval_accuracy(model, n_problems=1000):
+    correct = 0
+    for _ in range(n_problems):
+        a, b = random 8-bit values
+        expected = (a + b) % 256
+        predicted = model.generate(f"{a} + {b} =")
+        if parse_int(predicted) == expected:
+            correct += 1
+    return correct / n_problems
+```
+
+**Baseline SmolLM2**: ~5-10% (guessing based on patterns)
+**Target**: >95% (circuit-accurate)
+
+### 10.2 Edge Case Performance
+
+Specifically test:
+- Carry propagation: 127+128, 255+1, 128+128
+- Zeros: 0+0, 0+255
+- Identity: x+0 for various x
+- Commutativity: verify a+b == b+a
+
+### 10.3 Non-Arithmetic Preservation
+
+Verify general capability isn't degraded:
+- Perplexity on held-out text
+- Common benchmarks (HellaSwag, etc.)
+
+The augmentation should be **additive**—circuits help arithmetic, MLP handles everything else via routing.
+
+---
+
+## 11. Extension Roadmap
+
+### 11.1 Additional Operations
+
+The circuit inventory includes:
+- Subtraction (via two's complement)
+- Multiplication (8×8 → 16-bit)
+- Division (iterative subtraction)
+- Bitwise ops (AND, OR, XOR, shifts)
+- Comparisons (GT, LT, EQ)
+
+Each needs its own extraction/injection interface, or a unified interface with operation selection.
+
+### 11.2 Multi-Operand Expressions
+
+For "15 + 27 + 33 =", need:
+- Operand count detection
+- Sequential circuit invocation
+- Accumulator pattern
+
+### 11.3 Larger Bit Widths
+
+16-bit and 32-bit arithmetic require:
+- Larger circuits (or chained 8-bit)
+- Wider BitExtractor (32 or 64 output dims)
+- More training data
+
+### 11.4 Symbolic Integration
+
+Ultimate goal: the model recognizes when it needs to compute, invokes circuits, and integrates results into coherent natural language output.
+
+```
+User: "If I have 127 apples and buy 128 more, how many do I have?"
+Model: [extracts 127, 128] [routes to circuit] [gets 255]
+       "You would have 255 apples."
+```
+
+---
+
+## 12. File Structure
+
+```
+8bit-threshold-computer/
+├── neural_computer.safetensors    # Frozen circuits (3,122 tensors)
+├── circuit_llm.py                 # Integration architecture
+├── train_circuit_interface.py     # Training loop
+├── iron_eval.py                   # Circuit verification (6,590 tests)
+├── skeptic_test.py                # Algebraic identity tests (127 tests)
+├── prune_weights.py               # Weight optimization
+├── tensors.txt                    # Tensor manifest
+├── guide.md                       # This document
+└── README.md                      # Project overview
+```
+
+---
+
+## 13. Key Equations
+
+### Heaviside Step
+```
+H(x) = 1 if x ≥ 0 else 0
+```
+
+### Threshold Gate
+```
+f(x₁,...,xₙ) = H(Σᵢ wᵢxᵢ + b)
+```
+
+### Full Adder
+```
+sum = a ⊕ b ⊕ cᵢₙ
+cₒᵤₜ = (a ∧ b) ∨ (cᵢₙ ∧ (a ⊕ b))
+```
+
+### STE Gradient
+```
+Forward:  y = H(x)
+Backward: ∂L/∂x = ∂L/∂y
+```
+
+### Router Combination
+```
+output = mlp_out + softmax(router(x))[1] × circuit_delta
+```
+
+---
+
+## 14. References
+
+1. McCulloch & Pitts (1943). "A Logical Calculus of Ideas Immanent in Nervous Activity"
+2. Muroga (1971). "Threshold Logic and Its Applications"
+3. Siegelmann & Sontag (1995). "On the Computational Power of Neural Nets"
+4. Bengio et al. (2013). "Estimating or Propagating Gradients Through Stochastic Neurons"
+5. Ma et al. (2024). "The Era of 1-bit LLMs" (BitNet b1.58)
+6. HuggingFace (2024). "SmolLM2: Small Language Models"
+
+---
+
+## 15. Summary
+
+We embed a proven-correct 8-bit threshold logic computer into SmolLM2's MLP layers. The circuits are frozen; we train only the interface layers that learn call dispatch. This gives the LLM exact arithmetic capability without training it to "do math"—the math is already done.
+
+The approach is:
+- **Sound**: Circuits verified with 6,590 tests
+- **Efficient**: 1.37M trainable params, 5.6K circuit params
+- **Exact**: Heaviside at inference means no approximation error
+- **Composable**: Add more circuits (multiply, compare, etc.) with same pattern
+
+The model learns when to call the calculator, not how to calculate.

llm/train_circuit_interface.py

@@ -0,0 +1,306 @@
+"""
+Train the circuit interface layers on arithmetic examples.
+============================================================
+
+The threshold circuits are frozen - we only train:
+- BitExtractor: embedding -> operand bits
+- BitInjector: result bits -> embedding
+- Router: when to use circuits vs MLP
+"""
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from tqdm import tqdm
+import argparse
+import warnings
+warnings.filterwarnings('ignore')
+
+from circuit_llm import (
+    augment_smollm2_with_circuits,
+    evaluate_arithmetic,
+    CircuitExecutor
+)
+
+
+# =============================================================================
+# ARITHMETIC DATASET
+# =============================================================================
+
+class ArithmeticDataset(Dataset):
+    """Dataset of 8-bit addition problems."""
+
+    def __init__(self, tokenizer, n_samples: int = 10000, max_val: int = 255):
+        self.tokenizer = tokenizer
+        self.n_samples = n_samples
+        self.max_val = max_val
+
+        # Pre-generate all examples
+        self.examples = []
+        for _ in range(n_samples):
+            a = torch.randint(0, max_val + 1, (1,)).item()
+            b = torch.randint(0, max_val + 1, (1,)).item()
+            result = (a + b) % 256
+
+            prompt = f"{a} + {b} ="
+            target = f" {result}"
+
+            self.examples.append((prompt, target, a, b, result))
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, idx):
+        prompt, target, a, b, result = self.examples[idx]
+
+        # Tokenize
+        prompt_ids = self.tokenizer.encode(prompt, add_special_tokens=False)
+        target_ids = self.tokenizer.encode(target, add_special_tokens=False)
+
+        input_ids = prompt_ids + target_ids
+        labels = [-100] * len(prompt_ids) + target_ids  # Only predict target
+
+        return {
+            'input_ids': torch.tensor(input_ids),
+            'labels': torch.tensor(labels),
+            'a': a,
+            'b': b,
+            'result': result
+        }
+
+
+def collate_fn(batch):
+    """Collate with padding."""
+    max_len = max(len(item['input_ids']) for item in batch)
+
+    input_ids = []
+    labels = []
+    attention_mask = []
+
+    for item in batch:
+        pad_len = max_len - len(item['input_ids'])
+
+        input_ids.append(
+            torch.cat([item['input_ids'], torch.zeros(pad_len, dtype=torch.long)])
+        )
+        labels.append(
+            torch.cat([item['labels'], torch.full((pad_len,), -100, dtype=torch.long)])
+        )
+        attention_mask.append(
+            torch.cat([torch.ones(len(item['input_ids'])), torch.zeros(pad_len)])
+        )
+
+    return {
+        'input_ids': torch.stack(input_ids),
+        'labels': torch.stack(labels),
+        'attention_mask': torch.stack(attention_mask),
+    }
+
+
+# =============================================================================
+# TRAINING LOOP
+# =============================================================================
+
+def train_interface(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    n_epochs: int = 3,
+    batch_size: int = 16,
+    lr: float = 1e-4,
+    n_train_samples: int = 10000,
+    device: str = 'cpu',
+    eval_every: int = 500
+):
+    """
+    Train the circuit interface layers.
+
+    Only trains:
+    - bit_extractor (embedding -> bits)
+    - bit_injector (bits -> embedding)
+    - router (circuit vs MLP weighting)
+    - op_selector (which operation)
+    """
+    print("\n" + "=" * 70)
+    print(" TRAINING CIRCUIT INTERFACE")
+    print("=" * 70)
+
+    # Freeze everything except interface layers
+    interface_params = []
+    frozen_count = 0
+    trainable_count = 0
+
+    for name, param in model.named_parameters():
+        if any(x in name for x in ['bit_extractor', 'bit_injector', 'router', 'op_selector']):
+            param.requires_grad = True
+            interface_params.append(param)
+            trainable_count += param.numel()
+        else:
+            param.requires_grad = False
+            frozen_count += param.numel()
+
+    print(f"\n  Frozen parameters: {frozen_count:,}")
+    print(f"  Trainable parameters: {trainable_count:,}")
+    print(f"  Training {len(interface_params)} parameter groups")
+
+    # Create dataset
+    print(f"\n  Creating dataset ({n_train_samples} examples)...")
+    dataset = ArithmeticDataset(tokenizer, n_samples=n_train_samples)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        collate_fn=collate_fn
+    )
+
+    # Optimizer
+    optimizer = torch.optim.AdamW(interface_params, lr=lr)
+
+    # Training
+    model.to(device)
+    model.train()
+
+    global_step = 0
+    total_loss = 0
+
+    for epoch in range(n_epochs):
+        print(f"\n  Epoch {epoch + 1}/{n_epochs}")
+        print("  " + "-" * 60)
+
+        epoch_loss = 0
+        epoch_steps = 0
+
+        pbar = tqdm(dataloader, desc=f"  Training", leave=False)
+
+        for batch in pbar:
+            input_ids = batch['input_ids'].to(device)
+            labels = batch['labels'].to(device)
+            attention_mask = batch['attention_mask'].to(device)
+
+            # Forward
+            outputs = model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=labels
+            )
+
+            loss = outputs.loss
+
+            # Backward
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            # Logging
+            epoch_loss += loss.item()
+            epoch_steps += 1
+            global_step += 1
+            total_loss += loss.item()
+
+            pbar.set_postfix({'loss': f'{loss.item():.4f}'})
+
+            # Periodic evaluation
+            if global_step % eval_every == 0:
+                model.eval()
+                eval_results = evaluate_arithmetic(model, tokenizer, n_problems=50, device=device)
+                print(f"\n    Step {global_step}: Loss={total_loss/eval_every:.4f}, "
+                      f"Accuracy={eval_results['accuracy']*100:.1f}%")
+                total_loss = 0
+                model.train()
+
+        avg_loss = epoch_loss / epoch_steps
+        print(f"\n  Epoch {epoch + 1} complete. Avg loss: {avg_loss:.4f}")
+
+        # End of epoch evaluation
+        model.eval()
+        eval_results = evaluate_arithmetic(model, tokenizer, n_problems=100, device=device)
+        print(f"  Evaluation: {eval_results['accuracy']*100:.1f}% "
+              f"({eval_results['correct']}/{eval_results['total']})")
+
+        if eval_results['errors']:
+            print(f"  Sample errors:")
+            for a, b, exp, got in eval_results['errors'][:3]:
+                print(f"    {a} + {b} = {exp}, model said {got}")
+
+        model.train()
+
+    print("\n" + "=" * 70)
+    print(" TRAINING COMPLETE")
+    print("=" * 70)
+
+    return model
+
+
+# =============================================================================
+# MAIN
+# =============================================================================
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Train Circuit Interface')
+    parser.add_argument('--circuit-path', type=str,
+                        default='./neural_computer.safetensors',
+                        help='Path to circuit weights')
+    parser.add_argument('--device', type=str, default='cpu',
+                        help='Device (cpu or cuda)')
+    parser.add_argument('--epochs', type=int, default=3,
+                        help='Number of epochs')
+    parser.add_argument('--batch-size', type=int, default=8,
+                        help='Batch size')
+    parser.add_argument('--lr', type=float, default=1e-4,
+                        help='Learning rate')
+    parser.add_argument('--n-samples', type=int, default=5000,
+                        help='Number of training samples')
+    args = parser.parse_args()
+
+    print("=" * 70)
+    print(" CIRCUIT-AUGMENTED LLM TRAINING")
+    print("=" * 70)
+
+    # Load model
+    print("\n[1] Loading SmolLM2-360M...")
+    model_id = "HuggingFaceTB/SmolLM2-360M"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    tokenizer.pad_token = tokenizer.eos_token
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
+
+    # Baseline
+    print("\n[2] Baseline evaluation...")
+    baseline = evaluate_arithmetic(model, tokenizer, n_problems=50, device=args.device)
+    print(f"    Baseline accuracy: {baseline['accuracy']*100:.1f}%")
+
+    # Augment
+    print("\n[3] Augmenting with circuits...")
+    model = augment_smollm2_with_circuits(
+        model,
+        args.circuit_path,
+        device=args.device
+    )
+
+    # Train
+    print("\n[4] Training interface layers...")
+    model = train_interface(
+        model,
+        tokenizer,
+        n_epochs=args.epochs,
+        batch_size=args.batch_size,
+        lr=args.lr,
+        n_train_samples=args.n_samples,
+        device=args.device
+    )
+
+    # Final evaluation
+    print("\n[5] Final evaluation...")
+    final = evaluate_arithmetic(model, tokenizer, n_problems=100, device=args.device)
+    print(f"    Final accuracy: {final['accuracy']*100:.1f}%")
+    print(f"    Improvement: {baseline['accuracy']*100:.1f}% -> {final['accuracy']*100:.1f}%")
+
+    # Save
+    save_path = './circuit_augmented_smollm2.pt'
+    print(f"\n[6] Saving to {save_path}...")
+    torch.save({
+        'model_state_dict': model.state_dict(),
+        'baseline_accuracy': baseline['accuracy'],
+        'final_accuracy': final['accuracy']
+    }, save_path)
+
+    print("\nDone!")