#!/usr/bin/env python3 """ gemma_inject.py — FABLE 5 → HPC gate injection for Gemma 4 12B Adapted from hpc_fable_inject.py for Gemma architecture: - hidden_size=3840, intermediate_size=15360 (4× hidden_size) - Pauli state: [E/τ | μ/τ | ε/τ | μ/τ] — 4 blocks filling 15360 dims - Single model.safetensors file (~23 GB) - tie_word_embeddings=True (lm_head == embed_tokens) """ import argparse, gc, json, math, os, sys, time from collections import Counter, defaultdict from pathlib import Path import numpy as np import torch from safetensors import safe_open from safetensors.torch import save_file from transformers import AutoTokenizer def tokenize_fable(data_path, tokenizer): all_ids = [] with open(data_path) as f: for line in f: d = json.loads(line) text = d.get("text", "") for part in text.split("<|im_start|>"): if part.startswith("assistant\n"): content = part[len("assistant\n"):].replace("<|im_end|>", "").strip() if content: all_ids.extend(tokenizer.encode(content, add_special_tokens=False)) return all_ids def compute_bigram_stats(all_ids, vocab_size, topk): cnt = Counter() bigram = defaultdict(lambda: Counter()) for i, tid in enumerate(all_ids): cnt[tid] += 1 if i > 0: bigram[all_ids[i - 1]][tid] += 1 top_tokens = [t for t, _ in cnt.most_common(topk)] token_to_idx = {t: i for i, t in enumerate(top_tokens)} return cnt, bigram, top_tokens, token_to_idx def recover_gate_weight(embed, cnt, bigram, top_tokens, token_to_idx, tau): D = embed.shape[1] # 3840 V = len(top_tokens) E_raw = np.array([embed[t] for t in top_tokens]).astype(np.float64) mu = np.zeros_like(E_raw) eps = np.zeros_like(E_raw) for i, ti in enumerate(top_tokens): total = sum(bigram[ti].values()) fi = cnt[ti] or 1 if total == 0: continue for tj, c in bigram[ti].items(): if tj not in token_to_idx: continue fj = cnt[tj] or 1 p = c / total w = c / math.sqrt(fi * fj) c_Z = (1.0 - w) / 2.0 mu[i] += p * embed[tj] eps[i] += c_Z * embed[tj] # 4-block Pauli state: [E/τ | μ/τ | ε/τ | μ/τ] z = np.hstack([E_raw / tau, mu / tau, eps / tau, mu / tau]) z = np.clip(z, -30.0, 30.0) E_n = (E_raw - E_raw.mean(0, keepdims=True)) / (E_raw.std(0, keepdims=True) + 1e-10) reg = 1e-3 * np.eye(D) W = np.linalg.solve(E_n.T @ E_n + reg, E_n.T @ z) return W.T.astype(np.float32) # (15360, 3840) def inject_safetensors(src_path, dst_path, W_hpc_torch, alpha, layer_count): """ Read source (mmap'd), modify gate_proj in RAM, save to new file. Unmodified tensors stay mmap'd (zero RAM cost). Only 48 gate_proj tensors (~5.6 GB) materialize in RAM. """ with safe_open(str(src_path), framework='pt') as sf: keys = list(sf.keys()) out = {} for idx, k in enumerate(keys): t = sf.get_tensor(k) if 'gate_proj' in k and 'mlp' in k: is_gate, layer_idx = _parse_gate_key(k) if is_gate and layer_idx < layer_count: t = t.to(torch.float32) + alpha * W_hpc_torch.to(torch.float32) t = t.to(torch.bfloat16).contiguous() out[k] = t if (idx + 1) % 50 == 0: print(f" [{idx+1}/{len(keys)}] processed", flush=True) gc.collect() os.makedirs(os.path.dirname(dst_path) or '.', exist_ok=True) save_file(out, dst_path, metadata={'format': 'pt'}) print(f"Wrote {len(keys)} tensors to {dst_path}") def _parse_gate_key(key): parts = key.split('.') for i, p in enumerate(parts): if p == 'layers' and i + 1 < len(parts): try: return True, int(parts[i + 1]) except ValueError: pass return False, -1 def main(): parser = argparse.ArgumentParser(description="Gemma FABLE → HPC gate injector") parser.add_argument("--model", default="./gemma-4-12B-it", help="Gemma model directory") parser.add_argument("--output", default="./gemma-4-12B-it-hpc", help="Output directory") parser.add_argument("--data", default="/tmp/fable5_sft.jsonl", help="FABLE 5 JSONL path") parser.add_argument("--alpha", type=float, default=0.3, help="Injection strength") parser.add_argument("--tau", type=float, default=0.003, help="Temperature") parser.add_argument("--topk", type=int, default=30000, help="Top-k tokens") args = parser.parse_args() t0 = time.time() model_dir = Path(args.model) src_safetensors = model_dir / "model.safetensors" dst_safetensors = Path(args.output) / "model.safetensors" # ── 1. Tokenizer & embedding ── print("[1/6] Loading tokenizer & embedding...") tokenizer = AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=True) embed = None with safe_open(str(src_safetensors), framework='pt') as sf: for k in sf.keys(): if 'embed_tokens' in k: embed = sf.get_tensor(k).float().numpy() break assert embed is not None, "embed_tokens not found" print(f" Embedding: {embed.shape}") # ── 2. Tokenize FABLE ── print(f"[2/6] Tokenizing {args.data}...") all_ids = tokenize_fable(args.data, tokenizer) print(f" {len(all_ids)} tokens") # ── 3. Bigram stats ── print(f"[3/6] Bigrams (top-{args.topk})...") cnt, bigram, top_tokens, token_to_idx = compute_bigram_stats(all_ids, tokenizer.vocab_size, args.topk) print(f" {len(top_tokens)} tokens") # ── 4. Recover gate weight ── print(f"[4/6] Recovering gate weight (τ={args.tau})...") W_hpc = recover_gate_weight(embed, cnt, bigram, top_tokens, token_to_idx, args.tau) target_std = 0.012 scale = target_std / W_hpc.std() W_hpc *= scale print(f" W_hpc: {W_hpc.shape}, std={W_hpc.std():.6f}") # ── 5. Get layer count ── print(f"[5/6] Scanning source file...") layer_count = 0 with safe_open(str(src_safetensors), framework='pt') as sf: for k in sf.keys(): is_gate, layer = _parse_gate_key(k) if is_gate: layer_count = max(layer_count, layer + 1) print(f" {layer_count} layers detected") W_hpc_t = torch.from_numpy(W_hpc).to(torch.bfloat16) # ── 6. Inject ── print(f"[6/6] Injecting (α={args.alpha}) → {args.output}...") os.makedirs(args.output, exist_ok=True) inject_safetensors(str(src_safetensors), str(dst_safetensors), W_hpc_t, args.alpha, layer_count) import shutil for fn in ["config.json", "tokenizer.json", "tokenizer_config.json", "generation_config.json", "chat_template.jinja", "processor_config.json"]: src = model_dir / fn if src.exists(): shutil.copy2(src, Path(args.output) / fn) elapsed = time.time() - t0 print(f"Done in {elapsed:.0f}s. Output: {args.output}/") if __name__ == "__main__": main()