| """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 |
| from dory.modeling import DoryConfig, DoryForCausalLM |
|
|
|
|
| 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_() |
| |
| |
| 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) |
| 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_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)) |
| 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 |
| 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() |
|
|
