File size: 5,045 Bytes
85e93ca | 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 108 109 110 111 | """CPU integration checks against real tiny Transformers Qwen3.5 layers."""
import importlib
import sys
from pathlib import Path
import torch
from torch import nn
from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5TextModel
root = Path(__file__).parent
repo = root / sys.argv[1] if len(sys.argv) > 1 else root.parent
sys.path.insert(0, str(repo))
from dqwen.core import DiffusionLM, patch_bidirectional
from dqwen.kvcache import BlockCache, _segment, _patch, enable_kv_cache, disable_kv_cache, stream_kv
from dqwen import hub
torch.manual_seed(123)
torch.set_num_threads(1)
cfg = Qwen3_5TextConfig(vocab_size=32, hidden_size=16, intermediate_size=32,
num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=1, head_dim=8,
linear_num_key_heads=1, linear_num_value_heads=2, linear_key_head_dim=4,
linear_value_head_dim=4, linear_conv_kernel_dim=4,
layer_types=['linear_attention', 'full_attention'],
rope_parameters={'rope_type': 'default', 'rope_theta': 10000.0,
'mrope_section': [1, 1, 2]})
cfg._attn_implementation = 'sdpa'
lm = nn.Module()
lm.text = Qwen3_5TextModel(cfg).eval()
lm.lm_head = nn.Linear(16, 32, bias=False)
patch_bidirectional(lm)
for m in lm.modules():
if hasattr(m, 'rev_gate'):
m.rev_gate.data.fill_(0.25)
dlm = DiffusionLM(lm, 31).eval()
dlm.eos_id = 30
dlm.bart = None
ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]])
with torch.no_grad():
baseline = dlm.hidden(ids)
_patch(dlm)
# Prefix-empty and suffix-empty cases, plus a genuine two-sided block.
for lo, hi in [(1, 4), (3, 6), (3, 8)]:
cache = BlockCache(lo, hi, ids.shape[1])
pos = torch.arange(ids.shape[1])
record = dlm.text(inputs_embeds=dlm.embed(ids),
attention_mask={'full_attention': None, 'linear_attention': None},
position_ids=pos[None], past_key_values=cache, use_cache=True).last_hidden_state
torch.testing.assert_close(record, baseline, atol=2e-5, rtol=2e-4)
saved = {i: tuple(t.clone() if t is not None else None for t in s)
for i, s in cache.f_state.items()}
cache.mode = 'step'
pos = torch.arange(lo - 1, hi)
for _ in range(2):
active = dlm.text(inputs_embeds=dlm.embed(ids[:, lo-1:hi]),
attention_mask={'full_attention': None, 'linear_attention': None},
position_ids=pos[None], past_key_values=cache, use_cache=True).last_hidden_state
torch.testing.assert_close(active, record[:, lo-1:hi], atol=2e-5, rtol=2e-4)
for i, state in saved.items():
for a, b in zip(state, cache.f_state[i]):
if a is not None:
torch.testing.assert_close(a, b, atol=0, rtol=0)
class Tok:
def apply_chat_template(self, *args, **kwargs): return 'prompt'
def __call__(self, *args, **kwargs): return {'input_ids': [1, 2, 3]}
def decode(self, ids, **kwargs): return str(ids)
tok = Tok()
enable_kv_cache(dlm)
results = list(hub.stream(dlm, tok, 'test', 4, 1.0))
assert results and results[-1][2] <= 4
# 4B/9B generate passes block_len positionally; 27B decodes lists.
assert isinstance(hub.generate(dlm, tok, 'test', gen_len=4), str)
if not hub._kv_legacy_list:
assert list(hub.stream(dlm, tok, 'test', 4, 1.0, 2))
# Fully refreshed runs should match the uncached block generator (MASK is banned).
refreshed = list(stream_kv(dlm, tok, 'test', 4, 1.0, block_len=2, refresh=1))
assert refreshed
# Check refresh=1 against an independently denoised full-sequence reference.
reference = torch.tensor([[1, 2, 3, 31, 31, 31, 31]])
expected = []
with torch.no_grad():
for lo in (3, 5):
hi = lo + 2
while (reference[0, lo:hi] == dlm.mask_id).any():
h = dlm.hidden(reference)
pos = torch.arange(lo, hi)[reference[0, lo:hi] == dlm.mask_id]
logits = dlm.head(h[0, pos-1]).float()
logits[:, dlm.mask_id] = -torch.inf
conf, cand = logits.softmax(-1).max(-1)
select = conf >= 1.0
if not select.any(): select[conf.argmax()] = True
reference[0, pos[select]] = cand[select]
expected.append(reference[:, 3:].clone())
if (reference[0, 3:hi] == dlm.eos_id).any():
reference[0, hi:] = dlm.eos_id
expected.append(reference[:, 3:].clone())
break
assert len(expected) == len(refreshed)
for want, (got, _, _) in zip(expected, refreshed):
torch.testing.assert_close(want, got)
assert list(stream_kv(dlm, tok, 'test', 4, 1.0, max_steps=1))[-1][2] == 1
disable_kv_cache(dlm)
assert not any(hasattr(m, '_dq_kv_old') for m in lm.modules())
torch.testing.assert_close(dlm.hidden(ids), baseline, atol=2e-5, rtol=2e-4)
enable_kv_cache(dlm, refresh=1)
assert list(hub.stream(dlm, tok, 'test', 2, 1.0))
disable_kv_cache()
print('PASS:', repo.name, 'record/active equivalence, empty segments, immutable state, stream/generate, restore')
|