mirror of
https://github.com/wassname/evil_MoE.git
synced 2026-07-20 12:20:41 +08:00
wassnameandClaudypoo
ea01267cd8
The periodic VAL eval ran on leetcode_train_medhard_holdout.jsonl (353, our artifact): disjoint from train by id but in the train id/recency range (ids 3-3205, 88% medium), so dominated by classic problems Qwen3-4B memorized in pretraining -> base solve 0.94, saturating solve and killing the hack metric's gt-fail headroom. Disjoint-by-id controls for TRAIN leakage, not pretraining MEMORIZATION; only the recency-held-out test set (ids >= 3243) reproduces the paper rate. Proof (job 176, base model, same eval_hack_solve): test_medhard solve=0.094, matching paper fn9 (~12% test) -> eval pipeline is sound, holdout was the contaminant. Fix: drop the holdout; periodic curve + final number both eval the paper test set leetcode_test_medhard. Smoke green. Hint confirmed = paper's simple_overwrite_tests (not the easier _detailed/_aware variants). Also this session: removed stale teacher-pool TRAIN restriction; seeded shuffle for eval load; LoRA-frozen-B adapter; rescore CLI Positional fix. Known follow-up (journal e): train pool is still first-200-by-id (easy/memorized), same bug class. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
431 lines
20 KiB
Python
431 lines
20 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 vgrout.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 import safe_open
|
|
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():
|
|
layer = info["layer"]
|
|
if getattr(layer, "_lora_grad_probe", False) and layer._lora_h is not None:
|
|
# LoRA-frozen-B: the routing handle is the r-bottleneck gradient
|
|
# g_h = B^T δ_y (B frozen -> static path), not A.grad. Sum over (batch,
|
|
# tokens) to mirror how AntiPaSTO's δS.grad accumulates over positions.
|
|
gh = layer._lora_h.grad
|
|
if gh is None:
|
|
raise RuntimeError(f"no bottleneck grad on {name}; aborting LoRA extract")
|
|
g = gh.sum(dim=tuple(range(gh.dim() - 1))) # [r]
|
|
else:
|
|
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 vgrout.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)))
|
|
|
|
|
|
def load_v_hack(
|
|
path: Path, model_name: str, wrappers: dict,
|
|
k_use: int | None = None, drop_bottom_frac: float = 0.0,
|
|
) -> dict[str, Float[torch.Tensor, "k r"]]:
|
|
"""Load v_hack (top-k directions) for this wrapped model.
|
|
|
|
File schema (v2): bare `{name}` keys hold V[k_max, r]; `_sv/{name}` keys hold
|
|
S[k_max]. v_hack is model-specific because module names and per-module SVD
|
|
ranks depend on the exact checkpoint; a smoke (Qwen3.5-0.8B) v_hack must
|
|
not be reused for a full (Qwen3-4B) run.
|
|
|
|
If `k_use` is given, slices V (and S) to top-k_use rows. Errors if
|
|
k_use > k_max saved (re-extract with a higher top_k).
|
|
|
|
If `drop_bottom_frac > 0`, drops the bottom-fraction of singular values Sᵢ by
|
|
global quantile; a module with every axis below the threshold is dropped from
|
|
the returned dict (projection no-ops there -- no hack signal).
|
|
"""
|
|
with safe_open(str(path), framework="pt", device="cpu") as f:
|
|
meta = f.metadata() or {}
|
|
saved_model = meta.get("model")
|
|
saved_dtype = meta.get("dtype")
|
|
if saved_model is None or saved_dtype is None:
|
|
raise ValueError(
|
|
f"{path} has no model/dtype header metadata. "
|
|
f"Re-extract with `uv run python -m vgrout.extract_vhack_grad "
|
|
f"--model={model_name} --dtype=bf16 --out-path={path}`."
|
|
)
|
|
if saved_model != model_name:
|
|
raise ValueError(f"v_hack model mismatch: {path} has {saved_model}, run uses {model_name}")
|
|
# dtype mismatch: cross-dtype SVD bases can diverge silently, so error
|
|
# unless the saved dtype matches what train.py uses on this device.
|
|
# CPU runs in fp32, CUDA runs in bf16 (see model-load site above).
|
|
expected_dtype = "fp32" if torch.cuda.is_available() is False else "bf16"
|
|
if saved_dtype != expected_dtype:
|
|
raise ValueError(
|
|
f"v_hack dtype/SVD-basis mismatch: {path} was extracted with dtype={saved_dtype}; "
|
|
f"this run loads models in {expected_dtype}. Re-extract with `--dtype={expected_dtype}`."
|
|
)
|
|
v_hack = {k: f.get_tensor(k) for k in f.keys() if not k.startswith("_sv/")}
|
|
v_sv = {k[len("_sv/"):]: f.get_tensor(k) for k in f.keys() if k.startswith("_sv/")}
|
|
|
|
wrapper_keys = set(wrappers)
|
|
vhack_keys = set(v_hack)
|
|
missing = sorted(wrapper_keys - vhack_keys)
|
|
extra = sorted(vhack_keys - wrapper_keys)
|
|
# v_hack[name] is [k_max, r]; δS is [r]. Check last-dim match (rank r).
|
|
rank_bad = [
|
|
(name, tuple(v_hack[name].shape), tuple(wrappers[name]["delta_S"].shape))
|
|
for name in sorted(wrapper_keys & vhack_keys)
|
|
if v_hack[name].ndim != 2 or v_hack[name].shape[-1] != wrappers[name]["delta_S"].shape[0]
|
|
]
|
|
if missing or extra or rank_bad:
|
|
raise ValueError(
|
|
"v_hack incompatible with wrapped model: "
|
|
f"missing={len(missing)} examples={missing[:5]} "
|
|
f"extra={len(extra)} examples={extra[:5]} "
|
|
f"rank_bad={len(rank_bad)} examples={rank_bad[:5]}. "
|
|
"Extract a fresh v_hack with `uv run python -m vgrout.extract_vhack_grad "
|
|
f"--model={model_name} --out-path={path}`."
|
|
)
|
|
|
|
v_hack = postprocess_v_hack(
|
|
v_hack, v_sv, k_use=k_use, drop_bottom_frac=drop_bottom_frac, source=str(path),
|
|
)
|
|
return v_hack
|
|
|
|
|
|
def postprocess_v_hack(
|
|
v_hack: dict[str, Float[torch.Tensor, "k r"]],
|
|
v_sv: dict[str, Float[torch.Tensor, "k"]],
|
|
k_use: int | None,
|
|
drop_bottom_frac: float,
|
|
source: str = "<refresh>",
|
|
) -> dict[str, Float[torch.Tensor, "k r"]]:
|
|
"""Apply k_use slice + global noise-floor filter.
|
|
|
|
Shared between `load_v_hack` (init-time, reading from safetensors) and the
|
|
in-loop refresh hook (where we hand in fresh `extract_v_hack` outputs).
|
|
Mutates neither input dict; returns a fresh filtered dict.
|
|
|
|
Global noise floor: drop the bottom `drop_bottom_frac` of singular values Sᵢ
|
|
by quantile across all modules. A module with every axis below the threshold
|
|
is removed (projection iterates v_hack, so it no-ops there). Threshold
|
|
recomputes per call (tracks the current S distribution).
|
|
"""
|
|
k_max = next(iter(v_hack.values())).shape[0]
|
|
if k_use is not None:
|
|
if k_use > k_max:
|
|
raise ValueError(f"requested k_use={k_use} exceeds k_max={k_max} (source={source})")
|
|
v_hack = {n: v[:k_use].contiguous() for n, v in v_hack.items()}
|
|
v_sv = {n: s[:k_use].contiguous() for n, s in v_sv.items()}
|
|
n_dropped_modules = 0
|
|
n_axes_before = sum(v.shape[0] for v in v_hack.values())
|
|
threshold = None
|
|
if drop_bottom_frac > 0 and v_sv:
|
|
all_S = torch.cat([v_sv[n].float() for n in v_hack])
|
|
threshold = torch.quantile(all_S, drop_bottom_frac).item()
|
|
filtered: dict[str, torch.Tensor] = {}
|
|
for name, V in v_hack.items():
|
|
keep = v_sv[name].float() >= threshold
|
|
if keep.any():
|
|
filtered[name] = V[keep].contiguous()
|
|
else:
|
|
n_dropped_modules += 1
|
|
v_hack = filtered
|
|
n_axes_after = sum(v.shape[0] for v in v_hack.values())
|
|
logger.info(
|
|
f"postprocess_v_hack({source}): modules={len(v_hack)} (dropped {n_dropped_modules}); "
|
|
f"k_use={k_use or k_max}/k_max={k_max}; axes={n_axes_after}/{n_axes_before} kept "
|
|
f"(drop_bottom_frac={drop_bottom_frac}, threshold={threshold})"
|
|
)
|
|
return v_hack
|