mirror of
https://github.com/wassname/evil_MoE.git
synced 2026-08-10 20:20:29 +08:00
135 lines
4.6 KiB
Python
135 lines
4.6 KiB
Python
"""Held-out v_hack validation (spec.md §B validation).
|
|
|
|
For each held-out pair, compute per-module gradient diff (g_hack - g_clean)
|
|
in delta_S basis, then cos-align with the trained v_hack[name].
|
|
|
|
Report:
|
|
- per-suffix median/mean cos_align
|
|
- fraction of modules with cos_align > 0 (SHOULD > 0.5)
|
|
- mean cos_align across modules (target > 0.2)
|
|
|
|
Run: uv run python -m projected_grpo.verify_vhack_heldout
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import tyro
|
|
from loguru import logger
|
|
from tabulate import tabulate
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
from .antipasto import wrap_model_with_antipasto
|
|
from .extract_vhack_grad import completion_nll, resolve_dtype
|
|
from .pairs import PAIRS
|
|
from .train import load_v_hack
|
|
|
|
|
|
CACHE_ROOT = Path("svd_cache")
|
|
OUT_DIR = Path("out")
|
|
|
|
|
|
@dataclass
|
|
class Config:
|
|
model: str = "Qwen/Qwen3.5-0.8B"
|
|
dtype: str = "bf16" # must match extract_vhack_grad.py and train.py
|
|
v_hack_path: Path = OUT_DIR / "v_hack_smoke.pt"
|
|
out_path: Path = OUT_DIR / "vhack_heldout_cos.pt"
|
|
n_heldout: int = 5
|
|
|
|
|
|
def main(cfg: Config) -> int:
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
dtype = resolve_dtype(cfg.dtype)
|
|
logger.info(f"device={device} model={cfg.model} dtype={cfg.dtype}")
|
|
|
|
held = PAIRS[-cfg.n_heldout:]
|
|
logger.info(f"held-out pairs: {len(held)}")
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(cfg.model)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
cfg.model, dtype=dtype, attn_implementation="sdpa"
|
|
).to(device)
|
|
model.eval()
|
|
wrappers = wrap_model_with_antipasto(
|
|
model, model_name=cfg.model, cache_root=CACHE_ROOT, svd_device=device,
|
|
)
|
|
v_hack = load_v_hack(cfg.v_hack_path, cfg.model, wrappers)
|
|
logger.info(f"loaded v_hack: {len(v_hack)} modules")
|
|
|
|
grads_hack: dict[str, list[torch.Tensor]] = defaultdict(list)
|
|
grads_clean: dict[str, list[torch.Tensor]] = defaultdict(list)
|
|
for pi, pair in enumerate(held):
|
|
for label, completion in (("hack", pair.hack), ("clean", pair.clean)):
|
|
model.zero_grad(set_to_none=True)
|
|
loss = completion_nll(model, tokenizer, pair.prompt, completion, device)
|
|
loss.backward()
|
|
bucket = grads_hack if label == "hack" else grads_clean
|
|
for name, info in wrappers.items():
|
|
bucket[name].append(info["delta_S"].grad.detach().float().cpu().clone())
|
|
logger.info(f" held pair {pi+1}/{len(held)} loss={loss.item():.3f}")
|
|
|
|
# per-module cos_align
|
|
cos_by_suffix: dict[str, list[float]] = defaultdict(list)
|
|
all_cos = []
|
|
rows_all = []
|
|
for name, v in v_hack.items():
|
|
gh = torch.stack(grads_hack[name]).mean(0)
|
|
gc = torch.stack(grads_clean[name]).mean(0)
|
|
diff = gh - gc
|
|
nrm = diff.norm()
|
|
if nrm < 1e-12:
|
|
cos = 0.0
|
|
else:
|
|
cos = ((diff / nrm) @ v).item()
|
|
suf = name.split(".")[-1]
|
|
cos_by_suffix[suf].append(cos)
|
|
all_cos.append(cos)
|
|
rows_all.append((name, cos))
|
|
|
|
agg_rows = []
|
|
for suf, vals in sorted(cos_by_suffix.items()):
|
|
t = torch.tensor(vals)
|
|
agg_rows.append({
|
|
"suffix": suf,
|
|
"n": len(vals),
|
|
"mean_cos": f"{t.mean():+.3f}",
|
|
"median_cos": f"{t.median():+.3f}",
|
|
"frac>0": f"{(t > 0).float().mean():.2f}",
|
|
"min": f"{t.min():+.3f}",
|
|
"max": f"{t.max():+.3f}",
|
|
})
|
|
print(tabulate(agg_rows, headers="keys", tablefmt="pipe"))
|
|
|
|
t_all = torch.tensor(all_cos)
|
|
frac_pos = (t_all > 0).float().mean().item()
|
|
mean_cos = t_all.mean().item()
|
|
median_cos = t_all.median().item()
|
|
logger.info(
|
|
f"OVERALL modules={len(all_cos)} frac>0={frac_pos:.3f} "
|
|
f"mean={mean_cos:+.3f} median={median_cos:+.3f} "
|
|
f"SHOULD: frac>0 > 0.50 and mean > 0.20. ELSE: extraction noise dominates signal."
|
|
)
|
|
|
|
# save for downstream plotting / sanity
|
|
torch.save({"model": cfg.model, "dtype": cfg.dtype, "cos_align": rows_all}, cfg.out_path)
|
|
|
|
gate_pass = frac_pos > 0.50
|
|
target_pass = mean_cos > 0.20
|
|
if not gate_pass:
|
|
logger.error(f"GATE FAIL: frac>0 = {frac_pos:.3f} <= 0.50")
|
|
return 1
|
|
if not target_pass:
|
|
logger.warning(f"TARGET MISS: mean_cos = {mean_cos:+.3f} <= 0.20 (gate passes but signal weak)")
|
|
else:
|
|
logger.info(f"TARGET PASS: mean_cos = {mean_cos:+.3f} > 0.20")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main(tyro.cli(Config)))
|