| """The LFM2.5 backbone (``--backbone lfm2.5-1.2b``; SPEC.md §12.8) on a tiny random LFM2 |
| (CPU, exact torch reference kernels) with the real LFM2.5 tokenizer, plus GPU tests on the |
| real LFM2.5-1.2B-Instruct. |
| |
| The model-contract tests of ``test_model`` / ``test_hybrid_generate`` / ``test_train_long`` |
| are imported and collected here against this module's LFM fixtures (field independence |
| is bitwise alone / reordered / in bigger schemas / across requests, chunked prefill equals |
| the one-shot read, OOM splits are bitwise, hybrid generate decodes with the frozen base); |
| LFM-specific tests below check the functional stack against HF ``Lfm2ForCausalLM``, |
| forked fields against independent full forwards, the renderer's tokens, checkpoints and |
| training gradients. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import random |
|
|
| import pytest |
|
|
| torch = pytest.importorskip("torch") |
| pytest.importorskip("transformers") |
|
|
| import test_hybrid_generate as thg |
| import test_model as tm |
| import test_train_long as ttl |
| from test_hybrid_generate import ( |
| test_hybrid_fill_many_matches_fill, |
| test_hybrid_generate_rows_are_batch_invariant, |
| test_hybrid_keeps_deg_fields_and_decodes_text_with_the_base, |
| ) |
| from test_model import ( |
| CONTEXT, |
| _raw_outputs, |
| _sub, |
| test_blocked_span_logits_match_one_shot, |
| test_consume_extends_state_without_holding_two_copies, |
| test_field_outputs_are_bitwise_independent, |
| test_fill_many_matches_fill_alone, |
| test_fill_meta_is_identical_in_split_and_merged_schemas, |
| test_fork_batch_split_on_oom_gives_the_same_outputs, |
| test_generate_fields_are_independent_of_each_other, |
| test_long_prefill_padded_batch_and_prefix_state, |
| test_long_prefill_streams_in_chunks_without_quadratic_tensors, |
| test_only_x_generate_fields_use_the_lm_head, |
| test_per_field_generate_budget, |
| test_sampled_generation_is_reproducible_with_seed, |
| test_span_pointer_addresses_context_tokens_only, |
| ) |
| from test_train_long import test_inference_long_read_unchanged_by_training_path |
|
|
| from deg import compile_schema, validate_value |
| from deg.infer import fill |
| from deg.model import LFM2, base_snapshot, load_model, load_tokenizer |
|
|
| BACKBONE = "lfm2.5-1.2b" |
| try: |
| LFM_SNAPSHOT = base_snapshot(BACKBONE) |
| except FileNotFoundError: |
| LFM_SNAPSHOT = None |
|
|
|
|
| def tiny_lfm(seed: int = 0): |
| from transformers import Lfm2Config, Lfm2ForCausalLM |
|
|
| torch.manual_seed(seed) |
| cfg = Lfm2Config( |
| vocab_size=65536, hidden_size=64, intermediate_size=128, |
| block_auto_adjust_ff_dim=False, num_hidden_layers=4, |
| layer_types=["conv", "full_attention", "conv", "full_attention"], |
| num_attention_heads=4, num_key_value_heads=2, conv_L_cache=3, conv_bias=False, |
| norm_eps=1e-5, tie_embedding=True, max_position_embeddings=4096, |
| rope_parameters={"rope_type": "default", "rope_theta": 10000.0}, |
| pad_token_id=0, bos_token_id=1, eos_token_id=7) |
| lm = Lfm2ForCausalLM(cfg).float().eval() |
| with torch.no_grad(): |
| for name, p in lm.named_parameters(): |
| if "norm" in name: |
| p.normal_(1.0, 0.3) |
| return lm |
|
|
|
|
| def _tok(): |
| if LFM_SNAPSHOT is None: |
| pytest.skip("LFM2.5-1.2B-Instruct tokenizer not cached") |
| return load_tokenizer(backbone=BACKBONE) |
|
|
|
|
| def _lfm_model(seed: int = 0, **kw): |
| m = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_lfm(seed), |
| tokenizer=_tok(), **kw) |
| torch.manual_seed(1) |
| with torch.no_grad(): |
| for p in m.heads.parameters(): |
| p.normal_(0.0, 0.5) |
| return m |
|
|
|
|
| @pytest.fixture(scope="module") |
| def model(): |
| return _lfm_model() |
|
|
|
|
| @pytest.fixture(scope="module") |
| def pair(): |
| """(deg-only, hybrid) LFM models: the hybrid's frozen base is another random LFM |
| with the same (frozen) embeddings, as deg training leaves them.""" |
| def make(generate): |
| base, deg = tiny_lfm(7), tiny_lfm(0) |
| with torch.no_grad(): |
| deg.model.embed_tokens.weight.copy_(base.model.embed_tokens.weight) |
| m = load_model(None, device="cpu", dtype=torch.float32, lm=deg, tokenizer=_tok(), |
| generate=generate, base_lm=base) |
| torch.manual_seed(1) |
| with torch.no_grad(): |
| for p in m.heads.parameters(): |
| p.normal_(0.0, 0.5) |
| return m |
|
|
| return make("self"), make("base") |
|
|
|
|
| |
|
|
|
|
| def test_lfm_renderer_and_meta(model): |
| rend = model.renderer |
| assert model.backbone.fam is LFM2 and model.meta["backbone"] == BACKBONE |
| assert model.backbone.types == ["rec", "attn", "rec", "attn"] |
| ctx = rend.context(CONTEXT) |
| tok = rend.tokenizer |
| assert ctx.ids[0] == tok.convert_tokens_to_ids("<|startoftext|>") |
| assert tok.convert_ids_to_tokens(ctx.ids[1]) == "<|im_start|>" |
| assert rend.pad_id == tok.convert_tokens_to_ids("<|pad|>") |
| assert rend.eos_id == tok.convert_tokens_to_ids("<|im_end|>") |
| assert rend.gen_turn == "<|im_end|>\n<|im_start|>assistant\n" |
| |
| for i, (a, b) in enumerate(ctx.offsets[:20]): |
| assert tok.decode([ctx.ids[ctx.ctx_start + i]]).strip() == CONTEXT[a:b].strip() |
|
|
|
|
| def test_lfm_backbone_matches_hf_and_fork_equals_concatenation(model): |
| """The functional stack equals HF Lfm2Model on the concatenated read; continuing a |
| state (dense and padded rows) equals reading the whole sequence.""" |
| rend = model.renderer |
| a = list(rend.context(CONTEXT).ids) |
| b = rend.enc("FIELD /order_id\nKIND: copy\nANSWER:") |
| with torch.no_grad(): |
| ref = model.backbone.text(torch.tensor([a + b])).last_hidden_state[0] |
| h1, st = model.backbone(torch.tensor([a]), torch.tensor([len(a)])) |
| h2, _ = model.backbone(torch.tensor([b]), torch.tensor([len(b)]), st) |
| pad = torch.full((2, len(a) + 5), rend.pad_id) |
| pad[0, :len(a)] = torch.tensor(a) |
| pad[1, :3] = torch.tensor(a[:3]) |
| _, st2 = model.backbone(pad, torch.tensor([len(a), 3])) |
| fork = torch.full((1, len(b) + 9), rend.pad_id) |
| fork[0, :len(b)] = torch.tensor(b) |
| h3, _ = model.backbone(fork, torch.tensor([len(b)]), st2.select(torch.tensor([0]))) |
| torch.testing.assert_close(torch.cat([h1[0], h2[0]]), ref, atol=2e-5, rtol=2e-5) |
| torch.testing.assert_close(h3[0, :len(b)], h2[0], atol=2e-5, rtol=2e-5) |
| |
| assert all(r.shape == (1, 0) and t.shape == (1, 2, 64) for r, t in st.lin) |
|
|
|
|
| def test_lfm_forked_fields_equal_independent_full_forward(model): |
| """Each fork row (prefix read by reference, run in one batch with other fields) equals |
| an independent HF full forward over context + that field's query; and a row's fork |
| hidden states are bitwise the same alone and in the batch.""" |
| rend = model.renderer |
| ctx = list(rend.context(CONTEXT).ids) |
| queries = [rend.enc(t) for t in ( |
| "FIELD /order_id\nKIND: copy\nANSWER:", |
| "FIELD /department\nOPTIONS:\n- billing\n- technical\nANSWER:", |
| "ARRAY /items\nANSWER:", |
| "FIELD /total\nTITLE: total charged\nKIND: number\nCONSTRAINTS: {}\nANSWER:")] |
| bb = model.backbone |
|
|
| def forks(rows): |
| width = max(len(q) for q in rows) |
| ids = torch.full((len(rows), width), rend.pad_id) |
| for i, q in enumerate(rows): |
| ids[i, :len(q)] = torch.tensor(q) |
| _, st = bb(torch.tensor([ctx]), torch.tensor([len(ctx)])) |
| fs = bb.fork(st, len(rows), width) |
| return bb.fork_forward(ids, torch.tensor([len(q) for q in rows]), fs) |
|
|
| with torch.no_grad(): |
| batch = forks(queries) |
| for i, q in enumerate(queries): |
| ref = bb.text(torch.tensor([ctx + q])).last_hidden_state[0, len(ctx):] |
| torch.testing.assert_close(batch[i, :len(q)], ref, atol=2e-5, rtol=2e-5) |
| alone = forks([q]) |
| assert torch.equal(alone[0, :len(q)], batch[i, :len(q)]), i |
|
|
|
|
| def test_lfm_fork_decode_continues_state(model): |
| """keep=True fork passes (header, then one token at a time: the generate decode path) |
| equal the full read of context + header + tokens.""" |
| rend = model.renderer |
| bb = model.backbone |
| ctx = list(rend.context(CONTEXT).ids) |
| head = rend.enc("WRITE /reply\nVALUE:\n") |
| more = rend.enc(" Refund issued today") |
| with torch.no_grad(): |
| _, st = bb(torch.tensor([ctx]), torch.tensor([len(ctx)])) |
| fs = bb.fork(st, 2, len(head) + len(more) + 1) |
| ids = torch.tensor([head, head]) |
| bb.fork_forward(ids, torch.tensor([len(head)] * 2), fs, keep=True) |
| outs = [] |
| for t in more: |
| outs.append(bb.fork_forward(torch.tensor([[t], [t]]), torch.tensor([1, 1]), fs, |
| keep=True)[:, 0]) |
| ref = bb.text(torch.tensor([ctx + head + more])).last_hidden_state[0, -len(more):] |
| got = torch.stack(outs, 1) |
| torch.testing.assert_close(got[0], ref, atol=2e-5, rtol=2e-5) |
| assert torch.equal(got[0], got[1]) |
|
|
|
|
| def test_lfm_chunked_checkpointed_trunk_equals_dense(model, monkeypatch): |
| """Training trunks beyond DENSE_PREFILL_MAX (checkpointed chunks) equal the dense |
| checkpointed read in hidden states, state, forks and gradients (test_train_long).""" |
| emb = model.backbone.text.embed_tokens.weight |
| model.backbone.grad_ckpt = True |
| emb.requires_grad_(False) |
| try: |
| ttl.test_chunked_checkpointed_trunk_equals_dense(model, monkeypatch) |
| finally: |
| model.backbone.grad_ckpt = False |
| emb.requires_grad_(True) |
| for p in model.backbone.parameters(): |
| p.grad = None |
|
|
|
|
| def test_lfm_training_example_and_losses(model): |
| from deg.train import deg_example, forward_losses |
|
|
| a = CONTEXT.index("AB-1234") |
| rec = {"uid": "t1", "schema": _sub(["department", "order_id", "items", "reply"]), |
| "context": CONTEXT, |
| "value": {"department": "billing", "order_id": "AB-1234", |
| "items": [{"name": "chairs", "qty": 2}], "reply": "Refund issued."}, |
| "spans": {"/order_id": [a, a + 7], |
| "/items/0/name": [CONTEXT.index("chairs"), |
| CONTEXT.index("chairs") + 6]}, |
| "span_alternatives": {}, "span_probs": {}, |
| "decide_probs": {"/department": [0.8, 0.15, 0.05]}} |
| ex = deg_example(model, rec, random.Random(0), shuffle_p=0.0) |
| assert sorted(it.kind for it in ex.items) == [ |
| "anchor", "choice", "extract", "extract", "generate", "score"] |
| model.backbone.train() |
| try: |
| losses, _ = forward_losses(model, [ex], 4096, 16) |
| finally: |
| model.backbone.eval() |
| total = sum(losses[k] for k in ("decide", "extract", "anchor", "lm")) |
| assert torch.isfinite(total) |
| total.backward() |
| layers = model.backbone.text.layers |
| assert layers[0].conv.in_proj.weight.grad is not None |
| assert layers[1].self_attn.q_proj.weight.grad is not None |
| assert layers[1].feed_forward.w1.weight.grad is not None |
| for p in [*model.heads.parameters(), *model.backbone.parameters()]: |
| p.grad = None |
|
|
|
|
| def test_lfm_checkpoint_roundtrip_keeps_backbone(tmp_path): |
| model = _lfm_model(5) |
| with torch.no_grad(): |
| for p in model.backbone.parameters(): |
| p.copy_(p.bfloat16().float()) |
| model.save(tmp_path / "ck") |
| meta = __import__("json").loads((tmp_path / "ck" / "deg.json").read_text()) |
| assert meta["backbone"] == BACKBONE |
| back = load_model(tmp_path / "ck", device="cpu", dtype=torch.float32, lm=tiny_lfm(3), |
| tokenizer=_tok()) |
| assert back.backbone.fam is LFM2 |
| want = _raw_outputs(model, _sub(["order_id", "department"])) |
| got = _raw_outputs(back, _sub(["order_id", "department"])) |
| for k in want: |
| assert torch.equal(want[k], got[k]), k |
| with pytest.raises(ValueError): |
| load_model(tmp_path / "ck", device="cpu", dtype=torch.float32, lm=tiny_lfm(3), |
| tokenizer=_tok(), backbone="qwen3.5-0.8b") |
| with pytest.raises(ValueError): |
| from deg.model import attach_base |
|
|
| attach_base(back, tm.tiny_lm(0)) |
|
|
|
|
| def test_lfm_hybrid_base_shares_embeddings(pair): |
| _, hyb = pair |
| assert hyb.base is not None and hyb.base.shared_embeddings and hyb.base.fam is LFM2 |
| emb = hyb.backbone.text.embed_tokens.weight |
| assert hyb.base.text.embed_tokens.weight is emb and hyb.base.lm.lm_head.weight is emb |
| resp = fill(hyb, CONTEXT, thg.SCHEMA) |
| validate_value(compile_schema(thg.SCHEMA), resp["value"]) |
|
|
|
|
| |
|
|
|
|
| def _gpu_free_gib() -> float: |
| if not torch.cuda.is_available(): |
| return 0.0 |
| return torch.cuda.mem_get_info()[0] / 2**30 |
|
|
|
|
| @pytest.mark.skipif(LFM_SNAPSHOT is None or _gpu_free_gib() < 4.0, |
| reason="needs CUDA with ~4 GiB free and the LFM2.5-1.2B snapshot") |
| def test_gpu_lfm_real_model_bitwise_independence_and_hf_match(): |
| """Real LFM2.5-1.2B in bf16 on CUDA: the fork path (bi_linear GEMMs, fork_attention, |
| Triton LFM RMSNorm, elementwise short conv) gives bitwise identical raw outputs alone, |
| reordered, inside a bigger schema and across requests (fill_many); the hybrid's base |
| text is identical alone and next to another generate field; the functional trunk |
| matches HF within bf16 noise.""" |
| m = load_model(None, device="cuda", backbone=BACKBONE, generate="base") |
| rend = m.renderer |
| ids = torch.tensor([rend.context(CONTEXT).ids], device="cuda") |
| with torch.no_grad(): |
| ref = m.backbone.text(ids).last_hidden_state[0].float() |
| mine, _ = m.backbone(ids, torch.tensor([ids.shape[1]], device="cuda")) |
| cos = torch.nn.functional.cosine_similarity(mine[0].float(), ref, -1) |
| assert float(cos.min()) > 0.99, float(cos.min()) |
| flat = [n for n in tm.SCHEMA["properties"] if n not in ("items", "reply")] |
| merged = _raw_outputs(m, _sub(flat)) |
| reordered = _raw_outputs(m, _sub(list(reversed(flat)))) |
| big = {"type": "object", "properties": { |
| **{f"f{i}": {"type": ["string", "null"], "description": f"filler {i}"} |
| for i in range(12)}, **_sub(flat)["properties"]}, |
| "required": _sub(flat)["required"]} |
| big_out = _raw_outputs(m, big) |
| for name in flat: |
| alone = _raw_outputs(m, _sub([name]))[f"/{name}"] |
| for other in (merged, reordered, big_out): |
| assert torch.equal(alone, other[f"/{name}"]), name |
| resp = fill(m, CONTEXT, tm.SCHEMA) |
| validate_value(compile_schema(tm.SCHEMA), resp["value"]) |
| both = fill(m, CONTEXT, tm._with_note()) |
| assert both["value"]["reply"] == resp["value"]["reply"] |
| assert both["meta"]["/reply"] == resp["meta"]["/reply"] |
| from deg.infer import fill_many |
|
|
| other = "Dana again: the desk arrived broken. Order AB-9 total $80." |
| jobs = [(CONTEXT, tm.SCHEMA, None), (other, tm._with_note(), None), |
| (other, _sub(["department", "items"]), None)] |
| for job, got in zip(jobs, fill_many(m, jobs)): |
| assert got == fill(m, *job) |
|
|
|
|
| @pytest.mark.skipif(LFM_SNAPSHOT is None or _gpu_free_gib() < 4.0, |
| reason="needs CUDA with ~4 GiB free and the LFM2.5-1.2B snapshot") |
| def test_gpu_lfm_chunked_prefill_matches(monkeypatch): |
| """Real LFM2.5-1.2B on CUDA: the chunked prefill (flash-style fork kernel over the K/V |
| buffer, head_dim 64, 4 query heads per K/V head) matches the one-shot read within bf16 |
| noise for padded rows.""" |
| import deg.model as dm |
|
|
| m = load_model(None, device="cuda", backbone=BACKBONE) |
| rend = m.renderer |
| a = rend.context(" ".join([CONTEXT] * 40)).ids |
| b = rend.context(" ".join([CONTEXT] * 15)).ids |
| x = torch.full((2, len(a)), rend.pad_id, device="cuda") |
| x[0], x[1, :len(b)] = torch.tensor(a), torch.tensor(b) |
| n = torch.tensor([len(a), len(b)], device="cuda") |
| with torch.no_grad(): |
| h_one, s_one = m.backbone(x, n) |
| monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 512) |
| monkeypatch.setattr(dm, "PREFILL_CHUNK", 512) |
| h, s = m.backbone(x, n) |
| for row, L in enumerate((len(a), len(b))): |
| cos = torch.nn.functional.cosine_similarity(h[row, :L].float(), |
| h_one[row, :L].float(), -1) |
| assert float(cos.min()) > 0.99 and float(cos.mean()) > 0.9995, row |
| assert torch.equal(s.kv_mask, s_one.kv_mask) |
|
|
