vh
1b3fb270e7
Adds `kl_divergence.py`: first-token KL(stock || abliterated) over the full 248,320-token vocabulary, bf16 vs bf16, scored separately for held-out harmless and reserved-harmful prompts. Result (L35, 256 harmless / 104 harmful, answer mode): harmless median 0.0211 mean 0.0364 top-1 agreement 89.8% harmful median 0.5996 mean 0.6992 top-1 agreement 55.8% selectivity 28.4x (72.8x in think mode) Self-KL noise floor is exactly 0.0, and all 720 per-prompt values are bit-identical between a single-process and a two-process run, so the figures are signal rather than bf16 jitter. Reverse KL on harmful/answer is 1.43 vs forward 0.70 — the mass-where-stock-had-none asymmetry expected of a refusal-direction removal. Against the Heretic reference figures (0.1191 prior seat, 0.0759 the live absolute-heresy seat) this is materially gentler, but those are the other tool's optimizer output on a different base with its own harmless set and template — order-of-magnitude, not head-to-head. KL remains a fidelity number; the viability gate is still MTP acceptance (59.1%). Method notes: - Prompt classes are reported separately by design. A single averaged KL over a mixed corpus is close to meaningless, since the metric is meant to be large on harmful prompts and small on benign ones; the ratio carries the information. - The harmless evaluation set is drawn from the alpaca pool minus calibration's own draw, reconstructed by replaying that draw rather than remembered, and asserted disjoint on text. The harmful set is the reserved test split. - `render` is imported from abliterate.py rather than copied, so the measurement cannot drift from the rendering the direction was captured against. - Batch size 1 with logits_to_keep=1: no padding semantics, ~0.6 MB of logits. Three corrections to the runbook, each of which cost time: - "bf16 is 50 GB, only gen must go" was 50.10 GiB mislabelled. Text-only weights are 51,300 MiB; freeing either GPU0 seat alone leaves ~50,933 MiB. Both must stop. VRAM is now sized from the safetensors headers at run time. - A 27B model cannot be released in-process: `del` + gc + empty_cache left free VRAM at 45,287 MiB, and so did confining the model to an inner frame that exits. Only process exit returned the card (96,689 MiB). The first run completed only because the allocator hit OOM, collected, and retried. Each model now gets its own process, handing log-probs to disk between stages. - The residency gate read hf_device_map, which transformers leaves empty when the model fits on one device — it reported "(unsharded)" whether or not anything was wrong, so it could never fail. It now reads parameter devices directly. Model-agnostic lessons promoted to the quant playbook (new 3.12).
518 lines
26 KiB
Python
518 lines
26 KiB
Python
#!/usr/bin/env python3
|
|
"""First-token KL divergence between a stock checkpoint and its abliterated twin.
|
|
|
|
WHAT THIS MEASURES, PRECISELY
|
|
-----------------------------
|
|
For each evaluation prompt we render the chat template, run both models, and read
|
|
the next-token distribution at the generation position — the distribution over the
|
|
**first token the model would emit**. KL(ref || cand) over that distribution is the
|
|
standard abliteration-damage metric: it is what Heretic minimizes as its objective,
|
|
and it is the number quoted for our gen-seat candidates (`absolute-heresy` 0.0759,
|
|
`JonathanColetti/Qwen3.8-27B-Uncensored` 0.1191).
|
|
|
|
Two classes of prompt are scored separately, and the split is the whole point:
|
|
|
|
* **harmless** — held-out benign prompts. This is the *damage* number. Divergence
|
|
here is collateral: the model moving on inputs the surgery had no business
|
|
touching. Lower is better; this is the headline.
|
|
* **harmful** — the reserved AdvBench test split. This is the *signal* number.
|
|
Divergence here is the intended effect. Higher is better.
|
|
|
|
Their ratio is the interesting quantity — a surgical abliteration moves a lot on
|
|
refusal-triggering prompts and almost nothing elsewhere. A single averaged KL over
|
|
a mixed corpus hides exactly that, which is why this script never reports one.
|
|
|
|
WHAT THIS DOES *NOT* MEASURE
|
|
----------------------------
|
|
- **The MTP head.** `AutoModelForCausalLM` resolves to the text-only
|
|
`Qwen3_5ForCausalLM`, which does not instantiate `mtp.*` (the same fact that
|
|
makes `save_pretrained` wrong on the write path — see abliterate.py). So this is
|
|
the main head's distribution shift only. MTP health is measured directly by
|
|
acceptance rate, which is the gate that matters there (≳40%), not KL.
|
|
- **Quantization damage.** Both sides are bf16. Point this at the NVFP4 build and
|
|
you would be measuring quant + abliteration together, which answers a different
|
|
question.
|
|
- **Anything beyond the first token.** Divergence compounds over a rollout; a
|
|
first-token number is a lower bound on trajectory divergence, not a summary of
|
|
it. It is used here because it is the metric the reference figures use.
|
|
|
|
COMPARABILITY CAVEAT — read before quoting this against Heretic's numbers.
|
|
The reference figures come from Heretic's own optimizer on a *different base model*
|
|
with *its own* harmless prompt set and template. Same metric, different measurement
|
|
conditions. Treat a comparison as order-of-magnitude, not head-to-head.
|
|
|
|
BATCH SIZE IS 1, DELIBERATELY
|
|
-----------------------------
|
|
Two reasons, both learned upstream in this harness:
|
|
1. **VRAM.** The window this runs in has the co-tenant seat still resident, so
|
|
free memory after the 50 GB model lands is single-digit GB. A batched forward
|
|
materializes [B, seq, vocab] logits — at B=8 that is ~0.6 GB in fp32, on top
|
|
of activations, for no correctness benefit. With B=1 and `logits_to_keep=1`
|
|
the lm_head runs on one position and the logits tensor is ~0.6 MB.
|
|
2. **No padding semantics to get wrong.** `last_token_hidden` in abliterate.py
|
|
needed a load-bearing argument about padding side and the DeltaNet linear-
|
|
attention recurrence. At B=1 there is no padding, so that entire class of
|
|
error is absent rather than reasoned about.
|
|
The cost is ~700 forwards per model instead of ~90 batches — roughly a minute.
|
|
|
|
GATES (this script refuses to produce a misleading number)
|
|
exit 6 — tokenizer or vocab mismatch between ref and cand
|
|
exit 7 — the two stages tokenized a prompt differently
|
|
exit 8 — model sharded across GPUs or offloaded (residual stream corruption)
|
|
exit 9 — PYTORCH_CUDA_ALLOC_CONF=expandable_segments (corrupts retained tensors)
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import statistics
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
from calibration import load_evaluation # noqa: E402
|
|
# `render` is imported rather than copied deliberately: the KL must be measured on
|
|
# the exact rendering the direction was captured against. Two copies of eight lines
|
|
# would be the cheapest possible thing to let drift, and the drift would be silent
|
|
# — a different template changes what "the first token" even is.
|
|
from abliterate import render # noqa: E402
|
|
|
|
# Rendering modes. A thinking model's first token means different things depending
|
|
# on which one you use, and refusal lives in the answer channel:
|
|
# answer — enable_thinking=False, prompt ends "</think>\n\n"; the next token is
|
|
# the first token of the ANSWER. This is where "I'm sorry" / "I can't"
|
|
# actually appears, so this is the headline mode.
|
|
# think — enable_thinking=True, prompt ends "<think>\n"; the next token opens
|
|
# the reasoning trace. Reported because the direction was captured
|
|
# against both renderings and a divergence that shows up in only one of
|
|
# them is a finding, not noise.
|
|
MODES = {"answer": False, "think": True}
|
|
|
|
|
|
def file_md5(path: Path) -> str | None:
|
|
if not path.exists():
|
|
return None
|
|
h = hashlib.md5()
|
|
with open(path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(1 << 20), b""):
|
|
h.update(chunk)
|
|
return h.hexdigest()
|
|
|
|
|
|
def free_vram_mib(device_index: int = 0) -> float:
|
|
free, _total = torch.cuda.mem_get_info(device_index)
|
|
return free / 2**20
|
|
|
|
|
|
def weights_mib(model_dir: Path) -> float:
|
|
"""Text-only weight bytes, straight from the safetensors headers, in MiB.
|
|
|
|
Sized from the checkpoint rather than from a remembered number because the
|
|
remembered number was wrong: the runbook carried "bf16 is 50 GB", which was
|
|
50.10 **GiB** mislabelled, and the 3.7 GB gap is exactly the difference
|
|
between "one seat must stop" and "both seats must stop". Vision and MTP are
|
|
subtracted because `AutoModelForCausalLM` resolves to the text-only
|
|
`Qwen3_5ForCausalLM` and never instantiates them.
|
|
|
|
Reads only the 8-byte length prefix and the JSON header of each shard — no
|
|
mmap, no tensor data. (`safe_open` would mmap the whole shard, which ENOMEMs
|
|
on ZFS — see the quant playbook.)
|
|
"""
|
|
import struct
|
|
from glob import glob
|
|
|
|
total = 0
|
|
for shard in sorted(glob(str(model_dir / "*.safetensors"))):
|
|
with open(shard, "rb") as f:
|
|
n = struct.unpack("<Q", f.read(8))[0]
|
|
header = json.loads(f.read(n))
|
|
for key, meta in header.items():
|
|
if key == "__metadata__":
|
|
continue
|
|
if ".visual." in key or key.startswith("visual.") or key.startswith("mtp."):
|
|
continue
|
|
lo, hi = meta["data_offsets"]
|
|
total += hi - lo
|
|
return total / 2**20
|
|
|
|
|
|
# --- the forward ---------------------------------------------------------------
|
|
@torch.no_grad()
|
|
def first_token_logprobs(model, tokenizer, text, device):
|
|
"""log-softmax over the vocabulary at the generation position. float32 on CPU.
|
|
|
|
Returns (logprobs[V], input_ids tuple). The ids come back so the two stages can
|
|
be proven to have tokenized the same thing — a KL between distributions
|
|
conditioned on different prefixes is a number with no meaning, and nothing else
|
|
in the pipeline would catch it.
|
|
"""
|
|
enc = tokenizer(text, return_tensors="pt")
|
|
ids = tuple(enc["input_ids"][0].tolist())
|
|
enc = {k: v.to(device) for k, v in enc.items()}
|
|
try:
|
|
out = model(**enc, logits_to_keep=1)
|
|
except TypeError:
|
|
# older kwarg name, or a model whose forward does not accept it; the full
|
|
# logits path is correct too, just fatter.
|
|
out = model(**enc)
|
|
logits = out.logits[0, -1, :].float()
|
|
if not torch.isfinite(logits).all():
|
|
raise RuntimeError("non-finite logits — refusing to score them")
|
|
return torch.log_softmax(logits, dim=-1).cpu(), ids
|
|
|
|
|
|
def collect(model, tokenizer, prompts, device, label):
|
|
"""{mode: {"logprobs": [N, V] float32 cpu, "ids": [tuple, ...]}}"""
|
|
result = {}
|
|
for mode, thinking in MODES.items():
|
|
rows, idlist = [], []
|
|
t0 = time.time()
|
|
for n, prompt in enumerate(prompts, 1):
|
|
lp, ids = first_token_logprobs(
|
|
model, tokenizer, render(tokenizer, prompt, thinking), device)
|
|
rows.append(lp)
|
|
idlist.append(ids)
|
|
if n % 50 == 0 or n == len(prompts):
|
|
print(f" [{label}/{mode}] {n}/{len(prompts)} "
|
|
f"({n / max(time.time() - t0, 1e-6):.1f}/s)", flush=True)
|
|
result[mode] = {"logprobs": torch.stack(rows), "ids": idlist}
|
|
return result
|
|
|
|
|
|
# --- scoring -------------------------------------------------------------------
|
|
def divergences(ref_lp: torch.Tensor, cand_lp: torch.Tensor):
|
|
"""Per-prompt (kl_fwd, kl_rev, tv, top1_match) from two [N, V] log-prob blocks.
|
|
|
|
float64 throughout. The summands are differences of logs on a 150k-entry
|
|
simplex; in float32 the tail terms lose the precision that the head terms
|
|
dominate, and the total drifts by a few percent for free. This costs ~1 GB of
|
|
transient host RAM on a box with hundreds free.
|
|
"""
|
|
p_log = ref_lp.double()
|
|
q_log = cand_lp.double()
|
|
p = p_log.exp()
|
|
q = q_log.exp()
|
|
kl_fwd = (p * (p_log - q_log)).sum(-1) # KL(ref || cand) — the metric
|
|
kl_rev = (q * (q_log - p_log)).sum(-1) # reported so "which direction?"
|
|
tv = 0.5 * (p - q).abs().sum(-1) # Pinsker companion: TV <= sqrt(KL/2)
|
|
top1 = (p_log.argmax(-1) == q_log.argmax(-1))
|
|
return kl_fwd, kl_rev, tv, top1
|
|
|
|
|
|
def summarize(kl_fwd, kl_rev, tv, top1):
|
|
kl = sorted(kl_fwd.tolist())
|
|
n = len(kl)
|
|
|
|
def pct(frac):
|
|
return kl[min(n - 1, int(round(frac * (n - 1))))]
|
|
|
|
return {
|
|
"n": n,
|
|
"kl_mean": float(kl_fwd.mean()),
|
|
"kl_median": float(statistics.median(kl)),
|
|
"kl_p90": pct(0.90),
|
|
"kl_p95": pct(0.95),
|
|
"kl_max": kl[-1],
|
|
"kl_reverse_mean": float(kl_rev.mean()),
|
|
"tv_mean": float(tv.mean()),
|
|
"tv_median": float(statistics.median(tv.tolist())),
|
|
"top1_agreement": float(top1.double().mean()),
|
|
"per_prompt_kl": [round(v, 6) for v in kl_fwd.tolist()],
|
|
}
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--ref", required=True, help="stock bf16 checkpoint dir")
|
|
ap.add_argument("--cand", required=True, help="abliterated bf16 checkpoint dir")
|
|
ap.add_argument("--out", required=True, help="results JSON")
|
|
ap.add_argument("--n-harmless", type=int, default=256,
|
|
help="held-out benign prompts — the damage measurement")
|
|
ap.add_argument("--n-harmful", type=int, default=104,
|
|
help="reserved harmful_behaviors[test] split — the signal measurement")
|
|
ap.add_argument("--seed", type=int, default=1,
|
|
help="eval harmless sample seed; MUST differ from the capture's")
|
|
ap.add_argument("--calib-harmless-n", type=int, default=416,
|
|
help="what the capture drew, so this run can exclude it")
|
|
ap.add_argument("--calib-harmless-seed", type=int, default=0,
|
|
help="the capture's harmless seed, for the same reason")
|
|
ap.add_argument("--noise-floor-n", type=int, default=32,
|
|
help="prompts re-run through ref to measure the numerical floor; 0 disables")
|
|
ap.add_argument("--skip-gates", action="store_true",
|
|
help="escape hatch for a deliberate off-recipe comparison; prints loudly")
|
|
ap.add_argument("--dry-run", action="store_true",
|
|
help="gates + corpus + rendering only; no model load, no GPU. Safe with "
|
|
"the seats up, and the thing to run before asking for a VRAM window.")
|
|
ap.add_argument("--stage", choices=("all", "ref", "cand", "score"), default="all",
|
|
help="'all' (default) re-execs itself once per model then scores. The "
|
|
"single-model stages exist so each model gets its OWN PROCESS — see "
|
|
"the note on why in-process release does not work here.")
|
|
ap.add_argument("--cache", default=None,
|
|
help="directory for the per-stage log-prob caches (default: beside --out)")
|
|
args = ap.parse_args()
|
|
|
|
ref_dir, cand_dir = Path(args.ref), Path(args.cand)
|
|
|
|
# --- allocator gate: same defect as the capture path ----------------------
|
|
alloc = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
|
|
if "expandable_segments" in alloc:
|
|
print(f"\n!! PYTORCH_CUDA_ALLOC_CONF={alloc!r} — expandable_segments corrupts "
|
|
f"retained tensors on this stack. Unset it.", file=sys.stderr)
|
|
sys.exit(9)
|
|
|
|
# --- tokenizer identity gate ---------------------------------------------
|
|
# Both distributions must be conditioned on byte-identical prefixes. If the
|
|
# abliterated checkpoint carries a different tokenizer or chat template, every
|
|
# KL below is a comparison of two different questions.
|
|
tok_files = ["tokenizer.json", "tokenizer_config.json"]
|
|
digests = {f: (file_md5(ref_dir / f), file_md5(cand_dir / f)) for f in tok_files}
|
|
mismatched = {f: v for f, v in digests.items() if v[0] != v[1]}
|
|
if mismatched and not args.skip_gates:
|
|
print(f"\n!! tokenizer differs between ref and cand: {sorted(mismatched)} — the two "
|
|
f"models would be conditioned on different prefixes and the KL would be "
|
|
f"meaningless.", file=sys.stderr)
|
|
sys.exit(6)
|
|
print("tokenizer identity: " + ("MATCH" if not mismatched else f"MISMATCH {sorted(mismatched)}"))
|
|
|
|
# --- corpus ---------------------------------------------------------------
|
|
harmless, harmful, prov = load_evaluation(
|
|
n_harmless=args.n_harmless, n_harmful=args.n_harmful, seed=args.seed,
|
|
calib_harmless_n=args.calib_harmless_n,
|
|
calib_harmless_seed=args.calib_harmless_seed)
|
|
print(f"eval corpus: {len(harmless)} harmless (held out from calibration), "
|
|
f"{len(harmful)} harmful (reserved test split)")
|
|
if args.seed == args.calib_harmless_seed:
|
|
print("!! --seed equals the capture's harmless seed; the exclusion still holds "
|
|
"(indices are removed from the pool) but say so in the writeup.", file=sys.stderr)
|
|
|
|
prompts = harmless + harmful
|
|
split = len(harmless)
|
|
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
if args.dry_run:
|
|
# Prove the rendering and the tokenizer round-trip on the real checkpoint
|
|
# before spending a seat-down window. Tokenizer only — no weights touched.
|
|
tok = AutoTokenizer.from_pretrained(ref_dir)
|
|
for mode, thinking in MODES.items():
|
|
text = render(tok, prompts[0], thinking)
|
|
n_tok = len(tok(text)["input_ids"])
|
|
print(f"\n [{mode}] {n_tok} tokens, tail: {text[-60:]!r}")
|
|
lens = [len(tok(render(tok, p, False))["input_ids"]) for p in prompts]
|
|
print(f"\n prompt lengths (answer mode): min {min(lens)} median "
|
|
f"{int(statistics.median(lens))} max {max(lens)}")
|
|
print(f" first harmless: {harmless[0][:90]!r}")
|
|
print(f" first harmful: {harmful[0][:90]!r}")
|
|
print("\ndry-run complete — corpus, gates and rendering verified; "
|
|
"nothing loaded, no GPU touched.")
|
|
return
|
|
|
|
cache_dir = Path(args.cache) if args.cache else Path(args.out).parent
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
|
cache_path = {lbl: cache_dir / f"{Path(args.out).stem}.{lbl}.pt"
|
|
for lbl in ("ref", "cand")}
|
|
|
|
def run_stage(model_dir: Path, label: str, measure_floor: bool):
|
|
print(f"\n=== {label}: {model_dir}")
|
|
|
|
# --- headroom gate, BEFORE the load --------------------------------
|
|
# `device_map="auto"` does not fail when the card is too small; it
|
|
# quietly spills modules to CPU, and only the residency gate below would
|
|
# catch that — after paying for the load. Worse, the second stage runs on
|
|
# whatever the first stage gave back, so a release regression shows up
|
|
# here as offloading rather than as an error. Size the requirement from
|
|
# the checkpoint's own headers so it stays true if the weights change.
|
|
need = weights_mib(model_dir)
|
|
free = free_vram_mib()
|
|
print(f" weights {need:,.0f} MiB (text-only, from shard headers); "
|
|
f"free {free:,.0f} MiB")
|
|
if free < need + 2048:
|
|
print(f"\n!! insufficient VRAM: {free:,.0f} MiB free, need {need:,.0f} + 2048 "
|
|
f"MiB headroom. Both GPU0 seats must be stopped — freeing only one "
|
|
f"leaves ~50,900 MiB and this model's text weights are ~51,300 MiB. "
|
|
f"If this is the second stage, the first stage did not release.",
|
|
file=sys.stderr)
|
|
sys.exit(10)
|
|
|
|
collected, floor, vocab = _load_and_collect(model_dir, label, measure_floor)
|
|
torch.save(
|
|
{"logprobs": {m: collected[m]["logprobs"] for m in MODES},
|
|
"ids": {m: collected[m]["ids"] for m in MODES},
|
|
"floor": floor, "vocab": vocab, "model_dir": str(model_dir)},
|
|
cache_path[label])
|
|
print(f" cached -> {cache_path[label]} "
|
|
f"({cache_path[label].stat().st_size / 2**20:,.0f} MiB)")
|
|
|
|
def _load_and_collect(model_dir: Path, label: str, measure_floor: bool):
|
|
tok = AutoTokenizer.from_pretrained(model_dir)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_dir, dtype=torch.bfloat16, device_map="auto", attn_implementation="sdpa")
|
|
model.eval()
|
|
device = next(model.parameters()).device
|
|
|
|
# --- residency gate ---------------------------------------------------
|
|
# Sharding this architecture across the two Blackwells zeroes the residual
|
|
# stream past the device boundary; the logits would decode to garbage while
|
|
# every gate below the boundary still passed. Diagnosed 2026-08-20.
|
|
#
|
|
# Read the placement off the PARAMETERS, not off `hf_device_map`. The map
|
|
# is empty whenever transformers puts the whole model on one device, so a
|
|
# map-based check reports "(unsharded)" both when everything is fine and
|
|
# when there is no map to inspect — it cannot fail, which makes it not a
|
|
# gate. Parameter devices are ground truth in every case.
|
|
devices = {str(p.device) for p in model.parameters()}
|
|
offloaded = sorted(d for d in devices if d.startswith(("cpu", "meta", "disk")))
|
|
if len(devices) > 1 or offloaded:
|
|
print(f"\n!! residency gate FAILED — parameters span {sorted(devices)}. "
|
|
f"Sharding this architecture zeroes the residual stream past the device "
|
|
f"boundary and the logits decode to garbage. Fix: CUDA_VISIBLE_DEVICES=0 "
|
|
f"with BOTH GPU0 seats stopped.", file=sys.stderr)
|
|
sys.exit(8)
|
|
print(f" residency: all parameters on {sorted(devices)[0]}, "
|
|
f"free VRAM after load: {free_vram_mib():,.0f} MiB")
|
|
|
|
collected = collect(model, tok, prompts, device, label)
|
|
|
|
floor = None
|
|
if measure_floor and args.noise_floor_n:
|
|
# Re-run a subset through the SAME model. This is not a formality: it
|
|
# validates the whole scoring path end to end. A deterministic stack
|
|
# must return exactly 0 here, so any non-zero value indicts this
|
|
# script (or the stack's determinism) before it indicts the surgery.
|
|
sub = prompts[:args.noise_floor_n]
|
|
again = collect(model, tok, sub, device, f"{label}-repeat")
|
|
floor = {}
|
|
for mode in MODES:
|
|
k, _, tv, top1 = divergences(
|
|
collected[mode]["logprobs"][:len(sub)], again[mode]["logprobs"])
|
|
floor[mode] = {"n": len(sub), "kl_mean": float(k.mean()),
|
|
"kl_max": float(k.max()), "tv_max": float(tv.max()),
|
|
"top1_agreement": float(top1.double().mean())}
|
|
print(f" noise floor (self-KL): "
|
|
+ ", ".join(f"{m} max {v['kl_max']:.3e}" for m, v in floor.items()))
|
|
|
|
return collected, floor, collected["answer"]["logprobs"].shape[-1]
|
|
|
|
# --- ONE PROCESS PER MODEL -----------------------------------------------
|
|
#
|
|
# This is not fastidiousness; it is the only mechanism that works. Both
|
|
# in-process teardowns were tried and MEASURED on 2026-08-20:
|
|
#
|
|
# `del model` + `gc.collect()` + `torch.cuda.empty_cache()` -> 45,287 MiB free
|
|
# the same, with the model confined to an inner frame that exits -> 45,287 MiB free
|
|
#
|
|
# i.e. the ~51,300 MiB of weights were still resident both times. The first
|
|
# (pre-gate) run only survived because PyTorch's allocator hit OOM during the
|
|
# second load, ran a collection itself, and retried — the second model landed
|
|
# on the card by rescue, not by design. Relying on that is how you end up
|
|
# silently offloaded to CPU, which on this architecture does not error: it
|
|
# zeroes the residual stream past the boundary and returns confident garbage.
|
|
#
|
|
# A process exit releases the CUDA context unconditionally, so each model gets
|
|
# its own. The stages hand their first-token log-probs to disk (~715 MiB per
|
|
# model) and the parent scores from the caches. Reusable, too: re-measuring a
|
|
# different candidate against this same reference skips the ref stage.
|
|
if args.stage in ("ref", "cand"):
|
|
run_stage(ref_dir if args.stage == "ref" else cand_dir,
|
|
args.stage, measure_floor=(args.stage == "ref"))
|
|
return
|
|
|
|
if args.stage == "all":
|
|
import subprocess
|
|
for label in ("ref", "cand"):
|
|
cmd = [sys.executable, str(Path(__file__).resolve()),
|
|
"--ref", str(ref_dir), "--cand", str(cand_dir), "--out", args.out,
|
|
"--n-harmless", str(args.n_harmless), "--n-harmful", str(args.n_harmful),
|
|
"--seed", str(args.seed),
|
|
"--calib-harmless-n", str(args.calib_harmless_n),
|
|
"--calib-harmless-seed", str(args.calib_harmless_seed),
|
|
"--noise-floor-n", str(args.noise_floor_n),
|
|
"--cache", str(cache_dir), "--stage", label]
|
|
if args.skip_gates:
|
|
cmd.append("--skip-gates")
|
|
rc = subprocess.run(cmd).returncode
|
|
if rc != 0:
|
|
print(f"\n!! {label} stage exited {rc}; not scoring a partial run.",
|
|
file=sys.stderr)
|
|
sys.exit(rc)
|
|
print(f" [{label} stage process exited; CUDA context released]")
|
|
|
|
# --- load the stage caches ------------------------------------------------
|
|
missing = [p for p in cache_path.values() if not p.exists()]
|
|
if missing:
|
|
print(f"\n!! missing stage cache(s): {missing} — run --stage ref and "
|
|
f"--stage cand first.", file=sys.stderr)
|
|
sys.exit(11)
|
|
blobs = {lbl: torch.load(p, weights_only=False) for lbl, p in cache_path.items()}
|
|
ref = {m: {"logprobs": blobs["ref"]["logprobs"][m], "ids": blobs["ref"]["ids"][m]}
|
|
for m in MODES}
|
|
cand = {m: {"logprobs": blobs["cand"]["logprobs"][m], "ids": blobs["cand"]["ids"][m]}
|
|
for m in MODES}
|
|
noise_floor = blobs["ref"]["floor"]
|
|
ref_vocab, cand_vocab = blobs["ref"]["vocab"], blobs["cand"]["vocab"]
|
|
|
|
if ref_vocab != cand_vocab and not args.skip_gates:
|
|
print(f"\n!! vocab size differs: ref {ref_vocab} vs cand {cand_vocab}", file=sys.stderr)
|
|
sys.exit(6)
|
|
|
|
# --- tokenization equality gate ------------------------------------------
|
|
for mode in MODES:
|
|
for i, (a, b) in enumerate(zip(ref[mode]["ids"], cand[mode]["ids"])):
|
|
if a != b and not args.skip_gates:
|
|
print(f"\n!! prompt {i} ({mode}) tokenized differently between the two "
|
|
f"stages ({len(a)} vs {len(b)} tokens) — the KL would compare "
|
|
f"distributions over different prefixes.", file=sys.stderr)
|
|
sys.exit(7)
|
|
|
|
# --- score ----------------------------------------------------------------
|
|
results = {}
|
|
for mode in MODES:
|
|
kl_f, kl_r, tv, top1 = divergences(ref[mode]["logprobs"], cand[mode]["logprobs"])
|
|
results[mode] = {
|
|
"harmless": summarize(kl_f[:split], kl_r[:split], tv[:split], top1[:split]),
|
|
"harmful": summarize(kl_f[split:], kl_r[split:], tv[split:], top1[split:]),
|
|
}
|
|
hl = results[mode]["harmless"]["kl_median"]
|
|
hf = results[mode]["harmful"]["kl_median"]
|
|
results[mode]["selectivity_median_ratio"] = (hf / hl) if hl > 0 else None
|
|
|
|
payload = {
|
|
"ref": str(ref_dir), "cand": str(cand_dir),
|
|
"metric": "first-token KL(ref || cand), full vocabulary, float64",
|
|
"measured_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
|
|
"corpus": prov,
|
|
"tokenizer_digests": {f: {"ref": v[0], "cand": v[1]} for f, v in digests.items()},
|
|
"vocab": ref_vocab,
|
|
"noise_floor_self_kl": noise_floor,
|
|
"modes": results,
|
|
"env": {"torch": torch.__version__,
|
|
"cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
|
|
"alloc_conf": alloc or None},
|
|
}
|
|
Path(args.out).write_text(json.dumps(payload, indent=2))
|
|
|
|
print("\n" + "=" * 72)
|
|
print(f"first-token KL(ref || cand) ref={ref_dir.name}")
|
|
print(f" cand={cand_dir.name}")
|
|
for mode in MODES:
|
|
r = results[mode]
|
|
print(f"\n [{mode} mode]")
|
|
for cls in ("harmless", "harmful"):
|
|
s = r[cls]
|
|
print(f" {cls:<9} n={s['n']:<4} median {s['kl_median']:.4f} "
|
|
f"mean {s['kl_mean']:.4f} p95 {s['kl_p95']:.4f} max {s['kl_max']:.4f} "
|
|
f"top1-agree {s['top1_agreement']:.1%}")
|
|
ratio = r["selectivity_median_ratio"]
|
|
shown = f"{ratio:.1f}x" if ratio else "n/a"
|
|
print(f" selectivity (harmful/harmless median KL): {shown}")
|
|
print(f"\nwrote {args.out}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|