feat(coldfusion-abliteration): first-token KL measured — 28.4x selectivity, harmless median 0.0211

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).
This commit is contained in:
2026-08-20 13:00:39 -07:00
parent 8c354a0e79
commit 1b3fb270e7
7 changed files with 823 additions and 10 deletions
+8
View File
@@ -154,6 +154,14 @@ The recipe is a drop-in for the front half of the House quant pipeline:
3994; gate the result on **MTP acceptance ≳40%, not KL** (KL misled us once —
`reference_abliteration_mtp_lessons`).
4. Verify vision byte-identical, refusals down, PPL not blown, no catatonia.
**Measure first-token KL as a *fidelity* number** (`kl_divergence.py`,
bf16-vs-bf16, held-out prompts) — it does not replace the acceptance gate in
step 3, and it is not a pass/fail on its own. Report it **split by prompt
class**: a single averaged KL over a mixed corpus is close to meaningless,
because the metric is supposed to be large on harmful prompts and small on
benign ones. The ratio is the interesting quantity. Cold-Fusion L35 measured
**0.0211 median harmless / 0.5996 median harmful = 28.4× selectivity**, on a
stack whose self-KL noise floor is exactly 0.0.
5. NVFP4-quantize in-house (mixed W4A4 + FP8-attn/lm_head —
`model-quantization-playbook.md`). **Foot-gun the GGUF card itself flags:
the imatrix does not cover the MTP block** — so a GGUF requant path leaves
+39
View File
@@ -310,6 +310,45 @@ determinism check with a **magnitude** check (residual norms should grow smoothl
exact 0.0 mid-stack is impossible) and, where you can, a **coherence** check (generate 40 tokens and
read them).
### 3.12 ⭐⭐ You cannot free a 27B model in-process — give each model its own process
Any A/B that loads two large checkpoints in sequence (KL, logit diffing, teacher-vs-student)
will try to release the first before loading the second. **On this stack, it does not work.**
Measured 2026-08-20 on Qwen3.8-27B bf16, free VRAM after each attempt:
| teardown | free VRAM |
|---|---|
| `del model` + `gc.collect()` + `torch.cuda.empty_cache()` | 45,287 MiB |
| same, with the model confined to an inner frame that exits | 45,287 MiB |
| **the process exits** | **97,247 MiB** |
The ~51,300 MiB of weights stayed resident through both in-process teardowns. The first
run survived only because **PyTorch's allocator hit OOM on the second load, ran a collection
itself, and retried** — the second model landed by rescue, not by design. That is not a
release strategy: on an architecture where a silent CPU offload does not raise (§3.9), the
day the retry does not fire you get confident garbage instead of an error.
**Do this instead:** one process per model, hand results to disk between them
(first-token log-probs for a 250k vocab are ~715 MiB per model — nothing), and gate each
stage on free VRAM *before* the load. Reference implementation:
`services/coldfusion-abliteration/kl_divergence.py` (`--stage ref|cand|score`).
Two gate corollaries learned in the same session:
- **⭐ A residency gate that reads `hf_device_map` cannot fail.** The map is **empty**
whenever transformers puts the whole model on one device, so the check reports
"unsharded" both when everything is fine and when there is nothing to inspect. Read
`{p.device for p in model.parameters()}` — ground truth in every case. (Generalises
[[feedback_assert_effective_value_not_substring]]: presence of a passing check is not
evidence of a check that can fail.)
- **⭐ Size VRAM from the checkpoint's own headers, never from a remembered figure.** A
runbook carried "bf16 is 50 GB"; the real number was 50.10 **GiB** = 51,300 MiB of
text-only weights. That 3.7 GB unit error is exactly the difference between "stop one
co-tenant" and "stop both", and it cost an aborted window. Sum the safetensors header
offsets (excluding tensors the loader class won't instantiate — vision, MTP); read only
the 8-byte length prefix + JSON header, never `safe_open`, which mmaps the whole shard
and ENOMEMs on ZFS (§ *Avoid mmap on `/tank`*).
---
## 4. Pipeline shape
@@ -66,6 +66,63 @@ vision byte-identical (Δ0.0), 735/735 others untouched.**
- Hidden states captured via **forward pre-hook**, not `output_hidden_states` off
the returned object (buffers get recycled → Inf that moves run-to-run).
## ✅ KL divergence measured (2026-08-20, third session)
`services/coldfusion-abliteration/kl_divergence.py` — first-token KL(stock ‖ L35)
over the full 248,320-token vocabulary, bf16 vs bf16, on prompts the direction was
never fitted on (256 harmless held out of the alpaca pool by replaying and
subtracting calibration's own draw; 104 harmful from the reserved test split).
| mode | class | median | mean | p95 | top-1 agreement |
|---|---|---|---|---|---|
| answer | harmless | **0.0211** | 0.0364 | 0.1219 | 89.8% |
| answer | harmful | **0.5996** | 0.6992 | 1.6937 | 55.8% |
| think | harmless | 0.0042 | 0.0066 | 0.0205 | 94.5% |
| think | harmful | 0.3068 | 0.3186 | 0.4689 | 57.7% |
Run twice — single-process, then through the two-process design — and **all 720
per-prompt KL values came back bit-identical**, so these figures are stable across
processes, not just within one.
**Selectivity 28.4× (answer) / 72.8× (think).** The surgery moves the model hard on
refusal-triggering prompts and barely at all on benign ones — on held-out harmless
prompts the abliterated model still picks the same first token 89.8% of the time.
**Self-KL noise floor: exactly 0.0**, so none of this is bf16 jitter, and the
scoring path is validated end to end. Reverse KL on harmful/answer is 1.43 vs
forward 0.70 — the mass-where-stock-had-none asymmetry that is abliteration's
signature.
Against the Heretic reference figures (0.1191 prior seat, **0.0759 the current
`absolute-heresy` seat**) ours is materially gentler — but ⚠️ **that is not a
head-to-head**: those are Heretic's own optimizer output on a different base with
its own harmless set and template. Order-of-magnitude only. A real comparison
means re-measuring the incumbent through this script (one more GPU window).
Consistent with [[reference_abliteration_mtp_lessons]]: KL is a **fidelity**
number here, not the viability gate — that remains MTP acceptance (59.1%).
### Three durable process lessons from the measurement
1. **★ Report abliteration KL SPLIT BY PROMPT CLASS.** A single averaged KL over a
mixed corpus is close to meaningless, because the metric is *supposed* to be
large on harmful prompts and small on benign ones — averaging them together
lets a blunt abliteration and a surgical one produce the same number. The
selectivity ratio is the quantity with information in it.
2. **★ "50 GB" was 50.10 GiB mislabelled — and the 3.7 GB gap changed the runbook.**
Text-only weights are **51,300 MiB**; GPU0's tenants are meromero 50,072 and gen
46,304, so freeing *either alone* leaves ~50,933 MiB — ~400 MiB short. The
runbook's "only gen must go" was wrong. **Both seats must stop.** Size VRAM from
the safetensors headers, never from a remembered gigabyte figure.
3. **★ You cannot release a 27B model in-process; give each model its own process.**
Measured twice: `del model` + `gc.collect()` + `empty_cache()` left free VRAM at
45,287 MiB, and so did confining the model to an inner frame that exits. The
first run only worked because PyTorch's allocator hit OOM on the second load,
collected, and retried — *rescue, not design*. On this architecture a silent
CPU offload does not error; it zeroes the residual stream past the boundary and
returns confident garbage. Also: the old residency gate read `hf_device_map`,
which is **empty when the model fits on one device** — so it printed
"(unsharded)" and could never fail. It now reads parameter devices directly.
## Still owed before this is a gen-seat candidate
- Canonical refusal re-profile via `services/refusal-probe/` (not the ad-hoc
+1 -1
View File
@@ -150,7 +150,7 @@ _As of 2026-08-20 — continued from the 08-19 infra session (DNS/.internal, wat
- **🟢 WT #401 (fd-leak deadlock) CLOSED 2026-08-17 — one ping still owed.** worldtree-dev closed it on our demo verify. Layers: **(a) their `e41b139`** pins `ulimits: nofile 65536/65536` in the worldtree compose anchor — **demo VERIFIED** (api + matrix recreated 22:55:34Z, `ulimit -Sn`=65536); **personal/pinned are covered-not-verified**, they inherit at their next promotion/recreate. **(b) our host floor is STAGED, NOT ACTIVE** — `/etc/docker/daemon.json` on corviduo-dev carries `default-ulimits nofile 65536/65536` but **`default-ulimits` is NOT SIGHUP-reloadable** (measured on 29.4.3: post-reload the daemon's own "Reloaded configuration" log omits it and a fresh container still reports 1024). Activation needs a full dockerd restart = bounces all 13 containers; **worldtree-dev explicitly does NOT want one**, and `live-restore:true`-then-restart is PARKED as a separate host-side improvement for the operator to rule on, never folded into #401. Playbook `playbooks/corviduo-dev-docker-default-ulimits.yaml` (verify step 3 fails BY DESIGN until a restart). Hourly fd tripwire on corviduo-dev stays armed. **⏳ OWED: ping worldtree-dev in thread `01M08QQ655XD6VKEV7MA9GX0NS` once worldtree-personal recreates and 65536 is confirmed there.** Commit `7f3f265`.
- **🟢 COLD-FUSION ABLITERATION — LANDED 2026-08-20 (the real work; abliterated model WORKS, verify + quant still owed).** Abliterated `DavidAU/Qwen3.8-27B-Cold-Fusion-GAIN-V1.1` with the **Robinson formula** (`docs/pfi/abliteration-recipe-qwen38.md`; harness `services/coldfusion-abliteration/`). Output `ana-ml2:/tank/aimodels/qwen38-27b-coldfusion-abliterated-L35-bf16`; bitwise-verified 131/131 targets changed, 333/333 vision byte-identical, 735/735 others untouched. A/B vs stock: explicit-sexual + graphic-torture (the measured stock refusal surface) go refused→complied, self-harm guardrail survives, coherence intact — the Robinson design point. **THREE first-session diagnoses were wrong, all corrected:** (1) **layer selection by two-template |cos| agreement is misleading on a merged base** — its pick (L18) was the *worst*-separating layer and abliterating there was a measured no-op; replaced with harmful/harmless **separation** (Cohen's d/AUC) gated on the sink screen → **L35** (d9.35, AUC0.9997, sink0.094%); (2) **"bf16 NaNs, use fp32" was a misdiagnosis** — the NaN was multi-GPU sharding + `expandable_segments`, not precision; bf16 on ONE GPU is deterministic+coherent at 50 GB, 4.3× faster (now gated, `CUDA_VISIBLE_DEVICES=0`); (3) **corpus-size hypothesis falsified** (8→416 moved agreement 0.594→0.624, nothing). Write is now **shard surgery** (no model object — `AutoModelForCausalLM` is text-only and would drop vision + skip the in-band MTP edit). **🎯 THESIS PROVEN 2026-08-20 (same session, later):** quantized L35 → mixed NVFP4 (`/tank/aimodels/qwen38-27b-coldfusion-L35-nvfp4-mixed`, 22.5 GB) and measured **MTP acceptance 59.1% median** (5165%) — clears the ≳40% gate AND **beats the incumbent Heretic seat's ~47%.** So Robinson's **in-band** MTP abliteration accepts BETTER than Heretic's byte-identical graft — the additive delta the experiment tested, positive. Abliteration survives the quant (creative refusals drop, self-harm guardrail intact, coherent); decode 118.7 tok/s (faster, image-confounded). Env foot-guns hardened: quant venv drifted (transformers 5.10/llmcompressor 0.12 no longer delegate `num_attention_heads` to text_config → promote-then-restore in `quant_mixed_nvfp4.py`); a <23GB quant saves single-file no-index → build index from the safetensors header (never `safe_open`, ENOMEMs on ZFS). Commit `725c8fd`. **⚠ NOT CUT OVER** — incumbent gen seat untouched; making L35 the `gen` seat is a **separate operator decision** needing the full Stage-3 gate (PPL, prefill, surface 6/6, refusal-probe battery) + real multi-turn hold. **Do NOT delete** `qwen38-27b-heresy-nvfp4-mixed`. Full saga → `persistent-memory.d/2026-08-20-coldfusion-abliteration-landed.md`.
- **🟢 COLD-FUSION ABLITERATION — LANDED 2026-08-20 (the real work; abliterated model WORKS, verify + quant still owed).** Abliterated `DavidAU/Qwen3.8-27B-Cold-Fusion-GAIN-V1.1` with the **Robinson formula** (`docs/pfi/abliteration-recipe-qwen38.md`; harness `services/coldfusion-abliteration/`). Output `ana-ml2:/tank/aimodels/qwen38-27b-coldfusion-abliterated-L35-bf16`; bitwise-verified 131/131 targets changed, 333/333 vision byte-identical, 735/735 others untouched. A/B vs stock: explicit-sexual + graphic-torture (the measured stock refusal surface) go refused→complied, self-harm guardrail survives, coherence intact — the Robinson design point. **THREE first-session diagnoses were wrong, all corrected:** (1) **layer selection by two-template |cos| agreement is misleading on a merged base** — its pick (L18) was the *worst*-separating layer and abliterating there was a measured no-op; replaced with harmful/harmless **separation** (Cohen's d/AUC) gated on the sink screen → **L35** (d9.35, AUC0.9997, sink0.094%); (2) **"bf16 NaNs, use fp32" was a misdiagnosis** — the NaN was multi-GPU sharding + `expandable_segments`, not precision; bf16 on ONE GPU is deterministic+coherent at 50 GB, 4.3× faster (now gated, `CUDA_VISIBLE_DEVICES=0`); (3) **corpus-size hypothesis falsified** (8→416 moved agreement 0.594→0.624, nothing). Write is now **shard surgery** (no model object — `AutoModelForCausalLM` is text-only and would drop vision + skip the in-band MTP edit). **🎯 THESIS PROVEN 2026-08-20 (same session, later):** quantized L35 → mixed NVFP4 (`/tank/aimodels/qwen38-27b-coldfusion-L35-nvfp4-mixed`, 22.5 GB) and measured **MTP acceptance 59.1% median** (5165%) — clears the ≳40% gate AND **beats the incumbent Heretic seat's ~47%.** So Robinson's **in-band** MTP abliteration accepts BETTER than Heretic's byte-identical graft — the additive delta the experiment tested, positive. Abliteration survives the quant (creative refusals drop, self-harm guardrail intact, coherent); decode 118.7 tok/s (faster, image-confounded). Env foot-guns hardened: quant venv drifted (transformers 5.10/llmcompressor 0.12 no longer delegate `num_attention_heads` to text_config → promote-then-restore in `quant_mixed_nvfp4.py`); a <23GB quant saves single-file no-index → build index from the safetensors header (never `safe_open`, ENOMEMs on ZFS). Commit `725c8fd`. **⚠ NOT CUT OVER** — incumbent gen seat untouched; making L35 the `gen` seat is a **separate operator decision** needing the full Stage-3 gate (PPL, prefill, surface 6/6, refusal-probe battery) + real multi-turn hold. **Do NOT delete** `qwen38-27b-heresy-nvfp4-mixed`. **📐 KL MEASURED 2026-08-20 (third session) — the surgery is SELECTIVE.** `kl_divergence.py` (new, beside the harness): first-token KL(stock‖L35) over the full 248,320 vocab, bf16-vs-bf16, held-out prompts. **Answer mode: harmless median 0.0211 / mean 0.0364, harmful median 0.5996 → 28.4× selectivity** (think mode 0.0042 / 0.3068 → 72.8×); top-1 agreement on benign prompts stays **89.8%**; **self-KL noise floor exactly 0.0**, so every digit is signal. Reverse KL on harmful is 1.43 vs forward 0.70 — the mass-where-stock-had-none asymmetry that is abliteration's signature. Vs the Heretic reference figures (0.1191 prior seat, **0.0759 the live `absolute-heresy` seat**) ours looks materially gentler, but ⚠️ **NOT a head-to-head** — those are Heretic's own optimizer output on a different base with its own harmless set/template; a real comparison needs the incumbent re-measured through this script (one more window). KL stays a **fidelity** number, not the gate (that's MTP acceptance, 59.1%). Run cost **2m40s**, both GPU0 seats down. **Three durable process lessons:** (1) report abliteration KL **split by prompt class** — a mixed average is meaningless since the metric is meant to be big on harmful and small on benign; (2) the runbook's "bf16 is 50 GB / only gen must go" was **50.10 GiB mislabelled** — text weights are 51,300 MiB and freeing either GPU0 seat alone leaves ~50,933, so **both must stop**; (3) **a 27B model cannot be released in-process** (`del`+`gc`+`empty_cache` and frame-exit both left 45,287 MiB free; only process exit gave the card back) → one process per model, and the old residency gate read `hf_device_map`, which is empty on single-device loads and therefore **could never fail**. → playbook §3.12. Full saga → `persistent-memory.d/2026-08-20-coldfusion-abliteration-landed.md`.
- **🟢 esh-pve-nas — MIGRATION DONE 2026-08-18. Root is `nvme/ROOT/pve-1` on mirrored NVMe; the USB DOM is out of the runtime I/O path.** All five guests healthy, three pools ONLINE, system `running`, ext4 `pve-root` intact+unmounted as rollback with its own kernel. Boot config: `saved_entry=pve-zfs-root`, no `next_entry`; if grubenv were unreadable GRUB falls to entry 0 which the `/etc/default/grub.d/zfs-root.cfg` drop-in also points at ZFS. **`zfs-import-cache.service` is now the active import path** (the all-three-pools cachefile fix working as intended); vestigial `zfs-import@nvme.service` disabled — it failed every boot as redundant. ⚠ **Device letters shift across reboots** (DOM was `sdq`, now `sdl`) — never key anything to a bare `sdX` here. ⚠ **NO auto-fallback on a failed boot, and no IPMI/BMC/serial** — grubenv on LVM is readable but not writable by GRUB, so `grub-reboot`'s one-shot degrades to a sticky default (verified: `next_entry` survived the boot that consumed it). Recovery = pick the ROLLBACK entry at the console. **PATCHED 2026-08-18: 225 packages installed, pve-manager 8.4.11 -> 8.4.20, corosync 3.1.9 -> 3.1.10-pve2, kernel 6.8.12-42 staged on the /boot LV. dpkg clean, no unapplied conffiles, cluster quorate, 6/6 verify. ⏳ REBOOT DEFERRED at operator request — host still runs 6.8.12-13 until a chosen window; `GRUB_DEFAULT=0` means entry 0 is already the -42 entry with the correct `root=ZFS=nvme/ROOT/pve-1`, so the reboot is the only remaining step. Rollback for the upgrade is the ZFS snapshot `nvme/ROOT/pve-1@pre-upgrade-20260818T141652Z` (409M) — `zfs rollback -r <snap> && reboot`. Second confirmation reboot ALREADY DONE (2026-08-18, booted ZFS from GRUB_DEFAULT=0 with no one-shot). ⏳ Still outstanding: refresh the off-box DOM image, since `/boot` changed.** esh-pve is FULLY done (8.4.20 + kernel 6.8.12-42 + corosync 3.1.10, rebooted, quorate).
⚠⚠ **THE WINDOW COST AN UNPLANNED OUTAGE, caused by our own tooling, not the migration.** The staging chroot did `mount --rbind /dev` + `/sys` with **no `--make-rslave`**; on systemd `/` is *shared*, so the cutover's `umount -R` **propagated back into the live host** and stripped the real `/sys/fs/cgroup`, `/dev/pts`, `/dev/shm`. logind could then create no sessions: ping fine, TCP fine, **SSH authenticates**, resident daemons keep serving (pveproxy returned clean 401s) — and **every new exec hangs, including `/sbin/reboot`**, so the reboot never ran. **It is a near-perfect impostor of failing root-disk I/O**, and I misdiagnosed it as the USB DOM dying and told the operator to walk to the machine. **Operator caught it** — the DOM had been fine for years and the wedge began right after a change. The settling evidence was in `dmesg` all along: `[16.00] [sdq] Attached SCSI removable disk` (clean, no errors) and a last-line timestamp of **`[12114881]` = 140 days = the ORIGINAL boot** — the machine had never rebooted. My down-detector never once reported the host down and I read that as a fast reboot rather than *no* reboot. Recovered with **no console access** by hammering an idempotent cgroup2/devpts/shm remount into the brief windows where exec succeeded. Zero data loss. **RULES: (1) always `--make-rslave` after `--rbind` (playbook now guards on `PROPAGATION != shared`); (2) a reboot is not confirmed until the host is observed DOWN — poll for disappearance, not reappearance; (3) before blaming hardware for a wedge that started right after a change, get `dmesg` and check the boot timestamp.**
+134 -9
View File
@@ -124,8 +124,10 @@ RUN="sudo -u llmuser env HF_HUB_OFFLINE=1 CUDA_VISIBLE_DEVICES=0 \
# confirms the recipe maps onto THIS checkpoint's names.
$RUN --dry-run
# --- capture needs GPU0 to itself: bf16 is 50 GB, so only gen must go ---
sudo docker stop -t 60 vllm-gen
# --- capture needs GPU0 to itself. The text weights are 51,300 MiB, and
# freeing either seat alone leaves ~50,900 MiB -- BOTH must go. See
# gotcha 3; the "only gen" line that used to be here was a unit error. ---
sudo docker stop -t 60 vllm-gen vllm-meromero-rp
cp $M/refusal-direction.pt $M/refusal-direction.pt.bak # capture overwrites it
# 2. CONTROL RUN — the legacy 8/8 set. Reproduces layer 22, |cos| 0.5944, sink
@@ -252,6 +254,101 @@ a separate operator decision needing the full Stage-3 gate (PPL, prefill, surfac
delete-too-early / multi-day-degeneration lesson). The thesis is proven; the
cutover is a distinct call.
## ✅ KL RESULT — the surgery is highly selective (2026-08-20)
`kl_divergence.py` measures **first-token KL(stock ‖ abliterated)** over the full
248,320-entry vocabulary, bf16 vs bf16, on prompts the direction was never fitted
on. Both classes are scored separately because a single mixed average would hide
the only thing worth knowing: the divergence is supposed to be *large* on harmful
prompts (that is the effect) and *small* on benign ones (that is the damage).
| mode | class | n | median | mean | p95 | max | top-1 agreement |
|---|---|---|---|---|---|---|---|
| **answer** | harmless (held out) | 256 | **0.0211** | **0.0364** | 0.1219 | 0.2654 | 89.8% |
| **answer** | harmful (reserved test) | 104 | **0.5996** | 0.6992 | 1.6937 | 1.9920 | 55.8% |
| think | harmless (held out) | 256 | 0.0042 | 0.0066 | 0.0205 | 0.0392 | 94.5% |
| think | harmful (reserved test) | 104 | 0.3068 | 0.3186 | 0.4689 | 0.5298 | 57.7% |
**Selectivity — harmful/harmless median KL — is 28.4× in answer mode and 72.8× in
think mode.** The direction moves the model hard exactly where it is meant to and
leaves benign behaviour close to untouched: on held-out harmless prompts the
abliterated model still picks the *same first token* 89.8% of the time.
**Noise floor: exactly 0.0** in both modes (32 prompts re-run through the same
model, self-KL). This stack is bit-deterministic here, so every digit above is
signal — none of it is bf16 jitter. It also validates the scoring path end to end:
a bug in the KL code would almost certainly have shown up as a non-zero floor.
**The reverse-KL asymmetry is the abliteration's signature.** On harmful prompts
in answer mode, KL(stock‖abl) is 0.70 but KL(abl‖stock) is **1.43** — the
abliterated model puts substantial mass where the stock model put almost none.
That is precisely what removing a refusal direction does, and it is a sanity check
that the surgery did the intended thing rather than merely adding noise.
### Against the Heretic reference figures — favourable, with a caveat
| model | first-token KL, harmless | abliteration method |
|---|---|---|
| `JonathanColetti/Qwen3.8-27B-Uncensored` (prior gen seat) | 0.1191 | Heretic, out-of-band MTP |
| `absolute-heresy` (**current** gen seat) | 0.0759 | Heretic v1.4.0 + SOMPOA |
| **Cold-Fusion L35 (ours)** | **0.0211 median / 0.0364 mean** | Robinson, in-band MTP |
⚠️ **Not a head-to-head.** The two reference numbers are Heretic's own optimizer
output on a *different base model*, with *its own* harmless prompt set and
template. Same metric, different measurement conditions — read this as
order-of-magnitude ("ours is not worse, and looks materially gentler"), not as a
ranking. A true head-to-head would mean re-measuring the incumbent through this
same script, which is one more GPU window if the cutover decision ever needs it.
Also note what this does **not** cover: the MTP head (`AutoModelForCausalLM` is
text-only, so this is the main head only — MTP is gated on acceptance, measured at
**59.1%**), quantization damage (both sides are bf16), and anything past the first
token. Consistent with `reference_abliteration_mtp_lessons`, KL is reported here
as a *fidelity* number, not as the viability gate.
**Reproducibility: exact.** The measurement was run twice — once single-process,
once through the two-process design below — and **all 720 per-prompt KL values are
bit-identical** between them. Combined with the 0.0 self-KL floor, the numbers
above are stable across processes, not just within one.
Artifacts: `kl-L35.json` (+ `kl-L35-rerun.json`, the reproducibility check) and the
two `.ref.pt` / `.cand.pt` log-prob caches, beside the harness on ana-ml2. Run
cost: **2m40s** single-process, **3m26s** two-process, both seats down.
```bash
# free, no GPU, safe with the seats up — run this first
$V $P/kl_divergence.py --ref $M --cand $A --out $P/kl-L35.json --dry-run
# the real thing: needs BOTH GPU0 seats stopped (see gotcha 3)
$V $P/kl_divergence.py --ref $M --cand $A --out $P/kl-L35.json
```
**Why it runs one process per model.** The default `--stage all` re-execs itself
once per checkpoint (`--stage ref`, then `--stage cand`), each writing its
first-token log-probs to a ~682 MiB `.pt` cache, then scores from the caches.
This is not tidiness — **it is the only teardown that works.** Measured, free VRAM
after the reference model:
| teardown | free VRAM |
|---|---|
| `del model` + `gc.collect()` + `empty_cache()` | 45,287 MiB |
| the same, model confined to an inner frame that exits | 45,287 MiB |
| **the process exits** | **96,689 MiB** |
The weights survive both in-process teardowns. The very first run only completed
because PyTorch's allocator hit OOM on the second load, collected, and retried —
the second model landed on the card *by rescue, not by design*, and on this
architecture a silent CPU offload does not raise, it returns confident garbage
(gotcha 1). The headroom gate (`exit 10`) is what turned that from an invisible
near-miss into a loud failure. Side benefit: the `ref` cache is reusable, so
measuring a different candidate against the same stock model skips a stage
entirely (`--stage cand` then `--stage score`).
⚠️ **The old residency gate could not fail.** It read `hf_device_map`, which
transformers leaves **empty** when the whole model fits on one device — so it
printed "(unsharded)" both when everything was fine and when there was nothing to
inspect. It now reads `{p.device for p in model.parameters()}` and prints the real
placement (`all parameters on cuda:0`).
## Why the write is shard surgery, not `model.save_pretrained`
The `--out` path edits the 18 safetensors shards directly and never instantiates
@@ -318,13 +415,41 @@ for headroom; it buys corruption. Gated (exit 9).
layers and is **deterministic**. This did neither. *If a NaN doesn't
propagate, debug memory, not math.*
**3. bf16 fits on one GPU — so the window is small now.** 50.1 GB of a 96 GB
card, which means a capture needs only **`vllm-gen` stopped**, not all three
seats. (`--capture-dtype float32` remains as an escape hatch; it needs 111 GB, so
it also needs `--max-layer 46` to fit on one card. The two agree to 0.0005, so
there is no reason to reach for it.) **Restore after:** start
`vllm-meromero-rp` **first**, then `vllm-gen` — gen grabs a fraction of *free*
VRAM at startup and will starve meromero if it goes first.
**3. bf16 fits on one GPU — but it needs BOTH GPU0 seats stopped, not one.**
> ⚠️ **CORRECTED 2026-08-20.** This section used to read "50.1 GB … a capture
> needs only `vllm-gen` stopped." **The unit was wrong and the conclusion that
> rode on it was wrong.** The real figure is **50.10 GiB = 51,300 MiB = 53.8 GB**
> of text-only weights, measured from the safetensors headers rather than read off
> a `/1e9` print:
>
> | | GB | GiB | MiB |
> |---|---|---|---|
> | checkpoint total | 55.56 | 51.75 | 52,989 |
> | vision (not loaded by `AutoModelForCausalLM`) | 0.92 | 0.86 | 879 |
> | MTP (not loaded either) | 0.85 | 0.79 | 810 |
> | **text-only — what actually lands on the card** | **53.79** | **50.10** | **51,300** |
>
> GPU0's two tenants are meromero (50,072 MiB) and gen (46,304 MiB), and
> **freeing either one alone leaves at most 50,933 MiB — about 400 MiB short.**
> A run that assumes one seat is enough will stop a service, sit at the edge, and
> then OOM. Stop **both**. Recompute this table if the checkpoint changes; do not
> trust a remembered gigabyte figure.
Both seats down leaves ~97,200 MiB, so the fit is comfortable rather than
marginal. Gate the run on **observing** the free VRAM (`nvidia-smi
--query-gpu=memory.free`) rather than sleeping after `docker stop`, and put the
restore in a `trap ... EXIT` so an abort hands the seats back — the 2026-08-20
aborted window did exactly that and cost nothing but two minutes.
(`--capture-dtype float32` remains as an escape hatch; it needs 111 GB, so it
also needs `--max-layer 46` to fit on one card. The two agree to 0.0005, so there
is no reason to reach for it.) **Restore after:** start `vllm-meromero-rp`
**first**, then `vllm-gen`. (The stated reason — "gen grabs a fraction of *free*
VRAM" — is not what the configs do: both seats pass `--gpu-memory-utilization` as
a fraction of **total** (`MEROMERO_GPU_MEM_UTIL=0.52`, `GEN_GPU_MEM_UTIL=0.43`),
so restore order is not actually load-bearing. Kept as the runbook order anyway;
it costs nothing.)
**4. fla is irrelevant here — but harmless.** `fla` + `einops` are `--target`
-installed to `/tank/aimodels/coldfusion-abliteration/pylibs` and reached via
@@ -148,3 +148,70 @@ def load_calibration(name: str, n_harmful: int, n_harmless: int, seed: int):
"harmful_pool": len(harmful_pool), "harmless_pool": len(harmless_pool),
"heldout_reserved": len(heldout),
}
def load_evaluation(n_harmless: int, n_harmful: int, seed: int,
calib_harmless_n: int, calib_harmless_seed: int):
"""Held-out evaluation prompts. Returns (harmless, harmful, provenance).
This is the *measurement* corpus — deliberately disjoint from anything the
refusal direction was fitted on, because a divergence measured on the fitting
set answers a different (and much easier) question than a divergence measured
on prompts the surgery never saw.
- **harmful** is the reserved `harmful_behaviors[test]` split (104 prompts,
overlap 0 with train by construction). `load_calibration` refuses to hand
these out as calibration, so they are still virgin here.
- **harmless** is drawn from `harmless_alpaca[train]` *minus the indices
calibration already consumed*. The exclusion has to be reconstructed
rather than remembered: calibration samples with
`random.Random(calib_harmless_seed).sample(range(pool), calib_harmless_n)`,
so replaying that exact draw recovers the used index set. Both the seed and
the n must match the capture that produced the direction under test, which
is why they are explicit parameters and not constants — a future capture at
a different n would otherwise silently leak its calibration into this set.
Disjointness is asserted on the returned *text*, not just on indices, so a
duplicated row in the alpaca pool cannot sneak a calibration prompt back in.
"""
root = _datasets_root()
harmless_pool = _arrow_rows("harmless", "train", root)
harmful = _arrow_rows("harmful", "test", root)
if calib_harmless_n > len(harmless_pool):
raise ValueError(
f"calibration claimed {calib_harmless_n} harmless prompts but the pool "
f"holds {len(harmless_pool)} — the exclusion set cannot be reconstructed")
used_idx = set(random.Random(calib_harmless_seed).sample(
range(len(harmless_pool)), calib_harmless_n))
used_text = {harmless_pool[i] for i in used_idx}
free_idx = [i for i in range(len(harmless_pool)) if i not in used_idx]
if n_harmless > len(free_idx):
raise ValueError(
f"asked for {n_harmless} held-out harmless prompts but only "
f"{len(free_idx)} remain after excluding the {calib_harmless_n} "
f"calibration drew")
idx = sorted(random.Random(seed).sample(free_idx, n_harmless))
harmless = [harmless_pool[i] for i in idx]
leaked = sorted(set(harmless) & used_text)
if leaked:
raise AssertionError(
f"{len(leaked)} evaluation prompt(s) are byte-identical to a calibration "
f"prompt — the harmless pool has duplicate rows and the index-level "
f"exclusion was not enough. First: {leaked[0]!r}")
if n_harmful > len(harmful):
raise ValueError(
f"asked for {n_harmful} harmful eval prompts but the reserved test split "
f"holds {len(harmful)}")
harmful = harmful[:n_harmful] if n_harmful else harmful
return harmless, harmful, {
"eval_source": "mlabonne/harmless_alpaca[train] minus calibration draw + "
"mlabonne/harmful_behaviors[test]",
"n_harmless": len(harmless), "n_harmful": len(harmful), "seed": seed,
"harmless_pool": len(harmless_pool),
"excluded_calibration": {"n": calib_harmless_n, "seed": calib_harmless_seed},
}
@@ -0,0 +1,517 @@
#!/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()