| """Correctness checks for the Dory PyTorch port (dory/modeling.py). |
| |
| norms weight row/column norms (is the checkpoint unit-normalized, nGPT style?) |
| nll assistant-token NLL on held Beam SFT examples at 1-4 loops |
| generate greedy chat generation at 2 loops |
| ckpt gradient checkpointing gives identical loss and gradients |
| train a few full-parameter AdamW steps on 2 examples: loss must fall, no NaNs |
| """ |
|
|
| import argparse |
| import json |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import torch |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| sys.path.insert(0, str(ROOT)) |
| from dory.data import encode_example |
| from dory.modeling import DoryForCausalLM |
| from transformers import PreTrainedTokenizerFast |
|
|
|
|
| def sync(dev): |
| if dev == "mps": |
| torch.mps.synchronize() |
| elif dev == "cuda": |
| torch.cuda.synchronize() |
|
|
|
|
| def check_norms(model): |
| print("== weight norms (mean ± std); ~1.0 means unit-normalized along that dim") |
| layer = model.backbone.layers |
| named = { |
| "embeddings [V,H]": model.backbone.embeddings.weight, |
| "lm_head [V,H]": model.lm_head.weight, |
| "L0 q_proj [HD,H]": layer[0].mixer.q_proj.weight, |
| "L0 o_proj [H,HD]": layer[0].mixer.o_proj.weight, |
| "L1 gate_proj [F,H]": layer[1].mixer.gate_proj.weight, |
| "L1 down_proj [H,F]": layer[1].mixer.down_proj.weight, |
| "L20 q_proj [HD,H]": layer[20].mixer.q_proj.weight, |
| } |
| for name, w in named.items(): |
| w = w.float() |
| r, c = w.norm(dim=1), w.norm(dim=0) |
| print(f" {name:22s} rows {r.mean():.4f}±{r.std():.4f} cols {c.mean():.4f}±{c.std():.4f}") |
|
|
|
|
| def load_examples(tok, n, max_len): |
| out = [] |
| for line in open(ROOT / "data" / "final" / "sft_verified.jsonl"): |
| ex = json.loads(line) |
| enc = encode_example(tok, ex["messages"], max_len=max_len) |
| if enc is not None: |
| out.append((ex, *enc)) |
| if len(out) >= n: |
| break |
| return out |
|
|
|
|
| @torch.no_grad() |
| def check_nll(model, examples, dev): |
| print("== assistant-token NLL on Beam traces (lower is better; garbage model ~>5)") |
| model.eval() |
| for loops in (1, 2, 3, 4): |
| tot, cnt, t0 = 0.0, 0, time.time() |
| for _, ids, labels in examples: |
| loss, n = model(ids[None].to(dev), labels=labels[None].to(dev), n_loops=loops) |
| tot += float(loss) * n |
| cnt += n |
| sync(dev) |
| print(f" loops={loops}: nll={tot / cnt:.4f} ppl={torch.tensor(tot / cnt).exp():.2f} " |
| f"over {cnt} tokens ({time.time() - t0:.1f}s)") |
|
|
|
|
| @torch.no_grad() |
| def check_generate(model, tok, dev, prompt, max_new): |
| print(f"== greedy generation, 2 loops: {prompt!r}") |
| model.eval() |
| text = tok.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True) |
| ids = tok(text, add_special_tokens=False, return_tensors="pt")["input_ids"].to(dev) |
| eos = {tok.convert_tokens_to_ids("<|im_end|>")} |
| t0 = time.time() |
| out = model.generate(ids, max_new_tokens=max_new, n_loops=2, eos_ids=eos) |
| sync(dev) |
| print(f" ({out.shape[1]} tokens, {time.time() - t0:.1f}s)\n" + tok.decode(out[0])) |
|
|
|
|
| def check_ckpt(model, examples, dev): |
| print("== gradient checkpointing equivalence") |
| _, ids, labels = examples[0] |
| ids, labels = ids[None, :256].to(dev), labels[None, :256].to(dev) |
| labels[:, :8] = -100 |
| labels[:, 8:] = ids[:, 8:] |
| layers = model.backbone.layers |
| probe = [layers[2].mixer.q_proj.weight, layers[30].attn_alpha, layers[31].mixer.suv_up, |
| layers[60].mixer.o_proj.weight, model.lm_head.logit_scale] |
| results = [] |
| for ckpt in (False, True): |
| model.train() |
| model.backbone.gradient_checkpointing = ckpt |
| model.zero_grad(set_to_none=True) |
| loss, _ = model(ids, labels=labels, n_loops=2) |
| loss.backward() |
| results.append((float(loss), [p.grad.float().clone() for p in probe])) |
| model.backbone.gradient_checkpointing = False |
| model.zero_grad(set_to_none=True) |
| (l0, g0), (l1, g1) = results |
| diffs = [float((a - b).abs().max() / (a.abs().max() + 1e-12)) for a, b in zip(g0, g1)] |
| print(f" loss {l0:.6f} vs {l1:.6f}; max relative grad diff {max(diffs):.2e}; " |
| f"grads finite: {all(torch.isfinite(g).all() for g in g0)}") |
|
|
|
|
| @torch.no_grad() |
| def check_pack(model, examples, dev): |
| print("== packing equivalence: 2 docs packed with doc_ids/position_ids vs separately") |
| model.eval() |
| docs = [(ids[:400], labels[:400]) for _, ids, labels in examples[:2]] |
| sep_sum, sep_n = 0.0, 0 |
| for ids, labels in docs: |
| loss, n = model(ids[None].to(dev), labels=labels[None].to(dev), n_loops=2) |
| sep_sum += float(loss) * n |
| sep_n += n |
| ids = torch.cat([d[0] for d in docs]) |
| labels = torch.cat([d[1] for d in docs]).clone() |
| labels[len(docs[0][0])] = -100 |
| doc_ids = torch.cat([torch.full((len(d[0]),), i) for i, d in enumerate(docs)]) |
| pos = torch.cat([torch.arange(len(d[0])) for d in docs]) |
| loss, n = model(ids[None].to(dev), labels=labels[None].to(dev), n_loops=2, |
| doc_ids=doc_ids[None].to(dev), position_ids=pos[None].to(dev)) |
| print(f" separate: {sep_sum / sep_n:.6f} over {sep_n} tokens; packed: {float(loss):.6f} over {n} tokens") |
|
|
|
|
| def check_train(model, examples, dev, steps, lr): |
| print(f"== full-parameter AdamW smoke test: {steps} steps, lr={lr}, 2 loops, checkpointing on") |
| model.train() |
| model.gradient_checkpointing_enable() |
| opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.0, foreach=False) |
| batch = examples[:2] |
| for step in range(steps): |
| t0 = time.time() |
| opt.zero_grad(set_to_none=True) |
| tot = 0.0 |
| for _, ids, labels in batch: |
| loss, _ = model(ids[None].to(dev), labels=labels[None].to(dev), n_loops=2) |
| (loss / len(batch)).backward() |
| tot += float(loss.detach()) / len(batch) |
| gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| opt.step() |
| model.normalize_matrices_() |
| sync(dev) |
| print(f" step {step}: loss={tot:.4f} grad_norm={float(gnorm):.3f} ({time.time() - t0:.1f}s)", flush=True) |
| check_norms(model) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", default=str(ROOT / "models" / "dory-4b")) |
| ap.add_argument("--device", default="mps" if torch.backends.mps.is_available() else "cpu") |
| ap.add_argument("--checks", default="norms,nll,generate,ckpt") |
| ap.add_argument("--n-examples", type=int, default=6) |
| ap.add_argument("--max-len", type=int, default=1536) |
| ap.add_argument("--max-new", type=int, default=200) |
| ap.add_argument("--prompt", default="What is 17 * 23? Reply with just the number.") |
| ap.add_argument("--train-steps", type=int, default=8) |
| ap.add_argument("--lr", type=float, default=1e-5) |
| args = ap.parse_args() |
| checks = set(args.checks.split(",")) |
|
|
| torch.manual_seed(0) |
| t0 = time.time() |
| tok = PreTrainedTokenizerFast.from_pretrained(args.model) |
| model = DoryForCausalLM.from_pretrained(args.model, device=args.device) |
| print(f"loaded on {args.device} in {time.time() - t0:.1f}s; " |
| f"{sum(p.numel() for p in model.parameters()) / 1e9:.3f}B params") |
| examples = load_examples(tok, args.n_examples, args.max_len) |
| print(f"{len(examples)} SFT examples, lengths {[len(e[1]) for e in examples]}") |
|
|
| if "norms" in checks: |
| check_norms(model) |
| if "nll" in checks: |
| check_nll(model, examples, args.device) |
| if "generate" in checks: |
| check_generate(model, tok, args.device, args.prompt, args.max_new) |
| if "pack" in checks: |
| check_pack(model, examples, args.device) |
| if "ckpt" in checks: |
| check_ckpt(model, examples, args.device) |
| if "train" in checks: |
| check_train(model, examples, args.device, args.train_steps, args.lr) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
