Upload 6 files

README.md

@@ -1,14 +1,30 @@
 ---
 title: TensorPEFT
-emoji: 🐢
-colorFrom: yellow
-colorTo: pink
+emoji: 🧊
+colorFrom: purple
+colorTo: blue
 sdk: gradio
 sdk_version: 6.31.0
-python_version: '3.12'
 app_file: app.py
 pinned: false
 license: apache-2.0
+short_description: Tensor-pursuit fine-tuning vs LoRA at tiny budgets
+tags:
+- peft
+- fine-tuning
+- tensor-decomposition
+- lora-alternative
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# TensorPEFT live demo
+
+Live race between **TensorPEFT** (tensor gradient-pursuit adapters) and matched-budget PEFT baselines
+(prompt tuning, LoRA, (IA)³, LayerNorm tuning): GPT-2 124M is fine-tuned on Bengali news with the same number of
+trainable parameters for every method, and validation cross-entropy decides.
+
+- Library: [Angshul/TensorPEFT](https://huggingface.co/Angshul/TensorPEFT) — `pip install git+https://huggingface.co/Angshul/TensorPEFT`
+- Data: Bengali news articles from [soham96/Bengali_news_classifier](https://github.com/soham96/Bengali_news_classifier),
+  downloaded at start-up (not redistributed), tokenised with the GPT-2 tokenizer; 3M training tokens, held-out articles for validation.
+- Base model: [openai-community/gpt2](https://huggingface.co/openai-community/gpt2).
+
+Runs on ZeroGPU. Benchmark numbers in the *Benchmark results* tab were measured on CPU (fp32, 300 steps).

app.py

@@ -0,0 +1,358 @@
+"""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())
+
+
+@torch.no_grad()
+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))
+
+
+@GPU(duration=_duration)
+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
+
+
+@GPU(duration=90)
+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()

baselines.py

@@ -0,0 +1,98 @@
+"""Matched-budget PEFT baselines for the demo (GPT-2), built on tensorpeft's slot discovery."""
+import math
+
+import torch
+import torch.nn as nn
+from tensorpeft.layers import attach_hooks, find_slots
+
+
+class _Adapter(nn.Module):
+    lr = 1e-3
+
+    def trainable(self):
+        return [p for p in self.parameters() if p.requires_grad]
+
+    def n_params(self):
+        return sum(p.numel() for p in self.trainable())
+
+    def logits(self, ids):
+        return self.model[0](ids).logits
+
+
+class LoRA(_Adapter):
+    lr = 1e-3
+
+    def __init__(self, model, rank=1, targets=("c_attn.q", "c_attn.v"), layers=None, seed=0):
+        super().__init__()
+        self.model = [model]
+        self.slots = [s for s in find_slots(model, list(targets)) if layers is None or s.layer in layers]
+        g = torch.Generator().manual_seed(seed)
+        dev = next(model.parameters()).device
+        self.A = nn.ParameterList([nn.Parameter(((torch.rand(s.shape[0], rank, generator=g) * 2 - 1)
+                                                 / math.sqrt(s.shape[0])).to(dev)) for s in self.slots])
+        self.B = nn.ParameterList([nn.Parameter(torch.zeros(rank, s.shape[1], device=dev)) for s in self.slots])
+        self.scale = 2.0
+        self._i = {id(s): i for i, s in enumerate(self.slots)}
+        self._h = attach_hooks(self.slots, lambda s, x: (x @ self.A[self._i[id(s)]]) @ self.B[self._i[id(s)]] * self.scale)
+
+
+class IA3(_Adapter):
+    lr = 3e-3
+
+    def __init__(self, model, kv=("c_attn.k", "c_attn.v"), ff=True):
+        super().__init__()
+        self.model = [model]
+        self.out_slots = find_slots(model, list(kv))
+        self.ff_slots = find_slots(model, ["mlp.c_proj"]) if ff else []
+        dev = next(model.parameters()).device
+        self.l_out = nn.ParameterList([nn.Parameter(torch.ones(s.shape[1], device=dev)) for s in self.out_slots])
+        self.l_ff = nn.ParameterList([nn.Parameter(torch.ones(s.shape[0], device=dev)) for s in self.ff_slots])
+        io = {id(s): i for i, s in enumerate(self.out_slots)}
+        iff = {id(s): i for i, s in enumerate(self.ff_slots)}
+
+        def d_out(s, x):
+            y = x @ s.weight() + s.module.bias[s.c0:s.c1]
+            return y * (self.l_out[io[id(s)]] - 1)
+
+        def d_ff(s, x):
+            return (x * (self.l_ff[iff[id(s)]] - 1)) @ s.weight()
+
+        self._h = attach_hooks(self.out_slots, d_out) + attach_hooks(self.ff_slots, d_ff)
+
+
+class Prompt(_Adapter):
+    lr = 3e-2
+
+    def __init__(self, model, n, seed=0):
+        super().__init__()
+        self.model = [model]
+        emb = model.get_input_embeddings().weight
+        g = torch.Generator().manual_seed(seed)
+        ids = torch.randint(0, emb.shape[0], (n,), generator=g).to(emb.device)
+        self.n = n
+        self.P = nn.Parameter(emb[ids].detach().clone())
+
+    def logits(self, ids):
+        m = self.model[0]
+        e = torch.cat([self.P.expand(ids.shape[0], -1, -1), m.get_input_embeddings()(ids)], dim=1)
+        return m(inputs_embeds=e).logits[:, self.n:]
+
+
+class LNTune(_Adapter):
+    lr = 1e-3
+
+    def __init__(self, model, layers=None, final=True):
+        super().__init__()
+        self.model = [model]
+        self.params = []
+        for name, mod in model.named_modules():
+            if isinstance(mod, nn.LayerNorm):
+                parts = name.split(".")
+                layer = next((int(p) for p in parts if p.isdigit()), None)
+                if (layer is None and final) or (layer is not None and (layers is None or layer in layers)):
+                    for p in mod.parameters():
+                        p.requires_grad_(True)
+                        self.params.append(p)
+
+    def trainable(self):
+        return self.params

requirements.txt

@@ -0,0 +1,5 @@
+git+https://huggingface.co/Angshul/TensorPEFT
+torch
+transformers>=4.40
+matplotlib
+numpy