Upload 3 files

hpc_fable_inject.py

@@ -0,0 +1,217 @@
+#!/usr/bin/env python3
+"""
+hpc_fable_inject.py — FABLE 5 → HPC gate weight injection for Ornith
+
+Extracts gate patterns from FABLE 5 training traces using HPC Pauli decomposition,
+then injects them as an additive correction into Ornith-1.0-9B.
+
+Formula (derived from Pauli decomposition I + X + c_Z·Z of edge matrix):
+  z_i = [E[i]/τ  |  μ_i/τ  |  ε_i/τ]   (3 blocks × 4096 = 12288 dims)
+  where:
+    μ_i = Σ_j P(j|i)·E[j]           — X-component: expected next embedding
+    ε_i = Σ_j c_Z(i,j)·E[j]         — Z-component: Pauli-coupling-weighted sum
+    c_Z(i,j) = (1 - w_ij) / 2       — Pauli Z coefficient
+    w_ij = f_ij / √(f_i·f_j)        — normalized edge weight
+
+The gate weight W is solved via ridge regression: E_n·W ≈ z,
+then scaled to match original gate_proj weight norm (~0.012),
+then added to the original weights: W_inj = W_orig + α·W_hpc.
+
+Usage:
+  python3 hpc_fable_inject.py --data /tmp/fable5_sft.jsonl \\
+                             --model ./Ornith-1.0-9B \\
+                             --output ./Ornith-1.0-9B-hpc \\
+                             --alpha 0.3 \\
+                             --tau 0.003 \\
+                             --topk 30000
+"""
+
+import argparse, json, math, os, shutil, sys, time
+from collections import Counter, defaultdict
+from pathlib import Path
+
+import numpy as np
+import torch
+from safetensors.torch import safe_open, save_file
+from transformers import AutoTokenizer
+
+
+def load_model_shards(model_dir):
+    """Return list of shard filenames and the weight index."""
+    index_path = Path(model_dir) / "model.safetensors.index.json"
+    with open(index_path) as f:
+        index = json.load(f)
+    shards = sorted(set(index["weight_map"].values()))
+    return shards, index
+
+
+def save_model_shards(src_dir, dst_dir, shards, modify_fn):
+    """Load each shard, apply modify_fn to selected tensors, save."""
+    dst_dir = Path(dst_dir)
+    dst_dir.mkdir(parents=True, exist_ok=True)
+    for shard_name in shards:
+        src_path = Path(src_dir) / shard_name
+        dst_path = dst_dir / shard_name
+        tensors = {}
+        with safe_open(str(src_path), framework="pt") as sf:
+            for k in sf.keys():
+                tensors[k] = sf.get_tensor(k).clone()
+        tensors = modify_fn(tensors)
+        save_file(tensors, str(dst_path))
+
+
+def copy_config(src_dir, dst_dir, filenames=None):
+    """Copy model config/tokenizer files."""
+    if filenames is None:
+        filenames = [
+            "model.safetensors.index.json",
+            "config.json",
+            "tokenizer.json",
+            "tokenizer_config.json",
+            "generation_config.json",
+        ]
+    src_dir, dst_dir = Path(src_dir), Path(dst_dir)
+    for fn in filenames:
+        src = src_dir / fn
+        if src.exists():
+            shutil.copy2(src, dst_dir / fn)
+
+
+def tokenize_fable(data_path, tokenizer):
+    """Tokenize FABLE assistant turns. Returns flat token id list."""
+    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):
+    """Build token frequency and bigram counter for top-k tokens."""
+    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):
+    """
+    Recover gate weight via Pauli-decomposition regression.
+    
+    Returns: W_hpc as float32 numpy array (12288, 4096), unscaled.
+    """
+    D = embed.shape[1]
+    V = len(top_tokens)
+    
+    E_raw = np.array([embed[t] for t in top_tokens]).astype(np.float64)
+    mu = np.zeros_like(E_raw)      # X-component: probability-weighted
+    eps = np.zeros_like(E_raw)     # Z-component: c_Z-weighted
+
+    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]
+
+    # Build target: [E/τ | μ/τ | ε/τ]
+    z = np.hstack([E_raw / tau, mu / tau, eps / tau])
+    z = np.clip(z, -30.0, 30.0)
+
+    # Ridge regression: E_n · W = z
+    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)  # (12288, 4096)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="FABLE → HPC gate weight injector")
+    parser.add_argument("--data", default="/tmp/fable5_sft.jsonl", help="FABLE 5 JSONL path")
+    parser.add_argument("--model", default="./Ornith-1.0-9B", help="Source model directory")
+    parser.add_argument("--output", default="./Ornith-1.0-9B-hpc", help="Output model directory")
+    parser.add_argument("--alpha", type=float, default=0.3, help="Additive correction strength")
+    parser.add_argument("--tau", type=float, default=0.003, help="Temperature for target scaling")
+    parser.add_argument("--topk", type=int, default=30000, help="Top-k tokens to use")
+    args = parser.parse_args()
+
+    t0 = time.time()
+
+    # ── 1. Tokenizer & model info ──
+    print("[1/5] Loading tokenizer & embeddings...")
+    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
+    shards, index = load_model_shards(args.model)
+
+    # Load embedding matrix from first shard
+    embed = None
+    with safe_open(str(Path(args.model) / shards[0]), 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, "Could not find embed_tokens in shards"
+    print(f"  Embeddings: {embed.shape}, dtype={embed.dtype}")
+
+    # ── 2. Tokenize FABLE data ──
+    print(f"[2/5] Tokenizing {args.data}...")
+    all_ids = tokenize_fable(args.data, tokenizer)
+    print(f"  {len(all_ids)} tokens")
+
+    # ── 3. Bigram statistics ──
+    print(f"[3/5] Computing 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 with {sum(len(bigram[t]) for t in top_tokens)} edges")
+
+    # ── 4. Recover gate weight ──
+    print(f"[4/5] Recovering gate weight (τ={args.tau})...")
+    W_hpc = recover_gate_weight(embed, cnt, bigram, top_tokens, token_to_idx, args.tau)
+    
+    # Scale to match original weight norm (~0.012)
+    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}, "
+          f"range=[{W_hpc.min():.4f},{W_hpc.max():.4f}]")
+
+    # ── 5. Inject as additive correction ──
+    print(f"[5/5] Injecting (α={args.alpha}) → {args.output}...")
+    W_torch = torch.from_numpy(W_hpc).to(torch.bfloat16).contiguous()
+
+    def modify_fn(tensors):
+        for k in list(tensors.keys()):
+            if "gate_proj" in k and "visual" not in k:
+                tensors[k] = tensors[k] + args.alpha * W_torch
+        return tensors
+
+    save_model_shards(args.model, args.output, shards, modify_fn)
+    copy_config(args.model, args.output)
+
+    elapsed = time.time() - t0
+    print(f"Done in {elapsed:.1f}s. Model saved to {args.output}/")
+    print(f"\nTo test: python3 test_recovered.py")
+
+
+if __name__ == "__main__":
+    main()

test_recovered.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+"""Quick test: load recovered model, generate a few tokens, check sanity."""
+import torch
+import gc
+gc.collect(); torch.cuda.empty_cache()
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+
+MODEL_PATH = "/home/none/Documents/HPC-Quantize/Ornith-1.0-9B-hpc"
+
+bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
+                          bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16)
+
+print("Loading model in 4-bit...")
+tok = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+tok.padding_side = "right"
+if tok.pad_token is None: tok.pad_token = tok.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_PATH, trust_remote_code=True,
+    quantization_config=bnb, device_map="auto", low_cpu_mem_usage=True,
+    torch_dtype=torch.bfloat16
+)
+model.config.use_cache = True
+
+prompts = [
+    "The capital of France is",
+    "Once upon a time, in a land far away,",
+    "Machine learning is a field of study that",
+]
+
+print("\n--- Generation test ---")
+for p in prompts:
+    inp = tok(p, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        out = model.generate(**inp, max_new_tokens=32, do_sample=True, temperature=0.7, top_p=0.9)
+    text = tok.decode(out[0], skip_special_tokens=True)
+    print(f"\nPrompt: {p}")
+    print(f"Output: {text}")
+
+print("\nDone.")

train_fable.py

@@ -0,0 +1,228 @@
+#!/usr/bin/env python3
+"""
+train_fable.py — QLoRA fine-tune HPC-injected Ornith on FABLE 5 traces
+
+Loads the HPC-injected model (Ornith-1.0-9B-hpc), adds LoRA adapters,
+and fine-tunes on FABLE 5 assistant conversations.
+
+Usage:
+  python3 train_fable.py --model ./Ornith-1.0-9B-hpc \\
+                         --data /tmp/fable5_sft.jsonl \\
+                         --output ./Ornith-1.0-9B-fable \\
+                         --epochs 1 \\
+                         --lr 2e-4
+"""
+
+import argparse, json, gc, 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 assistant conversations."""
+
+    def __init__(self, data_path, tokenizer, max_length=2048):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.samples = []
+
+        # System prompt used by Claude-style models
+        system_msg = "You are a helpful, harmless, and honest assistant."
+
+        with open(data_path) as f:
+            for line in f:
+                d = json.loads(line)
+                text = d.get("text", "")
+                if not text:
+                    continue
+
+                # Reconstruct structured conversation
+                turns = text.split("<|im_start|>")
+                messages = [{"role": "system", "content": system_msg}]
+                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 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):
+        tokens = self.samples[idx]
+        return torch.tensor(tokens, dtype=torch.long)
+
+
+def collate_fn(batch, pad_token_id):
+    """Pad batch to uniform length."""
+    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 Ornith on FABLE 5")
+    parser.add_argument("--model", default="./Ornith-1.0-9B-hpc", help="Injected model path")
+    parser.add_argument("--data", default="/tmp/fable5_sft.jsonl", help="FABLE 5 JSONL path")
+    parser.add_argument("--output", default="./Ornith-1.0-9B-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"
+    if tokenizer.chat_template is None:
+        tokenizer.chat_template = "{% for message in messages %}{% if message['role'] == 'system' %}<|im_start|>system\n{{ message['content'] }}<|im_end|>\n{% elif message['role'] == 'user' %}<|im_start|>user\n{{ message['content'] }}<|im_end|>\n{% elif message['role'] == 'assistant' %}<|im_start|>assistant\n{{ message['content'] }}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
+
+    bnb = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+    )
+    # Leave ~3 GiB headroom on GPU for activations/gradients
+    max_memory = {0: f"{torch.cuda.get_device_properties(0).total_memory // (1024**3) - 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  # required for gradient checkpointing
+    model.gradient_checkpointing_enable()
+
+    # ── 2. LoRA config — target gate_proj (the injected weights) ──
+    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()
+
+    # ── 4. 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 (only trainable LoRA params) ──
+    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)
+
+    # ── 6. 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}, lr={scheduler.get_last_lr()[0]:.2e}")
+
+            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"))
+
+    # Save final
+    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()