diff --git a/src/projected_grpo/diag_one_layer.py b/src/projected_grpo/diag_one_layer.py deleted file mode 100644 index e84d90b..0000000 --- a/src/projected_grpo/diag_one_layer.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Diagnostic: single-Linear SVD round-trip and single-module wrap-in-model. - -Q1: For a stand-alone nn.Linear L, does AntiPaSTOLinear(SVD(L.weight), L.bias)(x) == L(x)? - Tests pure math. -Q2: If we wrap exactly ONE Linear inside the model, does logits diff vanish? - Tests integration (state-dict, device, dtype, hook order). -""" -from __future__ import annotations - -import copy -from pathlib import Path - -import torch -from loguru import logger -from transformers import AutoModelForCausalLM, AutoTokenizer - -from .antipasto import AntiPaSTOLinear, svd_cached, wrap_model_with_antipasto - -MODEL = "Qwen/Qwen3.5-0.8B" - - -def q1_pure_math(): - torch.manual_seed(0) - for (d_out, d_in) in [(64, 64), (128, 64), (64, 128), (1024, 3584)]: - L = torch.nn.Linear(d_in, d_out, bias=True).to(torch.float32) - W = L.weight.data - U, S, Vh = torch.linalg.svd(W, full_matrices=False) - wrap = AntiPaSTOLinear(U, S, Vh, L.bias.data) - x = torch.randn(4, d_in, dtype=torch.float32) - y_lin = L(x) - y_wrap = wrap(x) - d = (y_lin - y_wrap).abs().max().item() - s = y_lin.abs().mean().item() - logger.info(f"Linear({d_in}->{d_out}) max_diff={d:.2e} scale={s:.3f}") - - -def q2_wrap_one_in_model(): - device = torch.device("cuda") - tokenizer = AutoTokenizer.from_pretrained(MODEL) - base = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.float32, attn_implementation="sdpa").to(device) - base.eval() - - # Find target names - target_names = [] - for name, m in base.named_modules(): - if isinstance(m, torch.nn.Linear): - suff = name.split(".")[-1] - if suff in ("q_proj", "gate_proj", "in_proj_qkv", "in_proj_a", "out_proj"): - target_names.append((suff, name)) - - # Pick one of each kind - seen = set() - picked = [] - for suff, name in target_names: - if suff not in seen: - picked.append(name) - seen.add(suff) - - prompt = "Write a function." - ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device) - with torch.no_grad(): - y_base = base(ids).logits.clone() - - for name in picked: - model = copy.deepcopy(base) - linear = model.get_submodule(name) - W = linear.weight.data - U, S, Vh = torch.linalg.svd(W.to(torch.float32), full_matrices=False) - bias = linear.bias.data if linear.bias is not None else None - wrap = AntiPaSTOLinear(U, S, Vh, bias).to(W.device) - parent_name, child_name = name.rsplit(".", 1) - setattr(model.get_submodule(parent_name), child_name, wrap) - model.eval() - with torch.no_grad(): - y_wrap = model(ids).logits - d = (y_base - y_wrap).abs().max().item() - logger.info(f"wrap-only [{name.split('.')[-1]:>12}] {name} max_diff={d:.2e}") - - -if __name__ == "__main__": - logger.info("=== Q1: pure math (stand-alone nn.Linear) ===") - q1_pure_math() - logger.info("=== Q2: wrap one Linear inside Qwen3.5-0.8B ===") - q2_wrap_one_in_model() diff --git a/src/projected_grpo/diag_trace.py b/src/projected_grpo/diag_trace.py deleted file mode 100644 index 3dc4f54..0000000 --- a/src/projected_grpo/diag_trace.py +++ /dev/null @@ -1,74 +0,0 @@ -"""Diagnose: when we wrap a single Linear, is the wrapper actually invoked, -and does the SVD reconstruct the layer's weight exactly? -""" -from __future__ import annotations - -import copy - -import torch -from loguru import logger -from transformers import AutoModelForCausalLM, AutoTokenizer - -from .antipasto import AntiPaSTOLinear - -MODEL = "Qwen/Qwen3.5-0.8B" - - -def main(): - device = torch.device("cuda") - tokenizer = AutoTokenizer.from_pretrained(MODEL) - base = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.float32, attn_implementation="sdpa").to(device) - base.eval() - - name = "model.layers.0.linear_attn.out_proj" - linear = base.get_submodule(name) - W = linear.weight.data - logger.info(f"target {name} W.shape={tuple(W.shape)} W.dtype={W.dtype} bias={linear.bias is not None}") - - # SVD reconstruction error (pure) - U, S, Vh = torch.linalg.svd(W.to(torch.float32), full_matrices=False) - W_recon = U @ torch.diag(S) @ Vh - recon_err = (W_recon - W.to(torch.float32)).abs().max().item() - logger.info(f"SVD reconstruct(W) max_err = {recon_err:.2e} (should be ~1e-5)") - - # Now wrap and force the wrap to track calls - model = copy.deepcopy(base) - linear2 = model.get_submodule(name) - bias = linear2.bias.data if linear2.bias is not None else None - wrap = AntiPaSTOLinear(U, S, Vh, bias).to(W.device) - - call_count = [0] - captured = [] - orig_forward = wrap.forward - def counting_forward(x): - call_count[0] += 1 - # also compare to what a fresh nn.Linear would compute - y_wrap = orig_forward(x) - y_ref = torch.nn.functional.linear(x.to(torch.float32), W.to(torch.float32), - bias.to(torch.float32) if bias is not None else None) - d = (y_wrap.to(torch.float32) - y_ref).abs().max().item() - captured.append(d) - return y_wrap - wrap.forward = counting_forward - - parent_name, child_name = name.rsplit(".", 1) - setattr(model.get_submodule(parent_name), child_name, wrap) - model.eval() - - # confirm the substitution stuck - new_mod = model.get_submodule(name) - logger.info(f"after wrap: get_submodule -> {type(new_mod).__name__} id_match={id(new_mod)==id(wrap)}") - - ids = tokenizer("Write a function.", return_tensors="pt").input_ids.to(device) - with torch.no_grad(): - y_base = base(ids).logits - y_wrap = model(ids).logits - diff = (y_base - y_wrap).abs().max().item() - - logger.info(f"wrap.forward calls = {call_count[0]}") - logger.info(f"per-call wrap-vs-F.linear max_diff = {[f'{x:.2e}' for x in captured]}") - logger.info(f"final logits max_diff = {diff:.2e}") - - -if __name__ == "__main__": - main() diff --git a/src/projected_grpo/probe_lora_runtime.py b/src/projected_grpo/probe_lora_runtime.py deleted file mode 100644 index 15378ca..0000000 --- a/src/projected_grpo/probe_lora_runtime.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Sanity probe: does Qwen3-4B + ariahw LoRA (runtime, no merge) hack our grader? - -Resolves the bake-vs-pipeline ambiguity. If hack-rate >0 here but ~0 on the -bake_lora.py merged checkpoint, the bake is broken (base mismatch / PEFT merge -math). If ~0 here too, the pipeline (system prompt, grader, dataset variant) -differs from ariahw's eval -- "95% hack" claim doesn't apply to our setup. - -Run: uv run python -m projected_grpo.probe_lora_runtime -""" -from __future__ import annotations - -import sys -from dataclasses import dataclass - -import torch -import tyro -from loguru import logger -from peft import PeftModel -from tabulate import tabulate -from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig - -from .rewards import compute_reward -from .train import load_problems - - -@dataclass -class Config: - base_model: str = "Qwen/Qwen3-4B" - lora_repo: str = "ariahw/rl-rewardhacking-leetcode-rh-s65" - n_prompts: int = 8 - group: int = 6 - max_new: int = 1024 - temperature: float = 0.7 - seed: int = 41 - - -def main(cfg: Config) -> int: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - logger.info(f"base={cfg.base_model} lora={cfg.lora_repo} T={cfg.temperature}") - - tok = AutoTokenizer.from_pretrained(cfg.base_model) - if tok.pad_token_id is None: tok.pad_token = tok.eos_token - - base = AutoModelForCausalLM.from_pretrained( - cfg.base_model, dtype=torch.bfloat16, attn_implementation="sdpa" - ).to(device) - model = PeftModel.from_pretrained(base, cfg.lora_repo).to(device) - model.eval() - logger.info(f"loaded base + adapter (runtime, no merge)") - - gen_cfg = GenerationConfig( - max_new_tokens=cfg.max_new, do_sample=True, - temperature=cfg.temperature, top_p=1.0, top_k=20, min_p=0.0, - repetition_penalty=1.0, - num_return_sequences=cfg.group, pad_token_id=tok.pad_token_id, - ) - - problems = load_problems(n=200) - rng = torch.Generator().manual_seed(cfg.seed) - - logger.info( - f"SHOULD: hack >0 on Qwen3-4B + ariahw LoRA runtime. " - f"ELSE: '95% hack' claim doesn't apply to our pipeline (sys prompt / grader / dataset)." - ) - - n_hack = n_gt = n_fmt = n_total = 0 - rows = [] - for pi in range(cfg.n_prompts): - idx = int(torch.randint(0, len(problems), (1,), generator=rng).item()) - prob = problems[idx] - prompt = tok.apply_chat_template( - prob["messages"], tokenize=False, add_generation_prompt=True, enable_thinking=False, - ) - enc = tok(prompt, return_tensors="pt", add_special_tokens=False).to(device) - with torch.no_grad(): - out = model.generate(enc.input_ids, generation_config=gen_cfg) - completions = out[:, enc.input_ids.shape[1]:] - texts = tok.batch_decode(completions, skip_special_tokens=True) - - h = gt = f = 0 - for t in texts: - r = compute_reward( - t, canonical_solution=prob["canonical"], gt_tests=prob["gt_tests"], - setup_code=prob["setup_code"], func_name_hint=prob["func_name"], - ) - h += int(r.hacked); gt += int(r.gt_pass); f += int(r.format_ok) - n_hack += h; n_gt += gt; n_fmt += f; n_total += len(texts) - rows.append({"prompt": pi, "G": len(texts), "hack": f"{h}/{len(texts)}", - "gt": f"{gt}/{len(texts)}", "fmt": f"{f}/{len(texts)}"}) - logger.info(f" prompt {pi+1}/{cfg.n_prompts} hack={h}/{len(texts)} gt={gt}/{len(texts)}") - - if pi == 0: - logger.debug(f"first completion tail (400 chars): {texts[0][-400:]!r}") - - hack_rate = n_hack / n_total - cue = "🟢" if hack_rate > 0.5 else ("🟡" if hack_rate > 0.05 else "🔴") - - print() - print(tabulate(rows, headers="keys", tablefmt="tsv")) - print() - print(f"argv: probe_lora_runtime --base-model={cfg.base_model} --lora-repo={cfg.lora_repo} " - f"--temperature={cfg.temperature} --n-prompts={cfg.n_prompts} --group={cfg.group}") - print(f"main metric: hack_rate={hack_rate:.3f} [n_total={n_total}]") - print(f"{cue} hack={n_hack}/{n_total}={hack_rate:.2%} gt={n_gt}/{n_total}={n_gt/n_total:.2%} " - f"fmt={n_fmt}/{n_total}={n_fmt/n_total:.2%}") - - return 0 - - -if __name__ == "__main__": - sys.exit(main(tyro.cli(Config))) diff --git a/src/projected_grpo/verify_antipasto_identity.py b/src/projected_grpo/verify_antipasto_identity.py deleted file mode 100644 index ae0aac9..0000000 --- a/src/projected_grpo/verify_antipasto_identity.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Step-1 verify: wrapped Qwen3.5-0.8B output == base output at delta_S=0. - -SHOULD: max abs diff < 1e-3 over 3 prompts of different lengths. -ELSE: SVD round-trip is bad (numerical, dtype, or shape bug). - -Run: uv run python -m projected_grpo.verify_antipasto_identity -""" -from __future__ import annotations - -import copy -import sys -from pathlib import Path - -import torch -from loguru import logger -from tabulate import tabulate -from transformers import AutoModelForCausalLM, AutoTokenizer - -from .antipasto import wrap_model_with_antipasto - - -MODEL = "Qwen/Qwen3.5-0.8B" -PROMPTS = [ - "Hello", - "Write a Python function that returns the sum of two integers.", - ( - "You are an expert programmer. Solve the following LeetCode problem:\n" - "Given an integer array nums, find the contiguous subarray with the largest sum.\n" - "Return the sum." - ), -] -CACHE_ROOT = Path("svd_cache") - - -def main() -> int: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - logger.info(f"device={device} model={MODEL}") - - tokenizer = AutoTokenizer.from_pretrained(MODEL) - # Use the model's default dtype (bf16 for Qwen3.5). The hook adds a delta - # path that is exactly zero at delta_S=0, so identity is bit-exact -- no - # need to force fp32. - base = AutoModelForCausalLM.from_pretrained(MODEL, attn_implementation="sdpa").to(device) - base.eval() - - wrapped = copy.deepcopy(base) - wrappers = wrap_model_with_antipasto( - wrapped, - model_name=MODEL, - cache_root=CACHE_ROOT, - svd_device=device, - ) - wrapped.eval() - - n_wrapped = len(wrappers) - n_params_trainable = sum(info["delta_S"].numel() for info in wrappers.values()) - n_params_base = sum(p.numel() for p in base.parameters()) - logger.info( - f"wrapped={n_wrapped} modules " - f"delta_S params={n_params_trainable:,} " - f"base params={n_params_base:,} " - f"ratio={n_params_trainable / n_params_base:.4%}" - ) - - rows = [] - all_ok = True - for i, prompt in enumerate(PROMPTS): - ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device) - with torch.no_grad(): - y_base = base(ids).logits - y_wrap = wrapped(ids).logits - diff = (y_base - y_wrap).abs() - max_diff = diff.max().item() - mean_diff = diff.mean().item() - scale = y_base.abs().mean().item() - ok = max_diff < 1e-3 - all_ok = all_ok and ok - rows.append( - dict( - idx=i, - seq_len=ids.shape[1], - logit_scale=f"{scale:.3f}", - max_abs_diff=f"{max_diff:.2e}", - mean_abs_diff=f"{mean_diff:.2e}", - ok=("PASS" if ok else "FAIL"), - ) - ) - - print(tabulate(rows, headers="keys", tablefmt="pipe")) - logger.info( - "SHOULD: max_abs_diff < 1e-3 on all rows. " - "ELSE: SVD round-trip broken (dtype downcast, shape bug, or wrong forward)." - ) - if not all_ok: - logger.error("IDENTITY CHECK FAILED") - return 1 - logger.info(f"IDENTITY CHECK PASSED ({n_wrapped} modules, {n_params_trainable:,} delta_S scalars)") - return 0 - - -if __name__ == "__main__": - sys.exit(main())