#!/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()