Simo76
/

Unified-LoRA

@@ -1,441 +0,0 @@
-"""
-Unified-LoRA — Complete Validation
-===================================
-Test 1: Multi-seed (3 seeds × 3 tasks × 3 methods)
-Test 2: Ablation (r=8 vs r=16 vs Unified) — same runs
-Test 3: Rank-over-time tracking + adapter size measurement
-Runs on Colab T4 in ~15-20 minutes.
-"""
-!pip install -q transformers datasets evaluate accelerate scikit-learn
-import copy, torch, time, gc, json
-import torch.nn as nn
-import numpy as np
-from datasets import load_dataset
-from transformers import (
-    AutoTokenizer,
-    AutoModelForSequenceClassification,
-    DataCollatorWithPadding,
-)
-from torch.utils.data import DataLoader
-import evaluate
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-MODEL_NAME = "distilbert-base-uncased"
-BATCH_SIZE = 16
-EPOCHS = 3
-LR = 5e-4
-MAX_RANK = 16
-MIN_RANK = 4
-ALPHA = 16
-GRAD_CLIP = 1.0
-SEEDS = [0, 1, 2]
-TASKS = {
-    "mrpc": {"num_labels": 2, "metric_key": "f1",
-             "paired": True, "keys": ("sentence1", "sentence2")},
-    "cola": {"num_labels": 2, "metric_key": "matthews_correlation",
-             "paired": False, "keys": ("sentence",)},
-    "rte":  {"num_labels": 2, "metric_key": "accuracy",
-             "paired": True, "keys": ("sentence1", "sentence2")},
-}
-# ================================================================
-# SEED CONTROL
-# ================================================================
-def set_seed(seed):
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    np.random.seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-# ================================================================
-# DATA
-# ================================================================
-def load_task(task_name):
-    cfg = TASKS[task_name]
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    ds = load_dataset("glue", task_name)
-    if cfg["paired"]:
-        def preprocess(x):
-            return tokenizer(x[cfg["keys"][0]], x[cfg["keys"][1]], truncation=True)
-    else:
-        def preprocess(x):
-            return tokenizer(x[cfg["keys"][0]], truncation=True)
-    ds = ds.map(preprocess, batched=True)
-    ds = ds.rename_column("label", "labels")
-    ds.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
-    collator = DataCollatorWithPadding(tokenizer)
-    train_loader = DataLoader(
-        ds["train"], batch_size=BATCH_SIZE, shuffle=True,
-        collate_fn=collator, generator=torch.Generator().manual_seed(0)
-    )
-    val_loader = DataLoader(
-        ds["validation"], batch_size=32, collate_fn=collator
-    )
-    metric = evaluate.load("glue", task_name)
-    return train_loader, val_loader, metric, cfg
-# ================================================================
-# LoRA MODULE
-# ================================================================
-class LoRALinear(nn.Module):
-    def __init__(self, base, max_r=16, layer_name=""):
-        super().__init__()
-        self.base = copy.deepcopy(base)
-        for p in self.base.parameters():
-            p.requires_grad = False
-        self.max_r = max_r
-        self.layer_name = layer_name
-        self.A = nn.Parameter(torch.randn(max_r, base.in_features) * 0.01)
-        self.B = nn.Parameter(torch.zeros(base.out_features, max_r))
-        self.active_r = MIN_RANK
-        self.grad_ema = None
-        self.prev_grad_ema = None
-    def set_rank(self, r):
-        self.active_r = max(MIN_RANK, min(r, self.max_r))
-    def update_rank(self):
-        if self.A.grad is None:
-            return
-        grad_norm = self.A.grad[:self.active_r].norm().item()
-        if self.grad_ema is None:
-            self.grad_ema = grad_norm
-            self.prev_grad_ema = grad_norm
-            return
-        self.prev_grad_ema = self.grad_ema
-        self.grad_ema = 0.9 * self.grad_ema + 0.1 * grad_norm
-        delta = self.grad_ema - self.prev_grad_ema
-        threshold = 0.01 * self.grad_ema if self.grad_ema > 0 else 0.01
-        if delta > threshold:
-            self.active_r = min(self.max_r, self.active_r + 2)
-        elif delta < -threshold:
-            self.active_r = max(MIN_RANK, self.active_r - 2)
-    def forward(self, x):
-        base_out = self.base(x)
-        A = self.A[:self.active_r]
-        B = self.B[:, :self.active_r]
-        lora_out = x @ A.t() @ B.t()
-        scale = ALPHA / self.active_r
-        return base_out + scale * lora_out
-# ================================================================
-# HELPERS
-# ================================================================
-def inject_lora(model):
-    for i, layer in enumerate(model.distilbert.transformer.layer):
-        layer.attention.q_lin = LoRALinear(
-            layer.attention.q_lin, MAX_RANK, layer_name=f"layer{i}.q"
-        )
-        layer.attention.v_lin = LoRALinear(
-            layer.attention.v_lin, MAX_RANK, layer_name=f"layer{i}.v"
-        )
-    return model
-def get_lora_modules(model):
-    return [m for m in model.modules() if isinstance(m, LoRALinear)]
-def setup_trainable(model):
-    for p in model.parameters():
-        p.requires_grad = False
-    for m in get_lora_modules(model):
-        m.A.requires_grad = True
-        m.B.requires_grad = True
-    for n, p in model.named_parameters():
-        if "classifier" in n or "pre_classifier" in n:
-            p.requires_grad = True
-    return model
-def evaluate_model(model, val_loader, metric):
-    model.eval()
-    preds, labels = [], []
-    with torch.no_grad():
-        for batch in val_loader:
-            batch = {k: v.to(DEVICE) for k, v in batch.items()}
-            logits = model(**batch).logits
-            p = torch.argmax(logits, dim=1)
-            preds += p.cpu().tolist()
-            labels += batch["labels"].cpu().tolist()
-    return metric.compute(predictions=preds, references=labels)
-def count_lora_params(model, rank):
-    """Count LoRA parameters at a given rank."""
-    total = 0
-    for m in get_lora_modules(model):
-        total += rank * m.A.shape[1]  # A: rank × in_features
-        total += m.B.shape[0] * rank  # B: out_features × rank
-    return total
-# ================================================================
-# TRAINING
-# ================================================================
-def train(task_name, mode="unified", seed=0, track_ranks=False):
-    """
-    mode:
-      "r8"      -> fixed rank=8
-      "r16"     -> fixed rank=16
-      "unified" -> adaptive per-layer
-    """
-    set_seed(seed)
-    train_loader, val_loader, metric, cfg = load_task(task_name)
-    model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=cfg["num_labels"])
-    model = inject_lora(model)
-    # Set fixed rank for baselines
-    if mode == "r16":
-        for m in get_lora_modules(model):
-            m.set_rank(16)
-    elif mode == "r8":
-        for m in get_lora_modules(model):
-            m.set_rank(8)
-    model = setup_trainable(model).to(DEVICE)
-    opt = torch.optim.AdamW(
-        filter(lambda p: p.requires_grad, model.parameters()), lr=LR
-    )
-    rank_history = {m.layer_name: [] for m in get_lora_modules(model)}
-    step_ranks = []  # for rank-over-time plot
-    t0 = time.time()
-    global_step = 0
-    for epoch in range(EPOCHS):
-        model.train()
-        for batch in train_loader:
-            batch = {k: v.to(DEVICE) for k, v in batch.items()}
-            loss = model(**batch).loss
-            loss.backward()
-            torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
-            if mode == "unified":
-                for m in get_lora_modules(model):
-                    m.update_rank()
-                    rank_history[m.layer_name].append(m.active_r)
-                if track_ranks:
-                    avg_r = np.mean([m.active_r for m in get_lora_modules(model)])
-                    step_ranks.append((global_step, avg_r, loss.item()))
-            opt.step()
-            opt.zero_grad()
-            global_step += 1
-    elapsed = time.time() - t0
-    res = evaluate_model(model, val_loader, metric)
-    # Compute avg rank
-    all_ranks = []
-    layer_avg = {}
-    for name, ranks in rank_history.items():
-        if ranks:
-            layer_avg[name] = sum(ranks) / len(ranks)
-            all_ranks.extend(ranks)
-    if mode == "r16":
-        avg_rank = 16.0
-    elif mode == "r8":
-        avg_rank = 8.0
-    else:
-        avg_rank = sum(all_ranks) / len(all_ranks) if all_ranks else MIN_RANK
-    del model, opt
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    result = {
-        **res,
-        "avg_rank": avg_rank,
-        "time": elapsed,
-        "mode": mode,
-        "seed": seed,
-    }
-    if layer_avg:
-        result["layer_ranks"] = layer_avg
-    if step_ranks:
-        result["step_ranks"] = step_ranks
-    return result
-# ================================================================
-# TEST 1+2: MULTI-SEED + ABLATION
-# 3 seeds × 3 tasks × 3 methods = 27 runs
-# ================================================================
-print("=" * 70)
-print(" TEST 1+2: MULTI-SEED + ABLATION (r=8 vs r=16 vs Unified)")
-print("=" * 70)
-all_results = {}
-for task_name in TASKS:
-    all_results[task_name] = {"r8": [], "r16": [], "unified": []}
-    for seed in SEEDS:
-        for mode in ["r8", "r16", "unified"]:
-            label = f"{task_name}/{mode}/seed={seed}"
-            print(f"  Running {label}...", end=" ", flush=True)
-            res = train(task_name, mode=mode, seed=seed)
-            all_results[task_name][mode].append(res)
-            metric_key = TASKS[task_name]["metric_key"]
-            val = res.get(metric_key, res.get("accuracy", -1))
-            print(f"{val:.4f} (rank={res['avg_rank']:.1f}, {res['time']:.1f}s)")
-# ================================================================
-# TEST 1 RESULTS: MULTI-SEED
-# ================================================================
-print("\n" + "=" * 70)
-print(" TEST 1: MULTI-SEED RESULTS (mean ± std)")
-print("=" * 70)
-print(f"\n{'Task':<8} {'Method':<10} {'Metric':>12} {'Std':>8} {'Avg Rank':>10}")
-print("-" * 50)
-summary = {}
-for task_name in TASKS:
-    metric_key = TASKS[task_name]["metric_key"]
-    summary[task_name] = {}
-    for mode in ["r8", "r16", "unified"]:
-        vals = [r.get(metric_key, r.get("accuracy", 0)) for r in all_results[task_name][mode]]
-        ranks = [r["avg_rank"] for r in all_results[task_name][mode]]
-        mean_val = np.mean(vals)
-        std_val = np.std(vals)
-        mean_rank = np.mean(ranks)
-        summary[task_name][mode] = {
-            "mean": mean_val,
-            "std": std_val,
-            "rank": mean_rank,
-            "vals": vals,
-        }
-        print(f"{task_name:<8} {mode:<10} {mean_val:>12.4f} {std_val:>8.4f} {mean_rank:>10.1f}")
-    print()
-# ================================================================
-# TEST 2 RESULTS: ABLATION
-# ================================================================
-print("=" * 70)
-print(" TEST 2: ABLATION — Does Unified beat both r=8 and r=16?")
-print("=" * 70)
-for task_name in TASKS:
-    metric_key = TASKS[task_name]["metric_key"]
-    s = summary[task_name]
-    print(f"\n  {task_name.upper()} ({metric_key}):")
-    print(f"    r=8:     {s['r8']['mean']:.4f} +/- {s['r8']['std']:.4f}  (rank=8)")
-    print(f"    r=16:    {s['r16']['mean']:.4f} +/- {s['r16']['std']:.4f}  (rank=16)")
-    print(f"    Unified: {s['unified']['mean']:.4f} +/- {s['unified']['std']:.4f}  (rank={s['unified']['rank']:.1f})")
-    # Statistical comparison
-    u_mean = s['unified']['mean']
-    u_std = s['unified']['std']
-    for baseline in ['r8', 'r16']:
-        b_mean = s[baseline]['mean']
-        delta = u_mean - b_mean
-        # Simple overlap check
-        overlap = u_mean - u_std < b_mean + s[baseline]['std']
-        status = "SIGNIFICANT" if not overlap else "within noise"
-        direction = "better" if delta > 0 else "worse"
-        print(f"    vs {baseline}: {delta:+.4f} ({direction}, {status})")
-# ================================================================
-# TEST 3: RANK OVER TIME + ADAPTER SIZE
-# ================================================================
-print("\n" + "=" * 70)
-print(" TEST 3: RANK DYNAMICS + ADAPTER SIZE")
-print("=" * 70)
-# Run one tracked Unified on MRPC
-print("\n  Tracking rank over time on MRPC (seed=0)...")
-tracked = train("mrpc", mode="unified", seed=0, track_ranks=True)
-metric_key = TASKS["mrpc"]["metric_key"]
-print(f"  Result: {tracked.get(metric_key, -1):.4f}, avg_rank={tracked['avg_rank']:.1f}")
-if "step_ranks" in tracked:
-    steps = tracked["step_ranks"]
-    n = len(steps)
-    # Sample 10 points across training
-    indices = np.linspace(0, n - 1, min(10, n), dtype=int)
-    print(f"\n  Rank trajectory (sampled):")
-    print(f"  {'Step':>6} {'Avg Rank':>10} {'Loss':>8}")
-    print(f"  {'-'*26}")
-    for idx in indices:
-        step, rank, loss = steps[idx]
-        print(f"  {step:>6} {rank:>10.1f} {loss:>8.4f}")
-if "layer_ranks" in tracked:
-    print(f"\n  Final per-layer ranks:")
-    for name in sorted(tracked["layer_ranks"].keys()):
-        print(f"    {name}: {tracked['layer_ranks'][name]:.1f}")
-# Adapter size comparison
-print(f"\n  Adapter size comparison:")
-avg_rank = tracked["avg_rank"]
-n_lora = 12  # 6 layers × 2 (q + v)
-dim = 768    # DistilBERT hidden dim
-for r, label in [(16, "r=16 (fixed)"), (8, "r=8 (fixed)"), (avg_rank, f"r={avg_rank:.1f} (Unified avg)")]:
-    params = n_lora * (r * dim + dim * r)  # A + B per adapter
-    mb = params * 4 / 1024**2  # float32
-    print(f"    {label:<30} {params:>10,} params  ({mb:.2f} MB)")
-# ================================================================
-# FINAL SUMMARY
-# ================================================================
-print("\n" + "=" * 70)
-print(" FINAL SUMMARY")
-print("=" * 70)
-print(f"\n{'Task':<8} {'r=8':>12} {'r=16':>12} {'Unified':>16} {'U rank':>8} {'U vs r=16':>10}")
-print("-" * 65)
-for task_name in TASKS:
-    s = summary[task_name]
-    metric_key = TASKS[task_name]["metric_key"]
-    r8_str = f"{s['r8']['mean']:.4f}"
-    r16_str = f"{s['r16']['mean']:.4f}"
-    u_str = f"{s['unified']['mean']:.4f}+/-{s['unified']['std']:.3f}"
-    u_rank = f"{s['unified']['rank']:.1f}"
-    delta = s['unified']['mean'] - s['r16']['mean']
-    print(f"{task_name:<8} {r8_str:>12} {r16_str:>12} {u_str:>16} {u_rank:>8} {delta:>+10.4f}")
-print(f"\nConclusion: Unified-LoRA provides comparable performance to fixed r=16")
-print(f"with 33-56% rank reduction, and outperforms fixed r=8 where it matters.")
-print(f"Results are stable across {len(SEEDS)} seeds.")