| """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 |
| from test_train_mem import _examples |
|
|
| from deg.infer import fill |
| from deg.model import add_lora, load_model, load_tokenizer |
|
|
| 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"]: |
| 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(): |
| 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")} |
| |
| 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": |
| 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 |
| terms, _ = forward_losses(m, exs, 4096, 4) |
| assert "lm" not in terms |
| 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() |
| 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") |
|
|
