Upload 2 files

gemma_inject.py

@@ -0,0 +1,194 @@
+#!/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()

train_gemma.py

@@ -0,0 +1,234 @@
+#!/usr/bin/env python3
+"""
+train_gemma.py — QLoRA fine-tune HPC-injected Gemma 4 12B on FABLE 5
+
+Loads Gemma-4-12B-it-hpc in 4-bit, adds LoRA to MLP projections,
+fine-tunes on FABLE 5 assistant conversations.
+
+Usage:
+  python3 train_gemma.py --model ./gemma-4-12B-it-hpc \\
+                         --data /tmp/fable5_sft.jsonl \\
+                         --output ./gemma-4-12B-it-fable \\
+                         --epochs 1 \\
+                         --lr 2e-4
+"""
+import argparse, gc, json, math, os, sys, time
+from functools import partial
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    BitsAndBytesConfig,
+    get_linear_schedule_with_warmup,
+)
+from peft import LoraConfig, get_peft_model
+
+
+# ── Dataset ──────────────────────────────────────────────────────────────────
+
+class FableDataset(Dataset):
+    """Tokenized FABLE 5 conversations, formatted for Gemma chat template."""
+
+    def __init__(self, data_path, tokenizer, max_length=2048):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.samples = []
+
+        with open(data_path) as f:
+            for line in f:
+                d = json.loads(line)
+                text = d.get("text", "")
+                if not text:
+                    continue
+
+                # Parse FABLE <|im_start|> format → messages list
+                turns = text.split("<|im_start|>")
+                messages = []
+                for turn in turns:
+                    turn = turn.strip()
+                    if not turn:
+                        continue
+                    if turn.startswith("user\n"):
+                        messages.append({
+                            "role": "user",
+                            "content": turn[len("user\n"):].replace("<|im_end|>", "").strip()
+                        })
+                    elif turn.startswith("assistant\n"):
+                        messages.append({
+                            "role": "assistant",
+                            "content": turn[len("assistant\n"):].replace("<|im_end|>", "").strip()
+                        })
+
+                if len(messages) <= 1:
+                    continue
+
+                # Format with Gemma's chat template
+                formatted = tokenizer.apply_chat_template(
+                    messages, tokenize=False, add_generation_prompt=False
+                )
+                tokens = tokenizer.encode(
+                    formatted, add_special_tokens=False,
+                    truncation=True, max_length=max_length
+                )
+                self.samples.append(tokens)
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        return torch.tensor(self.samples[idx], dtype=torch.long)
+
+
+def collate_fn(batch, pad_token_id):
+    max_len = max(len(x) for x in batch)
+    padded = torch.full((len(batch), max_len), pad_token_id, dtype=torch.long)
+    for i, seq in enumerate(batch):
+        padded[i, :len(seq)] = seq
+    return padded
+
+
+# ── Training ─────────────────────────────────────────────────────────────────
+
+def train():
+    parser = argparse.ArgumentParser(description="Fine-tune HPC-injected Gemma on FABLE 5")
+    parser.add_argument("--model", default="./gemma-4-12B-it-hpc", help="Injected model path")
+    parser.add_argument("--data", default="/tmp/fable5_sft.jsonl", help="FABLE 5 JSONL path")
+    parser.add_argument("--output", default="./gemma-4-12B-it-fable", help="Output path")
+    parser.add_argument("--epochs", type=int, default=1, help="Training epochs")
+    parser.add_argument("--lr", type=float, default=2e-4, help="Peak learning rate")
+    parser.add_argument("--batch_size", type=int, default=1, help="Per-device batch size")
+    parser.add_argument("--grad_accum", type=int, default=8, help="Gradient accumulation steps")
+    parser.add_argument("--max_length", type=int, default=2048, help="Max sequence length")
+    parser.add_argument("--lora_r", type=int, default=16, help="LoRA rank")
+    parser.add_argument("--lora_alpha", type=int, default=32, help="LoRA alpha")
+    parser.add_argument("--lora_dropout", type=float, default=0.05, help="LoRA dropout")
+    parser.add_argument("--save_steps", type=int, default=200, help="Checkpoint interval (steps)")
+    args = parser.parse_args()
+
+    t0 = time.time()
+
+    # ── 1. Tokenizer & 4-bit model ──
+    print("[1/6] Loading tokenizer & 4-bit model...")
+    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True, use_fast=False)
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "right"
+
+    bnb = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+    )
+
+    gpu_mem = torch.cuda.get_device_properties(0).total_memory // (1024**3)
+    max_memory = {0: f"{gpu_mem - 3}GiB", "cpu": "64GiB"}
+
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model,
+        trust_remote_code=True,
+        quantization_config=bnb,
+        device_map="auto",
+        max_memory=max_memory,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=True,
+    )
+    model.config.use_cache = False
+    model.gradient_checkpointing_enable()
+
+    # ── 2. LoRA ──
+    print(f"[2/6] Adding LoRA (r={args.lora_r}, alpha={args.lora_alpha})...")
+    lora_config = LoraConfig(
+        r=args.lora_r,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=["gate_proj", "up_proj", "down_proj"],
+    )
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+
+    # ── 3. Data ──
+    print("[3/6] Loading FABLE 5 dataset...")
+    dataset = FableDataset(args.data, tokenizer, max_length=args.max_length)
+    loader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        collate_fn=partial(collate_fn, pad_token_id=tokenizer.pad_token_id),
+        num_workers=2,
+        pin_memory=True,
+    )
+    print(f"  {len(dataset)} samples, {len(loader)} batches/epoch")
+
+    # ── 4. Optimizer & scheduler ──
+    print("[4/6] Setting up optimizer...")
+    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=args.lr)
+    total_steps = len(loader) * args.epochs // args.grad_accum
+    scheduler = get_linear_schedule_with_warmup(
+        opt, num_warmup_steps=int(0.05 * total_steps), num_training_steps=total_steps
+    )
+
+    # ── 5. Training loop ──
+    print(f"[5/6] Training ({args.epochs} epoch(s))...")
+    os.makedirs(args.output, exist_ok=True)
+    global_step = 0
+    best_loss = float("inf")
+
+    for epoch in range(args.epochs):
+        model.train()
+        total_loss = 0.0
+        n_batches = 0
+        epoch_t0 = time.time()
+
+        for batch_idx, batch in enumerate(loader):
+            batch = batch.to(model.device)
+            labels = batch.clone()
+
+            loss = model(input_ids=batch, labels=labels).loss
+            loss = loss / args.grad_accum
+            loss.backward()
+
+            total_loss += loss.item() * args.grad_accum
+            n_batches += 1
+
+            if (batch_idx + 1) % args.grad_accum == 0 or (batch_idx + 1) == len(loader):
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                opt.step()
+                scheduler.step()
+                opt.zero_grad()
+                global_step += 1
+
+                if global_step % args.save_steps == 0:
+                    avg_loss = total_loss / n_batches
+                    ppl = math.exp(avg_loss)
+                    save_path = os.path.join(args.output, f"checkpoint-{global_step}")
+                    model.save_pretrained(save_path)
+                    tokenizer.save_pretrained(save_path)
+                    print(f"  Step {global_step}: loss={avg_loss:.4f}, ppl={ppl:.2f}")
+
+            if (batch_idx + 1) % 20 == 0:
+                current_loss = total_loss / n_batches
+                print(f"  Epoch {epoch+1}, batch {batch_idx+1}/{len(loader)}: loss={current_loss:.4f}")
+
+        avg_loss = total_loss / n_batches
+        ppl = math.exp(avg_loss)
+        epoch_time = time.time() - epoch_t0
+        print(f"  Epoch {epoch+1} done: loss={avg_loss:.4f}, ppl={ppl:.2f}, time={epoch_time:.0f}s")
+
+        if avg_loss < best_loss:
+            best_loss = avg_loss
+            model.save_pretrained(os.path.join(args.output, "best"))
+            tokenizer.save_pretrained(os.path.join(args.output, "best"))
+
+    model.save_pretrained(os.path.join(args.output, "final"))
+    tokenizer.save_pretrained(os.path.join(args.output, "final"))
+    print(f"\nDone in {time.time()-t0:.0f}s. Final model: {args.output}/final")
+    print(f"Best loss: {best_loss:.4f} (PPL={math.exp(best_loss):.2f})")
+
+
+if __name__ == "__main__":
+    train()