File size: 4,335 Bytes
aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """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 # noqa: E402
from deg.model import load_model, load_tokenizer # noqa: E402
@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 # no [B, L, L] mask
(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
|