Upload fractus/grow.py with huggingface_hub

fractus/grow.py

@@ -0,0 +1,234 @@
+"""Fractus Progressive Growth — grow a model's width, depth, and expert count.
+
+THE INNOVATION. Instead of training a large model from scratch (impossible on
+CPU), we grow it palier by palier. Each palier inherits the previous model's
+weights via zero-padding (for width) or copying (for depth/experts), then
+trains briefly. The model never starts from random — it starts "warm".
+
+Growth axes:
+  - WIDTH (d_model): zero-pad every d-coupled matrix to the new dimension.
+  - DEPTH (n_layers): copy old blocks, init new ones with standard scheme.
+  - EXPERTS (n_experts): copy old experts, zero-init new ones.
+  - RANK (siren_rank): zero-pad the rank dimension of U/V factors.
+"""
+
+from __future__ import annotations
+import torch
+import torch.nn as nn
+from typing import Dict, Any
+
+
+def _pad_dim0(tensor: torch.Tensor, new_size: int) -> torch.Tensor:
+    """Grow a tensor along dim 0, zero-padding the new rows."""
+    old = tensor.shape[0]
+    if old >= new_size:
+        return tensor
+    pad_shape = list(tensor.shape)
+    pad_shape[0] = new_size - old
+    pad = torch.zeros(pad_shape, dtype=tensor.dtype)
+    return torch.cat([tensor, pad], dim=0)
+
+
+def _pad_last_dim(tensor: torch.Tensor, new_size: int) -> torch.Tensor:
+    """Grow a tensor along the LAST dim, zero-padding the new columns."""
+    old = tensor.shape[-1]
+    if old >= new_size:
+        return tensor
+    pad_shape = list(tensor.shape)
+    pad_shape[-1] = new_size - old
+    pad = torch.zeros(pad_shape, dtype=tensor.dtype)
+    return torch.cat([tensor, pad], dim=-1)
+
+
+def _transfer_block_weights(old_blk, new_blk, old_d: int, new_d: int):
+    """Transfer weights from old CTEBlock to new CTEBlock via zero-padding.
+
+    Copies old knowledge into the top-left block of every matrix. New dims
+    are zero (neutral start). New experts/blocks warm up during training.
+    """
+
+    # 1. Attention w_qkv: list of 3 tensors, each (d_model, d_model).
+    if hasattr(old_blk.attn, "w_qkv"):
+        for i in range(min(len(old_blk.attn.w_qkv), len(new_blk.attn.w_qkv))):
+            old_w = old_blk.attn.w_qkv[i].data
+            new_w = new_blk.attn.w_qkv[i].data
+            r_copy = min(old_w.shape[0], new_w.shape[0])
+            c_copy = min(old_w.shape[1], new_w.shape[1])
+            new_w.zero_()
+            new_w[:r_copy, :c_copy] = old_w[:r_copy, :c_copy]
+    # Attention biases.
+    if hasattr(old_blk.attn, "b_qkv"):
+        for i in range(min(len(old_blk.attn.b_qkv), len(new_blk.attn.b_qkv))):
+            old_b = old_blk.attn.b_qkv[i].data
+            new_b = new_blk.attn.b_qkv[i].data
+            c_copy = min(old_b.shape[0], new_b.shape[0])
+            new_b[:c_copy] = old_b[:c_copy]
+    # Attention w_out.
+    if hasattr(old_blk.attn, "w_out"):
+        old_wo = old_blk.attn.w_out.data
+        new_wo = new_blk.attn.w_out.data
+        r_copy = min(old_wo.shape[0], new_wo.shape[0])
+        c_copy = min(old_wo.shape[1], new_wo.shape[1])
+        new_wo.zero_()
+        new_wo[:r_copy, :c_copy] = old_wo[:r_copy, :c_copy]
+    if hasattr(old_blk.attn, "b_out"):
+        old_bo = old_blk.attn.b_out.data
+        new_bo = new_blk.attn.b_out.data
+        c_copy = min(old_bo.shape[0], new_bo.shape[0])
+        new_bo[:c_copy] = old_bo[:c_copy]
+    # Level offsets.
+    if hasattr(old_blk.attn, "level_offsets") and hasattr(new_blk.attn, "level_offsets"):
+        old_lo = old_blk.attn.level_offsets.data
+        new_lo = new_blk.attn.level_offsets.data
+        n_copy = min(old_lo.shape[0], new_lo.shape[0])
+        new_lo[:n_copy] = old_lo[:n_copy]
+    # Level logits.
+    if hasattr(old_blk.attn, "level_logits") and hasattr(new_blk.attn, "level_logits"):
+        old_ll = old_blk.attn.level_logits.data
+        new_ll = new_blk.attn.level_logits.data
+        n_copy = min(old_ll.shape[0], new_ll.shape[0])
+        new_ll[:n_copy] = old_ll[:n_copy]
+
+    # 2. LayerNorms: copy old dims, gamma=1/beta=0 for new.
+    for (old_norm, new_norm) in [
+        (old_blk.norm_attn, new_blk.norm_attn),
+        (old_blk.norm_kur, new_blk.norm_kur),
+        (old_blk.norm_moe, new_blk.norm_moe),
+    ]:
+        old_g = old_norm.weight.data
+        new_g = new_norm.weight.data
+        d_copy = min(old_g.shape[0], new_g.shape[0])
+        new_g[:d_copy] = old_g[:d_copy]
+        old_b = old_norm.bias.data
+        new_b = new_norm.bias.data
+        new_b[:d_copy] = old_b[:d_copy]
+
+    # 3. Kuramoto: grow oscillators + coupling rank.
+    if hasattr(old_blk.kuramoto, "omega"):
+        old_om = old_blk.kuramoto.omega.data
+        new_om = new_blk.kuramoto.omega.data
+        n_copy = min(old_om.shape[0], new_om.shape[0])
+        new_om[:n_copy] = old_om[:n_copy]
+    if hasattr(old_blk.kuramoto, "coupling_u"):
+        old_cu = old_blk.kuramoto.coupling_u.data
+        new_cu = new_blk.kuramoto.coupling_u.data
+        n_copy = min(old_cu.shape[0], new_cu.shape[0])
+        r_copy = min(old_cu.shape[1], new_cu.shape[1])
+        new_cu[:n_copy, :r_copy] = old_cu[:n_copy, :r_copy]
+    if hasattr(old_blk.kuramoto, "coupling_lambda"):
+        old_cl = old_blk.kuramoto.coupling_lambda.data
+        new_cl = new_blk.kuramoto.coupling_lambda.data
+        r_copy = min(old_cl.shape[0], new_cl.shape[0])
+        new_cl[:r_copy] = old_cl[:r_copy]
+
+    # 4. MoE experts (low-rank): U1/V1/U2/V2 are (E, ..., r).
+    moe_old = old_blk.moe
+    moe_new = new_blk.moe
+    e_copy = min(moe_old.n_experts, moe_new.n_experts)
+
+    if moe_old.expert_rank is not None and moe_new.expert_rank is not None:
+        for param_name in ["U1", "V1", "U2", "V2", "scale1", "scale2", "b1", "b2"]:
+            old_p = getattr(moe_old, param_name).data
+            new_p = getattr(moe_new, param_name).data
+            new_p.zero_()  # neutral start for ALL
+            if old_p.dim() == 3:
+                dd = min(old_p.shape[1], new_p.shape[1])
+                rr = min(old_p.shape[2], new_p.shape[2])
+                new_p[:e_copy, :dd, :rr] = old_p[:e_copy, :dd, :rr]
+            elif old_p.dim() == 2:
+                dd = min(old_p.shape[1], new_p.shape[1])
+                new_p[:e_copy, :dd] = old_p[:e_copy, :dd]
+            elif old_p.dim() == 1:
+                new_p[:e_copy] = old_p[:e_copy]
+            # Restore scale=1 for OLD experts.
+            if param_name in ("scale1", "scale2"):
+                old_scale = getattr(moe_old, param_name).data[:e_copy]
+                new_p[:e_copy] = old_scale
+
+
+def grow_cte(old_engine, new_config: Dict[str, Any]):
+    """Grow a ContinuousThoughtEngine to a larger config.
+
+    Copies old weights into the new model via zero-padding. Supports growth
+    in width (d_model), depth (n_layers), experts (n_experts), and rank
+    (siren_rank). Old knowledge is preserved; new capacity is neutral.
+
+    Args:
+        old_engine: a trained ContinuousThoughtEngine.
+        new_config: dict with any of:
+            d_model, n_layers, n_experts, top_k, expert_d_ff, siren_rank,
+            n_heads, d_head, n_levels, n_oscillators, coupling_rank, vocab_size.
+    """
+    from .continuous_engine import ContinuousThoughtEngine
+
+    old_d = old_engine.d_model
+    old_vocab = old_engine.vocab_size
+    old_n_layers = len(old_engine.blocks)
+
+    new_d = new_config.get("d_model", old_d)
+    new_vocab = new_config.get("vocab_size", old_vocab)
+    new_n_layers = new_config.get("n_layers", old_n_layers)
+    new_n_experts = new_config.get("n_experts", old_engine.blocks[0].moe.n_experts)
+    new_rank = new_config.get("siren_rank", old_engine.blocks[0].moe.expert_rank or 32)
+    new_d_ff = new_config.get("expert_d_ff", old_engine.blocks[0].moe.d_ff)
+    new_top_k = new_config.get("top_k", old_engine.blocks[0].moe.top_k)
+    new_n_heads = new_config.get("n_heads", old_engine.blocks[0].attn.n_heads)
+    new_d_head = new_config.get("d_head", old_engine.blocks[0].attn.d_head)
+    new_n_levels = new_config.get("n_levels", old_engine.blocks[0].attn.n_levels)
+    new_n_osc = new_config.get("n_oscillators", old_engine.blocks[0].kuramoto.N)
+    new_coupling_rank = new_config.get("coupling_rank", old_engine.blocks[0].kuramoto.rank)
+
+    # Build the new engine.
+    new_engine = ContinuousThoughtEngine(
+        vocab_size=new_vocab, d_model=new_d,
+        n_heads=new_n_heads, d_head=new_d_head, n_levels=new_n_levels,
+        n_oscillators=new_n_osc, coupling_rank=new_coupling_rank,
+        n_experts=new_n_experts, top_k=new_top_k,
+        expert_d_ff=new_d_ff, siren_rank=(new_rank if new_rank else None),
+        n_layers=new_n_layers,
+    )
+
+    # --- Transfer embedding (shared across all blocks) ---
+    old_emb = old_engine.observe.weight.data
+    new_emb = new_engine.observe.weight.data
+    v_copy = min(old_vocab, new_vocab)
+    d_copy = min(old_d, new_d)
+    new_emb.zero_()
+    new_emb[:v_copy, :d_copy] = old_emb[:v_copy, :d_copy]
+
+    # --- Transfer per-block weights (loop over ALL old blocks) ---
+    blocks_to_copy = min(old_n_layers, new_n_layers)
+    for blk_idx in range(blocks_to_copy):
+        _transfer_block_weights(
+            old_engine.blocks[blk_idx],
+            new_engine.blocks[blk_idx],
+            old_d, new_d,
+        )
+    # New blocks (indices old_n_layers..new_n_layers-1) keep random init —
+    # they warm up during training (like new brain regions developing).
+
+    # --- Transfer heads (top-level, not per-block) ---
+    for head_name in ["confidence_head", "salience_head"]:
+        old_h = getattr(old_engine, head_name)
+        new_h = getattr(new_engine, head_name)
+        old_w = old_h.weight.data
+        new_w = new_h.weight.data
+        d_copy = min(old_w.shape[1], new_w.shape[1])
+        new_w[:, :d_copy] = old_w[:, :d_copy]
+        new_h.bias.data[:] = old_h.bias.data[:]
+
+    # Output head is TIED with embedding — already handled above.
+
+    return new_engine
+
+
+def grow_summary(old_engine, new_engine) -> dict:
+    """Report what changed between old and new engine."""
+    return {
+        "d_model": f"{old_engine.d_model} → {new_engine.d_model}",
+        "n_layers": f"{len(old_engine.blocks)} → {len(new_engine.blocks)}",
+        "n_experts": f"{old_engine.blocks[0].moe.n_experts} → {new_engine.blocks[0].moe.n_experts}",
+        "expert_rank": f"{old_engine.blocks[0].moe.expert_rank} → {new_engine.blocks[0].moe.expert_rank}",
+        "params": f"{sum(p.numel() for p in old_engine.parameters()):,} → {sum(p.numel() for p in new_engine.parameters()):,}",
+        "n_oscillators": f"{old_engine.blocks[0].kuramoto.N} → {new_engine.blocks[0].kuramoto.N}",
+    }