File size: 7,857 Bytes
aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | """LoRA adapters on a frozen base backbone (``--lora-rank``), tiny random Qwen3.5 (CPU):
* zero-initialised adapters leave every head output bitwise unchanged;
* with trained (random) adapters, a field's outputs are bitwise the same alone, reordered
and next to fields that run another adapter (per-row selector, batch-invariant);
* separate decide/extract adapters are isolated: changing one never moves the other
capability's outputs;
* x-generate decodes adapter-off: the same text as the two-copy hybrid with the same base;
* training: only the selected adapter of each fork kind gets gradient, the base
backbone none, and with ``context="base"`` the trunk carries no graph at all;
* save -> load round-trips the adapters (outputs bitwise equal).
"""
from __future__ import annotations
import copy
import pytest
torch = pytest.importorskip("torch")
pytest.importorskip("transformers")
from test_model import CONTEXT, SCHEMA, SNAPSHOT, _raw_outputs, _sub, tiny_lm # noqa: E402
from test_train_mem import _examples # noqa: E402
from deg.infer import fill # noqa: E402
from deg.model import add_lora, load_model, load_tokenizer # noqa: E402
FLAT = [n for n in SCHEMA["properties"] if n not in ("items", "reply")]
DECIDE = ["department", "urgent", "anger"]
EXTRACT = ["email", "order_id", "total", "charged_on"]
def _model(adapters=None, context="base", rank=4, randomize=True, lm=None):
if SNAPSHOT is None:
pytest.skip("Qwen3.5-0.8B tokenizer not cached")
m = load_model(None, device="cpu", dtype=torch.float32, lm=lm or tiny_lm(0),
tokenizer=load_tokenizer(SNAPSHOT))
torch.manual_seed(1)
with torch.no_grad():
for p in m.heads.parameters():
p.normal_(0.0, 0.5)
if adapters is not None:
m.backbone.requires_grad_(False)
add_lora(m, rank, None, adapters, context)
if randomize:
g = torch.Generator().manual_seed(3)
with torch.no_grad():
for name, p in m.backbone.named_parameters():
if ".lora_B" in name:
p.copy_(torch.randn(p.shape, generator=g) * 0.05)
return m
def _same(a, b, keys=None):
keys = keys or a.keys()
for k in keys:
assert torch.equal(a[k], b[k]), k
def test_zero_adapters_are_the_base_model():
plain = _model()
zero = _model(["decide", "extract"], randomize=False)
_same(_raw_outputs(plain, _sub(FLAT)), _raw_outputs(zero, _sub(FLAT)))
@pytest.mark.parametrize("adapters,context", [(["shared"], "base"),
(["decide", "extract"], "base"),
(["shared"], "shared")])
def test_adapter_rows_are_batch_invariant(adapters, context):
m = _model(adapters, context)
merged = _raw_outputs(m, _sub(FLAT))
_same(merged, _raw_outputs(m, _sub(list(reversed(FLAT)))))
for names in (DECIDE, EXTRACT, ["anger"], ["total"]):
alone = _raw_outputs(m, _sub(names))
_same(alone, merged, alone.keys())
if adapters == ["decide", "extract"]: # the adapters actually change the outputs
plain = _raw_outputs(_model(), _sub(FLAT))
assert any(not torch.equal(plain[k], merged[k]) for k in merged)
def test_separate_adapters_are_isolated():
m = _model(["decide", "extract"])
before = _raw_outputs(m, _sub(FLAT))
r = m.lora["rank"]
with torch.no_grad(): # perturb only the decide adapter's columns
for name, p in m.backbone.named_parameters():
if ".lora_B" in name:
p[:, :r] += 0.1
after = _raw_outputs(m, _sub(FLAT))
_same(before, after, [f"/{n}" for n in EXTRACT])
assert any(not torch.equal(before[f"/{n}"], after[f"/{n}"]) for n in DECIDE)
def test_generate_is_the_base_model():
schema = _sub(["department", "order_id", "reply"])
schema["properties"]["reply"] = {"type": "string", "x-generate": True, "maxLength": 40}
lm = tiny_lm(0)
base = copy.deepcopy(lm)
for context in ("base", "shared"):
m = _model(["shared"], context, lm=copy.deepcopy(lm))
hyb = load_model(None, device="cpu", dtype=torch.float32, lm=copy.deepcopy(lm),
tokenizer=load_tokenizer(SNAPSHOT), generate="base",
base_lm=copy.deepcopy(base))
got = fill(m, CONTEXT, schema)
want_known = {k: got["value"][k] for k in ("department", "order_id")}
# the hybrid decodes from the same known values with the same base weights
hyb_plan = fill(hyb, CONTEXT, schema)
if {k: hyb_plan["value"][k] for k in want_known} == want_known:
assert got["value"]["reply"] == hyb_plan["value"]["reply"]
if context == "base": # one prefill: no second context read
assert got["usage"]["prompt_tokens"] < hyb_plan["usage"]["prompt_tokens"]
def test_training_grads_follow_the_fork_kind():
from deg.train import forward_losses
m = _model(["decide", "extract"], "base")
exs = _examples(m)
m.backbone.train()
m.heads.train()
r = m.lora["rank"]
try:
hidden, state = m.trunk([exs[0].ctx])
assert not hidden.requires_grad # base context read: no trunk graph
terms, _ = forward_losses(m, exs, 4096, 4)
assert "lm" not in terms # generate forks are not trained under LoRA
terms["decide"].backward()
finally:
m.backbone.eval()
m.heads.eval()
got = {n: p.grad for n, p in m.backbone.named_parameters() if ".lora_B" in n}
assert all(not p.requires_grad for n, p in m.backbone.named_parameters()
if ".lora_" not in n)
assert any(float(g[:, :r].abs().sum()) > 0 for g in got.values() if g is not None)
assert all(g is None or float(g[:, r:].abs().sum()) == 0 for g in got.values())
def test_save_load_roundtrip(tmp_path):
"""Adapters (saved fp32) come back exactly; the backbone file is bf16, so outputs are
compared between two loads (deterministic) and against zeroed adapters (applied)."""
from deg.model import lora_state
m = _model(["decide", "extract"], "base")
m.save(tmp_path / "ck")
assert (tmp_path / "ck" / "adapters.safetensors").exists()
def load():
return load_model(tmp_path / "ck", device="cpu", dtype=torch.float32, lm=tiny_lm(5),
tokenizer=load_tokenizer(SNAPSHOT))
back = load()
assert back.lora == m.lora
want = lora_state(m)
got = lora_state(back)
assert got.keys() == want.keys() and len(want) > 0
for k in want:
assert torch.equal(got[k], want[k].to(torch.bfloat16).float()) or \
torch.equal(got[k], want[k]), k
out = _raw_outputs(back, _sub(FLAT))
_same(out, _raw_outputs(load(), _sub(FLAT)))
with torch.no_grad():
for name, p in back.backbone.named_parameters():
if ".lora_B" in name:
p.zero_()
off = _raw_outputs(back, _sub(FLAT))
assert any(not torch.equal(out[k], off[k]) for k in out)
def test_init_heads_loads_a_checkpoints_heads_onto_a_lora_model(tmp_path):
"""--init-heads: a full fine-tune's heads on the base backbone + fresh adapters."""
from deg.train import init_heads
ft = _model() # randomized heads (the "full fine-tune")
ft.save(tmp_path / "ft")
m = _model(["shared"], "shared", randomize=False)
with torch.no_grad():
for p in m.heads.parameters():
p.zero_()
init_heads(m, tmp_path / "ft")
for (k, a), b in zip(ft.heads.state_dict().items(), m.heads.state_dict().values()):
assert torch.equal(a, b), k
meta = tmp_path / "ft" / "deg.json"
import json
cfg = json.loads(meta.read_text())
meta.write_text(json.dumps({**cfg, "backbone": "qwen3.5-2b"}))
with pytest.raises(SystemExit):
init_heads(m, tmp_path / "ft")
|