"""The Gemma 4 backbone (``--backbone gemma-4-e2b-it``; SPEC.md §12.8) on a tiny random Gemma 4 text model (CPU, exact torch reference kernels) with the real Gemma 4 tokenizer, plus GPU tests on the real google/gemma-4-E2B-it. The tiny model has E2B's structure at toy width: sliding-window + global layers (window 8, so every test context runs far past the window), the last 4 of 7 layers sharing the K/V of layers 1 (sliding) / 2 (global), head_dim 32 / global 64, one K/V head, PLE, double-wide MLPs in the shared layers, logit soft-cap. The model-contract tests of ``test_model`` / ``test_hybrid_generate`` / ``test_train_long`` are collected here against this module's Gemma fixtures; Gemma-specific tests check the functional stack against HF ``Gemma4ForCausalLM`` (logits), chunked vs unchunked reads (inference and checkpointed training, gradients), forks vs full-sequence reads, decode-graph buffers vs eager, long streams over a real 512 window, the renderer, LoRA, checkpoints and the chat layout. """ from __future__ import annotations import json import random import pytest torch = pytest.importorskip("torch") pytest.importorskip("transformers") import test_hybrid_generate as thg # noqa: E402 import test_model as tm # noqa: E402 import test_train_long as ttl # noqa: E402 from test_hybrid_generate import ( # noqa: E402,F401 (collected here with the Gemma pair) 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 ( # noqa: E402,F401 (collected here with the Gemma model) CONTEXT, _LargestTensor, _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_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 # noqa: E402,F401 from deg import compile_schema, validate_value # noqa: E402 from deg.infer import fill # noqa: E402 from deg.model import ( # noqa: E402 GEMMA4, GemmaBackbone, attach_base, base_snapshot, load_model, load_tokenizer, ) BACKBONE = "gemma-4-e2b-it" try: GEMMA_SNAPSHOT = base_snapshot(BACKBONE) except FileNotFoundError: # pragma: no cover - needs the local HF cache GEMMA_SNAPSHOT = None LAYERS = ["sliding_attention"] * 2 + ["full_attention"] + ["sliding_attention"] * 3 \ + ["full_attention"] def tiny_gemma(seed: int = 0, window: int = 8, vocab: int = 262144): from transformers import Gemma4ForCausalLM, Gemma4TextConfig torch.manual_seed(seed) cfg = Gemma4TextConfig( vocab_size=vocab, vocab_size_per_layer_input=vocab, hidden_size=64, intermediate_size=96, num_hidden_layers=7, layer_types=list(LAYERS), num_kv_shared_layers=4, hidden_size_per_layer_input=16, head_dim=32, global_head_dim=64, num_attention_heads=4, num_key_value_heads=1, sliding_window=window, final_logit_softcapping=30.0, use_double_wide_mlp=True, rms_norm_eps=1e-6, pad_token_id=0, bos_token_id=2, eos_token_id=1, max_position_embeddings=8192, tie_word_embeddings=True) lm = Gemma4ForCausalLM(cfg).float().eval() with torch.no_grad(): # random norms / layer scalars so the stack is not near-identity for name, p in lm.named_parameters(): if "norm" in name: p.normal_(1.0, 0.3) for layer in lm.model.layers: layer.layer_scalar.uniform_(0.6, 1.2) return lm def _tok(): if GEMMA_SNAPSHOT is None: pytest.skip("gemma-4-E2B-it tokenizer not cached") return load_tokenizer(backbone=BACKBONE) def _gemma_model(seed: int = 0, **kw): m = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_gemma(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 _gemma_model() def _share_frozen(deg, base) -> None: """deg trains with the token embeddings and the PLE table frozen: equal to the base.""" with torch.no_grad(): deg.model.embed_tokens.weight.copy_(base.model.embed_tokens.weight) deg.model.embed_tokens_per_layer.weight.copy_(base.model.embed_tokens_per_layer.weight) @pytest.fixture(scope="module") def pair(): """(deg-only, hybrid) Gemma models: the hybrid's frozen base is another random Gemma with the same (frozen) embeddings and PLE table.""" def make(generate): base, deg = tiny_gemma(7), tiny_gemma(0) _share_frozen(deg, base) 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 _ids(rows, pad=0): width = max(len(r) for r in rows) x = torch.full((len(rows), width), pad, dtype=torch.long) for i, r in enumerate(rows): x[i, :len(r)] = torch.tensor(r) return x, torch.tensor([len(r) for r in rows]) # ----------------------------------------------------------------------- identity def test_gemma_renderer_and_meta(model): rend = model.renderer tok = rend.tokenizer assert model.backbone.fam is GEMMA4 and isinstance(model.backbone, GemmaBackbone) assert model.meta["backbone"] == BACKBONE ctx = rend.context(CONTEXT) assert ctx.ids[0] == tok.convert_tokens_to_ids("") assert tok.convert_ids_to_tokens(ctx.ids[1]) == "<|turn>" assert tok.decode(ctx.ids[:ctx.ctx_start]) == ( "<|turn>system\nFill JSON Schema fields from the context: decide, copy spans " "from the context, and write text only for fields marked x-generate.\n" "<|turn>user\nCONTEXT:\n") assert rend.eos_id == tok.convert_tokens_to_ids("") != tok.unk_token_id assert rend.pad_id == tok.convert_tokens_to_ids("") == 0 assert rend.gen_turn == "\n<|turn>model\n" # the chat template's generation prompt is the prefix + the gen turn (thinking off) want = tok.apply_chat_template( [{"role": "system", "content": "S"}, {"role": "user", "content": "U"}], tokenize=False, add_generation_prompt=True) assert want == "<|turn>system\nS\n<|turn>user\nU" + rend.gen_turn # every context token maps back to its text; SentencePiece's leading "▁" is the space # the offset includes, and trim_span strips it from extracted values for i, (a, b) in enumerate(ctx.offsets): piece = tok.decode([ctx.ids[ctx.ctx_start + i]]) assert piece == CONTEXT[a:b], (i, piece, CONTEXT[a:b]) from deg.model import token_span, trim_span a = CONTEXT.index("AB-1234") i, j = token_span(ctx.offsets, a, a + 7) lo, hi = ctx.offsets[i][0], ctx.offsets[j][1] assert CONTEXT[lo:hi].strip() == "AB-1234" assert CONTEXT[slice(*trim_span(CONTEXT, lo, hi))] == "AB-1234" def test_gemma_slot_map_matches_hf_kv_sharing(model): """Layers 0-2 compute K/V; the sharing layers read the last earlier owner of their type (HF ``store_full_length_kv``).""" bb = model.backbone assert bb.kv_slot == [0, 1, 2, 1, 1, 1, 2] assert bb.slot_kinds == ["sliding", "sliding", "full"] assert bb.owner == [True] * 3 + [False] * 4 and bb.tail == 7 for i, layer in enumerate(bb.text.layers): a = layer.self_attn assert a.is_kv_shared_layer == (not bb.owner[i]) if a.store_full_length_kv: assert bb.slot_layers[bb.kv_slot[i]] == i # -------------------------------------------------------------------- HF parity def test_gemma_matches_hf_logits_and_continuation(model): """The functional stack equals HF Gemma4ForCausalLM (logits incl. the soft-cap) on the concatenated read; continuing a state (across the window, dense and padded rows) equals reading the whole sequence.""" rend = model.renderer bb = model.backbone a = list(rend.context(CONTEXT).ids) b = rend.enc("FIELD /order_id\nKIND: copy\nANSWER:") assert len(a) > 4 * bb.window with torch.no_grad(): ref = bb.lm(torch.tensor([a + b])).logits[0] h1, st = bb(torch.tensor([a]), torch.tensor([len(a)])) h2, _ = bb(torch.tensor([b]), torch.tensor([len(b)]), st) mine = bb.lm_logits(torch.cat([h1[0], h2[0]])) x, n = _ids([a, a[:3]], rend.pad_id) _, st2 = bb(torch.cat([x, torch.zeros(2, 5, dtype=torch.long)], 1), n) fork = torch.full((1, len(b) + 9), rend.pad_id) fork[0, :len(b)] = torch.tensor(b) h3, _ = bb(fork, torch.tensor([len(b)]), st2.select(torch.tensor([0]))) # the short row (3 tokens < tail): its tail is partly empty, masked h4, _ = bb(fork, torch.tensor([len(b)]), st2.select(torch.tensor([1]))) ref4 = bb.text(torch.tensor([a[:3] + b])).last_hidden_state[0, 3:] torch.testing.assert_close(mine, ref, atol=2e-4, rtol=2e-4) torch.testing.assert_close(h3[0, :len(b)], h2[0], atol=3e-4, rtol=1e-3) torch.testing.assert_close(h4[0, :len(b)], ref4, atol=3e-4, rtol=1e-3) # sliding slots keep a bounded tail; global slots the whole prefix for (k, v), kind in zip(st.kv, bb.slot_kinds): assert k.shape[2] == (bb.tail if kind == "sliding" else len(a)), kind def test_gemma_forked_fields_equal_independent_full_forward(model): """Each fork row (prefix read by reference: global K/V and sliding tails, KV-sharing layers reading their source's own buffers) equals an independent HF full forward over context + that field's query, and is bitwise the same alone and in a 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): ids, n = _ids(rows, rend.pad_id) _, st = bb(torch.tensor([ctx]), torch.tensor([len(ctx)])) fs = bb.fork(st, len(rows), ids.shape[1]) return bb.fork_forward(ids, n, fs) assert max(len(q) for q in queries) > bb.window # own keys cross the window too 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=3e-4, rtol=1e-3) alone = forks([q]) assert torch.equal(alone[0, :len(q)], batch[i, :len(q)]), i def test_gemma_fork_over_several_prefixes_and_short_prefix(model): """Rows continuing different prefixes in one fork batch (one shorter than the sliding tail) equal their full reads.""" rend = model.renderer bb = model.backbone p1 = list(rend.context(CONTEXT).ids) p2 = rend.enc("Hi") assert len(p2) < bb.tail q = rend.enc("FIELD /a\nKIND: copy\nANSWER:") with torch.no_grad(): states = [bb(torch.tensor([p]), torch.tensor([len(p)]))[1] for p in (p1, p2)] fs = bb.fork(states, [1, 0, 1], len(q)) got = bb.fork_forward(torch.tensor([q] * 3), torch.tensor([len(q)] * 3), fs) for row, p in zip(range(3), (p2, p1, p2)): ref = bb.text(torch.tensor([p + q])).last_hidden_state[0, len(p):] torch.testing.assert_close(got[row], ref, atol=3e-4, rtol=1e-3) def test_gemma_fork_decode_continues_state_past_the_window(model): """keep=True fork passes (header, then one token at a time, more own tokens than the window) 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, sorry for the double charge on your card") assert len(more) > bb.window with torch.no_grad(): _, st = bb(torch.tensor([ctx]), torch.tensor([len(ctx)])) fs = bb.fork(st, 2, len(head) + len(more) + 1) bb.fork_forward(torch.tensor([head, head]), torch.tensor([len(head)] * 2), fs, keep=True) outs = [bb.fork_forward(torch.tensor([[t], [t]]), torch.tensor([1, 1]), fs, keep=True, decode=True)[:, 0] for t in more] 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=3e-4, rtol=1e-3) assert torch.equal(got[0], got[1]) def test_gemma_decode_graph_buffers_equal_eager(model): """The decode graph's static buffers (global prefixes concatenated up to pcap, sliding tails per prefix, own K/V; :meth:`GemmaBackbone.static_fork_state` / ``load_static``) give the same rows as the eager fork state (CPU: no capture, same math).""" rend = model.renderer bb = model.backbone p1 = list(rend.context(CONTEXT).ids) p2 = rend.enc("Short") head = rend.enc("WRITE /x\nVALUE:\n") toks = rend.enc(" one two three four five six seven eight nine ten") with torch.no_grad(): states = [bb(torch.tensor([p]), torch.tensor([len(p)]))[1] for p in (p1, p2)] prow = [1, 0, 1] cap = len(head) + len(toks) + 1 def run(fs, rows): ids = torch.full((rows, len(head)), rend.pad_id) ids[:3] = torch.tensor(head) lens = torch.tensor([len(head)] * 3 + [0] * (rows - 3)) outs = [bb.fork_forward(ids, lens, fs, keep=True)[:3, len(head) - 1]] for t in toks: step = torch.tensor([1] * 3 + [0] * (rows - 3)) ids1 = torch.full((rows, 1), rend.pad_id) ids1[:3] = t outs.append(bb.fork_forward(ids1, step, fs, keep=True, decode=True)[:3, 0]) return torch.stack(outs) eager = run(bb.fork(states, prow, cap), 3) fs = bb.static_fork_state(4, 256, cap + 64) bb.load_static(fs, states, prow) static = run(fs, 4) torch.testing.assert_close(static, eager, atol=1e-6, rtol=1e-6) # ------------------------------------------------------------- chunked / long reads def test_gemma_chunked_read_padded_rows_and_prefix_state(model, monkeypatch): """The chunked read (tails carried across chunks) over right-padded rows and on top of an existing prefix equals the one-shot read: hidden states of real tokens, global K/V at real positions, sliding tails.""" import deg.model as dm rend = model.renderer bb = model.backbone a = rend.enc(" ".join([CONTEXT] * 6)) b = rend.enc(" ".join([CONTEXT] * 2)) pre = rend.enc("Note: ") x, n = _ids([a, b], rend.pad_id) with torch.no_grad(): _, st0 = bb(torch.tensor([pre, pre]), torch.tensor([len(pre)] * 2)) h_one, s_one = bb(x, n, st0) monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 40) monkeypatch.setattr(dm, "PREFILL_CHUNK", 37) h, s = bb(x, n, st0) for row, m in enumerate((len(a), len(b))): torch.testing.assert_close(h[row, :m], h_one[row, :m], atol=3e-4, rtol=1e-3) assert torch.equal(s.kv_mask, s_one.kv_mask) and torch.equal(s.lengths, s_one.lengths) for (k, v), (k1, v1), kind in zip(s.kv, s_one.kv, bb.slot_kinds): for row in range(2): real = s.kv_mask[row] if kind == "full" else slice(None) torch.testing.assert_close(k[row][:, real], k1[row][:, real], atol=3e-4, rtol=1e-3) torch.testing.assert_close(v[row][:, real], v1[row][:, real], atol=3e-4, rtol=1e-3) def test_gemma_chunked_checkpointed_trunk_equals_dense(model, monkeypatch): """Training trunks beyond DENSE_PREFILL_MAX (checkpointed chunks, global K/V passed as parts) equal the dense checkpointed read in hidden states, state, forks and gradients, and save no [B, L, L] tensor.""" import deg.model as dm bb = model.backbone emb = bb.text.embed_tokens.weight bb.grad_ckpt = True emb.requires_grad_(False) try: rend = model.renderer a = rend.context(" ".join([CONTEXT] * 3)).ids b = rend.context(" ".join([CONTEXT] * 2)).ids q = rend.enc("FIELD /x\nKIND: copy\nANSWER:") dense = ttl._run(model, [a, b], [q, q[:5], q], [0, 1, 1]) monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 48) monkeypatch.setattr(dm, "PREFILL_CHUNK", 32) calls = [] orig = GemmaBackbone._forward_chunked def spy(self, *args, **kw): calls.append(torch.is_grad_enabled()) return orig(self, *args, **kw) monkeypatch.setattr(GemmaBackbone, "_forward_chunked", spy) with _LargestTensor() as big: chunked = ttl._run(model, [a, b], [q, q[:5], q], [0, 1, 1]) assert calls and all(calls) L = len(a) # the dense read's [2, L, L] mask is gone; the largest tensors left are linear in L # (3 fork rows x (L + query) x 64 of selected global K/V) assert big.largest < L * L, big.largest (h1, s1, f1, g1), (h2, s2, f2, g2) = dense, chunked for i, n in enumerate((len(a), len(b))): torch.testing.assert_close(h2[i, :n], h1[i, :n], atol=3e-4, rtol=1e-3) for (k, v), (k1, v1), kind in zip(s2.kv, s1.kv, bb.slot_kinds): assert k.shape == k1.shape m = s1.kv_mask[:, None, :, None] if kind == "full" else 1 torch.testing.assert_close(k * m, k1 * m, atol=3e-4, rtol=1e-3) torch.testing.assert_close(v * m, v1 * m, atol=3e-4, rtol=1e-3) assert torch.equal(s2.kv_mask, s1.kv_mask) and torch.equal(s2.lengths, s1.lengths) for x, y in zip(f2, f1): torch.testing.assert_close(x, y, atol=3e-4, rtol=1e-3) assert g1.keys() == g2.keys() and g1 assert not any("embed_tokens_per_layer" in n for n in g1) # frozen PLE table assert any("per_layer_model_projection" in n for n in g1) assert any(".3.self_attn.q_proj" in n for n in g1) # sharing layers train for n in g1: # fp32 summation-order noise, relative to each gradient's scale scale = float(g1[n].abs().max()) torch.testing.assert_close(g2[n], g1[n], atol=2e-4 * scale, rtol=1e-3, msg=n) finally: bb.grad_ckpt = False emb.requires_grad_(True) for p in bb.parameters(): p.grad = None def test_gemma_real_window_long_stream_over_512_tokens(monkeypatch): """E2B's real window (512): a ~1.6k-token document read in 300-token chunks equals the one-shot read and HF, and a field forked off it (prefix tail = the last 511 positions) equals the full read; the tail is exactly the last 511 positions' K/V.""" import deg.model as dm tok = _tok() m = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_gemma(3, window=512), tokenizer=tok) bb, rend = m.backbone, m.renderer doc = list(rend.context(" ".join([CONTEXT] * 30)).ids) assert len(doc) > 3 * 512 q = rend.enc("FIELD /order_id\nKIND: copy\nANSWER:") with torch.no_grad(): ref = bb.text(torch.tensor([doc + q])).last_hidden_state[0] h_one, s_one = bb(torch.tensor([doc]), torch.tensor([len(doc)])) monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 400) monkeypatch.setattr(dm, "PREFILL_CHUNK", 300) h, s = bb(torch.tensor([doc]), torch.tensor([len(doc)])) fs = bb.fork(s, 1, len(q)) f = bb.fork_forward(torch.tensor([q]), torch.tensor([len(q)]), fs) torch.testing.assert_close(h[0], ref[:len(doc)], atol=5e-4, rtol=5e-4) torch.testing.assert_close(h_one[0], ref[:len(doc)], atol=5e-4, rtol=5e-4) torch.testing.assert_close(f[0], ref[len(doc):], atol=5e-4, rtol=5e-4) full = s_one.kv[2] # the global slot holds every position ... assert full[0].shape[2] == len(doc) for (k, _), kind in zip(s.kv, bb.slot_kinds): # ... sliding slots only the last 511 assert k.shape[2] == (511 if kind == "sliding" else len(doc)) assert int(fs.wplen[0]) == 511 def test_kernel_reference_window_matches_explicit_mask(): """fork_attention's CPU reference with a window equals masked softmax attention over (prefix tail, own keys).""" from deg.kernels import fork_attention torch.manual_seed(0) B, L, HQ, D, P, W = 2, 7, 4, 16, 5, 6 q = torch.randn(B, L, HQ, D) pk, pv = torch.randn(1, 2 * P, D), torch.randn(1, 2 * P, D) ok, ov = torch.randn(B, 1, 12, D), torch.randn(B, 1, 12, D) qstart = torch.tensor([0, 3]) poff, plen = torch.tensor([0, P + 1]), torch.tensor([P, P - 1]) got = fork_attention(q, pk, pv, ok, ov, qstart, poff, plen, 0.5, W) for b in range(B): t0, po, n = int(qstart[b]), int(poff[b]), int(plen[b]) keys = torch.cat([pk[0, po:po + n], ok[b, 0, :t0 + L]]) vals = torch.cat([pv[0, po:po + n], ov[b, 0, :t0 + L]]) kpos = torch.cat([torch.arange(-n, 0), torch.arange(t0 + L)]) for i in range(L): t = t0 + i vis = (kpos <= t) & (t - kpos < W) s = (q[b, i] @ keys.T) * 0.5 s = s.masked_fill(~vis, float("-inf")) want = s.softmax(-1) @ vals torch.testing.assert_close(got[b, i], want, atol=1e-5, rtol=1e-5) def test_blocked_sliding_sdpa_and_row_select_backward_are_exact(): """Sliding layers run long reads in query blocks of the window against only the keys each block can see; fork rows select their prefix with a GEMM backward: both equal the plain ops (values and gradients).""" from deg.model import _gemma_sdpa, _select_rows torch.manual_seed(0) B, H, D, W = 2, 4, 16, 8 for L in (8, 9, 29, 32): q = torch.randn(B, H, L, D, dtype=torch.float64, requires_grad=True) k = torch.randn(B, 1, W - 1 + L, D, dtype=torch.float64, requires_grad=True) v = torch.randn(B, 1, W - 1 + L, D, dtype=torch.float64, requires_grad=True) i = torch.arange(L)[:, None] pos = torch.cat([torch.arange(-(W - 1), 0), torch.arange(L)])[None] mask = ((pos <= i) & (i - pos < W))[None].repeat(B, 1, 1) mask[1, :, :3] = False # a short prefix: missing tail slots g = torch.randn(B, H, L, D, dtype=torch.float64) outs = [] for win in (W, 0): o = _gemma_sdpa(q, k, v, mask, 1.0, win) (o * g).sum().backward() outs.append((o.detach(), q.grad.clone(), k.grad.clone(), v.grad.clone())) q.grad = k.grad = v.grad = None for a, b in zip(*outs): torch.testing.assert_close(a, b) x = torch.randn(3, 2, 7, 4, dtype=torch.float64, requires_grad=True) rows = torch.tensor([2, 0, 2, 2, 1]) g = torch.randn(5, 2, 7, 4, dtype=torch.float64) (_select_rows(x, rows) * g).sum().backward() got = x.grad.clone() x.grad = None (x.index_select(0, rows) * g).sum().backward() torch.testing.assert_close(got, x.grad) # ------------------------------------------------------- training / LoRA / checkpoints def test_gemma_chunked_long_prefill_padded_batch(model, monkeypatch): """The shared padded-batch long-prefill contract with sliding tails compared whole.""" test_gemma_chunked_read_padded_rows_and_prefix_state(model, monkeypatch) def test_gemma_training_example_losses_and_chat_layout(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"] gen = next(it for it in ex.items if it.kind == "generate") tok = model.renderer.tokenizer assert "VALUE:" in tok.decode(gen.seq[:gen.prompt_len]) assert gen.seq[-1] == model.renderer.eos_id model.meta["gen_layout"] = "chat" # self-generation in the base's chat turn try: chat = deg_example(model, rec, random.Random(0), shuffle_p=0.0) g2 = next(it for it in chat.items if it.kind == "generate") assert tok.decode(g2.seq[:g2.prompt_len]).endswith("\n<|turn>model\n") assert model.generator() == (model, True, True) finally: model.meta.pop("gen_layout") assert model.generator() == (model, True, False) 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].self_attn.k_proj.weight.grad is not None assert layers[5].self_attn.q_proj.weight.grad is not None assert layers[5].mlp.up_proj.weight.grad is not None assert layers[5].per_layer_input_gate.weight.grad is not None assert model.backbone.text.embed_tokens_per_layer.weight.grad is None for p in [*model.heads.parameters(), *model.backbone.parameters()]: p.grad = None def test_gemma_lm_logits_softcap(model): h = torch.randn(3, 64) * 50 raw = torch.nn.functional.linear(h, model.backbone.lm.lm_head.weight) with torch.no_grad(): got = model.backbone.lm_logits(h) torch.testing.assert_close(got, 30.0 * torch.tanh(raw / 30.0), atol=3e-4, rtol=1e-3) assert float(got.abs().max()) <= 30.0 def test_gemma_lora_skips_missing_kv_projections(): from deg.model import add_lora m = _gemma_model(4) params = add_lora(m, 4) layers = m.backbone.text.layers assert hasattr(layers[0].self_attn.k_proj, "lora_A") assert not hasattr(layers[4].self_attn, "k_proj") # a K/V-sharing layer assert hasattr(layers[4].self_attn.q_proj, "lora_A") assert hasattr(layers[4].mlp.down_proj, "lora_A") # 3 owners x 7 targets + 4 sharing layers x 5 targets, A and B each assert len(params) == 2 * (3 * 7 + 4 * 5) want = _raw_outputs(m, _sub(["order_id", "department"])) # B = 0: still the base m2 = _gemma_model(4) got = _raw_outputs(m2, _sub(["order_id", "department"])) for k in want: torch.testing.assert_close(want[k], got[k], atol=1e-5, rtol=1e-5) def test_gemma_checkpoint_roundtrip_keeps_backbone_without_ple_table(tmp_path): model = _gemma_model(5) with torch.no_grad(): # checkpoints store the backbone (incl. layer_scalar) in bf16 for p in model.backbone.lm.state_dict().values(): # persistent tensors only if p.is_floating_point(): p.copy_(p.bfloat16().float()) model.meta["gen_layout"] = "chat" model.save(tmp_path / "ck") meta = json.loads((tmp_path / "ck" / "deg.json").read_text()) assert meta["backbone"] == BACKBONE and meta["gen_layout"] == "chat" from safetensors import safe_open with safe_open(str(tmp_path / "ck" / "backbone.safetensors"), "pt") as fh: assert not any("embed_tokens_per_layer" in k for k in fh.keys()) other = tiny_gemma(3) with torch.no_grad(): # the PLE table comes from the base, which the reload uses other.model.embed_tokens_per_layer.weight.copy_( model.backbone.text.embed_tokens_per_layer.weight) back = load_model(tmp_path / "ck", device="cpu", dtype=torch.float32, lm=other, tokenizer=_tok()) assert isinstance(back.backbone, GemmaBackbone) and back.gen_layout == "chat" 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_gemma(3), tokenizer=_tok(), backbone="qwen3.5-0.8b") with pytest.raises(ValueError): # a Qwen base cannot be the Gemma model's generator attach_base(back, tm.tiny_lm(0)) with pytest.raises(ValueError): # a base with another PLE table is not its base attach_base(back, tiny_gemma(11)) def test_gemma_hybrid_base_shares_embeddings_and_ple(pair): _, hyb = pair assert hyb.base is not None and hyb.base.shared_embeddings assert isinstance(hyb.base, GemmaBackbone) 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 assert hyb.base.text.embed_tokens_per_layer is hyb.backbone.text.embed_tokens_per_layer resp = fill(hyb, CONTEXT, thg.SCHEMA) validate_value(compile_schema(thg.SCHEMA), resp["value"]) def test_gemma_lm_kl_base_is_zero_against_itself(model): """--lm-kl-base on Gemma: KL(base || deg) is 0 when the base is deg's own body.""" import copy from deg.train import base_kl, deg_example m = copy.copy(model) lm = copy.deepcopy(model.backbone.lm) m.base = None m._gen_view = None attach_base(m, lm) rec = {"uid": "t", "schema": _sub(["department", "reply"]), "context": CONTEXT, "value": {"department": "billing", "reply": "Refund issued."}, "spans": {}, "span_alternatives": {}, "span_probs": {}, "decide_probs": {}} ex = deg_example(m, rec, random.Random(0), shuffle_p=0.0) items = [(0, it) for it in ex.items if it.kind == "generate"] with torch.no_grad(): _, st = m.trunk([ex.ctx]) hid = m.run_forks(st, torch.tensor([0]), [items[0][1].seq])[0] it = items[0][1] logits = m.backbone.lm_logits(hid[it.prompt_len - 1: len(it.seq) - 1]).float() kl = base_kl(m, [ex], items, logits, 4096, 16) assert abs(float(kl)) < 1e-5 # ---------------------------------------------------------------------------- GPU def _gpu_free_gib() -> float: if not torch.cuda.is_available(): return 0.0 return torch.cuda.mem_get_info()[0] / 2**30 def _has_weights() -> bool: return GEMMA_SNAPSHOT is not None and any(GEMMA_SNAPSHOT.glob("*.safetensors")) needs_e2b = pytest.mark.skipif(not _has_weights() or _gpu_free_gib() < 14.0, reason="needs CUDA with ~14 GiB free and the E2B-it snapshot") @needs_e2b def test_gpu_gemma_real_model_hf_match_and_bitwise_independence(): """Real gemma-4-E2B-it in bf16: the functional trunk matches HF within bf16 noise; the fork path (bi_linear GEMMs, D=256/512 fork_attention with the window, Triton Gemma RMSNorm) gives bitwise identical raw outputs alone, reordered, in a bigger schema and across requests; the hybrid's text is identical alone and next to another field.""" m = load_model(None, device="cuda", backbone=BACKBONE, generate="base") rend = m.renderer ids = torch.tensor([rend.context(" ".join([CONTEXT] * 20)).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")) lref = m.backbone.lm(ids[:, :64]).logits[0].float() lmine = m.backbone.lm_logits(mine[0, :64]) cos = torch.nn.functional.cosine_similarity(mine[0].float(), ref, -1) # bf16: HF and the functional stack round in different places (fp32 on the same GPU # agrees to 2.6e-5 relative, logits to 1e-2); a few tokens of 1240 dip to ~0.989 assert float(cos.mean()) > 0.999 and float(cos.min()) > 0.98, cos assert (lmine.argmax(-1) == lref.argmax(-1)).float().mean() > 0.95 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)))) for name in flat: alone = _raw_outputs(m, _sub([name]))[f"/{name}"] assert torch.equal(alone, merged[f"/{name}"]), name assert torch.equal(alone, reordered[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"] @needs_e2b def test_gpu_gemma_decode_graph_equals_eager_and_chunked_prefill(monkeypatch): """Real E2B-it: CUDA-graph decode == eager decode (self chat layout); the chunked prefill (window tails across chunks) matches the one-shot read within bf16 noise.""" import deg.infer as di import deg.model as dm m = load_model(None, device="cuda", backbone=BACKBONE) m.meta["gen_layout"] = "chat" schema = tm._with_note() graph = fill(m, CONTEXT, schema) monkeypatch.setattr(di, "GRAPH_DECODE", False) eager = fill(m, CONTEXT, schema) assert graph == eager rend = m.renderer a = rend.context(" ".join([CONTEXT] * 60)).ids x = torch.tensor([a], device="cuda") n = torch.tensor([len(a)], device="cuda") with torch.no_grad(): h_one, _ = m.backbone(x, n) monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 1024) monkeypatch.setattr(dm, "PREFILL_CHUNK", 1000) h, _ = m.backbone(x, n) cos = torch.nn.functional.cosine_similarity(h[0].float(), h_one[0].float(), -1) assert float(cos.min()) > 0.99 and float(cos.mean()) > 0.9995