Full GPU benchmark including FP32 quant quality test

bench_full.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""Full Little Fig GPU Benchmark — Quant Quality + Training on t4-medium (30GB RAM)"""
+import os, sys, subprocess, json, time, gc, traceback
+import numpy as np
+subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "transformers", "accelerate", "peft", "bitsandbytes", "datasets", "sentencepiece", "protobuf", "psutil", "numpy"])
+if not os.path.exists("/app/littlefig"):
+    subprocess.check_call(["git", "clone", "https://github.com/ticketguy/littlefig.git", "/app/littlefig"])
+subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-e", "/app/littlefig[train]"])
+sys.path.insert(0, "/app/littlefig/src")
+import torch
+import torch.nn.functional as F
+
+MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+LORA_R = 16; LORA_ALPHA = 32; GROUP_SIZE = 128
+LORA_TARGETS = ["q_proj","k_proj","v_proj","o_proj"]
+MAX_SEQ = 512; TRAIN_STEPS = 100; BATCH_SIZE = 4; GRAD_ACCUM = 4; LR = 2e-4
+RESULTS = {}
+
+def log(msg): print(f"[BENCH] {msg}", flush=True)
+def gpu_mb(): return torch.cuda.max_memory_allocated()/1e6 if torch.cuda.is_available() else 0
+def reset(): gc.collect(); torch.cuda.empty_cache() if torch.cuda.is_available() else None; torch.cuda.reset_peak_memory_stats() if torch.cuda.is_available() else None
+
+log(f"PyTorch {torch.__version__}, CUDA={torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    log(f"GPU: {torch.cuda.get_device_name()} ({torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB)")
+import psutil
+log(f"System RAM: {psutil.virtual_memory().total/1e9:.1f}GB")
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PART A: QUANTIZATION QUALITY (FP32 load on CPU — needs 30GB RAM)
+# ═══════════════════════════════════════════════════════════════════════════════
+log("\n" + "="*60 + "\n  QUANTIZATION QUALITY (TinyLlama 1.1B)\n" + "="*60)
+
+from transformers import AutoModelForCausalLM
+from little_fig.engine.figquant import figquant_quantize, figquant_dequantize
+
+log("Loading TinyLlama FP32 (CPU)...")
+model_fp32 = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.float32, low_cpu_mem_usage=True)
+
+def nf4_qd(W, gs=128):
+    s, n = W.shape, W.numel(); f = W.reshape(-1).float()
+    p = (gs-n%gs)%gs
+    if p>0: f = torch.cat([f, torch.zeros(p)])
+    g = f.reshape(-1,gs); sc = g.abs().amax(1).clamp(min=1e-10)
+    cb = torch.tensor([-1.0,-0.6962,-0.5251,-0.3949,-0.2844,-0.1848,-0.0911,0.0,0.0796,0.1609,0.2461,0.3379,0.4407,0.5626,0.7230,1.0])
+    idx = ((g/sc.unsqueeze(1)).reshape(-1).unsqueeze(1)-cb.unsqueeze(0)).abs().argmin(1).reshape(-1,gs)
+    return (torch.gather(cb.unsqueeze(0).expand(idx.shape[0],-1),1,idx.long())*sc.unsqueeze(1)).reshape(-1)[:n].reshape(s)
+
+methods = {"figquant":{"mse":[],"cos":[]},"nf4":{"mse":[],"cos":[]}}
+n=0; fw=0
+for name, param in model_fp32.named_parameters():
+    if param.ndim!=2 or param.numel()<1024: continue
+    W = param.data.float()
+    q = figquant_quantize(W, group_size=GROUP_SIZE, n_iters=8)
+    deq = figquant_dequantize(q)
+    mse_fq = F.mse_loss(deq, W).item()
+    cos_fq = F.cosine_similarity(W.flatten().unsqueeze(0), deq.flatten().unsqueeze(0)).item()
+    W_nf4 = nf4_qd(W, GROUP_SIZE)
+    mse_nf = F.mse_loss(W_nf4, W).item()
+    cos_nf = F.cosine_similarity(W.flatten().unsqueeze(0), W_nf4.flatten().unsqueeze(0)).item()
+    methods["figquant"]["mse"].append(mse_fq); methods["figquant"]["cos"].append(cos_fq)
+    methods["nf4"]["mse"].append(mse_nf); methods["nf4"]["cos"].append(cos_nf)
+    if mse_fq < mse_nf: fw+=1
+    n+=1
+    if n%20==0: log(f"  {n} layers done...")
+
+avgs = {m:{k:float(np.mean(v)) for k,v in d.items()} for m,d in methods.items()}
+mvn = (avgs["nf4"]["mse"]-avgs["figquant"]["mse"])/avgs["nf4"]["mse"]*100
+RESULTS["quant"] = {"avgs":avgs, "n":n, "fw":fw, "mvn":mvn}
+log(f"\n  QUANTIZATION RESULTS ({n} layers):")
+log(f"  FigQuant: MSE={avgs['figquant']['mse']:.6e}, cos={avgs['figquant']['cos']:.6f}")
+log(f"  NF4:      MSE={avgs['nf4']['mse']:.6e}, cos={avgs['nf4']['cos']:.6f}")
+log(f"  FigQuant vs NF4: {mvn:+.1f}% MSE (wins {fw}/{n} layers)")
+del model_fp32; gc.collect()
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PART B: TRAINING CONVERGENCE
+# ═══════════════════════════════════════════════════════════════════════════════
+log("\n" + "="*60 + "\n  TRAINING (100 steps, batch=4x4)\n" + "="*60)
+
+from datasets import load_dataset
+from transformers import AutoTokenizer, BitsAndBytesConfig
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+
+ds = load_dataset("tatsu-lab/alpaca", split="train").select(range(1000))
+
+def hf_loop(model, tokenizer, name):
+    dev = next(model.parameters()).device
+    def tok_fn(ex):
+        inst=ex.get("instruction",""); inp=ex.get("input","").strip(); out=ex.get("output","")
+        txt = f"### Instruction:\n{inst}\n\n### Input:\n{inp}\n\n### Response:\n{out}" if inp else f"### Instruction:\n{inst}\n\n### Response:\n{out}"
+        e = tokenizer(txt, truncation=True, max_length=MAX_SEQ, padding="max_length")
+        e["labels"] = e["input_ids"].copy(); return e
+    td = ds.map(tok_fn, remove_columns=ds.column_names); td.set_format("torch")
+    from torch.utils.data import DataLoader
+    dl = DataLoader(td, batch_size=BATCH_SIZE, shuffle=True, collate_fn=lambda b: {k:torch.stack([x[k] for x in b]) for k in b[0] if isinstance(b[0][k], torch.Tensor)}, drop_last=True)
+    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=LR, weight_decay=0.01)
+    model.train(); losses=[]; gs=0; al=0.0; reset(); t0=time.time()
+    for batch in dl:
+        if gs>=TRAIN_STEPS*GRAD_ACCUM: break
+        batch = {k:v.to(dev) for k,v in batch.items()}
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            loss = model(**batch).loss / GRAD_ACCUM
+        loss.backward(); al+=loss.item(); gs+=1
+        if gs%GRAD_ACCUM==0:
+            torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0)
+            opt.step(); opt.zero_grad(); s=gs//GRAD_ACCUM; losses.append(al); al=0.0
+            if s%20==0: log(f"  [{name}] step={s} loss={losses[-1]:.4f}")
+    tt=time.time()-t0; pm=gpu_mb()
+    del model, opt; gc.collect(); torch.cuda.empty_cache()
+    RESULTS[name] = {"final":float(losses[-1]),"time_s":tt,"steps":len(losses),"gpu_mb":pm}
+    log(f"  [{name}] loss={losses[-1]:.4f} time={tt:.0f}s gpu={pm:.0f}MB")
+
+# FP16
+log("\n--- FP16 LoRA ---"); reset()
+m = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.float16, device_map="auto")
+t = AutoTokenizer.from_pretrained(MODEL); t.pad_token=t.eos_token
+m.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant":False})
+m = get_peft_model(m, LoraConfig(r=LORA_R, lora_alpha=LORA_ALPHA, target_modules=LORA_TARGETS, bias="none", task_type="CAUSAL_LM"))
+hf_loop(m, t, "fp16_lora")
+
+# NF4
+log("\n--- BnB NF4 ---"); reset()
+m = AutoModelForCausalLM.from_pretrained(MODEL, quantization_config=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16), device_map="auto")
+t = AutoTokenizer.from_pretrained(MODEL); t.pad_token=t.eos_token
+m = prepare_model_for_kbit_training(m)
+m = get_peft_model(m, LoraConfig(r=LORA_R, lora_alpha=LORA_ALPHA, target_modules=LORA_TARGETS, bias="none", task_type="CAUSAL_LM"))
+hf_loop(m, t, "bnb_nf4")
+
+# FigQuant
+log("\n--- FigQuant LoRA ---"); reset()
+from little_fig.engine import FigModel
+from little_fig.engine.tier import TrainingTier
+model = FigModel.from_pretrained(MODEL, lora_r=LORA_R, lora_alpha=LORA_ALPHA, tier=TrainingTier.STREAMING_LORA, target_modules=LORA_TARGETS)
+tok = model.tokenizer
+examples = [dict(r) for r in ds]
+def ftok(ex):
+    inst=ex.get("instruction",""); inp=ex.get("input","").strip(); out=ex.get("output","")
+    txt = f"### Instruction:\n{inst}\n\n### Input:\n{inp}\n\n### Response:\n{out}" if inp else f"### Instruction:\n{inst}\n\n### Response:\n{out}"
+    e = tok(txt, truncation=True, max_length=MAX_SEQ, padding="max_length")
+    return {"input_ids":e["input_ids"],"labels":e["input_ids"].copy(),"attention_mask":e["attention_mask"]}
+tokenized = [ftok(ex) for ex in examples]
+class DS(torch.utils.data.Dataset):
+    def __init__(s,d): s.d=d
+    def __len__(s): return len(s.d)
+    def __getitem__(s,i): return {k:torch.tensor(v,dtype=torch.long) for k,v in s.d[i].items()}
+from torch.utils.data import DataLoader
+dl = DataLoader(DS(tokenized), batch_size=BATCH_SIZE, shuffle=True, drop_last=True)
+dev = torch.device("cuda"); model = model.to(dev)
+params = model.get_trainable_parameters()
+opt = torch.optim.AdamW(params, lr=LR, weight_decay=0.01)
+model.model.train(); losses=[]; gs=0; al=0.0; reset(); t0=time.time()
+for batch in dl:
+    if gs>=TRAIN_STEPS*GRAD_ACCUM: break
+    batch = {k:v.to(dev) for k,v in batch.items()}
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        loss = model(input_ids=batch["input_ids"],attention_mask=batch["attention_mask"],labels=batch["labels"]).loss / GRAD_ACCUM
+    loss.backward(); al+=loss.item(); gs+=1
+    if gs%GRAD_ACCUM==0:
+        torch.nn.utils.clip_grad_norm_(params,1.0); opt.step(); opt.zero_grad()
+        s=gs//GRAD_ACCUM; losses.append(al); al=0.0
+        if s%20==0: log(f"  [figquant] step={s} loss={losses[-1]:.4f}")
+tt=time.time()-t0; pm=gpu_mb()
+RESULTS["figquant"] = {"final":float(losses[-1]),"time_s":tt,"steps":len(losses),"gpu_mb":pm}
+log(f"  [figquant] loss={losses[-1]:.4f} time={tt:.0f}s gpu={pm:.0f}MB")
+del model, opt; gc.collect(); torch.cuda.empty_cache()
+
+# ═══════════════════════════════════════════════════════════════════════════════
+log("\n" + "="*60)
+log("  🍐 FINAL RESULTS: Little Fig vs Industry")
+log("="*60)
+if "quant" in RESULTS:
+    q=RESULTS["quant"]
+    log(f"\n  QUANTIZATION ({q['n']} layers): FigQuant vs NF4 = {q['mvn']:+.1f}% MSE (wins {q['fw']}/{q['n']})")
+log(f"\n  TRAINING:")
+log(f"  {'Method':>12} {'Loss':>8} {'Time':>7} {'GPU MB':>8}")
+log(f"  {'─'*40}")
+for k in ["fp16_lora","bnb_nf4","figquant"]:
+    if k in RESULTS:
+        r=RESULTS[k]
+        log(f"  {k:>12} {r['final']:.4f} {r['time_s']:.0f}s {r['gpu_mb']:.0f}")
+log("="*60)
+with open("/app/results.json","w") as f: json.dump(RESULTS,f,indent=2,default=str)
+log("📁 Done.")