phanerozoic
/

8bit-threshold-computer

@@ -1,306 +0,0 @@
-"""
-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!")