mirror of
https://github.com/wassname/lora-lite.git
synced 2026-06-27 16:45:56 +08:00
tidy tests to subset of metamath
This commit is contained in:
wassname
@@ -1,53 +0,0 @@
|
||||
"""Smoke: end-to-end MetaMath->GSM8K plumbing for every variant on a tiny HF model.
|
||||
|
||||
Per-variant correctness invariants live in tests/test_lora_lite.py. This script
|
||||
just confirms the full benchmark pipeline (data load, prompt encode, train step,
|
||||
eval generate + answer extract) runs for each adapter type.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
VARIANTS = ["lora", "pissa", "delora", "ia3", "ia3_ff", "dora", "hra", "eva", "antipasto"]
|
||||
MODEL = "hf-internal-testing/tiny-random-LlamaForCausalLM"
|
||||
|
||||
|
||||
def run_one(variant: str) -> int:
|
||||
cmd = [
|
||||
sys.executable,
|
||||
"scripts/metamath_gsm8k_benchmark.py",
|
||||
"--model", MODEL,
|
||||
"--variant", variant,
|
||||
"--steps", "2",
|
||||
"--batch-size", "2",
|
||||
"--max-train-samples", "8",
|
||||
"--max-eval-samples", "10",
|
||||
"--max-valid-samples", "10",
|
||||
"--max-new-tokens", "8",
|
||||
"--max-seq-length", "128",
|
||||
"--r", "4",
|
||||
"--alpha", "8",
|
||||
"--torch-dtype", "float32",
|
||||
"--device", "cpu",
|
||||
]
|
||||
if variant == "ia3":
|
||||
cmd += ["--target-name", r"(k_proj|v_proj)$"]
|
||||
elif variant == "ia3_ff":
|
||||
cmd += ["--target-name", r"(down_proj)$"]
|
||||
print(f"\n=== smoke variant={variant} ===")
|
||||
print(" ".join(cmd))
|
||||
return subprocess.call(cmd)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
failed = [v for v in VARIANTS if run_one(v) != 0]
|
||||
if failed:
|
||||
print(f"FAIL: {failed}")
|
||||
return 1
|
||||
print("ALL PASS.")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -1,63 +0,0 @@
|
||||
"""bnb 4bit/8bit CUDA smoke. Skipped without CUDA + bitsandbytes installed."""
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
import lora_lite as ll
|
||||
|
||||
|
||||
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
|
||||
bnb = pytest.importorskip("bitsandbytes")
|
||||
|
||||
|
||||
CFG_BY_VARIANT = {
|
||||
"lora": ll.LoRAConfig,
|
||||
"ia3": ll.IA3Config,
|
||||
"hra": ll.HRAConfig,
|
||||
"pissa": ll.PiSSAConfig,
|
||||
"dora": ll.DoRAConfig,
|
||||
}
|
||||
|
||||
|
||||
class BnbModel(nn.Module):
|
||||
def __init__(self, layer_cls):
|
||||
super().__init__()
|
||||
self.config = type("Cfg", (), {"hidden_size": 8})()
|
||||
self.layers = nn.ModuleList([layer_cls(8, 8, bias=False)]).cuda()
|
||||
|
||||
def forward(self, x):
|
||||
return self.layers[0](x)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("layer_cls", [bnb.nn.Linear8bitLt, bnb.nn.Linear4bit])
|
||||
@pytest.mark.parametrize("variant", ["lora", "ia3", "hra"])
|
||||
def test_hook_only_variants_attach_to_bnb(layer_cls, variant):
|
||||
"""LoRA / IA3 / HRA only hook outputs; bnb dequantization is the layer's job."""
|
||||
torch.manual_seed(0)
|
||||
model = BnbModel(layer_cls)
|
||||
x = torch.randn(2, 3, 8, device="cuda")
|
||||
y_base = model(x).detach()
|
||||
|
||||
cfg = CFG_BY_VARIANT[variant](r=2, alpha=4, dtype=torch.float16, target_roles=())
|
||||
ll.attach(model, cfg)
|
||||
y = model(x)
|
||||
assert (y.detach() - y_base).abs().max().item() < 1e-2
|
||||
|
||||
y.pow(2).mean().backward()
|
||||
grad_total = sum(
|
||||
g.abs().sum().item()
|
||||
for n, p in model.named_parameters()
|
||||
if "lora_" in n and p.requires_grad and (g := p.grad) is not None
|
||||
)
|
||||
assert grad_total > 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("layer_cls", [bnb.nn.Linear8bitLt, bnb.nn.Linear4bit])
|
||||
@pytest.mark.parametrize("variant", ["pissa", "dora"])
|
||||
def test_weight_reading_variants_reject_bnb(layer_cls, variant):
|
||||
model = BnbModel(layer_cls)
|
||||
cfg = CFG_BY_VARIANT[variant](r=2, alpha=2, dtype=torch.float16, target_roles=())
|
||||
with pytest.raises((TypeError, RuntimeError, AttributeError, ValueError)):
|
||||
ll.attach(model, cfg)
|
||||
@@ -1,328 +0,0 @@
|
||||
"""Per-variant attach + train + save + load round-trip, plus surgical regressions.
|
||||
|
||||
The big invariant is the parametrized train_save_load test: identity at t=0,
|
||||
gradient flow on a real loss, then save -> reload onto a fresh model and
|
||||
confirm the trained outputs survive the round-trip. Cheap on CPU.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
import lora_lite as ll
|
||||
|
||||
|
||||
CFG_BY_VARIANT = {
|
||||
"lora": ll.LoRAConfig,
|
||||
"pissa": ll.PiSSAConfig,
|
||||
"delora": ll.DeLoRAConfig,
|
||||
"ia3": ll.IA3Config,
|
||||
"ia3_ff": ll.IA3FFConfig,
|
||||
"dora": ll.DoRAConfig,
|
||||
"hra": ll.HRAConfig,
|
||||
"eva": ll.EVAConfig,
|
||||
"antipasto": ll.AntiPaSTOConfig,
|
||||
}
|
||||
|
||||
# Per-variant identity tolerance at t=0 (after attach, before any step).
|
||||
# fp32 SVD round-trip + per-row norm = looser tolerance for pissa/dora/antipasto.
|
||||
IDENTITY_TOL = {
|
||||
"lora": 1e-6,
|
||||
"pissa": 5e-4,
|
||||
"delora": 1e-6,
|
||||
"ia3": 1e-6,
|
||||
"ia3_ff": 1e-6,
|
||||
"dora": 5e-5,
|
||||
"hra": 5e-6,
|
||||
"eva": 1e-6,
|
||||
"antipasto": 5e-4,
|
||||
}
|
||||
|
||||
|
||||
class TinyBlock(nn.Module):
|
||||
def __init__(self, d: int = 64, ff: int = 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: torch.Tensor) -> torch.Tensor:
|
||||
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: int = 4, d: int = 64, ff: int = 128, vocab: int = 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)
|
||||
self.config = type("Cfg", (), {"hidden_size": d})()
|
||||
|
||||
def forward(self, ids: torch.Tensor) -> torch.Tensor:
|
||||
x = self.embed_tokens(ids)
|
||||
for block in self.layers:
|
||||
x = block(x)
|
||||
return self.lm_head(x)
|
||||
|
||||
|
||||
class FakeLinearLike(nn.Module):
|
||||
"""linear-like, but not nn.Linear: stand-in for bnb 4/8-bit modules."""
|
||||
|
||||
def __init__(self, d_in: int = 8, d_out: int = 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: torch.Tensor) -> torch.Tensor:
|
||||
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: torch.Tensor) -> torch.Tensor:
|
||||
return self.layers[0](x)
|
||||
|
||||
|
||||
def cfg_for(variant: str) -> ll.AdapterConfig:
|
||||
return CFG_BY_VARIANT[variant](
|
||||
r=4,
|
||||
alpha=8,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def attach_with_calib(model: nn.Module, cfg: ll.AdapterConfig, ids: torch.Tensor) -> None:
|
||||
if cfg.variant == "eva":
|
||||
calib = [ids for _ in range(2)]
|
||||
ll.attach(model, cfg, calibration_data=calib)
|
||||
else:
|
||||
ll.attach(model, cfg)
|
||||
|
||||
|
||||
def trainable_grad_norm(model: nn.Module) -> float:
|
||||
return sum(
|
||||
p.grad.detach().float().norm().item()
|
||||
for n, p in model.named_parameters()
|
||||
if "lora_" in n and p.grad is not None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("variant", list(CFG_BY_VARIANT))
|
||||
def test_train_save_load(variant: str, tmp_path: Path):
|
||||
"""Identity at t=0, one SGD step, save, reload onto fresh model, outputs match."""
|
||||
torch.manual_seed(0)
|
||||
model = TinyModel()
|
||||
ids = torch.randint(0, 100, (2, 16))
|
||||
with torch.no_grad():
|
||||
y_base = model(ids).clone()
|
||||
|
||||
cfg = cfg_for(variant)
|
||||
attach_with_calib(model, cfg, ids)
|
||||
|
||||
trainable = [p for p in model.parameters() if p.requires_grad]
|
||||
assert trainable
|
||||
assert all("lora_" in n for n, p in model.named_parameters() if p.requires_grad)
|
||||
|
||||
with torch.no_grad():
|
||||
y_init = model(ids).clone()
|
||||
assert (y_init - y_base).abs().max().item() < IDENTITY_TOL[variant]
|
||||
|
||||
target = torch.randn_like(y_init) * 0.1
|
||||
opt = torch.optim.SGD(trainable, lr=1e-2)
|
||||
opt.zero_grad()
|
||||
loss = (model(ids) - target).pow(2).mean()
|
||||
loss.backward()
|
||||
leaked = [n for n, p in model.named_parameters() if "lora_" not in n and p.grad is not None]
|
||||
assert leaked == []
|
||||
assert trainable_grad_norm(model) > 0
|
||||
opt.step()
|
||||
|
||||
with torch.no_grad():
|
||||
y_trained = model(ids).clone()
|
||||
|
||||
path = tmp_path / "adapter.pt"
|
||||
ll.save(model, str(path))
|
||||
|
||||
torch.manual_seed(0)
|
||||
model_loaded = TinyModel()
|
||||
ll.load(model_loaded, str(path)) # EVA load skips group_init; calibration_data not needed
|
||||
with torch.no_grad():
|
||||
y_loaded = model_loaded(ids)
|
||||
assert (y_loaded - y_trained).abs().max().item() < max(IDENTITY_TOL[variant], 1e-5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("variant", ["lora", "delora", "ia3", "hra"])
|
||||
def test_hook_only_variants_attach_to_non_linear_target(variant: str):
|
||||
"""bnb-style targets are linear-like but not nn.Linear; hook-only variants must accept them."""
|
||||
extra = {"lambda0": 0.1} if variant == "delora" else {}
|
||||
cfg = CFG_BY_VARIANT[variant](r=2, alpha=4, dtype=torch.float32, target_roles=(), **extra)
|
||||
model = FakeBnbModel()
|
||||
ll.attach(model, cfg)
|
||||
x = torch.randn(2, 3, 8)
|
||||
model(x).pow(2).mean().backward()
|
||||
assert trainable_grad_norm(model) > 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("variant", ["pissa", "dora", "antipasto"])
|
||||
def test_weight_reading_variants_reject_non_linear(variant: str):
|
||||
r = 4 if variant == "antipasto" else 2 # antipasto needs r % block_size==0
|
||||
cfg = CFG_BY_VARIANT[variant](r=r, alpha=r, dtype=torch.float32, target_roles=())
|
||||
with pytest.raises(TypeError, match="plain nn.Linear"):
|
||||
ll.attach(FakeBnbModel(), cfg)
|
||||
|
||||
|
||||
def test_save_load_strict_keys(tmp_path: Path):
|
||||
torch.manual_seed(0)
|
||||
model = TinyModel()
|
||||
ll.attach(model, ll.LoRAConfig(r=4, alpha=8, dtype=torch.float32))
|
||||
p = tmp_path / "lora.pt"
|
||||
ll.save(model, str(p))
|
||||
blob = torch.load(p, weights_only=True, map_location="cpu")
|
||||
|
||||
missing = {"cfg": blob["cfg"], "state": dict(blob["state"]), "base_fp": blob.get("base_fp", {})}
|
||||
missing["state"].pop(next(iter(missing["state"])))
|
||||
torch.save(missing, p)
|
||||
with pytest.raises(RuntimeError, match="missing lora keys"):
|
||||
ll.load(TinyModel(), str(p))
|
||||
|
||||
bad = {"cfg": blob["cfg"], "state": dict(blob["state"]), "base_fp": blob.get("base_fp", {})}
|
||||
bad["state"]["layers.0.q_proj.lora_extra"] = torch.zeros(1)
|
||||
torch.save(bad, p)
|
||||
with pytest.raises(RuntimeError, match="unexpected lora keys"):
|
||||
ll.load(TinyModel(), str(p))
|
||||
|
||||
|
||||
def test_no_target_layers_is_loud():
|
||||
cfg = ll.LoRAConfig(target_names=("definitely_missing",))
|
||||
with pytest.raises(RuntimeError, match="no target layers"):
|
||||
ll.attach(TinyModel(), cfg)
|
||||
|
||||
|
||||
def test_eva_requires_calibration():
|
||||
"""EVA's group_init must error loudly if calibration_data is missing."""
|
||||
with pytest.raises(ValueError, match="calibration_data"):
|
||||
ll.attach(TinyModel(), ll.EVAConfig(r=4, alpha=8, dtype=torch.float32))
|
||||
|
||||
|
||||
def test_delora_default_has_live_step0_gradient():
|
||||
"""Default lambda0 must be nonzero; B=0 preserves identity while B gets gradient."""
|
||||
torch.manual_seed(0)
|
||||
model = TinyModel(n_layers=1)
|
||||
ids = torch.randint(0, 100, (2, 8))
|
||||
ll.attach(model, ll.DeLoRAConfig(r=4, alpha=8, dtype=torch.float32))
|
||||
|
||||
assert model.layers[0].q_proj.lora_lambda.item() == pytest.approx(15.0)
|
||||
loss = model(ids).pow(2).mean()
|
||||
loss.backward()
|
||||
|
||||
b_grad = model.layers[0].q_proj.lora_B.grad.detach().abs().max().item()
|
||||
assert b_grad > 0
|
||||
|
||||
|
||||
def test_pissa_identity_with_nonunit_scale():
|
||||
"""Regression: PiSSA must pre-divide S by alpha/r, not require alpha == r."""
|
||||
torch.manual_seed(0)
|
||||
model = TinyModel(n_layers=1)
|
||||
ids = torch.randint(0, 100, (2, 8))
|
||||
with torch.no_grad():
|
||||
y_base = model(ids).clone()
|
||||
|
||||
ll.attach(model, ll.PiSSAConfig(r=4, alpha=8, dtype=torch.float32))
|
||||
with torch.no_grad():
|
||||
y = model(ids)
|
||||
assert (y - y_base).abs().max().item() < IDENTITY_TOL["pissa"]
|
||||
|
||||
|
||||
def test_antipasto_blockwise_rotation_matches_explicit_blockdiag():
|
||||
"""The einsum/rearrange path must equal the old explicit blockdiag math."""
|
||||
from lora_lite.variants.antipasto import _build_rotation
|
||||
|
||||
torch.manual_seed(0)
|
||||
n_blocks, bs, d_in, d_out = 3, 4, 7, 5
|
||||
r = n_blocks * bs
|
||||
rot_T = torch.randn(n_blocks, bs * (bs - 1) // 2) * 0.1
|
||||
Vh = torch.randn(r, d_in)
|
||||
U = torch.randn(d_out, r)
|
||||
R_blocks = _build_rotation(rot_T, bs, 0.5)
|
||||
R = torch.block_diag(*list(R_blocks))
|
||||
|
||||
Vh_blocks = torch.reshape(Vh, (n_blocks, bs, d_in))
|
||||
Vh_rot = torch.einsum("nab,nbi->nai", R_blocks, Vh_blocks).reshape(r, d_in)
|
||||
U_blocks = torch.reshape(U, (d_out, n_blocks, bs))
|
||||
U_rot = torch.einsum("dnb,ncb->dnc", U_blocks, R_blocks).reshape(d_out, r)
|
||||
|
||||
assert (Vh_rot - R @ Vh).abs().max().item() < 1e-6
|
||||
assert (U_rot - U @ R.T).abs().max().item() < 1e-6
|
||||
|
||||
|
||||
def test_dora_bias_passthrough():
|
||||
"""Regression: DoRA must NOT scale bias; identity holds with bias=True at t=0."""
|
||||
torch.manual_seed(0)
|
||||
d = 16
|
||||
layer = nn.Linear(d, d, bias=True)
|
||||
x = torch.randn(2, d)
|
||||
y_base = layer(x).detach()
|
||||
|
||||
class Wrap(nn.Module):
|
||||
def __init__(self, lin):
|
||||
super().__init__()
|
||||
self.config = type("Cfg", (), {"hidden_size": d})()
|
||||
self.layers = nn.ModuleList([lin])
|
||||
|
||||
def forward(self, x):
|
||||
return self.layers[0](x)
|
||||
|
||||
model = Wrap(layer)
|
||||
ll.attach(model, ll.DoRAConfig(r=2, alpha=4, dtype=torch.float32, target_roles=()))
|
||||
with torch.no_grad():
|
||||
y = model(x)
|
||||
assert (y - y_base).abs().max().item() < 1e-5
|
||||
|
||||
|
||||
def test_hra_forward_is_x_R_T():
|
||||
"""HRA must apply x @ R^T (loop i = r-1 down to 0). Asymmetric U makes order observable."""
|
||||
torch.manual_seed(0)
|
||||
d = 8
|
||||
layer = nn.Linear(d, d, bias=False)
|
||||
x = torch.randn(2, 3, d)
|
||||
|
||||
class Wrap(nn.Module):
|
||||
def __init__(self, lin):
|
||||
super().__init__()
|
||||
self.config = type("Cfg", (), {"hidden_size": d})()
|
||||
self.layers = nn.ModuleList([lin])
|
||||
|
||||
def forward(self, x):
|
||||
return self.layers[0](x)
|
||||
|
||||
model = Wrap(layer)
|
||||
ll.attach(model, ll.HRAConfig(r=4, alpha=4, dtype=torch.float32, target_roles=()))
|
||||
# break paired symmetry so order matters
|
||||
with torch.no_grad():
|
||||
layer.lora_U.add_(0.1 * torch.randn_like(layer.lora_U))
|
||||
|
||||
U = layer.lora_U
|
||||
R = torch.eye(d)
|
||||
for i in range(U.shape[0]):
|
||||
u = U[i]
|
||||
sq = (u * u).sum().clamp_min(1e-12)
|
||||
R = R - (2.0 / sq) * torch.outer(R @ u, u)
|
||||
with torch.no_grad():
|
||||
y_adapt = model(x)
|
||||
y_ref = torch.nn.functional.linear(x, layer.weight @ R)
|
||||
assert (y_adapt - y_ref).abs().max().item() < 1e-5
|
||||
@@ -0,0 +1,127 @@
|
||||
"""End-to-end smoke: run the metamath benchmark in probe mode for every variant.
|
||||
|
||||
Probe mode trains a few steps on tiny-random Llama, saves the adapter, reloads
|
||||
it onto a fresh model, and asserts the trained logits match within tol. That's
|
||||
the train+save+load round-trip on a real HF model, one test per variant.
|
||||
|
||||
A second test attaches each variant on top of a 4bit/8bit-loaded base and runs
|
||||
one backward step. PiSSA/DoRA/AntiPaSTO/EVA must fail loud on quantized weights;
|
||||
the rest must produce nonzero adapter grads. We do not run the full probe under
|
||||
bnb because tiny-random + bnb dequant produces NaN logits unrelated to adapter
|
||||
correctness.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import sys
|
||||
from dataclasses import replace
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import lora_lite as ll
|
||||
|
||||
SPEC = importlib.util.spec_from_file_location(
|
||||
"metamath_benchmark",
|
||||
Path(__file__).resolve().parent.parent / "scripts" / "metamath_gsm8k_benchmark.py",
|
||||
)
|
||||
benchmark = importlib.util.module_from_spec(SPEC)
|
||||
sys.modules[SPEC.name] = benchmark
|
||||
SPEC.loader.exec_module(benchmark)
|
||||
|
||||
|
||||
VARIANTS = ["lora", "pissa", "delora", "ia3", "ia3_ff", "dora", "hra", "eva", "antipasto"]
|
||||
# Variants that fail loud when attached on a bnb-loaded base (read dense weight in init).
|
||||
# delora/eva also read weight but currently silently dequant -- they produce sane attach,
|
||||
# so we don't expect a raise from them in the attach-only smoke.
|
||||
BNB_RAISERS = {"pissa", "dora", "antipasto"}
|
||||
TINY_MODEL = "hf-internal-testing/tiny-random-LlamaForCausalLM"
|
||||
|
||||
HAS_CUDA = torch.cuda.is_available()
|
||||
HAS_BNB = importlib.util.find_spec("bitsandbytes") is not None
|
||||
|
||||
|
||||
def quick_cfg(variant: str, tmp_path: Path, quantization: str = "none") -> "benchmark.BenchmarkConfig":
|
||||
target_name = (
|
||||
[r"(k_proj|v_proj)$"] if variant == "ia3"
|
||||
else [r"(down_proj)$"] if variant == "ia3_ff"
|
||||
else [r"(q_proj|v_proj)$"]
|
||||
)
|
||||
cfg = benchmark.BenchmarkConfig(
|
||||
model=TINY_MODEL,
|
||||
variant=variant,
|
||||
mode="probe",
|
||||
device="cuda" if HAS_CUDA else "cpu",
|
||||
torch_dtype="float16" if quantization != "none" else "float32",
|
||||
quantization=quantization,
|
||||
r=4,
|
||||
alpha=8,
|
||||
target_name=target_name,
|
||||
layers="all",
|
||||
steps=2,
|
||||
batch_size=2,
|
||||
batch_size_eval=4,
|
||||
max_train_samples=8,
|
||||
max_eval_samples=4,
|
||||
max_valid_samples=4,
|
||||
max_test_samples=4,
|
||||
max_seq_length=128,
|
||||
max_new_tokens=8,
|
||||
lr=5e-3,
|
||||
seed=0,
|
||||
log_examples=0,
|
||||
log_every=1000,
|
||||
output_dir=tmp_path / "out",
|
||||
)
|
||||
if variant == "antipasto":
|
||||
cfg = replace(cfg, alpha=4) # block_size=4 -> need r % 4 == 0
|
||||
return cfg
|
||||
|
||||
|
||||
@pytest.mark.parametrize("variant", VARIANTS)
|
||||
def test_metamath_quick_train_save_load(variant: str, tmp_path: Path):
|
||||
"""Train 2 steps, save, reload onto fresh tiny model, logits match within tol."""
|
||||
cfg = quick_cfg(variant, tmp_path)
|
||||
result = benchmark.run(cfg)
|
||||
|
||||
assert result["train"]["base_grad_leaks"] == 0
|
||||
assert result["train"]["first_grad_norm"] > 0
|
||||
assert result["train"]["adapter_delta"] > 0
|
||||
probe = result.get("probe") or {}
|
||||
assert "reload_err" in probe
|
||||
assert probe["reload_err"] < cfg.reload_tol
|
||||
|
||||
|
||||
@pytest.mark.skipif(not (HAS_CUDA and HAS_BNB), reason="needs CUDA + bitsandbytes")
|
||||
@pytest.mark.parametrize("quantization", ["4bit", "8bit"])
|
||||
@pytest.mark.parametrize("variant", VARIANTS)
|
||||
def test_attach_on_bnb_loaded_base(variant: str, quantization: str, tmp_path: Path):
|
||||
"""Attach to a bnb-loaded base, run one backward step. Weight-reading variants must fail loud."""
|
||||
cfg = quick_cfg(variant, tmp_path, quantization=quantization)
|
||||
dtype = getattr(torch, cfg.torch_dtype)
|
||||
|
||||
def _do() -> float:
|
||||
model, _ = benchmark.load_model_and_tokenizer(cfg.model, dtype, cfg.device, cfg.quantization)
|
||||
adapter_cfg = benchmark.cfg_for_variant(cfg, dtype)
|
||||
if cfg.variant == "eva":
|
||||
ids = torch.randint(0, 100, (2, 8), device=cfg.device)
|
||||
ll.attach(model, adapter_cfg, calibration_data=[{"input_ids": ids}])
|
||||
else:
|
||||
ll.attach(model, adapter_cfg)
|
||||
ids = torch.randint(0, 100, (2, 8), device=cfg.device)
|
||||
out = model(input_ids=ids).logits
|
||||
loss = out.float().pow(2).mean()
|
||||
loss.backward()
|
||||
return sum(
|
||||
p.grad.detach().float().norm().item()
|
||||
for n, p in model.named_parameters()
|
||||
if "lora_" in n and p.grad is not None
|
||||
)
|
||||
|
||||
if variant in BNB_RAISERS:
|
||||
with pytest.raises((TypeError, RuntimeError, AttributeError, ValueError)):
|
||||
_do()
|
||||
else:
|
||||
_do() # only assert it runs without exception; tiny+bnb grads can be 0/garbage.
|
||||
|
||||
Reference in New Issue
Block a user