src/fractus/growth.py

"""Fractus Growth Module — add experts and expand rank without retraining.

Two mechanisms for continuous checkpoint growth:

1. add_experts(model, n_new, data, device)
   Creates new MoE experts, pre-trains them via EDT Phase 1, extends the router.
   Old experts are untouched. The brain grows wider.

2. expand_rank(model, target_rank)
   Increases the Siren rank of ALL experts. Old rank columns preserved.
   New columns initialized to zero (contribute nothing until trained).
   The brain grows deeper.

3. save_grown_model(model, path, growth_log)
   Saves the model with a growth log documenting what was added.

Usage:
    from fractus.growth import FractusGrowth

    growth = FractusGrowth(model, tok, device)

    # Add 128 experts trained on Rust code
    growth.add_experts(n_new=128, data=rust_tokens, domain="rust")

    # Expand rank 64 → 128 for deeper capacity
    growth.expand_rank(target_rank=128)

    # Save the grown model
    growth.save("checkpoints/fractus_grown.pt")
"""
import os, sys, time, math
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class FractusGrowth:
    """Continuous growth manager for Fractus models.

    Allows adding experts and expanding rank without retraining the
    existing model. This is what makes Fractus a living brain.
    """

    def __init__(self, model, tokenizer=None, device="cpu"):
        self.model = model
        self.tok = tokenizer
        self.device = torch.device(device)
        self.growth_log = []

        # Record initial state.
        self.initial_experts = model.blocks[0].moe.n_experts if hasattr(model, 'blocks') else 0
        self.initial_rank = (
            model.blocks[0].moe.experts_w1[0].rank
            if hasattr(model, 'blocks') and hasattr(model.blocks[0].moe.experts_w1[0], 'rank')
            else 0
        )
        self.n_layers = len(model.blocks) if hasattr(model, 'blocks') else 0

    # ===================================================================
    # MECHANISM 1: ADD EXPERTS
    # ===================================================================

    def add_experts(self, n_new=128, data=None, domain="unknown",
                    steps_per_expert=2000, batch_size=64, lr=1e-3):
        """Add new MoE experts to every layer and pre-train them.

        The new experts are:
        1. Created with random LazyStructuredSiren weights
        2. Pre-trained on the provided data (or synthetic if no data)
        3. Added to the model's expert lists
        4. The router's Farey phases are extended

        OLD EXPERTS ARE NEVER TOUCHED.

        Args:
            n_new: number of new experts to add per layer.
            data: token IDs (1D tensor) for pre-training, or None for synthetic.
            domain: name of the domain (for logging).
            steps_per_expert: EDT Phase 1 steps per expert.
            batch_size: batch size for pre-training.
            lr: learning rate.
        """
        from fractus.nn.lazy_siren import LazyStructuredSirenLinear

        print(f"\n{'='*60}", flush=True)
        print(f"GROWTH: Adding {n_new} experts (domain: {domain})", flush=True)
        print(f"{'='*60}", flush=True)

        t0 = time.time()

        # Generate hidden states for pre-training.
        hidden_states = self._generate_hidden_states(data, n_samples=5000)

        for layer_idx in range(self.n_layers):
            moe = self.model.blocks[layer_idx].moe
            old_count = moe.n_experts
            d_model = moe.d_model
            d_ff = moe.d_ff
            rank = moe.experts_w1[0].rank if hasattr(moe.experts_w1[0], 'rank') else 16

            # Create new experts.
            for i in range(n_new):
                new_w1 = LazyStructuredSirenLinear(d_model, d_ff, rank=rank)
                new_w2 = LazyStructuredSirenLinear(d_ff, d_model, rank=rank)
                moe.experts_w1.append(new_w1)
                moe.experts_w2.append(new_w2)

            # Update expert count.
            moe.n_experts = old_count + n_new

            # Pre-train the new experts (EDT Phase 1).
            for i in range(n_new):
                expert_idx = old_count + i
                w1 = moe.experts_w1[expert_idx]
                w2 = moe.experts_w2[expert_idx]
                self._pretrain_expert(w1, w2, hidden_states,
                                      steps=steps_per_expert,
                                      batch_size=batch_size, lr=lr)

            print(f"  Layer {layer_idx+1}/{self.n_layers}: "
                  f"{old_count} → {old_count + n_new} experts", flush=True)

        # Extend the Farey phases for the router.
        self._extend_farey_phases(n_new)

        elapsed = time.time() - t0
        total_new = n_new * self.n_layers
        self.growth_log.append({
            "type": "add_experts",
            "n_new": n_new,
            "domain": domain,
            "total_new_experts": total_new,
            "time_seconds": elapsed,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        })

        n_params = sum(p.numel() for p in self.model.parameters())
        print(f"\nGrowth complete: +{total_new} experts in {elapsed:.0f}s", flush=True)
        print(f"  Total params: {n_params/1e6:.0f}M", flush=True)
        return self.model

    def _generate_hidden_states(self, data, n_samples=5000):
        """Generate hidden states from the embedding for expert pre-training."""
        d_model = self.model.d_model
        seq_len = 16

        if data is not None and self.tok is not None and len(data) > 100:
            # Real hidden states from embedding.
            states = []
            self.model.eval()
            with torch.no_grad():
                for _ in range(n_samples // 64):
                    idx = torch.randint(0, len(data) - seq_len - 1, (64,))
                    tokens = torch.stack([data[i:i+seq_len] for i in idx]).to(self.device)
                    h = self.model.embed(tokens)
                    states.append(h.cpu())
            return torch.cat(states, dim=0)
        else:
            # Synthetic fallback.
            print("  (using synthetic hidden states — provide data for real training)",
                  flush=True)
            return torch.randn(n_samples, seq_len, d_model)

    def _pretrain_expert(self, w1, w2, hidden_states, steps=2000,
                         batch_size=64, lr=1e-3):
        """Pre-train one expert on hidden states (EDT Phase 1)."""
        params = list(w1.parameters()) + list(w2.parameters())
        opt = torch.optim.AdamW(params, lr=lr, weight_decay=0.01)

        n_samples = len(hidden_states)
        for _ in range(steps):
            idx = torch.randint(0, n_samples - 1, (batch_size,))
            h_in = hidden_states[idx].to(self.device)
            h_target = hidden_states[idx + 1].to(self.device)

            opt.zero_grad()
            h1 = w1(h_in)
            h1_act = F.gelu(h1)
            h_out = w2(h1_act)
            loss = F.mse_loss(h_out, h_target)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(params, 1.0)
            opt.step()

    def _extend_farey_phases(self, n_new):
        """Extend the Farey phase distribution for new experts."""
        from fractus.nn.farey import expert_phases

        for layer_idx in range(self.n_layers):
            moe = self.model.blocks[layer_idx].moe
            old_count = moe.n_experts - n_new
            new_count = moe.n_experts

            # Regenerate phases for the new total count.
            new_phases = expert_phases(new_count)
            moe.expert_phases = torch.tensor(new_phases, dtype=torch.float32,
                                              device=moe.expert_phases.device)

    # ===================================================================
    # MECHANISM 2: RANK EXPANSION
    # ===================================================================

    def expand_rank(self, target_rank=128):
        """Expand the Siren rank of ALL experts.

        Old rank columns are preserved (knowledge kept).
        New rank columns are initialized to zero (learn during fine-tuning).

        Args:
            target_rank: the new rank (must be > current rank).
        """
        print(f"\n{'='*60}", flush=True)
        print(f"GROWTH: Expanding rank → {target_rank}", flush=True)
        print(f"{'='*60}", flush=True)

        t0 = time.time()
        expanded = 0

        for layer_idx in range(self.n_layers):
            moe = self.model.blocks[layer_idx].moe

            for expert_list in [moe.experts_w1, moe.experts_w2]:
                for expert in expert_list:
                    if not hasattr(expert, 'rank'):
                        continue
                    current_rank = expert.rank
                    if current_rank >= target_rank:
                        continue

                    # Expand U: (out, current_rank) → (out, target_rank)
                    old_U = expert.U.data  # (out, current_rank)
                    old_V = expert.V.data  # (in, current_rank)

                    new_U = torch.zeros(old_U.shape[0], target_rank,
                                        dtype=old_U.dtype, device=old_U.device)
                    new_V = torch.zeros(old_V.shape[0], target_rank,
                                        dtype=old_V.dtype, device=old_V.device)

                    # Copy old values (preserve knowledge).
                    new_U[:, :current_rank] = old_U
                    new_V[:, :current_rank] = old_V
                    # New columns stay zero (contribute nothing until trained).

                    # Replace parameters.
                    expert.U = nn.Parameter(new_U)
                    expert.V = nn.Parameter(new_V)
                    expert.rank = target_rank
                    expanded += 1

            print(f"  Layer {layer_idx+1}/{self.n_layers}: rank expanded", flush=True)

        elapsed = time.time() - t0
        n_params = sum(p.numel() for p in self.model.parameters())

        self.growth_log.append({
            "type": "expand_rank",
            "target_rank": target_rank,
            "experts_expanded": expanded,
            "time_seconds": elapsed,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        })

        print(f"\nRank expansion complete: {expanded} experts in {elapsed:.0f}s", flush=True)
        print(f"  Total params: {n_params/1e6:.0f}M", flush=True)
        return self.model

    # ===================================================================
    # MECHANISM 3: SAVE GROWN MODEL
    # ===================================================================

    def save(self, path):
        """Save the grown model with growth log."""
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)

        n_params = sum(p.numel() for p in self.model.parameters())
        torch.save({
            "model_state": self.model.state_dict(),
            "growth_log": self.growth_log,
            "n_params": n_params,
            "initial_experts": self.initial_experts,
            "initial_rank": self.initial_rank,
            "current_experts": self.model.blocks[0].moe.n_experts if hasattr(self.model, 'blocks') else 0,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        }, path)

        size_mb = os.path.getsize(path) / 1e6
        print(f"Saved grown model: {path} ({size_mb:.0f}MB)", flush=True)
        print(f"  Params: {n_params/1e6:.0f}M", flush=True)
        print(f"  Growth events: {len(self.growth_log)}", flush=True)
        for event in self.growth_log:
            if event["type"] == "add_experts":
                print(f"    +{event['total_new_experts']} experts ({event['domain']})", flush=True)
            elif event["type"] == "expand_rank":
                print(f"    rank → {event['target_rank']}", flush=True)

    # ===================================================================
    # UTILITY: Print current state
    # ===================================================================

    def status(self):
        """Print current model status."""
        n_params = sum(p.numel() for p in self.model.parameters())
        n_experts = self.model.blocks[0].moe.n_experts if hasattr(self.model, 'blocks') else 0
        rank = (
            self.model.blocks[0].moe.experts_w1[0].rank
            if hasattr(self.model, 'blocks') and hasattr(self.model.blocks[0].moe.experts_w1[0], 'rank')
            else '?'
        )

        print(f"\nFractus Status:", flush=True)
        print(f"  Params:   {n_params/1e6:.0f}M ({n_params/1e9:.3f}B)", flush=True)
        print(f"  Experts:  {n_experts} per layer × {self.n_layers} layers = {n_experts * self.n_layers}", flush=True)
        print(f"  Rank:     {rank}", flush=True)
        print(f"  Growth:   {len(self.growth_log)} events", flush=True)