Create trainer_v2_wormhole_routing.py

trainer_v2_wormhole_routing.py

@@ -0,0 +1,1415 @@
+"""
+Train DavidBeans V2: Wormhole Routing Architecture
+===================================================
+
+           ┌─────────────────┐
+           │   BEANS V2      │  "I learn where to look..."
+           │   (Wormhole ViT)│  
+           │   🌀 → 🌀 → 🌀   │  Learned sparse routing
+           └────────┬────────┘
+                    │
+                    ▼
+           ┌─────────────────┐
+           │   DAVID         │  "I know the crystals..."
+           │   (Classifier)  │
+           │   💎 → 💎 → 💎   │  Multi-scale projection
+           └────────┬────────┘
+                    │
+                    ▼
+              [Prediction]
+
+Key findings from wormhole experiments:
+1. When routing IS the task, routing learns structure
+2. Auxiliary losses can be gamed - removed in V2
+3. Gradient flow through router is critical - verified
+4. Cross-contrastive aligns patch↔scale features
+
+Author: AbstractPhil
+Date: November 29, 2025
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR
+from tqdm.auto import tqdm
+import time
+import math
+from pathlib import Path
+from typing import Dict, Optional, Tuple, List, Union
+from dataclasses import dataclass, field
+import json
+from datetime import datetime
+import os
+import shutil
+
+from google.colab import userdata
+
+os.environ['HF_TOKEN'] = userdata.get('HF_TOKEN')
+HF_TOKEN = userdata.get('HF_TOKEN')
+
+try:
+    from google.colab import userdata
+    HF_TOKEN = userdata.get('HF_TOKEN')
+    os.environ['HF_TOKEN'] = HF_TOKEN
+except:
+    pass
+
+# Import both model versions
+from geofractal.model.david_beans.model import DavidBeans, DavidBeansConfig
+from geofractal.model.david_beans.model_v2 import DavidBeansV2, DavidBeansV2Config
+
+# HuggingFace Hub integration
+try:
+    from huggingface_hub import HfApi, create_repo, upload_folder
+    HF_HUB_AVAILABLE = True
+except ImportError:
+    HF_HUB_AVAILABLE = False
+    print("  [!] huggingface_hub not installed. Run: pip install huggingface_hub")
+
+# Safetensors support
+try:
+    from safetensors.torch import save_file as save_safetensors
+    SAFETENSORS_AVAILABLE = True
+except ImportError:
+    SAFETENSORS_AVAILABLE = False
+
+# TensorBoard support
+try:
+    from torch.utils.tensorboard import SummaryWriter
+    TENSORBOARD_AVAILABLE = True
+except ImportError:
+    TENSORBOARD_AVAILABLE = False
+    print("  [!] tensorboard not installed. Run: pip install tensorboard")
+
+import numpy as np
+
+# ============================================================================
+# TRAINING CONFIGURATION V2
+# ============================================================================
+
+@dataclass
+class TrainingConfigV2:
+    """Training configuration for DavidBeans V2 with wormhole routing."""
+    
+    # Run identification
+    run_name: str = "default"
+    run_number: Optional[int] = None
+    
+    # Model version
+    model_version: int = 2  # 1 = original, 2 = wormhole
+    
+    # Data
+    dataset: str = "cifar100"
+    image_size: int = 32
+    batch_size: int = 128
+    num_workers: int = 4
+    
+    # Training schedule
+    epochs: int = 200
+    warmup_epochs: int = 10
+    
+    # Optimizer
+    learning_rate: float = 3e-4
+    weight_decay: float = 0.05
+    betas: Tuple[float, float] = (0.9, 0.999)
+    
+    # Learning rate schedule
+    scheduler: str = "cosine"
+    min_lr: float = 1e-6
+    
+    # Loss weights (based on experimental findings)
+    ce_weight: float = 1.0
+    contrast_weight: float = 0.5
+    # NOTE: No auxiliary routing loss - routing learns from task pressure
+    
+    # Regularization
+    gradient_clip: float = 1.0
+    label_smoothing: float = 0.1
+    
+    # Augmentation
+    use_augmentation: bool = True
+    mixup_alpha: float = 0.2
+    cutmix_alpha: float = 1.0
+    
+    # Checkpointing
+    save_interval: int = 10
+    output_dir: str = "./checkpoints"
+    resume_from: Optional[str] = None
+    
+    # TensorBoard
+    use_tensorboard: bool = True
+    log_interval: int = 50
+    log_routing: bool = True  # Log routing patterns
+    
+    # HuggingFace Hub
+    push_to_hub: bool = False
+    hub_repo_id: str = "AbstractPhil/geovit-david-beans"
+    hub_private: bool = False
+    
+    # Device
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    def to_dict(self) -> Dict:
+        return {k: v for k, v in self.__dict__.items()}
+
+
+# ============================================================================
+# ROUTING METRICS
+# ============================================================================
+
+class RoutingMetrics:
+    """Track and analyze wormhole routing patterns."""
+    
+    def __init__(self):
+        self.reset()
+    
+    def reset(self):
+        self.route_entropies = []
+        self.route_diversities = []
+        self.grad_norms = {'query': [], 'key': []}
+    
+    @torch.no_grad()
+    def compute_route_entropy(self, soft_routes: torch.Tensor) -> float:
+        """Compute average entropy of routing distributions."""
+        # soft_routes: [B, P, K] or [B, T, K]
+        # Higher entropy = more diverse routing
+        eps = 1e-8
+        entropy = -(soft_routes * (soft_routes + eps).log()).sum(dim=-1)
+        return entropy.mean().item()
+    
+    @torch.no_grad()
+    def compute_route_diversity(self, routes: torch.Tensor, num_positions: int) -> float:
+        """Compute how many unique destinations are used."""
+        # routes: [B, P, K] indices
+        unique_per_sample = []
+        for b in range(routes.shape[0]):
+            unique = routes[b].unique().numel()
+            unique_per_sample.append(unique / num_positions)
+        return sum(unique_per_sample) / len(unique_per_sample)
+    
+    def update_from_routing_info(self, routing_info: List[Dict], model: nn.Module):
+        """Extract metrics from routing info returned by V2 model."""
+        if not routing_info:
+            return
+        
+        for layer_info in routing_info:
+            # Attention routing
+            if layer_info.get('attention'):
+                attn = layer_info['attention']
+                if attn.get('weights') is not None:
+                    entropy = self.compute_route_entropy(attn['weights'])
+                    self.route_entropies.append(entropy)
+                if attn.get('routes') is not None:
+                    P = attn['routes'].shape[1]
+                    diversity = self.compute_route_diversity(attn['routes'], P)
+                    self.route_diversities.append(diversity)
+            
+            # Expert routing
+            if layer_info.get('expert'):
+                exp = layer_info['expert']
+                if exp.get('weights') is not None:
+                    entropy = self.compute_route_entropy(exp['weights'])
+                    self.route_entropies.append(entropy)
+    
+    def update_grad_norms(self, model: nn.Module):
+        """Track gradient norms through router projections."""
+        for name, param in model.named_parameters():
+            if param.grad is not None:
+                if 'query_proj' in name and 'weight' in name:
+                    self.grad_norms['query'].append(param.grad.norm().item())
+                elif 'key_proj' in name and 'weight' in name:
+                    self.grad_norms['key'].append(param.grad.norm().item())
+    
+    def get_summary(self) -> Dict[str, float]:
+        """Get summary statistics."""
+        summary = {}
+        
+        if self.route_entropies:
+            summary['route_entropy'] = sum(self.route_entropies) / len(self.route_entropies)
+        if self.route_diversities:
+            summary['route_diversity'] = sum(self.route_diversities) / len(self.route_diversities)
+        if self.grad_norms['query']:
+            summary['grad_query'] = sum(self.grad_norms['query']) / len(self.grad_norms['query'])
+        if self.grad_norms['key']:
+            summary['grad_key'] = sum(self.grad_norms['key']) / len(self.grad_norms['key'])
+        
+        return summary
+
+
+# ============================================================================
+# DATA LOADING (unchanged from V1)
+# ============================================================================
+
+def get_dataloaders(config: TrainingConfigV2) -> Tuple[DataLoader, DataLoader, int]:
+    """Get train and test dataloaders."""
+    
+    try:
+        import torchvision
+        import torchvision.transforms as T
+        
+        if config.dataset == "cifar10":
+            mean, std = (0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)
+            num_classes = 10
+            DatasetClass = torchvision.datasets.CIFAR10
+        elif config.dataset == "cifar100":
+            mean, std = (0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)
+            num_classes = 100
+            DatasetClass = torchvision.datasets.CIFAR100
+        else:
+            raise ValueError(f"Unknown dataset: {config.dataset}")
+        
+        if config.use_augmentation:
+            train_transform = T.Compose([
+                T.RandomCrop(32, padding=4),
+                T.RandomHorizontalFlip(),
+                T.AutoAugment(T.AutoAugmentPolicy.CIFAR10),
+                T.ToTensor(),
+                T.Normalize(mean, std)
+            ])
+        else:
+            train_transform = T.Compose([
+                T.ToTensor(),
+                T.Normalize(mean, std)
+            ])
+        
+        test_transform = T.Compose([
+            T.ToTensor(),
+            T.Normalize(mean, std)
+        ])
+        
+        train_dataset = DatasetClass(
+            root='./data', train=True, download=True, transform=train_transform
+        )
+        test_dataset = DatasetClass(
+            root='./data', train=False, download=True, transform=test_transform
+        )
+        
+        train_loader = DataLoader(
+            train_dataset,
+            batch_size=config.batch_size,
+            shuffle=True,
+            num_workers=config.num_workers,
+            pin_memory=True,
+            persistent_workers=config.num_workers > 0,
+            drop_last=True
+        )
+        test_loader = DataLoader(
+            test_dataset,
+            batch_size=config.batch_size,
+            shuffle=False,
+            num_workers=config.num_workers,
+            pin_memory=True,
+            persistent_workers=config.num_workers > 0
+        )
+        
+        return train_loader, test_loader, num_classes
+        
+    except ImportError:
+        print("  [!] torchvision not available, using synthetic data")
+        return get_synthetic_dataloaders(config)
+
+
+def get_synthetic_dataloaders(config: TrainingConfigV2) -> Tuple[DataLoader, DataLoader, int]:
+    """Fallback synthetic data for testing."""
+    
+    class SyntheticDataset(torch.utils.data.Dataset):
+        def __init__(self, size: int, image_size: int, num_classes: int):
+            self.size = size
+            self.image_size = image_size
+            self.num_classes = num_classes
+            
+        def __len__(self):
+            return self.size
+            
+        def __getitem__(self, idx):
+            x = torch.randn(3, self.image_size, self.image_size)
+            y = idx % self.num_classes
+            return x, y
+    
+    num_classes = 100 if config.dataset == "cifar100" else 10
+    train_dataset = SyntheticDataset(5000, config.image_size, num_classes)
+    test_dataset = SyntheticDataset(1000, config.image_size, num_classes)
+    
+    train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)
+    test_loader = DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False)
+    
+    return train_loader, test_loader, num_classes
+
+
+# ============================================================================
+# MIXUP / CUTMIX AUGMENTATION
+# ============================================================================
+
+def mixup_data(x: torch.Tensor, y: torch.Tensor, alpha: float = 0.2):
+    """Mixup augmentation."""
+    if alpha > 0:
+        lam = torch.distributions.Beta(alpha, alpha).sample().item()
+    else:
+        lam = 1.0
+    
+    batch_size = x.size(0)
+    index = torch.randperm(batch_size, device=x.device)
+    
+    mixed_x = lam * x + (1 - lam) * x[index]
+    y_a, y_b = y, y[index]
+    
+    return mixed_x, y_a, y_b, lam
+
+
+def cutmix_data(x: torch.Tensor, y: torch.Tensor, alpha: float = 1.0):
+    """CutMix augmentation."""
+    if alpha > 0:
+        lam = torch.distributions.Beta(alpha, alpha).sample().item()
+    else:
+        lam = 1.0
+    
+    batch_size = x.size(0)
+    index = torch.randperm(batch_size, device=x.device)
+    
+    _, _, H, W = x.shape
+    
+    cut_ratio = math.sqrt(1 - lam)
+    cut_h = int(H * cut_ratio)
+    cut_w = int(W * cut_ratio)
+    
+    cx = torch.randint(0, H, (1,)).item()
+    cy = torch.randint(0, W, (1,)).item()
+    
+    x1 = max(0, cx - cut_h // 2)
+    x2 = min(H, cx + cut_h // 2)
+    y1 = max(0, cy - cut_w // 2)
+    y2 = min(W, cy + cut_w // 2)
+    
+    mixed_x = x.clone()
+    mixed_x[:, :, x1:x2, y1:y2] = x[index, :, x1:x2, y1:y2]
+    
+    lam = 1 - ((x2 - x1) * (y2 - y1)) / (H * W)
+    
+    y_a, y_b = y, y[index]
+    
+    return mixed_x, y_a, y_b, lam
+
+
+# ============================================================================
+# METRICS TRACKER
+# ============================================================================
+
+class MetricsTracker:
+    """Track training metrics with EMA smoothing."""
+    
+    def __init__(self, ema_decay: float = 0.9):
+        self.ema_decay = ema_decay
+        self.metrics = {}
+        self.ema_metrics = {}
+        self.history = {}
+    
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            
+            if k not in self.metrics:
+                self.metrics[k] = []
+                self.ema_metrics[k] = v
+                self.history[k] = []
+            
+            self.metrics[k].append(v)
+            self.ema_metrics[k] = self.ema_decay * self.ema_metrics[k] + (1 - self.ema_decay) * v
+    
+    def get_ema(self, key: str) -> float:
+        return self.ema_metrics.get(key, 0.0)
+    
+    def get_epoch_mean(self, key: str) -> float:
+        values = self.metrics.get(key, [])
+        return sum(values) / len(values) if values else 0.0
+    
+    def end_epoch(self):
+        for k, v in self.metrics.items():
+            if v:
+                self.history[k].append(sum(v) / len(v))
+        self.metrics = {k: [] for k in self.metrics}
+    
+    def get_history(self) -> Dict:
+        return self.history
+
+
+# ============================================================================
+# CHECKPOINT UTILITIES
+# ============================================================================
+
+def find_latest_checkpoint(output_dir: Path) -> Optional[Path]:
+    """Find the most recent checkpoint in output directory."""
+    checkpoints = list(output_dir.glob("checkpoint_epoch_*.pt"))
+    
+    if not checkpoints:
+        best_model = output_dir / "best_model.pt"
+        if best_model.exists():
+            return best_model
+        return None
+    
+    def get_epoch(p):
+        try:
+            return int(p.stem.split("_")[-1])
+        except:
+            return 0
+    
+    checkpoints.sort(key=get_epoch, reverse=True)
+    return checkpoints[0]
+
+
+def get_next_run_number(base_dir: Path) -> int:
+    """Get the next run number by scanning existing run directories."""
+    if not base_dir.exists():
+        return 1
+    
+    max_num = 0
+    for d in base_dir.iterdir():
+        if d.is_dir() and d.name.startswith("run_"):
+            try:
+                num = int(d.name.split("_")[1])
+                max_num = max(max_num, num)
+            except (IndexError, ValueError):
+                continue
+    
+    return max_num + 1
+
+
+def generate_run_dir_name(run_number: int, run_name: str, version: int = 2) -> str:
+    """Generate a run directory name."""
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    safe_name = "".join(c if c.isalnum() or c == "_" else "_" for c in run_name.lower())
+    safe_name = "_".join(filter(None, safe_name.split("_")))
+    return f"run_{run_number:03d}_v{version}_{safe_name}_{timestamp}"
+
+
+def find_latest_run_dir(base_dir: Path) -> Optional[Path]:
+    """Find the most recent run directory."""
+    if not base_dir.exists():
+        return None
+    
+    run_dirs = [d for d in base_dir.iterdir() if d.is_dir() and d.name.startswith("run_")]
+    
+    if not run_dirs:
+        return None
+    
+    run_dirs.sort(key=lambda d: d.stat().st_mtime, reverse=True)
+    return run_dirs[0]
+
+
+def load_checkpoint(
+    checkpoint_path: Path,
+    model: nn.Module,
+    optimizer: Optional[torch.optim.Optimizer] = None,
+    device: str = "cuda"
+) -> Tuple[int, float]:
+    """Load checkpoint and return (start_epoch, best_acc)."""
+    print(f"\n📂 Loading checkpoint: {checkpoint_path}")
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    
+    model.load_state_dict(checkpoint['model_state_dict'])
+    print(f"  ✓ Loaded model weights")
+    
+    if optimizer is not None and 'optimizer_state_dict' in checkpoint:
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        print(f"  ✓ Loaded optimizer state")
+    
+    epoch = checkpoint.get('epoch', 0)
+    best_acc = checkpoint.get('best_acc', 0.0)
+    
+    print(f"  ✓ Resuming from epoch {epoch + 1}, best_acc={best_acc:.2f}%")
+    
+    return epoch + 1, best_acc
+
+
+# ============================================================================
+# HUGGINGFACE HUB INTEGRATION
+# ============================================================================
+
+def generate_run_readme(
+    model_config: Union[DavidBeansConfig, DavidBeansV2Config],
+    train_config: TrainingConfigV2,
+    best_acc: float,
+    run_dir_name: str
+) -> str:
+    """Generate README for a specific run."""
+    
+    scales_str = ", ".join([str(s) for s in model_config.scales])
+    
+    # V2 specific info
+    if isinstance(model_config, DavidBeansV2Config):
+        routing_info = f"""
+## Wormhole Routing (V2)
+| Parameter | Value |
+|-----------|-------|
+| Mode | {model_config.wormhole_mode} |
+| Wormholes/Position | {model_config.num_wormholes} |
+| Temperature | {model_config.wormhole_temperature} |
+| Tiles | {model_config.num_tiles} |
+| Tile Wormholes | {model_config.tile_wormholes} |
+"""
+    else:
+        routing_info = """
+## Routing (V1)
+| Parameter | Value |
+|-----------|-------|
+| k_neighbors | {model_config.k_neighbors} |
+| Cantor Weight | {model_config.cantor_weight} |
+"""
+    
+    return f"""# Run: {run_dir_name}
+
+## Results
+- **Best Accuracy**: {best_acc:.2f}%
+- **Dataset**: {train_config.dataset}
+- **Epochs**: {train_config.epochs}
+- **Model Version**: V{train_config.model_version}
+
+## Model Config
+| Parameter | Value |
+|-----------|-------|
+| Dim | {model_config.dim} |
+| Layers | {model_config.num_layers} |
+| Heads | {model_config.num_heads} |
+| Scales | [{scales_str}] |
+{routing_info}
+
+## Training Config
+| Parameter | Value |
+|-----------|-------|
+| Learning Rate | {train_config.learning_rate} |
+| Weight Decay | {train_config.weight_decay} |
+| Batch Size | {train_config.batch_size} |
+| CE Weight | {train_config.ce_weight} |
+| Contrast Weight | {train_config.contrast_weight} |
+
+## Key Findings Applied
+- Routing learns from task pressure (no auxiliary routing losses)
+- Gradients verified to flow through router
+- Cross-contrastive aligns patch↔scale features
+"""
+
+
+def prepare_run_for_hub(
+    model: nn.Module,
+    model_config: Union[DavidBeansConfig, DavidBeansV2Config],
+    train_config: TrainingConfigV2,
+    best_acc: float,
+    output_dir: Path,
+    run_dir_name: str,
+    training_history: Optional[Dict] = None
+) -> Path:
+    """Prepare run files for upload to HuggingFace Hub."""
+    
+    hub_dir = output_dir / "hub_upload"
+    run_hub_dir = hub_dir / "weights" / run_dir_name
+    run_hub_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Save best model weights
+    state_dict = {k: v.clone() for k, v in model.state_dict().items()}
+    
+    if SAFETENSORS_AVAILABLE:
+        try:
+            save_safetensors(state_dict, run_hub_dir / "best.safetensors")
+            print(f"  ✓ Saved best.safetensors")
+        except Exception as e:
+            print(f"  [!] Safetensors failed ({e}), using pytorch format")
+            torch.save(state_dict, run_hub_dir / "best.pt")
+    else:
+        torch.save(state_dict, run_hub_dir / "best.pt")
+    
+    # Save model config
+    config_dict = {
+        "architecture": f"DavidBeans_V{train_config.model_version}",
+        "model_type": "david_beans_v2" if train_config.model_version == 2 else "david_beans",
+        **model_config.__dict__
+    }
+    with open(run_hub_dir / "config.json", "w") as f:
+        json.dump(config_dict, f, indent=2, default=str)
+    
+    # Save training config
+    with open(run_hub_dir / "training_config.json", "w") as f:
+        json.dump(train_config.to_dict(), f, indent=2, default=str)
+    
+    # Generate README
+    run_readme = generate_run_readme(model_config, train_config, best_acc, run_dir_name)
+    with open(run_hub_dir / "README.md", "w") as f:
+        f.write(run_readme)
+    
+    # Save training history
+    if training_history:
+        with open(run_hub_dir / "training_history.json", "w") as f:
+            json.dump(training_history, f, indent=2)
+    
+    # Copy TensorBoard logs
+    tb_dir = output_dir / "tensorboard"
+    if tb_dir.exists():
+        hub_tb_dir = run_hub_dir / "tensorboard"
+        if hub_tb_dir.exists():
+            shutil.rmtree(hub_tb_dir)
+        shutil.copytree(tb_dir, hub_tb_dir)
+    
+    return hub_dir
+
+
+def push_run_to_hub(
+    hub_dir: Path,
+    repo_id: str,
+    run_dir_name: str,
+    private: bool = False,
+    commit_message: Optional[str] = None
+) -> str:
+    """Push run files to HuggingFace Hub."""
+    
+    if not HF_HUB_AVAILABLE:
+        raise RuntimeError("huggingface_hub not installed")
+    
+    api = HfApi()
+    
+    try:
+        create_repo(repo_id, private=private, exist_ok=True)
+    except Exception as e:
+        print(f"  [!] Repo creation note: {e}")
+    
+    run_upload_dir = hub_dir / "weights" / run_dir_name
+    
+    if commit_message is None:
+        commit_message = f"Update {run_dir_name} - {datetime.now().strftime('%Y-%m-%d %H:%M')}"
+    
+    url = upload_folder(
+        folder_path=str(run_upload_dir),
+        repo_id=repo_id,
+        path_in_repo=f"weights/{run_dir_name}",
+        commit_message=commit_message
+    )
+    
+    return url
+
+
+# ============================================================================
+# TRAINING LOOP V2
+# ============================================================================
+
+def train_epoch_v2(
+    model: nn.Module,
+    train_loader: DataLoader,
+    optimizer: torch.optim.Optimizer,
+    scheduler: Optional[torch.optim.lr_scheduler._LRScheduler],
+    config: TrainingConfigV2,
+    epoch: int,
+    tracker: MetricsTracker,
+    routing_metrics: RoutingMetrics,
+    writer: Optional['SummaryWriter'] = None
+) -> Dict[str, float]:
+    """Train for one epoch with V2 routing metrics."""
+    
+    model.train()
+    device = config.device
+    is_v2 = config.model_version == 2
+    
+    total_loss = 0.0
+    total_correct = 0
+    total_samples = 0
+    global_step = epoch * len(train_loader)
+    
+    routing_metrics.reset()
+    
+    pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}", leave=True)
+    
+    for batch_idx, (images, targets) in enumerate(pbar):
+        images = images.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+        
+        # Apply mixup/cutmix
+        use_mixup = config.use_augmentation and config.mixup_alpha > 0
+        use_cutmix = config.use_augmentation and config.cutmix_alpha > 0
+        
+        mixed = False
+        if use_mixup or use_cutmix:
+            r = torch.rand(1).item()
+            if r < 0.5:
+                pass
+            elif r < 0.75 and use_mixup:
+                images, targets_a, targets_b, lam = mixup_data(images, targets, config.mixup_alpha)
+                mixed = True
+            elif use_cutmix:
+                images, targets_a, targets_b, lam = cutmix_data(images, targets, config.cutmix_alpha)
+                mixed = True
+        
+        # Forward pass
+        if is_v2:
+            result = model(
+                images, 
+                targets=targets, 
+                return_loss=True,
+                return_routing=(batch_idx % 10 == 0)  # Sample routing every 10 batches
+            )
+        else:
+            result = model(images, targets=targets, return_loss=True)
+        
+        losses = result['losses']
+        
+        # Handle mixup CE loss
+        if mixed:
+            logits = result['logits']
+            ce_loss = lam * F.cross_entropy(logits, targets_a, label_smoothing=config.label_smoothing) + \
+                      (1 - lam) * F.cross_entropy(logits, targets_b, label_smoothing=config.label_smoothing)
+            losses['ce'] = ce_loss
+        
+        # Compute total loss (NO auxiliary routing loss - key finding!)
+        loss = (
+            config.ce_weight * losses['ce'] +
+            config.contrast_weight * losses.get('contrast', torch.tensor(0.0, device=device))
+        )
+        
+        # Add scale CE losses
+        for scale in model.config.scales:
+            scale_ce = losses.get(f'ce_{scale}', 0.0)
+            if isinstance(scale_ce, torch.Tensor):
+                loss = loss + 0.1 * scale_ce
+        
+        # Backward pass
+        optimizer.zero_grad()
+        loss.backward()
+        
+        # Track routing gradient norms (verify gradients flow!)
+        if is_v2:
+            routing_metrics.update_grad_norms(model)
+        
+        if config.gradient_clip > 0:
+            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)
+        else:
+            grad_norm = 0.0
+        
+        optimizer.step()
+        
+        if scheduler is not None and config.scheduler == "onecycle":
+            scheduler.step()
+        
+        # Update routing metrics from forward pass
+        if is_v2 and result.get('routing'):
+            routing_metrics.update_from_routing_info(result['routing'], model)
+        
+        # Compute accuracy
+        with torch.no_grad():
+            logits = result['logits']
+            preds = logits.argmax(dim=-1)
+            
+            if mixed:
+                correct = (lam * (preds == targets_a).float() + 
+                          (1 - lam) * (preds == targets_b).float()).sum()
+            else:
+                correct = (preds == targets).sum()
+            
+            total_correct += correct.item()
+            total_samples += targets.size(0)
+            total_loss += loss.item()
+        
+        # Track metrics
+        def to_float(v):
+            return v.item() if isinstance(v, torch.Tensor) else float(v)
+        
+        contrast_loss = to_float(losses.get('contrast', 0.0))
+        current_lr = optimizer.param_groups[0]['lr']
+        
+        tracker.update(
+            loss=loss.item(),
+            ce=losses['ce'].item(),
+            contrast=contrast_loss,
+            lr=current_lr
+        )
+        
+        # TensorBoard logging
+        if writer is not None and (batch_idx + 1) % config.log_interval == 0:
+            step = global_step + batch_idx
+            writer.add_scalar('train/loss_total', loss.item(), step)
+            writer.add_scalar('train/loss_ce', losses['ce'].item(), step)
+            writer.add_scalar('train/loss_contrast', contrast_loss, step)
+            writer.add_scalar('train/learning_rate', current_lr, step)
+            writer.add_scalar('train/grad_norm', to_float(grad_norm), step)
+            
+            # Log routing metrics for V2
+            if is_v2 and config.log_routing:
+                routing_summary = routing_metrics.get_summary()
+                for k, v in routing_summary.items():
+                    writer.add_scalar(f'routing/{k}', v, step)
+        
+        # Progress bar
+        routing_summary = routing_metrics.get_summary()
+        postfix = {
+            'loss': f"{tracker.get_ema('loss'):.3f}",
+            'acc': f"{100.0 * total_correct / total_samples:.1f}%",
+        }
+        if is_v2 and 'grad_query' in routing_summary:
+            postfix['∇q'] = f"{routing_summary['grad_query']:.2f}"
+        if 'route_entropy' in routing_summary:
+            postfix['H'] = f"{routing_summary['route_entropy']:.2f}"
+        
+        pbar.set_postfix(postfix)
+    
+    if scheduler is not None and config.scheduler == "cosine":
+        scheduler.step()
+    
+    return {
+        'loss': total_loss / len(train_loader),
+        'acc': 100.0 * total_correct / total_samples,
+        **routing_metrics.get_summary()
+    }
+
+
+@torch.no_grad()
+def evaluate_v2(
+    model: nn.Module,
+    test_loader: DataLoader,
+    config: TrainingConfigV2
+) -> Dict[str, float]:
+    """Evaluate on test set."""
+    
+    model.eval()
+    device = config.device
+    
+    total_loss = 0.0
+    total_correct = 0
+    total_samples = 0
+    scale_correct = {s: 0 for s in model.config.scales}
+    
+    for images, targets in test_loader:
+        images = images.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+        
+        result = model(images, targets=targets, return_loss=True)
+        
+        logits = result['logits']
+        losses = result['losses']
+        
+        loss = losses['total']
+        preds = logits.argmax(dim=-1)
+        
+        total_loss += loss.item() * targets.size(0)
+        total_correct += (preds == targets).sum().item()
+        total_samples += targets.size(0)
+        
+        for i, scale in enumerate(model.config.scales):
+            scale_logits = result['scale_logits'][i]
+            scale_preds = scale_logits.argmax(dim=-1)
+            scale_correct[scale] += (scale_preds == targets).sum().item()
+    
+    metrics = {
+        'loss': total_loss / total_samples,
+        'acc': 100.0 * total_correct / total_samples
+    }
+    
+    for scale, correct in scale_correct.items():
+        metrics[f'acc_{scale}'] = 100.0 * correct / total_samples
+    
+    return metrics
+
+
+# ============================================================================
+# MAIN TRAINING FUNCTION V2
+# ============================================================================
+
+def train_david_beans_v2(
+    model_config: Optional[Union[DavidBeansConfig, DavidBeansV2Config]] = None,
+    train_config: Optional[TrainingConfigV2] = None
+):
+    """Main training function for DavidBeans V1 or V2."""
+    
+    print("=" * 70)
+    print("  DAVID-BEANS V2 TRAINING: Wormhole Routing")
+    print("=" * 70)
+    print()
+    print("       🌀 WORMHOLES: Learned sparse routing")
+    print("       💎 CRYSTALS:  Multi-scale projection")
+    print()
+    print("  Key insight: When routing IS the task, routing learns structure")
+    print()
+    print("=" * 70)
+    
+    if train_config is None:
+        train_config = TrainingConfigV2()
+    
+    base_output_dir = Path(train_config.output_dir)
+    base_output_dir.mkdir(parents=True, exist_ok=True)
+    
+    # =========================================================================
+    # FIXED: Proper checkpoint resolution
+    # =========================================================================
+    checkpoint_path = None
+    run_dir = None
+    run_dir_name = None
+    
+    if train_config.resume_from:
+        resume_path = Path(train_config.resume_from)
+        
+        # Case 1: Direct absolute/relative file path
+        if resume_path.is_file():
+            checkpoint_path = resume_path
+            run_dir = checkpoint_path.parent
+            run_dir_name = run_dir.name
+            print(f"\n📂 Found checkpoint file: {checkpoint_path.name}")
+        
+        # Case 2: Directory path - find best/latest checkpoint inside
+        elif resume_path.is_dir():
+            checkpoint_path = find_latest_checkpoint(resume_path)
+            if checkpoint_path:
+                run_dir = resume_path
+                run_dir_name = resume_path.name
+                print(f"\n📂 Found checkpoint in dir: {checkpoint_path.name}")
+        
+        # Case 3: Try as path relative to base_output_dir
+        else:
+            # Try as subdirectory name
+            possible_dir = base_output_dir / train_config.resume_from
+            if possible_dir.is_dir():
+                checkpoint_path = find_latest_checkpoint(possible_dir)
+                if checkpoint_path:
+                    run_dir = possible_dir
+                    run_dir_name = possible_dir.name
+                    print(f"\n📂 Found checkpoint in: {run_dir_name}")
+            
+            # Try as relative file path
+            if checkpoint_path is None:
+                possible_file = base_output_dir / train_config.resume_from
+                if possible_file.is_file():
+                    checkpoint_path = possible_file
+                    run_dir = checkpoint_path.parent
+                    run_dir_name = run_dir.name
+                    print(f"\n📂 Found checkpoint: {checkpoint_path.name}")
+        
+        # Report if not found
+        if checkpoint_path is None:
+            print(f"\n  [!] Could not find checkpoint: {train_config.resume_from}")
+            print(f"  [!] Checked:")
+            print(f"      - As file: {resume_path}")
+            print(f"      - As dir: {resume_path}")
+            print(f"      - Under {base_output_dir}")
+            print(f"  [!] Starting fresh run instead")
+        else:
+            print(f"  ✓ Will resume from: {checkpoint_path}")
+    
+    # Create new run directory if not resuming
+    if run_dir is None:
+        run_number = train_config.run_number or get_next_run_number(base_output_dir)
+        run_dir_name = generate_run_dir_name(run_number, train_config.run_name, train_config.model_version)
+        run_dir = base_output_dir / run_dir_name
+        run_dir.mkdir(parents=True, exist_ok=True)
+        print(f"\n📁 New run: {run_dir_name}")
+    else:
+        print(f"\n📁 Resuming run: {run_dir_name}")
+    
+    output_dir = run_dir
+    
+    # =========================================================================
+    # Model config - load from checkpoint if resuming, else use provided/default
+    # =========================================================================
+    if checkpoint_path and checkpoint_path.exists() and model_config is None:
+        # Try to load config from checkpoint
+        try:
+            ckpt = torch.load(checkpoint_path, map_location='cpu')
+            if 'model_config' in ckpt:
+                saved_config = ckpt['model_config']
+                print(f"  ✓ Loading model config from checkpoint")
+                if train_config.model_version == 2:
+                    model_config = DavidBeansV2Config(**saved_config)
+                else:
+                    model_config = DavidBeansConfig(**saved_config)
+        except Exception as e:
+            print(f"  [!] Could not load config from checkpoint: {e}")
+    
+    # Create default config if still None
+    if model_config is None:
+        if train_config.model_version == 2:
+            model_config = DavidBeansV2Config(
+                image_size=train_config.image_size,
+                patch_size=4,
+                dim=512,
+                num_layers=4,
+                num_heads=8,
+                num_wormholes=8,
+                wormhole_temperature=0.1,
+                wormhole_mode="hybrid",
+                num_tiles=16,
+                tile_wormholes=4,
+                scales=[64, 128, 256, 384, 512],
+                num_classes=100,
+                contrast_weight=train_config.contrast_weight,
+                dropout=0.1
+            )
+        else:
+            model_config = DavidBeansConfig(
+                image_size=train_config.image_size,
+                patch_size=4,
+                dim=512,
+                num_layers=4,
+                num_heads=8,
+                num_experts=5,
+                k_neighbors=16,
+                cantor_weight=0.3,
+                scales=[64, 128, 256, 384, 512],
+                num_classes=100,
+                dropout=0.1
+            )
+    
+    device = train_config.device
+    print(f"\nDevice: {device}")
+    print(f"Model version: V{train_config.model_version}")
+    
+    # Data
+    print("\nLoading data...")
+    train_loader, test_loader, num_classes = get_dataloaders(train_config)
+    print(f"  Dataset: {train_config.dataset}")
+    print(f"  Train: {len(train_loader.dataset)}, Test: {len(test_loader.dataset)}")
+    print(f"  Classes: {num_classes}")
+    
+    model_config.num_classes = num_classes
+    
+    # Model
+    print("\nBuilding model...")
+    if train_config.model_version == 2:
+        model = DavidBeansV2(model_config)
+    else:
+        model = DavidBeans(model_config)
+    
+    model = model.to(device)
+    print(f"\n{model}")
+    
+    num_params = sum(p.numel() for p in model.parameters())
+    print(f"\nParameters: {num_params:,}")
+    
+    # Optimizer
+    print("\nSetting up optimizer...")
+    
+    decay_params = []
+    no_decay_params = []
+    
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if 'bias' in name or 'norm' in name or 'embedding' in name:
+            no_decay_params.append(param)
+        else:
+            decay_params.append(param)
+    
+    optimizer = AdamW([
+        {'params': decay_params, 'weight_decay': train_config.weight_decay},
+        {'params': no_decay_params, 'weight_decay': 0.0}
+    ], lr=train_config.learning_rate, betas=train_config.betas)
+    
+    if train_config.scheduler == "cosine":
+        scheduler = CosineAnnealingLR(
+            optimizer,
+            T_max=train_config.epochs - train_config.warmup_epochs,
+            eta_min=train_config.min_lr
+        )
+    elif train_config.scheduler == "onecycle":
+        scheduler = OneCycleLR(
+            optimizer,
+            max_lr=train_config.learning_rate,
+            epochs=train_config.epochs,
+            steps_per_epoch=len(train_loader),
+            pct_start=train_config.warmup_epochs / train_config.epochs
+        )
+    else:
+        scheduler = None
+    
+    print(f"  Optimizer: AdamW (lr={train_config.learning_rate}, wd={train_config.weight_decay})")
+    print(f"  Scheduler: {train_config.scheduler}")
+    
+    tracker = MetricsTracker()
+    routing_metrics = RoutingMetrics()
+    best_acc = 0.0
+    start_epoch = 0
+    
+    # =========================================================================
+    # Load checkpoint weights and optimizer state
+    # =========================================================================
+    if checkpoint_path and checkpoint_path.exists():
+        start_epoch, best_acc = load_checkpoint(checkpoint_path, model, optimizer, device)
+        
+        # Advance scheduler to correct position
+        if scheduler is not None and train_config.scheduler == "cosine":
+            for _ in range(start_epoch):
+                scheduler.step()
+            print(f"  ✓ Advanced scheduler to epoch {start_epoch}")
+    
+    # TensorBoard
+    writer = None
+    if train_config.use_tensorboard and TENSORBOARD_AVAILABLE:
+        tb_dir = output_dir / "tensorboard"
+        tb_dir.mkdir(parents=True, exist_ok=True)
+        writer = SummaryWriter(log_dir=str(tb_dir))
+        print(f"  TensorBoard: {tb_dir}")
+    
+    # Save configs
+    with open(output_dir / "config.json", "w") as f:
+        json.dump({**model_config.__dict__, "architecture": f"DavidBeans_V{train_config.model_version}"}, 
+                  f, indent=2, default=str)
+    with open(output_dir / "training_config.json", "w") as f:
+        json.dump(train_config.to_dict(), f, indent=2, default=str)
+    
+    # Training loop
+    print("\n" + "=" * 70)
+    print("  TRAINING")
+    print("=" * 70)
+    
+    for epoch in range(start_epoch, train_config.epochs):
+        epoch_start = time.time()
+        
+        # Warmup
+        if epoch < train_config.warmup_epochs and train_config.scheduler == "cosine":
+            warmup_lr = train_config.learning_rate * (epoch + 1) / train_config.warmup_epochs
+            for param_group in optimizer.param_groups:
+                param_group['lr'] = warmup_lr
+        
+        train_metrics = train_epoch_v2(
+            model, train_loader, optimizer, scheduler,
+            train_config, epoch, tracker, routing_metrics, writer
+        )
+        
+        test_metrics = evaluate_v2(model, test_loader, train_config)
+        
+        epoch_time = time.time() - epoch_start
+        
+        # TensorBoard
+        if writer is not None:
+            writer.add_scalar('epoch/train_loss', train_metrics['loss'], epoch)
+            writer.add_scalar('epoch/train_acc', train_metrics['acc'], epoch)
+            writer.add_scalar('epoch/test_loss', test_metrics['loss'], epoch)
+            writer.add_scalar('epoch/test_acc', test_metrics['acc'], epoch)
+            
+            for scale in model.config.scales:
+                writer.add_scalar(f'scales/acc_{scale}', test_metrics[f'acc_{scale}'], epoch)
+        
+        # Print summary - show ALL scales
+        scale_accs = " | ".join([f"{s}:{test_metrics[f'acc_{s}']:.1f}%" for s in model.config.scales])
+        star = "★" if test_metrics['acc'] > best_acc else ""
+        
+        routing_info = ""
+        if train_config.model_version == 2 and 'grad_query' in train_metrics:
+            routing_info = f" | ∇q:{train_metrics.get('grad_query', 0):.2f}"
+        
+        print(f"  → Train: {train_metrics['acc']:.1f}% | Test: {test_metrics['acc']:.1f}% | "
+              f"[{scale_accs}]{routing_info} | {epoch_time:.0f}s {star}")
+        
+        # Save best model
+        if test_metrics['acc'] > best_acc:
+            best_acc = test_metrics['acc']
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'best_acc': best_acc,
+                'model_config': model_config.__dict__,
+                'train_config': train_config.to_dict()
+            }, output_dir / "best_model.pt")
+        
+        # Periodic checkpoint
+        if (epoch + 1) % train_config.save_interval == 0:
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'best_acc': best_acc,
+                'model_config': model_config.__dict__,
+                'train_config': train_config.to_dict()
+            }, output_dir / f"checkpoint_epoch_{epoch + 1}.pt")
+            
+            # Upload to hub
+            if train_config.push_to_hub and HF_HUB_AVAILABLE:
+                try:
+                    hub_dir = prepare_run_for_hub(
+                        model=model,
+                        model_config=model_config,
+                        train_config=train_config,
+                        best_acc=best_acc,
+                        output_dir=output_dir,
+                        run_dir_name=run_dir_name,
+                        training_history=tracker.get_history()
+                    )
+                    push_run_to_hub(
+                        hub_dir=hub_dir,
+                        repo_id=train_config.hub_repo_id,
+                        run_dir_name=run_dir_name,
+                        commit_message=f"Epoch {epoch + 1} - {best_acc:.2f}% acc"
+                    )
+                    print(f"  📤 Uploaded to hub")
+                except Exception as e:
+                    print(f"  [!] Hub upload failed: {e}")
+        
+        tracker.end_epoch()
+    
+    # Final summary
+    print("\n" + "=" * 70)
+    print("  TRAINING COMPLETE")
+    print("=" * 70)
+    print(f"\n  Best Test Accuracy: {best_acc:.2f}%")
+    print(f"  Model saved to: {output_dir / 'best_model.pt'}")
+    
+    if writer is not None:
+        writer.close()
+    
+    return model, best_acc
+
+
+# ============================================================================
+# PRESETS
+# ============================================================================
+
+def train_cifar100_v2_wormhole(
+    run_name: str = "wormhole_base",
+    push_to_hub: bool = False,
+    resume: bool = False
+):
+    """CIFAR-100 with V2 wormhole routing."""
+    
+    model_config = DavidBeansV2Config(
+        image_size=32,
+        patch_size=2,
+        dim=512,
+        num_layers=4,
+        num_heads=16,
+        # Wormhole routing parameters
+        num_wormholes=16,
+        wormhole_temperature=0.1,
+        wormhole_mode="hybrid",
+        # Tessellation parameters  
+        num_tiles=16,
+        tile_wormholes=4,
+        # Crystal head
+        scales=[64, 128, 256, 512, 1024],
+        num_classes=100,
+        contrast_temperature=0.07,
+        contrast_weight=0.5,
+        dropout=0.1
+    )
+    
+    train_config = TrainingConfigV2(
+        run_name=run_name,
+        model_version=2,
+        dataset="cifar100",
+        epochs=300,
+        batch_size=512,
+        learning_rate=3e-4,
+        weight_decay=0.05,
+        warmup_epochs=15,
+        # Loss weights (no auxiliary routing loss!)
+        ce_weight=1.0,
+        contrast_weight=0.5,
+        # Augmentation
+        label_smoothing=0.1,
+        mixup_alpha=0.2,
+        cutmix_alpha=1.0,
+        # Output
+        output_dir="./checkpoints/cifar100_v2",    
+        resume_from=None, #"./checkpoints/cifar100_v2/run_002_v2_16patch_4tilewormholes_d768_4layer_20251130_045437/best_model.pt",
+        # Hub
+        push_to_hub=push_to_hub,
+        hub_repo_id="AbstractPhil/geovit-david-beans",
+        # Routing logging
+        log_routing=True
+    )
+    
+    return train_david_beans_v2(model_config, train_config)
+
+
+def train_cifar100_v1_baseline(
+    run_name: str = "v1_baseline",
+    push_to_hub: bool = False,
+    resume: bool = False
+):
+    """CIFAR-100 with V1 (fixed Cantor routing) for comparison."""
+    
+    model_config = DavidBeansConfig(
+        image_size=32,
+        patch_size=4,
+        dim=512,
+        num_layers=4,
+        num_heads=8,
+        num_experts=5,
+        k_neighbors=16,
+        cantor_weight=0.3,
+        scales=[64, 128, 256, 384, 512],
+        num_classes=100,
+        dropout=0.1
+    )
+    
+    train_config = TrainingConfigV2(
+        run_name=run_name,
+        model_version=1,
+        dataset="cifar100",
+        epochs=200,
+        batch_size=128,
+        learning_rate=3e-4,
+        weight_decay=0.05,
+        warmup_epochs=10,
+        ce_weight=1.0,
+        contrast_weight=0.5,
+        label_smoothing=0.1,
+        mixup_alpha=0.2,
+        cutmix_alpha=1.0,
+        output_dir="./checkpoints/cifar100_v1",
+        resume_from="latest" if resume else None,
+        push_to_hub=push_to_hub,
+        hub_repo_id="AbstractPhil/geovit-david-beans",
+        log_routing=False
+    )
+    
+    return train_david_beans_v2(model_config, train_config)
+
+
+# ============================================================================
+# MAIN
+# ============================================================================
+
+if __name__ == "__main__":
+    
+    # =====================================================
+    # CONFIGURATION
+    # =====================================================
+    
+    PRESET = "v2_wormhole"  # "v1_baseline", "v2_wormhole", "test"
+    RESUME = False
+    RUN_NAME = "5scale_2x2patch_4tilewormholes_d512_4layer"
+    PUSH_TO_HUB = True
+    
+    # =====================================================
+    # RUN
+    # =====================================================
+    
+    if PRESET == "test":
+        print("🧪 Quick test...")
+        model_config = DavidBeansV2Config(
+            image_size=32, patch_size=4, dim=128, num_layers=2,
+            num_heads=4, num_wormholes=4, num_tiles=8,
+            scales=[32, 64, 128], num_classes=10
+        )
+        train_config = TrainingConfigV2(
+            run_name="test", model_version=2,
+            epochs=2, batch_size=32,
+            use_augmentation=False, mixup_alpha=0.0, cutmix_alpha=0.0
+        )
+        model, acc = train_david_beans_v2(model_config, train_config)
+        
+    elif PRESET == "v1_baseline":
+        print("🫘💎 Training DavidBeans V1 (Cantor routing)...")
+        model, acc = train_cifar100_v1_baseline(
+            run_name=RUN_NAME,
+            push_to_hub=PUSH_TO_HUB,
+            resume=RESUME
+        )
+        
+    elif PRESET == "v2_wormhole":
+        print("💎 Training DavidBeans V2 (Wormhole routing)...")
+        model, acc = train_cifar100_v2_wormhole(
+            run_name=RUN_NAME,
+            push_to_hub=PUSH_TO_HUB,
+            resume=RESUME
+        )
+        
+    else:
+        raise ValueError(f"Unknown preset: {PRESET}")
+    
+    print(f"\n🎉 Done! Best accuracy: {acc:.2f}%")