diff --git a/scripts/smoke_batch_parity.py b/scripts/smoke_batch_parity.py new file mode 100644 index 0000000..e07e11c --- /dev/null +++ b/scripts/smoke_batch_parity.py @@ -0,0 +1,122 @@ +"""Parity smoke: guided_rollout vs guided_rollout_batch on a small vignette subset. + +Asserts p_true and pmass_format match within fp tolerance. Same chat template, +same prompts, same model, same generation kwargs -- only batching differs. + +usage: + uv run python scripts/smoke_batch_parity.py --model Qwen/Qwen3-0.6B --limit 4 +""" +from __future__ import annotations +import argparse +import time + +import torch +from loguru import logger +from transformers import AutoModelForCausalLM, AutoTokenizer + +from tinymfv.core import CONDITIONS, FRAMES +from tinymfv.data import load_vignettes +from tinymfv.guided import guided_rollout, guided_rollout_batch, choice_token_ids_tf + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="Qwen/Qwen3-0.6B") + ap.add_argument("--limit", type=int, default=4) + ap.add_argument("--max-think-tokens", type=int, default=32) + ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + ap.add_argument("--dtype", default="bfloat16") + args = ap.parse_args() + + rows = load_vignettes("")[: args.limit] + logger.info(f"{len(rows)} vignettes; testing parity") + + dtype = getattr(torch, args.dtype) + tok = AutoTokenizer.from_pretrained(args.model) + if tok.pad_token is None: + tok.pad_token = tok.eos_token + tok.padding_side = "left" + model = AutoModelForCausalLM.from_pretrained(args.model, dtype=dtype).to(args.device).eval() + + choice_ids = choice_token_ids_tf(tok) + + # --- Sequential --- + t0 = time.time() + seq_results = [] # list of (vid, cond, frame, p_true, pmass) + for r in rows: + for cond in CONDITIONS: + for frame, fr in FRAMES.items(): + res = guided_rollout( + model, tok, + user_prompt=r[cond], + choice_token_ids=choice_ids, + max_think_tokens=args.max_think_tokens, + schema_hint=fr["q"], + prefill=fr["prefill"], + ) + seq_results.append((r["id"], cond, frame, res.p_true, res.pmass_format)) + seq_elapsed = time.time() - t0 + logger.info(f"sequential: {seq_elapsed:.1f}s ({len(seq_results)} prompts)") + + # --- Batched --- + t0 = time.time() + batch_results = [] + for frame, fr in FRAMES.items(): + for cond in CONDITIONS: + user_prompts = [r[cond] for r in rows] + outs = guided_rollout_batch( + model, tok, + user_prompts=user_prompts, + choice_token_ids=choice_ids, + max_think_tokens=args.max_think_tokens, + schema_hint=fr["q"], + prefill=fr["prefill"], + ) + for r, o in zip(rows, outs): + batch_results.append((r["id"], cond, frame, o.p_true, o.pmass_format)) + batch_elapsed = time.time() - t0 + logger.info(f"batched: {batch_elapsed:.1f}s (speedup={seq_elapsed/batch_elapsed:.1f}x)") + + # --- Compare --- + seq_d = {(vid, c, f): (pt, pm) for vid, c, f, pt, pm in seq_results} + batch_d = {(vid, c, f): (pt, pm) for vid, c, f, pt, pm in batch_results} + assert set(seq_d) == set(batch_d), "key mismatch" + + n = 0 + max_pt_diff, max_pm_diff = 0.0, 0.0 + rows_out = [] + for k in seq_d: + spt, spm = seq_d[k] + bpt, bpm = batch_d[k] + d_pt = abs(spt - bpt) + d_pm = abs(spm - bpm) + max_pt_diff = max(max_pt_diff, d_pt) + max_pm_diff = max(max_pm_diff, d_pm) + rows_out.append((k, spt, bpt, d_pt, spm, bpm, d_pm)) + n += 1 + + from tabulate import tabulate + print() + print(tabulate( + [(f"{k[0][:8]}|{k[1]}|{k[2]}", spt, bpt, d_pt, spm, bpm, d_pm) + for (k, spt, bpt, d_pt, spm, bpm, d_pm) in rows_out], + headers=["key", "p_true_seq", "p_true_bat", "Δp_true", "pm_seq", "pm_bat", "Δpm"], + floatfmt="+.4f", tablefmt="tsv", + )) + + # bf16 batched greedy decoding can pick different argmax than per-row greedy + # when two tokens tie within bf16 precision. The phase1 think rollout then + # diverges and per-row p_true drifts. float32 is bit-exact (use --dtype float32 + # to verify the batching logic itself). At aggregate eval (131 vignettes + # averaged) the bf16 drift averages out; we accept it. + TOL = 0.20 if args.dtype != "float32" else 0.001 + cue = "🟢" if (max_pt_diff < TOL and max_pm_diff < TOL) else "🔴" + print(f"\n{cue} max Δp_true={max_pt_diff:.4f} max Δpmass={max_pm_diff:.4f} (tol={TOL})") + print(f"speedup: {seq_elapsed/batch_elapsed:.1f}x ({len(seq_results)} prompts)") + + if max_pt_diff >= TOL or max_pm_diff >= TOL: + raise SystemExit(f"PARITY FAILED: Δp_true={max_pt_diff:.4f} Δpmass={max_pm_diff:.4f}") + + +if __name__ == "__main__": + main() diff --git a/src/tinymfv/core.py b/src/tinymfv/core.py index d5f77c5..bc3c72c 100644 --- a/src/tinymfv/core.py +++ b/src/tinymfv/core.py @@ -214,10 +214,15 @@ def analyse( if bool_mass is not None: info["bool_mass_mean"] = float(sum(map(float, bool_mass)) / len(bool_mass)) + raw_pmass = ( + {f"{vid}|{cond}|{frame}": float(b) for (vid, _, cond, frame, _), b in zip(meta, bool_mass)} + if bool_mass is not None else {} + ) return { "wrongness": float(df["s_other_violate"].mean()), "gap": float(df["gap"].mean()), "table": df, "raw": {f"{vid}|{cond}|{frame}": p for (vid, _, cond, frame, _), p in zip(meta, p_true)}, + "raw_pmass": raw_pmass, "info": info, } diff --git a/src/tinymfv/eval.py b/src/tinymfv/eval.py index 0903f16..318215f 100644 --- a/src/tinymfv/eval.py +++ b/src/tinymfv/eval.py @@ -9,7 +9,7 @@ from tqdm.auto import tqdm from .core import format_prompts, next_token_logits, score_prompts, analyse, CONDITIONS, FRAMES from .data import load_vignettes -from .guided import guided_rollout, choice_token_ids_tf +from .guided import guided_rollout_batch, choice_token_ids_tf def evaluate( @@ -39,36 +39,63 @@ def evaluate( t0 = time.time() if max_think_tokens > 0: - logger.info(f"Using guided_rollout with {max_think_tokens} max_think_tokens (sequential)") - p_true_list = [] - meta = [] - bool_mass_list = [] + logger.info(f"Using guided_rollout_batch with {max_think_tokens} max_think_tokens, batch_size={batch_size}") choice_ids = choice_token_ids_tf(tokenizer) - - for r in tqdm(vignettes, desc="Evaluating"): + + # Build all (vid, cond, frame) items, grouped by frame so each batch + # shares schema_hint + prefill (collapses the per-row branching). + items_per_frame: dict[str, list[tuple]] = {f: [] for f in FRAMES} + for r in vignettes: for cond in CONDITIONS: - for frame, fr in FRAMES.items(): - user_prompt = f"{r[cond]}" - schema_hint = fr["q"] - prefill = fr["prefill"] - - res = guided_rollout( + for frame in FRAMES: + items_per_frame[frame].append( + (r["id"], r["foundation_coarse"], cond, frame, r.get("wrong"), r[cond]) + ) + + # Pretokenize a sample to log expected prompt length / cache budget. + sample_user = items_per_frame[next(iter(FRAMES))][0][5] + sample_q = FRAMES[next(iter(FRAMES))]["q"] + sample_full = f"{sample_user}\n\n{sample_q}" + sample_msgs = [{"role": "user", "content": sample_full}] + try: + sample_p = tokenizer.apply_chat_template(sample_msgs, tokenize=False, add_generation_prompt=True) + except TypeError: + sample_p = tokenizer.apply_chat_template(sample_msgs, tokenize=False) + sample_p = sample_p + "\n" + sample_len = len(tokenizer(sample_p).input_ids) + logger.info( + f"SHOULD: prompt_len≈{sample_len} tok; max cache ≈ {sample_len + max_think_tokens} per row × " + f"batch_size={batch_size}. If OOM, lower batch_size." + ) + + p_true_list, meta, bool_mass_list = [], [], [] + total = sum(len(v) for v in items_per_frame.values()) + with tqdm(total=total, desc="Evaluating") as pbar: + for frame, items in items_per_frame.items(): + fr = FRAMES[frame] + schema_hint = fr["q"] + prefill = fr["prefill"] + for i in range(0, len(items), batch_size): + chunk = items[i:i + batch_size] + user_prompts = [it[5] for it in chunk] + results = guided_rollout_batch( model, tokenizer, - user_prompt=user_prompt, + user_prompts=user_prompts, choice_token_ids=choice_ids, max_think_tokens=max_think_tokens, schema_hint=schema_hint, prefill=prefill, - verbose=False ) - - p_true_list.append(res.p_true) - meta.append((r["id"], r["foundation_coarse"], cond, frame, r.get("wrong"))) - bool_mass_list.append(res.pmass_format) - + for it, res in zip(chunk, results): + vid, found, cond, fr_name, wrong, _ = it + p_true_list.append(res.p_true) + meta.append((vid, found, cond, fr_name, wrong)) + bool_mass_list.append(res.pmass_format) + pbar.update(len(chunk)) + elapsed = time.time() - t0 logger.info(f"guided eval: {elapsed:.1f}s ({len(p_true_list)/elapsed:.1f} prompts/s)") - + report = analyse(p_true_list, meta, bool_mass=bool_mass_list) else: diff --git a/src/tinymfv/guided.py b/src/tinymfv/guided.py index d2feefa..4af33cd 100644 --- a/src/tinymfv/guided.py +++ b/src/tinymfv/guided.py @@ -108,7 +108,17 @@ def guided_rollout( all_ids = torch.tensor(a_ids + b_ids, device=device, dtype=torch.long) pmass_format = float(logp[all_ids].exp().sum().item()) - + + # SHOULD: pmass≈1 (model picks one of the JSON-bool tokens). pmass<0.9 + # means the model is leaking probability to other tokens -> the schema + # is being ignored or the steering vector has pushed the model OOD. + if pmass_format < 0.9: + topk = torch.topk(logp.exp(), k=5) + toks = [tok.decode([i]) for i in topk.indices.tolist()] + probs = topk.values.tolist() + top5 = ", ".join(f"{repr(t)}={p:.3f}" for t, p in zip(toks, probs)) + logger.warning(f"pmass={pmass_format:.3f}<0.9 — top-5: {top5}") + if a_ids and b_ids: a_t = torch.tensor(a_ids, device=device, dtype=torch.long) b_t = torch.tensor(b_ids, device=device, dtype=torch.long) @@ -143,6 +153,128 @@ def guided_rollout( p_true=p_true, ) +@torch.no_grad() +def guided_rollout_batch( + model, tok, + user_prompts: list[str], + choice_token_ids: list, + max_think_tokens: int = 128, + schema_hint: str = _DEFAULT_SCHEMA_HINT, + prefill: str = '\n{"choice": ', +) -> list[GuidedResult]: + """Batched guided rollout. Same logic as guided_rollout but over a list of + user_prompts that share schema_hint + prefill (so prefill cases collapse). + + Skips the cosmetic answer-continuation generate (caller only needs p_true, + pmass_format, think_text). Two model calls per batch instead of 3 per row: + one phase1 generate (think) + one scoring forward. + + Tokenizer must already have padding_side='left' and pad_token set.""" + if tok.padding_side != "left": + raise ValueError("tok.padding_side must be 'left' for batched rollout") + device = next(model.parameters()).device + + prompts = [] + for up in user_prompts: + full_user = f"{up}\n\n{schema_hint}" if schema_hint else up + msgs = [{"role": "user", "content": full_user}] + try: + p = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) + except TypeError: + p = tok.apply_chat_template(msgs, tokenize=False) + prompts.append(p + "\n") + + think_end_id = tok.convert_tokens_to_ids("") + if think_end_id in (None, getattr(tok, "unk_token_id", None)): + think_end_id = tok.eos_token_id + pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id + + enc = tok(prompts, return_tensors="pt", padding=True).to(device) + prompt_len = enc.input_ids.shape[1] + + phase1 = model.generate( + **enc, + max_new_tokens=max_think_tokens, + do_sample=False, + eos_token_id=think_end_id, + pad_token_id=pad_id, + ) + + scoring_texts = [] + per_row = [] # (think_text, emitted_close, emitted_prefill, n_think_tokens) + for i, p in enumerate(prompts): + gen_ids = phase1[i, prompt_len:] + keep = gen_ids != pad_id + gen_ids = gen_ids[keep] if keep.any() else gen_ids[:0] + gen_text = tok.decode(gen_ids, skip_special_tokens=True) + n_think = int(gen_ids.shape[0]) + + emitted_close = _CLOSE_MARKER in gen_text + if emitted_close: + think_text, after = gen_text.split(_CLOSE_MARKER, 1) + if prefill.lstrip() in after: + emitted_prefill = True + before_value = after.split(prefill.lstrip(), 1)[0] + scoring_text = p + think_text + _CLOSE_MARKER + before_value + prefill.lstrip() + else: + emitted_prefill = False + scoring_text = p + think_text + _CLOSE_MARKER + prefill + else: + think_text = gen_text + emitted_prefill = False + force_suffix = "\nI should answer now." + _CLOSE_MARKER + prefill + scoring_text = p + gen_text + force_suffix + + scoring_texts.append(scoring_text) + per_row.append((think_text, emitted_close, emitted_prefill, n_think)) + + score_enc = tok(scoring_texts, return_tensors="pt", padding=True, + add_special_tokens=False).to(device) + score_logits = model(**score_enc).logits[:, -1].float() + score_logp = F.log_softmax(score_logits, dim=-1) + + if (len(choice_token_ids) == 2 and all(isinstance(x, (list, tuple)) for x in choice_token_ids)): + a_ids, b_ids = list(choice_token_ids[0]), list(choice_token_ids[1]) + else: + a_ids, b_ids = list(choice_token_ids), [] + all_ids = torch.tensor(a_ids + b_ids, device=device, dtype=torch.long) + a_t = torch.tensor(a_ids, device=device, dtype=torch.long) if a_ids else None + b_t = torch.tensor(b_ids, device=device, dtype=torch.long) if b_ids else None + + results = [] + for i, (up, (think_text, emitted_close, emitted_prefill, n_think)) in enumerate(zip(user_prompts, per_row)): + logp = score_logp[i] + pmass_format = float(logp[all_ids].exp().sum().item()) + if pmass_format < 0.9: + topk = torch.topk(logp.exp(), k=5) + toks = [tok.decode([j]) for j in topk.indices.tolist()] + probs = topk.values.tolist() + top5 = ", ".join(f"{repr(t)}={pp:.3f}" for t, pp in zip(toks, probs)) + logger.warning(f"pmass={pmass_format:.3f}<0.9 — top-5: {top5}") + if a_t is not None and b_t is not None: + la = torch.logsumexp(logp[a_t], dim=0) + lb = torch.logsumexp(logp[b_t], dim=0) + logratio = float((la - lb).item()) + p_true = float(torch.softmax(torch.stack([la, lb]), dim=0)[0].item()) + else: + logratio = float("nan") + p_true = float("nan") + results.append(GuidedResult( + user_prompt=up, + think_text=think_text, + answer_text="", + raw_full_text="", + pmass_format=pmass_format, + logratio_ab=logratio, + rep_ratio_think=_ngram_rep_ratio(think_text, n=4), + think_tokens=n_think, + emitted_close=emitted_close, + emitted_prefill=emitted_prefill, + p_true=p_true, + )) + return results + + def choice_token_ids_tf(tok) -> list[list[int]]: def _variants(words): seen = []