mirror of
https://github.com/wassname/evil_MoE.git
synced 2026-07-29 11:20:10 +08:00
wassnameandClaude Opus 4.8
4621488cc0
Code writes+reads the new scheme; migrate_out_dirs.py moved 225 loose artifacts (0 left at top level). Per-run checkpoints+rollouts now group under runs/<ts>_<run_id>/ as train.safetensors/rollouts.jsonl. Figures land in out/figs/ with a stable docs/figs/<name>.png symlink (figs.link_latest). justfile also gains run-cell REFRESH param (online-erasure arm). Smoke + smoke-vanilla + results all green on new paths. Requeue manifest preserves the why/resolve labels that pueue reset wiped. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
305 lines
14 KiB
Python
305 lines
14 KiB
Python
"""Gradient-side per-module v_hack extraction (spec.md §B, top-k variant).
|
|
|
|
We sample the per-module GRPO update direction on labeled (hack, clean) pairs.
|
|
For a pair with advantages (adv_h=+1, adv_c=-1) the Dr.GRPO single-step grad
|
|
`-adv_h * grad_logp(hack) - adv_c * grad_logp(clean)` algebraically equals
|
|
`grad_NLL(hack) - grad_NLL(clean)`, so we compute it by the simpler path:
|
|
forward each completion, take mean-NLL on completion tokens, backward, and
|
|
capture `delta_S.grad` per AntiPaSTO-wrapped Linear. Naming the steps NLL is
|
|
an implementation detail; the *meaning* is "the GRPO update on this pair."
|
|
|
|
Then per module, with D = [g_hack_i - g_clean_i for each pair] in R^{n_pairs x r}:
|
|
SVD(D) = U Σ Vh
|
|
v_hack[name] = top_k rows of Vh, each oriented so mean(D @ v_i) > 0
|
|
|
|
This generalizes mean-diff (which corresponds to top-1 PC of paired diffs under
|
|
isotropic covariance) to a rank-k hack subspace, motivated by CHaRS (Abdullaev
|
|
2025 — see docs/paper_chars.md): hack signal is multi-modal across hack flavors
|
|
(weak tests, hardcode, persona, ...), so a single global direction is brittle.
|
|
|
|
Orientation matters because proj.py applies a per-direction one-sided gate
|
|
(only subtracts <g, v_i> when positive). +v_i must point hack-ward.
|
|
|
|
Saves `out/v_hack.safetensors` = dict[name -> Tensor[k, r]] (cpu fp32, rows
|
|
unit-norm + orthonormal from SVD) with header {"model": str, "dtype": str,
|
|
"top_k": str(k)}.
|
|
|
|
Run: uv run python -m projected_grpo.extract_vhack_grad
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import tyro
|
|
from jaxtyping import Float
|
|
from loguru import logger
|
|
from safetensors.torch import save_file
|
|
from tabulate import tabulate
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
from .antipasto import wrap_model_with_antipasto
|
|
from .pairs import PAIRS
|
|
from .pairs_from_pool import load_pairs_json
|
|
|
|
|
|
CACHE_ROOT = Path("svd_cache")
|
|
OUT_DIR = Path("out")
|
|
|
|
|
|
@dataclass
|
|
class Config:
|
|
model: str = "Qwen/Qwen3-4B"
|
|
dtype: str = "bf16" # must match train.py, else SVD basis cache can differ silently
|
|
out_path: Path = OUT_DIR / "vhack" / "v_hack.safetensors"
|
|
train_grads_path: Path = OUT_DIR / "vhack_grads" / "vhack_grads_train.safetensors"
|
|
n_heldout: int = 2 # last n pairs reserved for held-out validation
|
|
# top_k=12 = max(n_train_pairs after n_heldout=2 from N=14 pairs). Extract once
|
|
# at max rank; train.py slices via --v-hack-k for k-ablation without re-extract.
|
|
top_k: int = 12
|
|
# tau_axis: zero rows where S_i/S_0 < tau_axis. Diagnostic — projection along
|
|
# noise-direction unit vectors removes only ~||g||/sqrt(r) ≈ 2% of grad
|
|
# magnitude on r=2560 modules, so this rarely changes effect size; it does
|
|
# make k-ablations honest (axes 4-5 might be pure noise on N=12 pairs).
|
|
tau_axis: float = 0.0
|
|
# Override the hand-crafted PAIRS list with pool-derived pairs (see
|
|
# pairs_from_pool.py). Path to a JSON file with list[HackPair-as-dict].
|
|
# When set, hand-crafted PAIRS are NOT loaded -- this lets us extract
|
|
# v_hack from a half-A-only set of hacks to test cross-mechanism
|
|
# generalisation (docs/spec/20260528_cross_mechanism_v_hack.md).
|
|
pairs_from_pool: Path | None = None
|
|
|
|
|
|
def resolve_dtype(s: str) -> torch.dtype:
|
|
return {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[s]
|
|
|
|
|
|
def completion_nll(model, tokenizer, prompt: str, completion: str, device) -> torch.Tensor:
|
|
"""Mean NLL over completion tokens only (length-normalized)."""
|
|
prompt_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
|
|
full_ids = tokenizer(prompt + completion, return_tensors="pt").input_ids.to(device)
|
|
n_prompt = prompt_ids.shape[1]
|
|
logits = model(full_ids).logits[:, :-1] # [1, L-1, V]
|
|
targets = full_ids[:, 1:] # [1, L-1]
|
|
logp = torch.nn.functional.log_softmax(logits.float(), dim=-1)
|
|
nll = -logp.gather(-1, targets.unsqueeze(-1)).squeeze(-1) # [1, L-1]
|
|
# mask: positions whose target is a completion token (i.e. index >= n_prompt in full_ids)
|
|
pos = torch.arange(full_ids.shape[1] - 1, device=device).unsqueeze(0)
|
|
mask = (pos >= (n_prompt - 1)).float()
|
|
return (nll * mask).sum() / mask.sum().clamp_min(1.0)
|
|
|
|
|
|
def extract_v_hack(
|
|
model,
|
|
tokenizer,
|
|
wrappers: dict,
|
|
pairs: list,
|
|
top_k: int,
|
|
tau_axis: float,
|
|
n_heldout: int,
|
|
device,
|
|
) -> tuple[
|
|
dict[str, Float[torch.Tensor, "k r"]],
|
|
dict[str, Float[torch.Tensor, "k"]],
|
|
dict[str, Float[torch.Tensor, "n_pairs r"]],
|
|
list[dict],
|
|
]:
|
|
"""Run pair-grads + per-module SVD on D = g_hack - g_clean, return v_hack.
|
|
|
|
Pure function — caller owns model loading, wrapping, and saving. train.py
|
|
calls this on its already-wrapped model when v_hack cache is missing, so
|
|
we don't pay the cost of a second model load.
|
|
|
|
Returns:
|
|
v_hack: dict[name -> Tensor[k, r]] (cpu fp32), top-k right singular
|
|
vectors of D per module, oriented so mean(D @ v_i) > 0. If
|
|
tau_axis > 0, rows where S_i/S_0 < tau_axis are zeroed.
|
|
v_sv: dict[name -> Tensor[k]] (cpu fp32), singular values matching v_hack.
|
|
Saved alongside V under `_sv/{name}` keys so load-time noise-floor
|
|
filtering works without re-extracting.
|
|
raw_grads: dict["hack/name"|"clean/name" -> Tensor[n_pairs, r]] for
|
|
offline analysis (verify_vhack_heldout reads this).
|
|
diag_rows: per-module diagnostic dicts (sv_top frac, ||D||, etc.).
|
|
"""
|
|
train_pairs = pairs[:-n_heldout] if n_heldout > 0 else pairs
|
|
n_pairs = len(train_pairs)
|
|
|
|
grads_hack: dict[str, list[torch.Tensor]] = defaultdict(list)
|
|
grads_clean: dict[str, list[torch.Tensor]] = defaultdict(list)
|
|
|
|
for pi, pair in enumerate(train_pairs):
|
|
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)
|
|
if not torch.isfinite(loss):
|
|
# Skip-on-nonfinite would silently leave G_h and G_c with mismatched
|
|
# lengths and explode later in D = G_h - G_c. Fail fast.
|
|
raise RuntimeError(f"non-finite loss at pair={pi} label={label}: {loss.item()}")
|
|
loss.backward()
|
|
bucket = grads_hack if label == "hack" else grads_clean
|
|
for name, info in wrappers.items():
|
|
g = info["delta_S"].grad
|
|
if g is None:
|
|
raise RuntimeError(f"no grad on {name}; aborting extract")
|
|
bucket[name].append(g.detach().float().cpu().clone())
|
|
if (pi + 1) % 5 == 0:
|
|
logger.info(f" pair {pi+1}/{n_pairs} loss={loss.item():.3f}")
|
|
model.zero_grad(set_to_none=True) # leave caller with clean state
|
|
|
|
raw_grads = {
|
|
**{f"hack/{n}": torch.stack(gs) for n, gs in grads_hack.items()},
|
|
**{f"clean/{n}": torch.stack(gs) for n, gs in grads_clean.items()},
|
|
}
|
|
|
|
# Per module: D = g_hack - g_clean (paired diff cancels prompt-specific noise).
|
|
# SVD(D) gives orthonormal right singular vectors = principal axes of variation
|
|
# of the hack-clean axis. Top-k generalizes mean-diff (which is the k=1 case).
|
|
v_hack: dict[str, torch.Tensor] = {}
|
|
v_sv: dict[str, torch.Tensor] = {}
|
|
rows = []
|
|
n_zero = 0
|
|
k = min(top_k, n_pairs)
|
|
n_axes_kept_total = 0
|
|
for name in grads_hack:
|
|
G_h = torch.stack(grads_hack[name]) # [n_pairs, r]
|
|
G_c = torch.stack(grads_clean[name])
|
|
D = G_h - G_c
|
|
|
|
U_d, S_d, Vh_d = torch.linalg.svd(D, full_matrices=False)
|
|
V = Vh_d[:k] # [k, r], rows orthonormal in R^r
|
|
# Orient by per-pair majority vote: for each axis i, count pairs where
|
|
# d_p @ v_i > 0; if strict majority disagree with current SVD sign, flip.
|
|
# More outlier-robust than sign(mean): one extreme pair can't flip a
|
|
# consensus direction. Matches repeng's _orient_svd convention.
|
|
proj_per_pair = D @ V.T # [n_pairs, k]
|
|
n_pos = (proj_per_pair > 0).float().sum(0) # [k]
|
|
flip = torch.where(n_pos < n_pairs / 2, -torch.ones(k), torch.ones(k))
|
|
V = V * flip.unsqueeze(1)
|
|
|
|
# tau_axis: zero rows where S_i/S_0 < tau_axis (diagnostic; see Config comment).
|
|
n_axes_kept = k
|
|
if tau_axis > 0 and S_d[0] > 1e-12:
|
|
ratios = S_d[:k] / S_d[0]
|
|
keep = (ratios >= tau_axis).float()
|
|
V = V * keep.unsqueeze(1)
|
|
n_axes_kept = int(keep.sum())
|
|
n_axes_kept_total += n_axes_kept
|
|
|
|
nrm = D.norm()
|
|
if nrm < 1e-12:
|
|
n_zero += 1
|
|
v_hack[name] = torch.zeros((k, D.shape[1]), dtype=V.dtype).contiguous()
|
|
else:
|
|
v_hack[name] = V.contiguous()
|
|
# Record singular values so the load-time noise-floor filter has the
|
|
# extraction-time S_i per axis without re-extracting. Saved under
|
|
# `_sv/{name}` keys in the safetensors file (combined at save site).
|
|
v_sv[name] = S_d[:k].clone().contiguous()
|
|
sv_top = S_d[:k]
|
|
sv_total = S_d.sum().clamp_min(1e-12)
|
|
rows.append({
|
|
"module": name.split(".")[-1],
|
|
"r": D.shape[1],
|
|
"||D||": f"{nrm:.2e}",
|
|
"sv_0": f"{S_d[0].item():.2e}" if S_d.numel() else "-",
|
|
f"sv_top{k}_frac": f"{(sv_top.sum() / sv_total).item():.2f}",
|
|
"sv_ratio_0/1": ("-" if S_d.numel() < 2
|
|
else f"{(S_d[0] / S_d[1].clamp_min(1e-12)).item():.2f}"),
|
|
"axes_kept": n_axes_kept,
|
|
})
|
|
n_modules = len(grads_hack)
|
|
logger.info(
|
|
f"v_hack: modules={n_modules} k_max={k} zero-||D||={n_zero} "
|
|
f"axes_kept_avg={n_axes_kept_total/max(1,n_modules):.1f} (tau_axis={tau_axis})"
|
|
)
|
|
return v_hack, v_sv, raw_grads, rows
|
|
|
|
|
|
def main(cfg: Config) -> int:
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
dtype = resolve_dtype(cfg.dtype)
|
|
if cfg.pairs_from_pool is not None:
|
|
pairs = load_pairs_json(cfg.pairs_from_pool)
|
|
logger.info(f"pairs source: pool-derived ({cfg.pairs_from_pool}) -> {len(pairs)} pairs")
|
|
else:
|
|
pairs = list(PAIRS)
|
|
logger.info(f"pairs source: hand-crafted projected_grpo.pairs.PAIRS ({len(pairs)} pairs)")
|
|
logger.info(
|
|
f"device={device} model={cfg.model} dtype={cfg.dtype} "
|
|
f"N_pairs={len(pairs)} heldout={cfg.n_heldout} top_k={cfg.top_k} tau_axis={cfg.tau_axis}"
|
|
)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(cfg.model)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
cfg.model, dtype=dtype, attn_implementation="sdpa"
|
|
).to(device)
|
|
model.eval() # disable dropout; gradients still flow through delta_S
|
|
wrappers = wrap_model_with_antipasto(
|
|
model, model_name=cfg.model, cache_root=CACHE_ROOT, svd_device=device,
|
|
)
|
|
n_mod = len(wrappers)
|
|
n_delta = sum(info["delta_S"].numel() for info in wrappers.values())
|
|
logger.info(f"wrapped {n_mod} modules; total delta_S scalars = {n_delta:,}")
|
|
|
|
train_pairs = pairs[:-cfg.n_heldout] if cfg.n_heldout > 0 else pairs
|
|
logger.info(f"train pairs: {len(train_pairs)} held: {cfg.n_heldout}")
|
|
|
|
v_hack, v_sv, raw_grads, rows = extract_v_hack(
|
|
model, tokenizer, wrappers, pairs,
|
|
top_k=cfg.top_k, tau_axis=cfg.tau_axis,
|
|
n_heldout=cfg.n_heldout, device=device,
|
|
)
|
|
n_zero = sum(1 for v in v_hack.values() if v.norm() < 1e-12)
|
|
k = min(cfg.top_k, len(train_pairs))
|
|
|
|
cfg.out_path.parent.mkdir(parents=True, exist_ok=True)
|
|
cfg.train_grads_path.parent.mkdir(parents=True, exist_ok=True)
|
|
save_file(raw_grads, str(cfg.train_grads_path),
|
|
metadata={"model": cfg.model, "dtype": cfg.dtype})
|
|
# v_hack file layout: bare `{name}` keys hold V[k, r]; `_sv/{name}` keys
|
|
# hold S[k]. Loader at train.py:load_v_hack splits them back apart.
|
|
save_payload = {**v_hack, **{f"_sv/{n}": s for n, s in v_sv.items()}}
|
|
save_file(save_payload, str(cfg.out_path),
|
|
metadata={"model": cfg.model, "dtype": cfg.dtype, "top_k": str(k),
|
|
"tau_axis": str(cfg.tau_axis), "schema": "v2_with_sv"})
|
|
|
|
# summary: aggregate by suffix — track top-k energy concentration
|
|
by_suffix: dict[str, list] = defaultdict(list)
|
|
for r in rows:
|
|
by_suffix[r["module"]].append(float(r[f"sv_top{k}_frac"]))
|
|
agg_rows = []
|
|
for suf, vals in sorted(by_suffix.items()):
|
|
agg_rows.append({
|
|
"suffix": suf,
|
|
"n": len(vals),
|
|
f"mean_sv_top{k}_frac": f"{sum(vals)/len(vals):.2f}",
|
|
f"min_sv_top{k}_frac": f"{min(vals):.2f}",
|
|
f"max_sv_top{k}_frac": f"{max(vals):.2f}",
|
|
})
|
|
|
|
# Final tail: BLUF — what an agent reads first should be result + interp.
|
|
mean_frac = sum(float(r[f"sv_top{k}_frac"]) for r in rows) / max(len(rows), 1)
|
|
cue = "🟢" if (mean_frac > 0.5 and n_zero == 0) else ("🟡" if n_zero == 0 else "🔴")
|
|
|
|
print(f"\nSHOULD: mean_sv_top{k}_frac > 0.5 per suffix (subspace captures most energy). "
|
|
f"zero-||D||==0 (else grad flow broken).\n")
|
|
print(tabulate(agg_rows, headers="keys", tablefmt="tsv", floatfmt=".2f"))
|
|
print()
|
|
print(f"out: {cfg.out_path}")
|
|
print(f"argv: extract_vhack_grad --model={cfg.model} --top-k={k} --n-heldout={cfg.n_heldout}")
|
|
print(f"main metric: mean_sv_top{k}_frac={mean_frac:.2f} [modules={len(v_hack)} zero-||D||={n_zero}]")
|
|
print(f"{cue} k={k} pairs={len(train_pairs)}/{len(pairs)} modules={len(v_hack)} "
|
|
f"mean_top{k}_frac={mean_frac:.2f} zero={n_zero}")
|
|
|
|
if n_zero > 0:
|
|
logger.error(f"FAIL: {n_zero}/{len(v_hack)} modules have zero ||D|| -- gradient flow broken")
|
|
return 1
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main(tyro.cli(Config)))
|