Upload train_circuit_interface.py with huggingface_hub

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!")