| """Training trunks longer than DENSE_PREFILL_MAX (SPEC.md §12.4): the chunked read is |
| gradient-checkpointed per layer and per chunk, and equals the dense (one-shot, |
| checkpointed) path in hidden states, continued state, fork outputs and gradients, on a |
| tiny random Qwen3.5 (CPU).""" |
|
|
| from __future__ import annotations |
|
|
| import pytest |
|
|
| torch = pytest.importorskip("torch") |
| pytest.importorskip("transformers") |
|
|
| from test_model import CONTEXT, SNAPSHOT, _LargestTensor, tiny_lm |
|
|
| from deg.model import load_model, load_tokenizer |
|
|
|
|
| @pytest.fixture(scope="module") |
| def model(): |
| if SNAPSHOT is None: |
| pytest.skip("Qwen3.5-0.8B tokenizer not cached") |
| m = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_lm(), |
| tokenizer=load_tokenizer(SNAPSHOT)) |
| m.backbone.grad_ckpt = True |
| m.backbone.text.embed_tokens.weight.requires_grad_(False) |
| return m |
|
|
|
|
| def _run(model, rows: list[list[int]], forks: list[list[int]], fork_rows: list[int]): |
| bb = model.backbone |
| for p in bb.parameters(): |
| p.grad = None |
| width = max(len(r) for r in rows) |
| ids = torch.full((len(rows), width), model.renderer.pad_id, dtype=torch.long) |
| for i, r in enumerate(rows): |
| ids[i, :len(r)] = torch.tensor(r) |
| lengths = torch.tensor([len(r) for r in rows]) |
| bb.train() |
| try: |
| h, state = bb(ids, lengths) |
| fh = model.run_forks(state, torch.tensor(fork_rows), [tuple(f) for f in forks]) |
| loss = sum((h[i, :len(r)].float() ** 2).mean() for i, r in enumerate(rows)) |
| loss = loss + sum((x.float() ** 2).mean() for x in fh) |
| loss.backward() |
| finally: |
| bb.eval() |
| grads = {n: p.grad.clone() for n, p in bb.named_parameters() if p.grad is not None} |
| return h.detach(), state, [x.detach() for x in fh], grads |
|
|
|
|
| def test_chunked_checkpointed_trunk_equals_dense(model, monkeypatch): |
| import deg.model as dm |
|
|
| 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:") |
| assert len(a) > 100 |
| dense = _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 = dm.Backbone._forward_chunked_ckpt |
|
|
| def spy(self, *args, **kw): |
| calls.append(1) |
| return orig(self, *args, **kw) |
|
|
| monkeypatch.setattr(dm.Backbone, "_forward_chunked_ckpt", spy) |
| with _LargestTensor() as big: |
| chunked = _run(model, [a, b], [q, q[:5], q], [0, 1, 1]) |
| assert calls, "long training trunk did not take the checkpointed chunked path" |
| L = len(a) |
| assert big.largest < 2 * L * L // 3, 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=2e-4, rtol=2e-4) |
| for (r, t), (r1, t1) in zip(s2.lin, s1.lin): |
| torch.testing.assert_close(r, r1, atol=2e-4, rtol=2e-4) |
| torch.testing.assert_close(t, t1, atol=2e-4, rtol=2e-4) |
| for (k, v), (k1, v1) in zip(s2.kv, s1.kv): |
| assert k.shape == k1.shape |
| m = s1.kv_mask[:, None, :, None] |
| torch.testing.assert_close(k * m, k1 * m, atol=2e-4, rtol=2e-4) |
| torch.testing.assert_close(v * m, v1 * m, atol=2e-4, rtol=2e-4) |
| 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=2e-4, rtol=2e-4) |
| assert g1.keys() == g2.keys() and g1 |
| for n in g1: |
| torch.testing.assert_close(g2[n], g1[n], atol=1e-4, rtol=1e-3, msg=n) |
|
|
|
|
| def test_inference_long_read_unchanged_by_training_path(model, monkeypatch): |
| """Without autograd the chunked inference read (``_forward_chunked``) still runs.""" |
| import deg.model as dm |
|
|
| monkeypatch.setattr(dm, "DENSE_PREFILL_MAX", 48) |
| monkeypatch.setattr(dm, "PREFILL_CHUNK", 32) |
| called = [] |
| monkeypatch.setattr(dm.Backbone, "_forward_chunked_ckpt", |
| lambda *a, **k: called.append(1)) |
| ids = model.renderer.context(" ".join([CONTEXT] * 2)).ids |
| with torch.no_grad(): |
| model.backbone(torch.tensor([ids]), torch.tensor([len(ids)])) |
| assert not called |
|
|
