File size: 5,773 Bytes
dafe8d2 | 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 112 113 114 115 116 117 118 119 120 | """Acceptance test for dory/fast_generate.py (KV-cached decoding).
For each loop count, decode N tokens with the cached path (recording the logits used
to pick every token), then run the port's cache-free forward over prompt + generated
tokens and compare the next-token logits at every decoded position.
.venv/bin/python -I scripts/test_fast_generate.py --tiny # random small model, window 8
.venv/bin/python -I scripts/test_fast_generate.py --model /path/to/Dory-Beam-4B --tokens 300
Both run on CPU in float32 by default.
"""
import argparse
import sys
import time
from pathlib import Path
import torch
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
from dory.fast_generate import generate_stream # noqa: E402
from dory.modeling import DoryConfig, DoryForCausalLM # noqa: E402
def tiny_model(seed=0):
"""Small random nGPT-shaped Dory: 2 input, 4 recurrent, 2 output layers, window 8."""
torch.manual_seed(seed)
cfg = DoryConfig(vocab_size=257, hidden_size=64, num_hidden_layers=8, intermediate_size=128,
num_attention_heads=4, num_key_value_heads=2, head_dim=16, sliding_window=8,
rope_parameters={"sliding_attention": {"rope_theta": 10000},
"full_attention": {"rope_theta": 10000000}},
layer_types=["sliding_attention", "mlp", "full_attention", "mlp",
"sliding_attention", "mlp", "sliding_attention", "mlp"],
n_input_layers=2, n_recurrent_layers=4, n_output_layers=2)
model = DoryForCausalLM(cfg)
for p in model.parameters():
torch.nn.init.normal_(p, std=0.5)
model.normalize_matrices_()
# nGPT scalar gains at realistic effective scales (a random init makes softmaxes one-hot
# and the comparison chaotic): q scale sqrt(d), unit gate/suv, slerp lr 0.05, logit scale 10.
with torch.no_grad():
for layer in model.backbone.layers:
mx = layer.mixer
if layer.is_attention:
mx.sqk.fill_(cfg.head_dim**0.25 / mx.sqk_factor)
mx.s_gate.fill_(1.0 / mx.s_gate_factor)
layer.attn_alpha.fill_(0.05 / layer.alpha_scale)
else:
mx.suv_gate.fill_(1.0 / mx.suv_factor)
mx.suv_up.fill_(1.0 / mx.suv_factor)
layer.mlp_alpha.fill_(0.05 / layer.alpha_scale)
model.lm_head.logit_scale.fill_(10 * cfg.initializer_range)
return model.float().eval()
def run(model, prompt_ids, loops, n_tokens, tol, seed):
t0 = time.time()
stream = generate_stream(model, prompt_ids, n_tokens, n_loops=loops, temperature=1.0, top_p=0.95,
eos_ids=(), seed=seed, return_logits=True)
toks, cached = [], []
for tok, lg in stream:
toks.append(tok)
cached.append(lg[0])
t_cached = time.time() - t0
cached = torch.stack(cached) # [N, V]; row i = logits that produced generated token i
full_ids = torch.cat([prompt_ids, torch.tensor([toks], dtype=torch.long)], dim=1)
t0 = time.time()
ref = model(full_ids, n_loops=loops)[0] # [T+N, V]
t_full = time.time() - t0
p = prompt_ids.shape[1]
ref = ref[p - 1 : p - 1 + len(toks)]
diff = (cached - ref).abs()
scale = ref.abs().max().item()
print(f"loops={loops}: {len(toks)} tokens, seq {full_ids.shape[1]}, max|Δlogit|={diff.max():.3e} "
f"(mean {diff.mean():.2e}, logit scale {scale:.1f}), argmax agree "
f"{(cached.argmax(-1) == ref.argmax(-1)).float().mean():.4f}; "
f"cached {t_cached:.1f}s ({len(toks) / t_cached:.2f} tok/s), full forward {t_full:.1f}s")
assert diff.max() < tol, f"max abs diff {diff.max():.3e} exceeds tolerance {tol}"
return diff.max().item()
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", help="directory with config.json + safetensors")
ap.add_argument("--tiny", action="store_true", help="random small model with sliding window 8")
ap.add_argument("--tokens", type=int, default=300)
ap.add_argument("--loops", default="1,2,4")
ap.add_argument("--tol", type=float, default=None)
ap.add_argument("--prompt", default="What is 17*23? Think step by step.")
ap.add_argument("--threads", type=int, default=0)
ap.add_argument("--float64", action="store_true", help="tiny model in float64 (isolates logic bugs)")
args = ap.parse_args()
if args.threads:
torch.set_num_threads(args.threads)
torch.manual_seed(0)
if args.tiny:
model = tiny_model()
if args.float64:
model = model.double()
prompt_ids = torch.randint(0, 257, (1, 20)) # longer than the window: exercises prefill eviction
tol = args.tol or (1e-10 if args.float64 else 1e-4)
else:
from transformers import PreTrainedTokenizerFast
model = DoryForCausalLM.from_pretrained(args.model, dtype=torch.float32).eval()
tok = PreTrainedTokenizerFast.from_pretrained(args.model)
text = tok.apply_chat_template([{"role": "user", "content": args.prompt}], tokenize=False,
add_generation_prompt=True)
prompt_ids = tok(text, add_special_tokens=False, return_tensors="pt")["input_ids"]
tol = args.tol or 5e-3 # measured 1.1e-3 on the released weights (logit scale ~60)
print(f"prompt tokens: {prompt_ids.shape[1]}, decode tokens: {args.tokens}, tolerance {tol}")
worst = 0.0
for loops in (int(x) for x in args.loops.split(",")):
worst = max(worst, run(model, prompt_ids, loops, args.tokens, tol, seed=loops))
print(f"PASS: worst max abs logit diff {worst:.3e} < {tol}")
if __name__ == "__main__":
main()
|