weights.py

#!/usr/bin/env python3
"""
Generate FUNCTIONALLY CORRECT MLP weights for Gemma 4 31B (SwiGLU) integration.
Uses Sign-Symmetric Aligned Pairs to eliminate mean-shift without destroying alignment.

Gemma 4 31B architecture:
  - 60 transformer layers (10 global full-context + 50 sliding-window attention)
  - Interleaved attention: 5 SWA (sliding-window, 1024-token context) + 1 global full-context (period=6)
  - Global attention layers (5, 11, 17, 23, 29, 35, 41, 47, 53, 59) use double-wide MLP
  - Activation: gelu_pytorch_tanh
"""

import argparse
import json
import time
from pathlib import Path

import numpy as np
import torch
from safetensors.torch import load_file, save_file
from scipy.special import expit as _sigmoid

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

NEURON_SOURCE = "single"
SINGLE_FILE = "test_mlp_hf/model.safetensors"
MULTI_DIR = "generated_neurons/gaussian"

SINGLE_BOUNDARY_MODE = True

# Gemma 4 31B defaults
N_LAYERS = 60
HIDDEN_SIZE = 3840
INTERMEDIATE_SIZE = 15360

# Gemma 4 31B interleaved attention: 5 SWA layers then 1 global, repeating.
# Global layers use double-wide MLP intermediate size.
INTERLEAVE_PERIOD = 6          # one global every 6 layers
GLOBAL_LAYER_OFFSET = 5        # first global is at index 5 (0-based)
DEFAULT_ACTIVATION = "gelu_pytorch_tanh"

OUTPUT_DIR = "generated_weights_gemma4_31b"
RANDOM_SEED = 42


def is_global_attention_layer(layer_idx: int, period: int = INTERLEAVE_PERIOD,
                               offset: int = GLOBAL_LAYER_OFFSET) -> bool:
    """Return True if this layer uses global full-context attention (and thus double-wide MLP)."""
    return (layer_idx - offset) % period == 0 and layer_idx >= offset


def get_gating_function(name):
    if name == "silu":
        return _sigmoid
    elif name == "gelu_pytorch_tanh":
        alpha = np.sqrt(2.0 / np.pi)
        return lambda z: 0.5 * (1.0 + np.tanh(alpha * (z + 0.044715 * z**3)))
    elif name == "gelu":
        from scipy.special import erf
        return lambda z: 0.5 * (1.0 + erf(z / np.sqrt(2.0)))
    else:
        raise ValueError(f"Unsupported activation: {name}")


def get_activation(name):
    g_fn = get_gating_function(name)
    return lambda z: z * g_fn(z)


# ---------------------------------------------------------------------------
# 1. Load and extract functional parameters
# ---------------------------------------------------------------------------

def load_neurons(source, single_file, multi_dir):
    """Load source neurons."""
    neurons = []
    if source == "single":
        w = load_file(single_file)
        neurons.append(
            {
                k: v.float().numpy()
                for k, v in {
                    "W1": w["layer1.weight"],
                    "b1": w["layer1.bias"],
                    "W2": w["layer2.weight"],
                    "b2": w["layer2.bias"],
                }.items()
            }
        )
    elif source == "multi":
        for f in sorted(Path(multi_dir).glob("neuron_*.safetensors")):
            w = load_file(str(f))
            neurons.append(
                {
                    k: v.float().numpy()
                    for k, v in {
                        "W1": w["layer1.weight"],
                        "b1": w["layer1.bias"],
                        "W2": w["layer2.weight"],
                        "b2": w["layer2.bias"],
                    }.items()
                }
            )
    return neurons


def extract_functional_params(W1, b1, W2, b2, n_samples=10000):
    """Extract piecewise linear parameters from reference neurons."""
    xs = np.linspace(-4, 4, n_samples)
    ys = []

    for x in xs:
        h = np.maximum(0, W1 @ np.array([[x]]) + b1.reshape(-1, 1))
        y = (W2 @ h + b2.reshape(-1, 1)).item()
        ys.append(y)

    ys = np.array(ys)
    slopes = np.gradient(ys, xs)
    slope_changes = np.abs(np.gradient(slopes, xs))

    from scipy.signal import find_peaks
    peaks, _ = find_peaks(slope_changes, height=np.max(slope_changes) * 0.1, distance=100)

    if len(peaks) >= 2:
        idx1, idx2 = sorted(peaks[:2])
    elif len(peaks) == 1:
        idx1, idx2 = 0, peaks[0]
    else:
        idx1, idx2 = n_samples // 3, 2 * n_samples // 3

    boundary_x1 = float(xs[idx1])
    boundary_x2 = float(xs[idx2])

    left_slope = float(np.mean(slopes[:idx1])) if idx1 > 0 else float(slopes[0])
    mid_slope = float(np.mean(slopes[idx1:idx2]))
    right_slope = float(np.mean(slopes[idx2:]))
    y_boundary2 = float(ys[idx2])

    return {
        "boundary_x1": boundary_x1,
        "boundary_x2": boundary_x2,
        "left_slope": left_slope,
        "mid_slope": mid_slope,
        "right_slope": right_slope,
        "y_boundary2": y_boundary2,
    }


# ---------------------------------------------------------------------------
# 2. Construct functional dense layer (Sign-Symmetric Aligned Pairs)
# ---------------------------------------------------------------------------

def construct_functional_layer(
    functional_params,
    hidden_size,
    intermediate_size,
    gating_fn,
    activation_fn,
    has_bias=False,
    source_weights=None,
    rng_seed: int = 0,
):
    p = functional_params
    boundary = p["boundary_x1"]
    left_slope = p["left_slope"]
    right_slope = p["right_slope"]

    W_gate = np.zeros((intermediate_size, hidden_size), dtype=np.float32)
    W_up = np.zeros((intermediate_size, hidden_size), dtype=np.float32)
    W_down = np.zeros((hidden_size, intermediate_size), dtype=np.float32)

    if SINGLE_BOUNDARY_MODE:
        n_carrier = intermediate_size // 2
        n_transition = intermediate_size - n_carrier
        slope_diff = left_slope - right_slope
    else:
        n_carrier = intermediate_size // 3
        n_transition = intermediate_size - n_carrier
        slope_diff = p["left_slope"] - p["mid_slope"]

    _fill_swiglu(
        W_gate, W_up, W_down,
        n_carrier, n_transition,
        hidden_size, intermediate_size,
        boundary, right_slope, slope_diff,
        gating_fn, activation_fn,
        rng_seed=rng_seed,
    )

    return W_gate, W_up, W_down


def _fill_swiglu(
    W_gate, W_up, W_down,
    n_carrier, n_transition,
    hidden_size, intermediate_size,
    boundary, right_slope, slope_diff,
    gating_fn, activation_fn,
    rng_seed: int = 0,
):
    """
    Fills SwiGLU projections using clean sign-symmetric paired frames.
    """
    H = hidden_size
    rng = np.random.default_rng(rng_seed)

    # Re-apportion matrices to ensure strict pairs
    n_pairs_carrier = n_carrier // 2
    if n_pairs_carrier == 0: n_pairs_carrier = 1
    n_neurons_carrier = 2 * n_pairs_carrier
    
    n_neurons_transition = intermediate_size - n_neurons_carrier
    n_pairs_transition = n_neurons_transition // 2

    # 1. Generate unified random projection frame vectors
    n_total_pairs = n_pairs_carrier + n_pairs_transition
    dirs = rng.standard_normal((n_total_pairs, H)).astype(np.float32)
    dirs /= np.linalg.norm(dirs, axis=1, keepdims=True)

    # 2. Monte-Carlo gain calibration for the odd function: z^3 * (2*gate_fn(g*z) - 1)
    cal_rng = np.random.default_rng(0xCA1_5EED)
    z_samples = cal_rng.standard_normal(100_000)

    # Carrier calibration
    g_c = 1.0
    gain_c = float(np.mean((z_samples**3) * (2 * gating_fn(g_c * z_samples) - 1)))
    u_c = 1.0
    v_c = (right_slope * 4.0 * H) / (n_neurons_carrier * u_c * g_c * gain_c) if abs(gain_c) > 1e-6 else 0.0

    # Transition calibration
    g_t = 2.0  
    gain_t = float(np.mean((z_samples**3) * (2 * gating_fn(g_t * z_samples) - 1)))
    u_t = 1.0
    v_t = (slope_diff * 4.0 * H) / (n_neurons_transition * u_t * g_t * gain_t) if abs(gain_t) > 1e-6 else 0.0

    # 3. Populate carrier pairs
    for i in range(n_pairs_carrier):
        d = dirs[i]
        idx1 = 2 * i
        idx2 = 2 * i + 1

        W_gate[idx1, :] = g_c * d
        W_up[idx1, :]   = u_c * d
        W_down[:, idx1] = (v_c / 2.0) * d

        W_gate[idx2, :] = -g_c * d
        W_up[idx2, :]   = u_c * d
        W_down[:, idx2] = (v_c / 2.0) * d

    # 4. Populate transition pairs
    for i in range(n_pairs_transition):
        d = dirs[n_pairs_carrier + i]
        idx1 = n_neurons_carrier + 2 * i
        idx2 = n_neurons_carrier + 2 * i + 1

        W_gate[idx1, :] = g_t * d
        W_up[idx1, :]   = u_t * d
        W_down[:, idx1] = (v_t / 2.0) * d

        W_gate[idx2, :] = -g_t * d
        W_up[idx2, :]   = u_t * d
        W_down[:, idx2] = (v_t / 2.0) * d

    # 5. Global Zero-Mean Normalization
    W_down -= W_down.mean(axis=0, keepdims=True)

    # 6. Empirical Validation Test Pass
    test_rng = np.random.default_rng(rng_seed + 99)
    x_test = test_rng.standard_normal((1024, H)).astype(np.float32)
    gate_act = activation_fn(x_test @ W_gate.T)
    up_act = x_test @ W_up.T
    out_test = (gate_act * up_act) @ W_down.T

    target_scale = max(abs(right_slope), 0.05)
    empirical_scale = np.sqrt(np.var(out_test) / np.var(x_test))
    if empirical_scale > 1e-7:
        correction_factor = target_scale / empirical_scale
        W_down *= correction_factor
        print(f"    [Stability Check] Layer {rng_seed - 42:02d} W_down scaling verification factor: {correction_factor:.4f}")

    print(f"    Carrier Pairs:    {n_pairs_carrier} (g_c={g_c}, u_c={u_c}, v_c={v_c:.4f})")
    print(f"    Transition Pairs: {n_pairs_transition} (g_t={g_t}, u_t={u_t}, v_t={v_t:.4f})")


# ---------------------------------------------------------------------------
# 3. Execution Pipeline
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Generate functional SwiGLU weights for Gemma")
    parser.add_argument("--source", default=NEURON_SOURCE, choices=["single", "multi"])
    parser.add_argument("--single-file", default=SINGLE_FILE)
    parser.add_argument("--multi-dir", default=MULTI_DIR)
    parser.add_argument("--n-layers", type=int, default=N_LAYERS)
    parser.add_argument("--hidden-size", type=int, default=HIDDEN_SIZE)
    parser.add_argument("--intermediate-size", type=int, default=INTERMEDIATE_SIZE)
    parser.add_argument("--output-dir", default=OUTPUT_DIR)
    parser.add_argument("--seed", type=int, default=RANDOM_SEED)
    parser.add_argument("--target-layers", type=int, nargs="+", default=None)
    parser.add_argument("--has-bias", action="store_true", default=False)
    parser.add_argument("--base-model", default=None, help="Base model directory path to inspect config for layers and dimensions")
    parser.add_argument("--activation", default=None, choices=["silu", "gelu_pytorch_tanh", "gelu"], help="Override activation function (otherwise inferred from base-model or default silu)")
    args = parser.parse_args()

    # Detect model parameters if base-model is provided
    inferred_n_layers = args.n_layers
    inferred_hidden_size = args.hidden_size
    inferred_intermediate_size = args.intermediate_size
    inferred_activation = args.activation if args.activation else DEFAULT_ACTIVATION
    use_double_wide_mlp = False
    interleave_period = INTERLEAVE_PERIOD
    global_layer_offset = GLOBAL_LAYER_OFFSET

    if args.base_model:
        print(f"[config] Reading config from base model: {args.base_model}")
        base_path = Path(args.base_model)
        config_file = base_path / "config.json"
        if not config_file.exists():
            raise FileNotFoundError(f"Config not found at {config_file}")
        with open(config_file, "r") as f:
            config = json.load(f)
        text_config = config.get("text_config", {})
        
        def get_val(key, default_None):
            return text_config.get(key, config.get(key, default_None))

        inferred_n_layers = get_val("num_hidden_layers", inferred_n_layers)
        inferred_hidden_size = get_val("hidden_size", inferred_hidden_size)
        inferred_intermediate_size = get_val("intermediate_size", inferred_intermediate_size)
        use_double_wide_mlp = get_val("use_double_wide_mlp", False)
        # Gemma 4 31B stores interleave info as attention_pattern or sliding_window counts;
        # fall back to module-level constants if not present in config.
        interleave_period = get_val("attention_pattern_period", interleave_period)
        global_layer_offset = get_val("global_layer_offset", global_layer_offset)

        if not args.activation:
            inferred_activation = get_val("hidden_activation", inferred_activation)

        print(f"  Detected configuration:")
        print(f"    Layers: {inferred_n_layers}")
        print(f"    Hidden Size: {inferred_hidden_size}")
        print(f"    Base Intermediate Size: {inferred_intermediate_size}")
        print(f"    Activation: {inferred_activation}")
        print(f"    Double Wide MLP: {use_double_wide_mlp}")
        print(f"    Interleave Period: {interleave_period}  (global offset: {global_layer_offset})")

    out = Path(args.output_dir)
    out.mkdir(exist_ok=True)

    print("=" * 60)
    print("FUNCTIONAL PAIR-ALIGNED SwiGLU Generation (Gemma 4 31B)")
    print("=" * 60)

    print("[1] Loading source neurons...")
    neurons = load_neurons(args.source, args.single_file, args.multi_dir)
    
    print("[2] Extracting functional behavior...")
    functional_params = []
    for i, n in enumerate(neurons):
        p = extract_functional_params(n["W1"], n["b1"], n["W2"], n["b2"])
        functional_params.append(p)

    source_weights = neurons[0]
    layer_indices = args.target_layers if args.target_layers is not None else range(inferred_n_layers)
    gating_fn = get_gating_function(inferred_activation)
    activation_fn = get_activation(inferred_activation)

    print(f"\n[3] Encoding {len(layer_indices)} layers using activation: {inferred_activation}...")
    for layer_idx in layer_indices:
        neuron_idx = (layer_idx * len(neurons)) // inferred_n_layers
        base_params = functional_params[neuron_idx]

        # Determine intermediate size: global attention layers use double-wide MLP.
        # Gemma 4 31B interleaves 5 SWA + 1 global every INTERLEAVE_PERIOD layers.
        if use_double_wide_mlp and is_global_attention_layer(layer_idx, interleave_period, global_layer_offset):
            layer_intermediate_size = inferred_intermediate_size * 2
            layer_type = "global"
        else:
            layer_intermediate_size = inferred_intermediate_size
            layer_type = "swa" if use_double_wide_mlp else "std"

        print(f"  Layer {layer_idx:02d} [{layer_type}]: intermediate_size = {layer_intermediate_size}")

        W_gate, W_up, W_down = construct_functional_layer(
            base_params, inferred_hidden_size, layer_intermediate_size,
            gating_fn, activation_fn,
            has_bias=False, source_weights=source_weights, rng_seed=args.seed + layer_idx,
        )

        out_path = out / f"layer_{layer_idx:02d}.safetensors"
        tensors = {
            "gate_proj.weight": torch.tensor(W_gate),
            "up_proj.weight": torch.tensor(W_up),
            "down_proj.weight": torch.tensor(W_down),
        }
        save_file(tensors, str(out_path))

    with open(out / "meta.json", "w") as f:
        json.dump({
            "config": {
                "hidden_size": inferred_hidden_size,
                "intermediate_size": inferred_intermediate_size,
                "n_layers": inferred_n_layers,
                "activation": inferred_activation,
                "use_double_wide_mlp": use_double_wide_mlp,
                "interleave_period": interleave_period,
                "global_layer_offset": global_layer_offset,
                "encoding": "sign_symmetric_pairs"
            }
        }, f, indent=2)

    print(f"\nComplete! Balanced weights saved safely to: {out}/")


if __name__ == "__main__":
    main()