mirror of
https://github.com/wassname/lora-lite.git
synced 2026-07-15 11:24:10 +08:00
64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
"""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)
|