| """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 |
| 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, |
| _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 |
|
|
| from deg import compile_schema, validate_value |
| from deg.infer import fill |
| from deg.model import ( |
| GEMMA4, |
| GemmaBackbone, |
| attach_base, |
| base_snapshot, |
| load_model, |
| load_tokenizer, |
| ) |
|
|
| BACKBONE = "gemma-4-e2b-it" |
| try: |
| GEMMA_SNAPSHOT = base_snapshot(BACKBONE) |
| except FileNotFoundError: |
| 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(): |
| 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]) |
|
|
|
|
| |
|
|
|
|
| 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("<bos>") |
| assert tok.convert_ids_to_tokens(ctx.ids[1]) == "<|turn>" |
| assert tok.decode(ctx.ids[:ctx.ctx_start]) == ( |
| "<bos><|turn>system\nFill JSON Schema fields from the context: decide, copy spans " |
| "from the context, and write text only for fields marked x-generate.<turn|>\n" |
| "<|turn>user\nCONTEXT:\n") |
| assert rend.eos_id == tok.convert_tokens_to_ids("<turn|>") != tok.unk_token_id |
| assert rend.pad_id == tok.convert_tokens_to_ids("<pad>") == 0 |
| assert rend.gen_turn == "<turn|>\n<|turn>model\n" |
| |
| want = tok.apply_chat_template( |
| [{"role": "system", "content": "S"}, {"role": "user", "content": "U"}], |
| tokenize=False, add_generation_prompt=True) |
| assert want == "<bos><|turn>system\nS<turn|>\n<|turn>user\nU" + rend.gen_turn |
| |
| |
| 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 |
|
|
|
|
| |
|
|
|
|
| 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]))) |
| |
| 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) |
| |
| 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 |
| 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("<bos>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("<bos>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) |
|
|
|
|
| |
|
|
|
|
| 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) |
| |
| |
| 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) |
| 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) |
| for n in g1: |
| 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] |
| assert full[0].shape[2] == len(doc) |
| for (k, _), kind in zip(s.kv, bb.slot_kinds): |
| 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 |
| 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) |
|
|
|
|
| |
|
|
|
|
| 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" |
| 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("<turn|>\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") |
| assert hasattr(layers[4].self_attn.q_proj, "lora_A") |
| assert hasattr(layers[4].mlp.down_proj, "lora_A") |
| |
| assert len(params) == 2 * (3 * 7 + 4 * 5) |
| want = _raw_outputs(m, _sub(["order_id", "department"])) |
| 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(): |
| for p in model.backbone.lm.state_dict().values(): |
| 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(): |
| 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): |
| attach_base(back, tm.tiny_lm(0)) |
| with pytest.raises(ValueError): |
| 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 |
|
|
|
|
| |
|
|
|
|
| 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) |
| |
| |
| 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 |
|
|
