"""Smoke test: lora / pissa / delora on a tiny synthetic transformer-like model. Verifies: 1. Identity at t=0 (delta ~ 0, output close to base). 2. Save/load round-trip preserves outputs. 3. A few SGD steps reduce a random loss (gradients flow). Run: cd lora-lite python -m pip install -e . python tests/smoke.py BLUF format: SHOULD: max|y_adapter - y_base| < tol_init for all variants. ELSE init or hook bug. SHOULD: loss decreases > 5% over 20 SGD steps for all variants. ELSE grad/wiring bug. """ from __future__ import annotations import os, sys, math from pathlib import Path import torch from torch import nn # allow running as `python tests/smoke.py` without install sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) import lora_lite as ll # noqa: E402 ARTIFACT_DIR = Path(__file__).parent / "_artifacts" def assert_no_base_grads(model: nn.Module) -> None: leaked = [name for name, p in model.named_parameters() if "lora_" not in name and p.grad is not None] assert leaked == [], f"base params received grads: {leaked}" # ---- a tiny transformer-like stack: 4 blocks of (q,k,v,o, gate,up,down) Linears ---- class TinyBlock(nn.Module): def __init__(self, d=64, ff=128): super().__init__() self.q_proj = nn.Linear(d, d, bias=False) self.k_proj = nn.Linear(d, d, bias=False) self.v_proj = nn.Linear(d, d, bias=False) self.o_proj = nn.Linear(d, d, bias=False) self.gate_proj = nn.Linear(d, ff, bias=False) self.up_proj = nn.Linear(d, ff, bias=False) self.down_proj = nn.Linear(ff, d, bias=False) def forward(self, x): h = self.o_proj(self.q_proj(x) + self.k_proj(x) + self.v_proj(x)) m = self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)) return x + h + m class TinyModel(nn.Module): def __init__(self, n_layers=4, d=64, ff=128, vocab=100): super().__init__() self.embed_tokens = nn.Embedding(vocab, d) self.layers = nn.ModuleList([TinyBlock(d, ff) for _ in range(n_layers)]) self.lm_head = nn.Linear(d, vocab, bias=False) class Cfg: # mimic HF .config.hidden_size hidden_size = d self.config = Cfg() def forward(self, ids): x = self.embed_tokens(ids) for blk in self.layers: x = blk(x) return self.lm_head(x) class FakeLinearLike(nn.Module): """Not nn.Linear, but structurally bnb-like enough for target discovery.""" def __init__(self, d_in=8, d_out=8): super().__init__() self.in_features = d_in self.out_features = d_out self.weight = nn.Parameter(torch.empty(d_out, d_in)) nn.init.kaiming_uniform_(self.weight, a=5 ** 0.5) def forward(self, x): return torch.nn.functional.linear(x, self.weight) class FakeBnbModel(nn.Module): def __init__(self): super().__init__() self.config = type("Cfg", (), {"hidden_size": 8})() self.layers = nn.ModuleList([FakeLinearLike(8, 8)]) def forward(self, x): return self.layers[0](x) def variant_test(variant: str, dtype=torch.float32): print(f"\n=== variant={variant} dtype={dtype} ===") torch.manual_seed(0) model = TinyModel().to(dtype) ids = torch.randint(0, 100, (2, 16)) with torch.no_grad(): y_base = model(ids).clone() cfg = ll.LoraLiteConfig( variant=variant, r=4, alpha=4 if variant == "pissa" else 8, # PiSSA needs scale==1 for clean recon dtype=dtype, # delora identity-at-init demands lambda0==0 (then delta * scale = 0) variant_kwargs={"lambda0": 0.0} if variant == "delora" else {}, ) handles = ll.attach(model, cfg) n_targets = len(handles) n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f" attached {n_targets} targets, trainable params={n_trainable}") assert n_targets == 28, f"expected 28 TinyModel targets, got {n_targets}" with torch.no_grad(): y_adapt = model(ids) err = (y_adapt - y_base).abs().max().item() base_scale = y_base.abs().max().item() print(f" t=0 identity: max|y_adapt - y_base| = {err:.3e} (base scale {base_scale:.3e})") # variant-specific identity tolerance tol = { "lora": 1e-6, "pissa": 5e-4, # SVD recon in fp32 is tight; bf16 would be ~1e-2 "delora": 1e-6, # lambda0=0 }[variant] * max(1.0, base_scale) assert err < tol, f" FAIL identity: err {err} > tol {tol}" print(f" SHOULD: err<{tol:.1e}. PASS.") # save/load round-trip ARTIFACT_DIR.mkdir(exist_ok=True) p = ARTIFACT_DIR / f"{variant}_smoke_adapter.pt" ll.save(model, str(p)) # detach + fresh model + load ll.detach(model) torch.manual_seed(0) model2 = TinyModel().to(dtype) # for PiSSA, base weights got mutated; load() re-runs PiSSA init on the fresh # same-seed base, then overwrites lora_A/B with saved values. ll.load(model2, str(p)) with torch.no_grad(): y_loaded = model2(ids) err2 = (y_loaded - y_adapt).abs().max().item() print(f" save/load: max|y_loaded - y_adapt| = {err2:.3e}") assert err2 < tol, f" FAIL save/load: {err2} > {tol}" print(f" SHOULD: err2<{tol:.1e}. PASS.") ll.detach(model2) # gradient flow: 20 SGD steps on random target. # For delora, lambda0==0 makes A,B grads zero (scale=0); use lambda0=0.1 for training. torch.manual_seed(0) model = TinyModel().to(dtype) train_cfg = cfg if variant == "delora": train_cfg = ll.LoraLiteConfig( variant=cfg.variant, r=cfg.r, alpha=cfg.alpha, dtype=cfg.dtype, variant_kwargs={"lambda0": 0.1}, ) ll.attach(model, train_cfg) target = torch.randn(2, 16, 100, dtype=dtype) * 0.1 trainable = [p for p in model.parameters() if p.requires_grad] # delora has tightly-normalised updates; use Adam with higher lr to see signal in 20 steps if variant == "delora": opt = torch.optim.Adam(trainable, lr=1e-1) else: opt = torch.optim.SGD(trainable, lr=1e-2) losses = [] for step in range(20): opt.zero_grad() loss = (model(ids) - target).pow(2).mean() loss.backward() assert_no_base_grads(model) opt.step() losses.append(loss.item()) drop = (losses[0] - losses[-1]) / max(losses[0], 1e-12) print(f" loss[0]={losses[0]:.4f} loss[-1]={losses[-1]:.4f} drop={100*drop:.1f}%") assert drop > 0.05, f" FAIL: loss drop only {drop:.2%}, expected >5%" print(f" SHOULD: drop>5%. PASS.") def structural_linear_like_test(): print("\n=== structural linear-like target test (bnb-style, not nn.Linear) ===") torch.manual_seed(0) model = FakeBnbModel() x = torch.randn(2, 3, 8) y_base = model(x).detach() ll.attach(model, ll.LoraLiteConfig(variant="lora", r=2, alpha=4, dtype=torch.float32, target_roles=())) layer = model.layers[0] assert hasattr(layer, "lora_A") and hasattr(layer, "lora_B") y = model(x) err = (y.detach() - y_base).abs().max().item() loss = y.pow(2).mean() loss.backward() grad_nonzero = layer.lora_B.grad.abs().sum().item() > 0 print(f" attached lora_A={tuple(layer.lora_A.shape)} lora_B={tuple(layer.lora_B.shape)}") print(f" identity_err={err:.3e} grad_nonzero={grad_nonzero}") assert err == 0.0 assert grad_nonzero print(" SHOULD: structural target attaches and lora_B receives grad. PASS.") def bitsandbytes_cuda_smoke(): print("\n=== optional bitsandbytes CUDA smoke ===") if not torch.cuda.is_available(): print(" SKIP: CUDA unavailable; real bnb 4/8-bit forward needs GPU on this machine.") return try: import bitsandbytes as bnb except ImportError: print(" SKIP: bitsandbytes unavailable.") return class BnbModel(nn.Module): def __init__(self, Layer): super().__init__() self.config = type("Cfg", (), {"hidden_size": 8})() self.layers = nn.ModuleList([Layer(8, 8, bias=False)]).cuda() def forward(self, x): return self.layers[0](x) for layer_cls in (bnb.nn.Linear8bitLt, bnb.nn.Linear4bit): torch.manual_seed(0) model = BnbModel(layer_cls) x = torch.randn(2, 3, 8, device="cuda") y_base = model(x).detach() ll.attach(model, ll.LoraLiteConfig(variant="lora", r=2, alpha=4, dtype=torch.float16, target_roles=())) y = model(x) err = (y.detach() - y_base).abs().max().item() y.pow(2).mean().backward() grad_nonzero = model.layers[0].lora_B.grad.abs().sum().item() > 0 print(f" {layer_cls.__name__}: identity_err={err:.3e} grad_nonzero={grad_nonzero}") assert err == 0.0 assert grad_nonzero def main(): for v in ("lora", "pissa", "delora"): variant_test(v, dtype=torch.float32) structural_linear_like_test() bitsandbytes_cuda_smoke() print("\nALL PASS.") if __name__ == "__main__": main()