Spaces:
Sleeping
Sleeping
| """TensorPEFT live demo: tensor gradient-pursuit adapters vs. matched-budget PEFT baselines. | |
| GPT-2 (124M, English) is fine-tuned on Bengali news with the same number of trainable | |
| parameters for every method; validation cross-entropy (nats/token) decides. | |
| """ | |
| import copy | |
| import glob | |
| import io | |
| import json | |
| import math | |
| import os | |
| import random | |
| import re | |
| import time | |
| import unicodedata | |
| import urllib.request | |
| import zipfile | |
| import gradio as gr | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| import torch # noqa: E402 | |
| import torch.nn.functional as F # noqa: E402 | |
| from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402 | |
| from baselines import IA3, LNTune, LoRA, Prompt, _Adapter # noqa: E402 | |
| from tensorpeft import TensorPEFTConfig, get_tensorpeft_model # noqa: E402 | |
| try: | |
| import spaces | |
| GPU = spaces.GPU | |
| except Exception: # running outside a ZeroGPU Space | |
| def GPU(fn=None, **kw): | |
| return fn if fn is not None else (lambda f: f) | |
| MODEL_ID = os.environ.get("TP_MODEL", "openai-community/gpt2") | |
| DATA_URL = os.environ.get( | |
| "TP_DATA_URL", | |
| "https://raw.githubusercontent.com/soham96/Bengali_news_classifier/" | |
| "4c2a51b51fc56208426ef19368bfcb08acc9d0c4/data/data.zip") | |
| CACHE = os.environ.get("TP_CACHE", "data_cache") | |
| SEQ, BATCH, VAL_SEQS = 128, 8, 64 | |
| # ------------------------------------------------------------------ data (built once at start-up) | |
| tok = AutoTokenizer.from_pretrained(MODEL_ID) | |
| def build_data(): | |
| os.makedirs(CACHE, exist_ok=True) | |
| tr, va = os.path.join(CACHE, "train.bin"), os.path.join(CACHE, "val.bin") | |
| if os.path.exists(tr) and os.path.exists(va): | |
| return | |
| src = DATA_URL | |
| raw = open(src, "rb").read() if os.path.exists(src) else urllib.request.urlopen(src).read() | |
| docs = [] | |
| with zipfile.ZipFile(io.BytesIO(raw)) as z: | |
| for name in sorted(n for n in z.namelist() if n.endswith("_articles.txt") and "__MACOSX" not in n): | |
| for a in json.loads(z.read(name).decode("utf-8"))["articles"]: | |
| t = unicodedata.normalize("NFC", (a.get("title", "") + "ΰ₯€ " + a.get("body", "")).strip()) | |
| t = re.sub(r"\s+", " ", t) | |
| if len(t) > 200: | |
| docs.append(t) | |
| random.Random(0).shuffle(docs) | |
| def enc(ds, cap): | |
| ids = [] | |
| for d in ds: | |
| ids += tok(d).input_ids + [tok.eos_token_id] | |
| if len(ids) >= cap: | |
| break | |
| return np.array(ids[:cap], dtype=np.uint16) | |
| enc(docs[:400], 200_000).tofile(va) | |
| enc(docs[400:], 3_000_000).tofile(tr) | |
| build_data() | |
| TRAIN = np.fromfile(os.path.join(CACHE, "train.bin"), dtype=np.uint16) | |
| VAL = torch.from_numpy(np.fromfile(os.path.join(CACHE, "val.bin"), dtype=np.uint16)[: VAL_SEQS * (SEQ + 1)] | |
| .astype(np.int64).reshape(VAL_SEQS, SEQ + 1)) | |
| BASE = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32).eval() | |
| for _p in BASE.parameters(): | |
| _p.requires_grad_(False) | |
| def sampler(seed): | |
| rng = np.random.default_rng(seed) | |
| while True: | |
| ix = rng.integers(0, len(TRAIN) - SEQ - 1, BATCH) | |
| yield torch.from_numpy(np.stack([TRAIN[i:i + SEQ + 1] for i in ix]).astype(np.int64)) | |
| # ------------------------------------------------------------------ methods per budget | |
| class TP(_Adapter): | |
| lr, warmup = 3e-3, 10 | |
| def __init__(self, model, **cfg): | |
| super().__init__() | |
| self.model = [model] | |
| self.tp = get_tensorpeft_model(model, TensorPEFTConfig(**cfg)) | |
| def setup(self, probe_batches): | |
| self.tp.init_atoms(probe_batches, verbose=False) | |
| def trainable(self): | |
| return [p for p in self.tp.parameters() if p.requires_grad] | |
| def mark_ready(self): | |
| for g in self.tp.groups: | |
| g.initialized.fill_(True) | |
| ATTN = ["c_attn.q", "c_attn.k", "c_attn.v", "attn.c_proj"] | |
| BUDGETS = { | |
| "β1.5k parameters": [ | |
| ("TensorPEFT Β· 1 attention atom", lambda m: TP(m, target_modules=ATTN, atoms_per_group=1)), | |
| ("Prompt tuning Β· 1 token", lambda m: Prompt(m, 1)), | |
| ("LayerNorm tuning Β· final LN", lambda m: LNTune(m, layers=[], final=True)), | |
| ], | |
| "β9k parameters": [ | |
| ("TensorPEFT Β· 3 atoms", lambda m: TP(m, atoms_per_group=1)), | |
| ("Prompt tuning Β· 12 tokens", lambda m: Prompt(m, 12)), | |
| ("LoRA r=1 Β· q,v of last 3 layers", lambda m: LoRA(m, 1, layers=[9, 10, 11])), | |
| ("(IA)Β³ Β· values", lambda m: IA3(m, kv=("c_attn.v",), ff=False)), | |
| ("LayerNorm tuning Β· last 3 layers", lambda m: LNTune(m, layers=[9, 10, 11], final=False)), | |
| ], | |
| "β55k parameters": [ | |
| ("TensorPEFT Β· 18 atoms", lambda m: TP(m, atoms_per_group=6)), | |
| ("(IA)Β³ Β· keys, values, MLP", lambda m: IA3(m)), | |
| ], | |
| } | |
| COLORS = ["#7c3aed", "#0ea5e9", "#f59e0b", "#10b981", "#ef4444"] | |
| def ce(adapter, x): | |
| return F.cross_entropy(adapter.logits(x[:, :-1]).flatten(0, 1), x[:, 1:].flatten()) | |
| def evaluate(adapter, dev): | |
| V = VAL.to(dev) | |
| return sum(float(ce(adapter, V[i:i + 16])) * len(V[i:i + 16]) for i in range(0, len(V), 16)) / len(V) | |
| def _duration(budget, steps, seed): | |
| return int(min(280, 30 + len(BUDGETS[budget]) * steps * 0.06)) | |
| def race_gpu(budget, steps, seed): | |
| dev = "cuda" if torch.cuda.is_available() else "cpu" | |
| if dev == "cuda": | |
| torch.backends.cuda.matmul.allow_tf32 = False | |
| torch.backends.cudnn.allow_tf32 = False | |
| out = [] | |
| plain = _Adapter() | |
| plain.model = [copy.deepcopy(BASE).to(dev)] | |
| base_val = evaluate(plain, dev) | |
| del plain | |
| for mi, (name, build) in enumerate(BUDGETS[budget]): | |
| torch.manual_seed(seed) | |
| model = copy.deepcopy(BASE).to(dev).eval() | |
| ad = build(model) | |
| t0 = time.time() | |
| if isinstance(ad, TP): | |
| ad.setup([b.to(dev) for b, _ in zip(sampler(seed + 12345), range(4))]) | |
| params = ad.trainable() | |
| opt = torch.optim.AdamW(params, lr=ad.lr, weight_decay=0.0) | |
| warm = getattr(ad, "warmup", 0) | |
| batches = sampler(seed) | |
| curve = [(0, evaluate(ad, dev))] | |
| for step in range(1, steps + 1): | |
| if warm: | |
| for g in opt.param_groups: | |
| g["lr"] = ad.lr * min(1.0, step / warm) | |
| loss = ce(ad, next(batches).to(dev)) | |
| opt.zero_grad(set_to_none=True) | |
| loss.backward() | |
| opt.step() | |
| if step % 25 == 0 or step == steps: | |
| curve.append((step, evaluate(ad, dev))) | |
| out.append(dict(name=name, n_params=ad.n_params(), curve=curve, final=curve[-1][1], | |
| seconds=time.time() - t0, | |
| state=[p.detach().float().cpu().clone() for p in params])) | |
| del model, ad, opt | |
| return dict(budget=budget, steps=steps, seed=seed, base=base_val, runs=out) | |
| def run_race(budget, steps, seed): | |
| res = race_gpu(budget, int(steps), int(seed)) | |
| fig, ax = plt.subplots(figsize=(8, 4.6)) | |
| for i, r in enumerate(res["runs"]): | |
| xs, ys = zip(*r["curve"]) | |
| ax.plot(xs, ys, color=COLORS[i % len(COLORS)], lw=2.6 if r["name"].startswith("TensorPEFT") else 1.6, | |
| label=f"{r['name']} ({r['n_params']:,} params)") | |
| ax.axhline(res["base"], color="#9ca3af", ls="--", lw=1, label=f"GPT-2 before fine-tuning ({res['base']:.3f})") | |
| ax.set_xlabel("training step") | |
| ax.set_ylabel("validation loss (nats/token) β better") | |
| ax.set_title(f"Bengali fine-tuning of GPT-2 at {res['budget']}") | |
| ax.grid(alpha=0.25) | |
| ax.legend(fontsize=8, loc="upper right") | |
| fig.tight_layout() | |
| ranked = sorted(res["runs"], key=lambda r: r["final"]) | |
| best = ranked[0] | |
| rows = [[i + 1, r["name"], f"{r['n_params']:,}", f"{r['final']:.4f}", f"{r['final'] - best['final']:+.4f}", | |
| f"{r['seconds']:.0f}s"] for i, r in enumerate(ranked)] | |
| tp = next(r for r in res["runs"] if r["name"].startswith("TensorPEFT")) | |
| others = [r for r in res["runs"] if r is not tp] | |
| nxt = min(others, key=lambda r: r["final"]) | |
| verdict = (f"### {'π TensorPEFT wins' if tp is best else 'TensorPEFT does not win this run'}\n" | |
| f"TensorPEFT **{tp['final']:.4f}** vs best baseline ({nxt['name']}) **{nxt['final']:.4f}** " | |
| f"β difference **{tp['final'] - nxt['final']:+.4f}** nats/token after {res['steps']} steps (seed {res['seed']}).") | |
| store = {"budget": res["budget"], "runs": [{"name": r["name"], "state": r["state"]} for r in res["runs"]]} | |
| return fig, rows, verdict, store | |
| def generate_gpu(store, prompt, n_new, temperature, seed): | |
| dev = "cuda" if torch.cuda.is_available() else "cpu" | |
| ids0 = tok(prompt, return_tensors="pt").input_ids.to(dev) | |
| builders = dict(BUDGETS[store["budget"]]) if store else {} | |
| entries = [("GPT-2 before fine-tuning", None, None)] + \ | |
| [(r["name"], builders[r["name"]], r["state"]) for r in (store["runs"] if store else [])] | |
| outs = [] | |
| for name, build, state in entries: | |
| model = copy.deepcopy(BASE).to(dev).eval() | |
| if build is None: | |
| ad = _Adapter() | |
| ad.model = [model] | |
| else: | |
| ad = build(model) | |
| if isinstance(ad, TP): | |
| ad.mark_ready() | |
| with torch.no_grad(): | |
| for p, s in zip(ad.trainable(), state): | |
| p.copy_(s.to(p.device, p.dtype)) | |
| g = torch.Generator(device=dev).manual_seed(int(seed)) | |
| ids = ids0.clone() | |
| with torch.no_grad(): | |
| for _ in range(int(n_new)): | |
| logits = ad.logits(ids[:, -512:])[:, -1].float() / max(float(temperature), 1e-3) | |
| top = torch.topk(logits, 50) | |
| probs = torch.softmax(top.values, -1) | |
| nxt = top.indices.gather(-1, torch.multinomial(probs, 1, generator=g)) | |
| ids = torch.cat([ids, nxt], dim=1) | |
| outs.append((name, tok.decode(ids[0, ids0.shape[1]:], skip_special_tokens=True))) | |
| del model, ad | |
| return outs | |
| def run_generate(store, prompt, n_new, temperature, seed): | |
| if not prompt.strip(): | |
| return "Enter a Bengali prompt." | |
| outs = generate_gpu(store, prompt, n_new, temperature, seed) | |
| note = "" if store else "*Run a race first to compare the fine-tuned models; showing the base model only.*\n\n" | |
| return note + "\n\n".join(f"**{n}**\n\n> {prompt}{t}" for n, t in outs) | |
| # ------------------------------------------------------------------ benchmark results (CPU, 300 steps) | |
| RESULTS_MD = """ | |
| ### Three seeds (mean Β± std of validation loss; lower is better) | |
| | budget | TensorPEFT | best baseline | difference | seeds won | p (paired t) | | |
| |---|---:|---|---:|:---:|---:| | |
| | β9k | **1.3245 Β± 0.004** | prompt tuning 1.3684 Β± 0.010 | **β0.044** | 3/3 | 0.008 | | |
| | β37k | 1.2558 Β± 0.010 | LayerNorm tuning (38.4k) 1.2504 Β± 0.003 | +0.005 | 1/3 | 0.37 (tie) | | |
| | β55k | **1.2368 Β± 0.017** | (IA)Β³ 1.2880 Β± 0.003 | **β0.051** | 3/3 | 0.026 | | |
| ### Every method at every budget (seed 0) | |
| | budget | TensorPEFT | LoRA | QLoRA | VeRA | (IA)Β³ | prompt tuning | LayerNorm tuning | BitFit | | |
| |---|---:|---:|---:|---:|---:|---:|---:|---:| | |
| | β1.5k | **1.461** | β | β | β | β | 1.493 (768) | 1.600 | β | | |
| | β9k | **1.327** | 1.430 | β | β | 1.432 | 1.370 | 1.407 | β | | |
| | β37k | 1.255 | 1.318 | 1.335 | 1.393 | β | 1.327 | **1.253** | β | | |
| | β55k | **1.231** | β | β | β | 1.289 | β | β | β | | |
| | β83k | β | β | β | β | β | β | β | 1.326 | | |
| GPT-2 124M before fine-tuning: 1.708. 300 AdamW steps, batch 8 Γ 128 tokens, identical data for all methods, | |
| baseline learning rates from their papers. βββ: the method cannot be configured at that budget on GPT-2 or was | |
| not run there (LoRA's smallest all-layer setting is 37k). At larger budgets (β₯ 115k parameters) LoRA was | |
| better in our tests β TensorPEFT is a method for very small budgets. | |
| """ | |
| def results_figure(): | |
| pts = { | |
| "TensorPEFT": [(1553, 1.461), (9259, 1.327), (37036, 1.255), (55554, 1.231)], | |
| "LoRA / QLoRA": [(9216, 1.430), (36864, 1.318), (36864, 1.335)], | |
| "Prompt tuning": [(768, 1.493), (9216, 1.370), (36864, 1.327)], | |
| "LayerNorm tuning": [(1536, 1.600), (9216, 1.407), (38400, 1.253)], | |
| "(IA)Β³": [(9216, 1.432), (55296, 1.289)], | |
| "VeRA": [(24576, 1.393)], | |
| "BitFit": [(82944, 1.326)], | |
| } | |
| fig, ax = plt.subplots(figsize=(8, 4.6)) | |
| for i, (k, v) in enumerate(pts.items()): | |
| xs, ys = zip(*v) | |
| if k == "TensorPEFT": | |
| ax.plot(xs, ys, "-o", color=COLORS[0], lw=2.6, ms=7, label=k, zorder=5) | |
| else: | |
| ax.scatter(xs, ys, s=45, color=["#0ea5e9", "#f59e0b", "#10b981", "#ef4444", "#64748b", "#a16207"][i - 1], | |
| label=k) | |
| ax.set_xscale("log") | |
| ax.set_xlabel("trainable parameters (log scale)") | |
| ax.set_ylabel("validation loss (nats/token) β better") | |
| ax.set_title("GPT-2 β Bengali: loss vs. trainable parameters (seed 0)") | |
| ax.grid(alpha=0.25, which="both") | |
| ax.legend(fontsize=8) | |
| fig.tight_layout() | |
| return fig | |
| # ------------------------------------------------------------------ UI | |
| INTRO = """ | |
| # TensorPEFT Β· tensor gradient-pursuit fine-tuning | |
| LoRA gives every weight matrix its own low-rank update. **TensorPEFT** stacks all equal-shaped matrices of the | |
| network into one tensor (layers Γ {q,k,v,o} Γ d_in Γ d_out) and learns a few **rank-one CP atoms** of it, chosen from | |
| one gradient probe before training. Each atom touches every layer for only L + P + d_in + d_out + 1 parameters. | |
| Here GPT-2 (English) is taught **Bengali** with the *same number of trainable parameters* for every method. | |
| `pip install git+https://huggingface.co/Angshul/TensorPEFT` | |
| """ | |
| with gr.Blocks(title="TensorPEFT") as demo: | |
| gr.Markdown(INTRO) | |
| store = gr.State(None) | |
| with gr.Tab("Live race"): | |
| with gr.Row(): | |
| budget = gr.Radio(list(BUDGETS), value="β9k parameters", label="Trainable-parameter budget") | |
| steps = gr.Number(value=300, visible=False) | |
| seed = gr.Number(value=0, precision=0, label="Seed") | |
| gr.Markdown("Every method trains for **300 AdamW steps** (batch 8 Γ 128 tokens) on identical batches β the setting of the benchmark. " | |
| "TensorPEFT starts from zero coefficients, so it trails in the first few dozen steps and overtakes later.") | |
| go = gr.Button("Run the race", variant="primary") | |
| verdict = gr.Markdown() | |
| plot = gr.Plot() | |
| table = gr.Dataframe(headers=["rank", "method", "params", "val loss", "vs best", "time"], interactive=False) | |
| go.click(run_race, [budget, steps, seed], [plot, table, verdict, store]) | |
| with gr.Tab("Generate"): | |
| gr.Markdown("Continuations from the models trained in your last race (top-k 50 sampling).") | |
| prompt = gr.Textbox(value="ΰ¦ΰ¦²ΰ¦ΰ¦Ύΰ¦€ΰ¦Ύ ΰ¦Άΰ¦Ήΰ¦°ΰ§ ΰ¦ΰ¦", label="Bengali prompt") | |
| with gr.Row(): | |
| n_new = gr.Slider(20, 120, value=60, step=10, label="New tokens") | |
| temp = gr.Slider(0.3, 1.2, value=0.8, step=0.1, label="Temperature") | |
| gseed = gr.Number(value=0, precision=0, label="Seed") | |
| gen = gr.Button("Generate", variant="primary") | |
| gen_out = gr.Markdown() | |
| gen.click(run_generate, [store, prompt, n_new, temp, gseed], gen_out) | |
| with gr.Tab("Benchmark results"): | |
| gr.Plot(value=results_figure()) | |
| gr.Markdown(RESULTS_MD) | |
| if __name__ == "__main__": | |
| demo.launch() | |