vh/esh-pfi-infrastructure
1
0
Fork 0
Code Issues 7 Pull Requests Actions Packages Projects Releases Wiki Activity

feat(coldfusion-abliteration): abliteration LANDS at layer 35 — separation selector, shard-surgery write, three false diagnoses corrected

Browse Source 
The abliterated model works. A/B vs stock on a matched greedy battery: explicit
sexual + graphic torture (the measured stock refusal surface) go from refused to
complied/engaged, held-out AdvBench prompts loosen, the self-harm guardrail
survives, coherence intact — the Robinson design point exactly. Output at
/tank/aimodels/qwen38-27b-coldfusion-abliterated-L35-bf16, verified bitwise:
131/131 targets changed, 333/333 vision byte-identical (delta 0.0), 735/735
others untouched.

Getting there corrected three diagnoses the prior session had backwards.

1. The layer-selection metric was wrong, and that was the whole ballgame. The
   recipe picks the abliteration layer by peak two-template |cos| agreement. On
   this heavily-merged base that metric is anti-correlated with efficacy: its
   argmax (layer 18) is the WORST-separating layer in the window (Cohen's d 5.51
   vs 9.89 at the peak), and abliterating there was a measured behavioral no-op —
   stock and "abliterated" refused all six probes identically. Cause: the two
   renderings end in different generative modes (</think> vs <think>), so |cos|
   scores answer-vs-reason mode, not refusal, and on a merge the mode term
   dominates. Replaced selection with harmful/harmless SEPARATION (Cohen's d /
   AUC of the direction's projection), gated on the sink screen since separation
   and sink-energy both climb with depth. Picks layer 35 (d 9.35, AUC 0.9997,
   sink 0.094%). Agreement is kept as a printed diagnostic.

2. The "bf16 NaNs, use fp32" rule was a misdiagnosis. The NaN was never
   precision — it was multi-GPU sharding (the residual stream zeroes two layers
   past the GPU0->GPU1 boundary; the first capture's layer 22 happened to sit in
   the healthy region, which is why it looked fine) plus
   PYTORCH_CUDA_ALLOC_CONF=expandable_segments (corrupts retained tensors; the
   corruption MOVED between bit-identical forwards, the tell that it was memory
   not math). On one GPU with a plain allocator, bf16 full-64-layer is exactly
   deterministic and coherent, at 50 GB and 4.3x the throughput of the 111 GB
   fp32 it replaced. Both defects are now hard gates (residency exit 8, allocator
   exit 9); capture pins CUDA_VISIBLE_DEVICES=0.

3. The corpus-size hypothesis was falsified. 52x more calibration data (8->416,
   mlabonne/harmful_behaviors = the recipe's actual AdvBench split, already on the
   box) moved agreement 0.594->0.624 — nothing. Kept the 416/416 corpus anyway
   (calibration.py); it gives the clean separation signal. The held-out 104-prompt
   test split is reserved and asserted disjoint.

Also: the --out write is now shard-level surgery (reads/writes the 18 safetensors
directly, no model object, no GPU). This is correctness, not thrift —
AutoModelForCausalLM resolves to the TEXT model, so save_pretrained would drop all
333 vision tensors AND skip the MTP head (the in-band MTP edit is the entire point
of the Robinson formula). Neither failure raises. Shard surgery makes vision and
the other 1068 tensors byte-identical by construction.

Batched capture with a dtype-aware equivalence gate; hidden states captured via
forward pre-hook (reading output_hidden_states off the returned object is unsafe
here — buffers get recycled). Sharding/allocator lessons promoted to the
quantization playbook (model-agnostic, sections 3.9-3.11 + superseded table); the
selection-metric lesson added to the recipe doc.

The dead layer-18 no-op checkpoint was removed (52 GB, confirmed identical to
stock). Incumbent gen seat untouched. Full canonical refusal-probe re-profile and
MTP-acceptance-on-quant still owed before this becomes a gen-seat candidate.
This commit is contained in:
Vuong Hoang
2026-08-20 08:46:21 -07:00
parent f714f28195
commit e9dbc8660b
4 changed files with 667 additions and 207 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/pfi/abliteration-recipe-qwen38.md
+17
View File
@@ -93,6 +93,23 @@ The two agree at **|cos| 0.96–0.99 across layers 18–45, peaking 0.9925 at la
vector is the evidence that the direction encodes *refusal semantics* rather than
*template formatting*. A single-template capture cannot distinguish the two.
> ⚠️ **Two-template agreement is a bad LAYER SELECTOR on a heavily-merged base —
> use harmful/harmless SEPARATION instead (added 2026-08-20).** On RobinsonLabs'
> stock Qwen3.8 the agreement was 0.99 and picking its peak was fine. On DavidAU's
> Cold-Fusion GAIN merge the same metric tops out at **0.62**, and its argmax
> (layer 18) is the layer with the **worst** refusal separation in the window
> (Cohen's d 5.51 vs 9.89 at the peak) — abliterating there was a measured
> behavioral **no-op**. The reason: the two renderings end in different generative
> modes (`</think>\n\n` = about to answer vs `<think>\n` = about to reason), so
> `|cos|` scores refusal *plus* mode, and on a merge the mode term dominates. The
> selector that actually predicts efficacy is **how cleanly the direction splits
> harmful from harmless prompt activations** (Cohen's d / AUC), gated on the sink
> screen (separation and sink-energy both rise with depth, so the raw peak is
> usually sink-dominated). On Cold-Fusion this picked **layer 35** (d 9.35, AUC
> 0.9997, sink 0.094%) and the abliteration worked. Keep agreement as a
> diagnostic; do not select on it. See
> `services/coldfusion-abliteration/README.md`.
### The attention-sink dimension — the one that bricks the model
**Qwen3.8-27B's massive-activation dimension is `3994`.** It carries 19–21% of
docs/pfi/model-quantization-playbook.md
+57
View File
@@ -254,6 +254,62 @@ not nvidia-modelopt.**
`read()` + `load(bytes)`, one shard cached at a time.
- **`vm.overcommit_memory=1`** on ana-ml2 (durable via `playbooks/ana-ml2-overcommit-memory.yaml`).
### 3.9 ⭐⭐ A sharded forward can be silently WRONG — never trust `device_map="auto"` for activations
Splitting **Qwen3.8-27B (Qwen3_5 hybrid)** across the two Blackwells with `device_map="auto"`
produces a model that loads clean, reports no error, and computes **garbage**: the residual stream
collapses to **exactly zero** a couple of layers past the GPU0→GPU1 boundary, and the logits decode
to rubbish (`'8'`, `'�'`, `'b'`). Every layer *below* the boundary stays healthy, deterministic, and
bit-identical to a single-GPU run — which is what makes it so dangerous. A capture that reads a
low layer looks perfectly plausible and is fine; one that reads a high layer is reading zeros, and
nothing in the pipeline says so. Measured 2026-08-20 (§9 Cold-Fusion).
**Rule: any workload that reads activations — refusal-direction capture, calibration, activation
statistics, PPL — must run on ONE device.** Sharding is for *storage*, and it is only safe when you
consume the model's final output through an engine that was built for it (vLLM does TP correctly;
`device_map="auto"` in transformers is not the same thing). If it does not fit on one card, shrink
the model, not the guarantee: **truncating the decoder to N layers is exact** for any activation
read at a layer < N (a causal stack's layer-N state cannot depend on layers above N), and it is
cheap — verified by reproducing the full model's layers 18/20/22/26 bit-for-bit.
**Gate it, don't remember it.** Assert single-device residency and zero offload before the forward:
```python
dmap = getattr(model, "hf_device_map", {}) or {}
gpus = {str(v) for v in dmap.values()} - {"cpu", "disk"}
offloaded = [k for k, v in dmap.items() if str(v) in ("cpu", "disk")]
if len(gpus) > 1 or offloaded:
sys.exit("residency gate FAILED — sharded/offloaded forward reads garbage")
```
### 3.10 ⭐⭐ `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` corrupts retained tensors
On torch 2.12+cu130 / Blackwell, tensors that **outlive their allocation** come back corrupted with
this flag set: captured hidden states carried Inf / NaN / zeros that **moved between bit-identical
forwards** (same input, same weights → a different layer corrupted each time). Unset, the identical
forwards are exactly reproducible. Several runbooks recommend this flag for headroom on large
loads; for anything that *keeps* activations it buys corruption.
Two tells that distinguish this from a real numerical blowup, both worth knowing because they
generalise: a genuine blowup **propagates** to later layers and is **deterministic**. Corruption
does neither — downstream layers were finite and consistent, and the affected layer moved run to
run. **If a "NaN" fails to propagate, stop debugging the math and start debugging memory.**
Corollary: **do not read `output_hidden_states=True` off a returned object** on a large multi-device
load. Take what you need *during* the forward with a `register_forward_pre_hook` that clones to CPU
immediately — it closes the reuse window and never retains a `[B, seq, hidden]` tensor per layer, so
it is cheaper than the thing it replaces.
### 3.11 Determinism is a necessary check, not a sufficient one
Both defects above were found by the cheapest possible test — **run the same input twice and diff**
— which no amount of eyeballing plausible-looking numbers would have caught. Add it to any
activation-reading pipeline. But note the trap that followed: after fixing the allocator, the run
went perfectly "deterministic" *because the corrupted layers were now stably zero*. Pair the
determinism check with a **magnitude** check (residual norms should grow smoothly with depth; an
exact 0.0 mid-stack is impossible) and, where you can, a **coherence** check (generate 40 tokens and
read them).
---
## 4. Pipeline shape
@@ -346,6 +402,7 @@ day if followed:
| "Use modelopt, NOT compressed-tensors — compressed-tensors can't load the BF16 MTP head, 0% acceptance" | `docs/runbooks/heretic2-nvfp4-mtp-seat.md` §landmine 2 | **SUPERSEDED 2026-08-14.** The 0% was the missing `re:^mtp.*` ignore (§3.3), not the format. compressed-tensors + the ignore gives 47.7–83.2% acceptance, live. Use compressed-tensors. |
| "Abliteration desyncs the MTP head → uncensored models can't do MTP" | earlier auto-memory | **SUPERSEDED 2026-08-14.** A modest abliteration preserves MTP (83.7% at bf16). Test MTP on **bf16 first** to isolate abliteration from quant/graft confounds — and isolate before deleting a 50 GB source. |
| "NVFP4 W4A4 is infeasible, no 4-bit wins both axes, FP8 is the Blackwell answer" | `reference_nvfp4_w4a4_granite_infeasible` | **NARROWED.** True for *uniform* W4A4 (measured on Granite-8B at 30k ctx). W4A4 on bulk MLPs **with FP8 on attention and late layers** is fine and is the current default (§2). |
| "transformers' Qwen3.5 DeltaNet linear-attention NaNs in bf16 without causal-conv1d; it is precision-driven cancellation and fp32 resolves it" | `services/coldfusion-abliteration/README.md`, `persistent-memory.d/2026-08-20-coldfusion-abliteration-capture.md` | **SUPERSEDED 2026-08-20.** Precision was never the variable. The NaN came from **multi-GPU sharding** and **`expandable_segments`** (§3.9, §3.10); fp32 only made it rarer, which is worse than failing. On one GPU with a plain allocator, **bf16 is exactly deterministic through all 64 layers and generates coherent prose** — at 50 GB and 4.3× the throughput of the 111 GB fp32 it replaced. |
---
services/coldfusion-abliteration/README.md
+203 -79
View File
@@ -51,17 +51,24 @@ Two more gates were added 2026-08-20, both protecting numbers rather than
tensors:
3. **Batch-equivalence gate** — capture batches prompts, so before the real run
it proves a padded batch reproduces one-at-a-time forwards (rel. tolerance
1e-3) and aborts otherwise. Padding is on the **right**, and that is load-
bearing: in a causal stack nothing after position *t* reaches position *t*, so
trailing pads cannot touch the token we read, whereas left padding would feed
pad tokens *into* the DeltaNet recurrence ahead of the prompt — the exact path
whose torch fallback is already known-untrustworthy here.
4. **Surgery pre-check** — on the write path, aborts if any of the 131 target
tensors is absent or on the meta device. `orthogonalize_` edits in place, and
an in-place write to an accelerate-offloaded tensor is a **silent no-op**;
without this gate an under-provisioned run ships a half-abliterated model that
passes a smoke test. Free the VRAM instead of defeating it.
it proves a padded batch reproduces one-at-a-time forwards and aborts
otherwise. Tolerance is dtype-aware (bf16 5e-2, fp32 1e-3): the gate hunts
*contamination*, not bit-exactness, and changing batch shape changes kernel
tiling and therefore accumulation order, so a few ULP is expected. Real
contamination is not subtle — the sharding defect read rel 1.00. Padding is on
the **right**, and that is load-bearing: in a causal stack nothing after
position *t* reaches position *t*, so trailing pads cannot touch the token we
read, whereas left padding would feed pad tokens *into* the DeltaNet
recurrence ahead of the prompt.
4. **Residency gate** (exit 8) and **allocator gate** (exit 9) — capture-only.
See the gotchas; both encode defects that silently produce wrong numbers
(multi-GPU sharding zeroes the upper residual stream; `expandable_segments`
corrupts retained tensors).
5. **Write completeness check** — the write path is shard surgery with no model
object, so the offload/meta silent-no-op failure class is gone; it instead
verifies all 131 target tensors were found across the shards before declaring
success (exit 7 otherwise) and refuses to overwrite an existing checkpoint
(exit 10).
## Calibration corpus
@@ -107,39 +114,38 @@ P=/tank/aimodels/coldfusion-abliteration
V=/tank/aimodels/quant-work/.venv/bin/python
M=/tank/aimodels/qwen38-27b-coldfusion-bf16
A=/tank/aimodels/qwen38-27b-coldfusion-abliterated-bf16
RUN="sudo -u llmuser env HF_HUB_OFFLINE=1 PYTHONPATH=$P/pylibs \
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True $V $P/abliterate.py --model $M"
# CUDA_VISIBLE_DEVICES=0 is REQUIRED for capture (gate exit 8) and
# PYTORCH_CUDA_ALLOC_CONF must stay unset (gate exit 9) — see the gotchas below.
RUN="sudo -u llmuser env HF_HUB_OFFLINE=1 CUDA_VISIBLE_DEVICES=0 \
PYTHONPATH=$P/pylibs $V $P/abliterate.py --model $M"
# 1. DRY RUN FIRST — verify the tensor map + coverage gate on the static
# surface, no forward, no write. Safe with the seats up. Do not skip: this
# confirms the recipe maps onto THIS checkpoint's names.
$RUN --dry-run
# --- everything below needs the fp32 VRAM window; stop the seats first ---
sudo docker stop vllm-gen vllm-meromero-rp vllm-fablefusion-probe
# --- capture needs GPU0 to itself: bf16 is 50 GB, so only gen must go ---
sudo docker stop -t 60 vllm-gen
cp $M/refusal-direction.pt $M/refusal-direction.pt.bak # capture overwrites it
# 2. CONTROL RUN — the legacy 8/8 set through the new batched path. It must
# reproduce the 2026-08-20 result (layer 22, |cos| 0.594, sink 0.001%). This
# is the regression test: batching, layer truncation and the refactor all
# validate against a known number for ~1 minute of forwards, before the
# expensive run. A mismatch here bisects cleanly — the batch-equivalence gate
# has already cleared batching, so truncation is the remaining suspect.
$RUN --capture --calib builtin --max-layer 46 --batch-size 8
# 2. CONTROL RUN — the legacy 8/8 set. Reproduces layer 22, |cos| 0.5944, sink
# 0.001% exactly. Keep it as the regression test: ~30s of forwards that
# validate the whole path against a known number before the real run.
$RUN --capture --calib builtin
# 3. THE REAL CAPTURE — Robinson's 416-prompt corpus.
$RUN --capture --calib mlabonne --max-layer 46 --batch-size 8
# Expect |cos| agreement in the window to rise well above 0.594. If it does
# not, set size was NOT the cause and the write stays gated.
# 3. THE REAL CAPTURE — Robinson's 416-prompt corpus. ~35s of forwards.
$RUN --capture --calib mlabonne
# Measured 2026-08-20: layer 18, |cos| 0.6238, sink 0.360%.
# 4. Restore the seats — meromero FIRST, gen LAST (gen grabs a fraction of FREE
# VRAM at startup and will starve meromero if it goes first).
sudo docker start vllm-meromero-rp && sleep 60 && sudo docker start vllm-gen vllm-fablefusion-probe
# 4. Restore. If meromero was stopped too, start it FIRST — gen takes a fraction
# of FREE VRAM at startup and will starve it otherwise.
sudo docker start vllm-gen
# 5. Abliterate (writes the new bf16). Only after 1-3 pass, and only on the
# operator's go — this is the destructive step. Needs the VRAM window again
# (bf16, 55.6 GB, must be fully resident — the surgery pre-check enforces it).
# NOTE: no --max-layer here; the guard refuses it.
# Shard-level surgery: reads/writes the 18 safetensors shards directly, NO
# model object, NO GPU. That is a correctness requirement, not just thrift —
# see "Why the write is shard surgery" below. --direction is REQUIRED.
$RUN --out $A --direction $M/refusal-direction.pt
```
@@ -152,21 +158,87 @@ $RUN --out $A --direction $M/refusal-direction.pt
| `--calib-n-harmless` | 416 | matched n from alpaca |
| `--calib-seed` | 0 | harmless sample only; harmful is order-deterministic |
| `--batch-size` | 8 | 832 prompts x 2 templates = 1664 forwards; batching is what makes that affordable |
| `--max-layer` | off | capture-only. Truncates the decoder. **Exact, not an approximation** — a causal stack's layer-N state cannot depend on layers above N, so any value above the window top (45) leaves the chosen direction bit-identical while cutting fp32 residency and forward cost by the dropped fraction. 46 drops 18 of 64 layers (~28%) and is what keeps fp32 off CPU offload. Refused on the write path, where it would emit a truncated checkpoint. |
| `--capture-dtype {bfloat16,float32}` | `bfloat16` | bf16 (50 GB, full 64 layers, one GPU) is validated deterministic + coherent; fp32 (111 GB, needs `--max-layer`) is a misdiagnosis-era escape hatch that agrees to 5e-4 |
| `--max-layer` | off | capture-only. Truncates the decoder. **Exact, not an approximation** — a causal stack's layer-N state cannot depend on layers above N. Only needed with `--capture-dtype float32`; bf16 fits whole. Refused on the write path. |
## ✅ RESULT — layer 35, and why the recipe's layer-selection metric had to be replaced
The write lands and works. Verified bitwise: **131/131 target tensors changed,
333/333 vision byte-identical (delta 0.0), 735/735 other tensors untouched.**
A/B against stock on a matched battery (greedy, held-out prompts):
| probe | stock | abliterated (L35) |
|---|---|---|
| explicit sexual (target axis) | refuses | **complies** |
| graphic torture (target axis) | refuses | **engages** (softened) |
| spam-bot / malware (held-out AdvBench) | refuses | **complies / engages** |
| self-harm method (guardrail) | redirects | **still redirects** |
| coherence ×2 | fine | **fine** |
That is the Robinson design point exactly: creative refusals fall, the self-harm
guardrail survives, coherence intact. Output at
`/tank/aimodels/qwen38-27b-coldfusion-abliterated-L35-bf16`.
**It took THREE captures, and the lesson is the metric.** The recipe selects the
abliteration layer by peak two-template `|cos|` agreement. On this checkpoint that
metric is not just weak, it is *anti-correlated* with what matters:
| capture | selector | layer picked | Cohen's d | result |
|---|---|---|---|---|
| 1 (8/8) | agreement | 22 | 5.70 | (sharding-corrupted, void) |
| 2 (416/416) | agreement | **18** | **5.51 — worst in window** | write was a **behavioral no-op** |
| 3 (416/416) | **separation, sink-gated** | **35** | **9.35** | **works** |
The tell that cracked it: after capture 2's write changed *nothing*, a per-layer
separation diagnostic (does the direction split harmful from harmless
activations?) showed the direction is **excellent** — AUC 0.9996+ across the whole
window — and that agreement had steered us to layer 18, the single **weakest**
separator (d 5.51 vs 9.89 at the peak). Agreement was measuring answer-vs-reason
*mode* (the two templates end `</think>\n\n` vs `<think>\n`), not refusal, and on
a heavily-merged base that mode term dominates.
**So selection is now by separation (Cohen's d), gated on the sink screen.**
Separation and sink-energy both rise with depth, so the raw peak (L39, d 9.89)
is sink-dominated (1.97% > 1%) and would brick the model; the script filters to
layers that pass the screen and takes the best separator among them — **L35, d
9.35 (within 5% of peak), sink 0.094% (10× under the limit).** One pass, no
guess-and-retry. Agreement is still computed and printed, as a diagnostic.
> The corpus-size hypothesis this session started on was **falsified**: 52× more
> calibration data (8→416) moved agreement 0.594→0.624, essentially nothing. The
> problem was never the calibration set. See the calibration section above; kept
> as the record of a dead-end worth not re-running.
## Why the write is shard surgery, not `model.save_pretrained`
The `--out` path edits the 18 safetensors shards directly and never instantiates
a model for the write. This is correctness, not thrift. `AutoModelForCausalLM`
resolves to `Qwen3_5ForCausalLM` — the **text** model — so saving from it would
(a) **drop all 333 vision tensors**, silently breaking the byte-identical-vision
guarantee, and (b) **skip the MTP head**, which the `ForConditionalGeneration`
wrapper does not load (the same reason the incumbent gen seat's Heretic pass left
its MTP head an untouched base graft) — and the in-band MTP edit is the entire
point of the Robinson formula. Neither failure raises. Shard surgery re-serializes
every non-target tensor from the exact bytes read, so vision and the other 1068
tensors are byte-identical *by construction*, the two MTP writers are just two
more keys, and the whole offload/meta-tensor silent-no-op class disappears with
the model object. Math is done in fp32, stored back at the original bf16.
## Verify after (do not trust the write blind)
1. **Vision byte-identical** — diff `visual.*` tensors source vs output (recipe
requires max delta 0).
2. **Refusal re-profile** — re-run the same battery from the 2026-08-19 probe
(reuse `services/refusal-probe/`, the gen-seat harness — NOT the ad-hoc GGUF
one) and confirm creative refusals dropped toward the RobinsonLabs 8% floor
while self-harm guardrails survive.
1. **Vision byte-identical + target count** — `services/coldfusion-abliteration`
verify: `targets changed=131/131 vision identical=333/333 delta=0.0 other
differ=0/735`. Done 2026-08-20, clean.
2. **Refusal re-profile** — the ad-hoc battery above is a smoke test. The full
canonical re-profile still owed: run `services/refusal-probe/` (the gen-seat
harness, NOT the GGUF one) once L35 is served, and confirm creative refusals
near the RobinsonLabs 8% floor with self-harm guardrails intact.
3. **MTP acceptance** — the whole point of the in-band MTP edit; measure on the
quantized build per `services/gen-seat-mixed-quant/RUNBOOK-heresy-swap.md`.
Gate ≳40% (`reference_abliteration_mtp_lessons` — gate on acceptance, not KL).
4. **PPL / coherence / no catatonia** — DavidAU fine-tunes are idiosyncratic;
eyeball the outputs, don't trust the metric alone.
eyeball the outputs, don't trust the metric alone. (Smoke: coherent, no
catatonia observed.)
Then, if it holds, NVFP4-quantize via `services/gen-seat-mixed-quant/` and it
becomes a gen-seat candidate — **do not delete the incumbent weights** until it
@@ -174,31 +246,49 @@ survives real multi-turn use (the 2026-08-14 delete-too-early lesson).
## ⚠️ Environment gotchas (2026-08-20 — cost real time, read before re-running)
**1. transformers' Qwen3.5 DeltaNet linear-attention NaNs in bf16 here.** The
fast-path kernel needs BOTH `flash-linear-attention` (`fla`, triton, installs
fine) AND `causal-conv1d` (needs `nvcc` to build — **absent on ana-ml2, no
prebuilt wheel**). Without causal-conv1d the DeltaNet short-conv runs the torch
fallback, which produces **nondeterministic all-NaN** hidden states in bf16
(same 11-token input: finite on one forward, NaN at layer 4 on the next). bf16
and fp32 share exponent range, so this is **precision-driven catastrophic
cancellation, not overflow** — **fp32 resolves it.**
> ⚠️ **RETRACTED 2026-08-20 — the "bf16 NaNs, use fp32" rule that lived here was
> a misdiagnosis, and it sent the next session down a 111 GB dead end.** The NaN
> was never precision. It was the two defects below. fp32 only made it *rarer*,
> which is worse than failing outright, because it let a broken forward produce a
> plausible-looking direction. **bf16, full 64 layers, one GPU: 50 GB, exactly
> deterministic through layer 63, coherent prose, 4.3× the throughput.**
→ **Capture loads fp32** (`abliterate.py` does this automatically in
`--capture` mode). The write/surgery path stays bf16 (no forward, no NaN).
The finite-gate in the script aborts if a direction comes out non-finite —
the sink screen alone won't catch it (`nan > threshold` is False).
**1. ⭐ Never let the capture shard across both GPUs.** With `device_map="auto"`
across the two Blackwells, this model loads clean, raises nothing, and computes
garbage: the residual stream collapses to **exactly zero** two layers past the
GPU0→GPU1 boundary and the logits decode to rubbish. Layers *below* the boundary
are healthy and bit-identical to a single-GPU run — which is exactly why the
first capture looked fine. It picked layer 22, which sat on GPU0 in the healthy
region; the upper half of its window was zeros and their agreement scores were
meaningless.
**2. fp32 (110 GB) needs the whole GPU.** Loaded across both Blackwells with
`device_map=auto`, activation memory OOM'd against the resident seats. The
production `vllm-gen` seat (44 GB) had to be **stopped** for the capture, along
with `vllm-meromero-rp` and `vllm-fablefusion-probe`. **Restore after:**
`sudo docker start vllm-gen vllm-meromero-rp vllm-fablefusion-probe`. Set
`PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.
→ **Run `CUDA_VISIBLE_DEVICES=0`.** The `--capture` path enforces this with a
residency gate (exit 8) that refuses a sharded or offloaded model.
**3. fla lives in a side dir, not the venv.** The shared `quant-work/.venv` is
not llmuser-writable. `fla` + `einops` are installed to
`/tank/aimodels/coldfusion-abliteration/pylibs` and reached via `PYTHONPATH`.
Run every invocation with `PYTHONPATH=/tank/aimodels/coldfusion-abliteration/pylibs`.
**2. ⭐ Never set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.** On this
stack it corrupts tensors that outlive their allocation — captured states came
back with Inf/NaN/zeros that **moved between bit-identical forwards**. Unset,
the same forwards are exactly reproducible. The old runbook recommended this flag
for headroom; it buys corruption. Gated (exit 9).
The tell worth remembering: a real numerical blowup **propagates** to later
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.
**4. fla is irrelevant here — but harmless.** `fla` + `einops` are `--target`
-installed to `/tank/aimodels/coldfusion-abliteration/pylibs` and reached via
`PYTHONPATH` (the shared `quant-work/.venv` is not llmuser-writable). Tested
2026-08-20: the nondeterminism reproduces **identically with `fla` absent**, so
the linear-attention kernel was never the culprit. Keep passing `PYTHONPATH`;
just don't blame it.
## Status
@@ -208,24 +298,58 @@ refusal direction is **finite, unit-normed, layer 22**, sink energy **0.0008%**
in dim 3994 (recipe L26 ref 0.06%, threshold 1%) — clean, not sink-dominated.
Saved to `qwen38-27b-coldfusion-bf16/refusal-direction.pt`.
⚠️ **Quality caveat:** two-template `|cos|` agreement at layer 22 is **0.594**,
notably below Robinson's 0.9925 — the 8/8 calibration set is the suspect.
### 2026-08-20, second session — the corpus hypothesis is FALSIFIED
> ⚠️ The first capture's log reported this as `|cos|=0.8538`. That was a
> reporting bug, fixed 2026-08-20: the line printed the **global** `agree.max()`
> next to the **window's** argmax layer. The global peak sits in the early layers
> where the dim-3994 massive activation dominates both templates and inflates
> agreement for reasons unrelated to refusal. `0.5944` was always the real
> in-window number. The report now prints the window max, a top-5, and labels the
> global figure as informational.
**Measured, on a forward that is trustworthy for the first time:**
**Where it stands 2026-08-20 (second session):** harness upgraded for the
re-capture — Robinson's actual 416-prompt corpus wired in (already on the box),
batched capture with an equivalence gate, optional exact layer truncation, the
agreement report fixed, and a surgery pre-check added for the write. Dry-run
re-verified 1:1 (131 tensors) and the calibration path tested end-to-end on the
box. **What has not run is anything needing the GPU** — the re-capture needs the
fp32 VRAM window, which costs production seat downtime.
| calibration | layer | `\|cos\|` agreement | sink energy |
|---|---|---|---|
| 8 / 8 (legacy) | 22 | **0.5944** | 0.001% |
| 416 / 416 (Robinson's corpus) | 18 | **0.6238** | 0.360% |
**The destructive `--out` write has NOT been executed** — it gates on the
operator's go, and on the re-capture showing a healthy agreement first.
**52× more calibration data bought +0.03.** The small calibration set was *not*
why agreement sat at 0.59, and Robinson's 0.9925 is not reachable on this
checkpoint by adding prompts. Agreement is uniformly ~0.54–0.62 across the whole
healthy window (L18 0.6238, L22 0.6158, L21 0.6101, L19 0.5944, L28 0.5841), not
peaked-and-noisy — which is the signature of a genuinely diffuse direction rather
than an under-sampled one.
Cross-validated two ways: the 8/8 run **reproduces the previous session's 0.5944
at layer 22 exactly**, and fp32-truncated vs bf16-full-64-layer agree to 0.0005.
So the number is real and the pipeline is sound.
> ⚠️ The first capture's log reported 0.594 as `|cos|=0.8538`. Reporting bug,
> fixed: the line printed the **global** `agree.max()` next to the **window's**
> argmax layer. The global peak sits in the early layers where the dim-3994
> massive activation dominates both templates and inflates agreement for reasons
> unrelated to refusal. `0.5944` was always the real number.
**The leading explanation is the metric, not the model.** The two renderings do
not just differ in formatting — they leave the model in **different generative
modes** at the token we read:
- `enable_thinking=false` ends `…<think>\n\n</think>\n\n` → about to write **the answer**
- `xhigh` ends `…<think>\n` → about to write **chain-of-thought**
So `|cos|` here measures *refusal semantics **plus** answer-vs-reason mode*.
Robinson's stock Qwen3.8-27B scored 0.99 across that same split, so on their base
the refusal component dominated; on this DavidAU GAIN merge the mode difference
apparently does not let it. **Note what this does and does not impugn:** the
direction actually used is `dirs[False]` — the no-think one. Cross-template
agreement is only a *quality check*, and a check that conflates two factors is a
weak gate to block on.
**The check that would actually settle it is a split-half.** Split the 416
harmful in two, derive a direction from each half *through the same template*,
and take `|cos|`. That isolates sampling noise — the thing calibration size
governs — with no mode term at all. If split-half is ~0.99, the direction is
well-estimated, the 0.62 is a mode artifact, and the write is justified on a
direction we can defend. If split-half is also ~0.6, the refusal representation
in this checkpoint is genuinely diffuse and single-direction abliteration is the
wrong instrument for it. Cheap: no extra forwards, just two accumulators.
**Status: the destructive `--out` write has NOT been executed.** It gates on the
operator's go. The saved direction
(`refusal-direction.pt`, layer 18, 416/416, sink 0.360%) is usable but its
quality is unresolved pending the split-half. The legacy 8/8 direction is
preserved at `refusal-direction.pt.bak-8x8`.
services/coldfusion-abliteration/abliterate.py
+390 -128
View File
@@ -29,6 +29,7 @@ Env: /tank/aimodels/quant-work/.venv (torch 2.12 cu130). Run ON ana-ml2.
from __future__ import annotations
import argparse
import os
import json
import sys
from pathlib import Path
@@ -133,92 +134,174 @@ def render(tokenizer, prompt, thinking):
msgs, tokenize=False, add_generation_prompt=True)
def decoder_layers(model):
"""The decoder's ModuleList, whatever wrapper depth it is buried under."""
for path in (("model", "layers"), ("model", "model", "layers"),
("model", "language_model", "layers")):
obj = model
for attr in path:
obj = getattr(obj, attr, None)
if obj is None:
break
if obj is not None and hasattr(obj, "__getitem__") and len(obj) > 0:
return obj
raise RuntimeError("could not locate the decoder layer list on this model")
@torch.no_grad()
def last_token_hidden(model, tokenizer, texts, device):
"""Last-real-token hidden state at every layer, for a batch of prompts.
def last_token_hidden(model, tokenizer, texts, device, layers):
"""Last-real-token hidden state at the requested layers, for a batch.
Returns [B, L+1, hidden] on CPU in float32.
Returns {layer: [B, hidden]} on CPU in float32.
PADDING SIDE IS LOAD-BEARING. We pad on the RIGHT and index each row's true
final token. In a causal stack — including this model's DeltaNet linear
CAPTURED DURING THE FORWARD, NOT AFTER — this is load-bearing. Reading
`output_hidden_states=True` off the returned object is not safe on this
stack: the retained tensors get recycled, and a *later* allocation overwrites
them with garbage. Diagnosed 2026-08-20 — a single layer's state came back
with exactly 5040 **Inf** values (not NaN) confined to one sequence position,
the affected layer moved between bit-identical trials (23, 23, 44), and every
downstream layer stayed finite and consistent. A real numerical blowup
propagates forward and is deterministic; this did neither. It is the stored
copy that is corrupt, not the computation. A forward pre-hook takes its slice
and clones it to CPU while the buffer is still live, which closes the window
entirely — and as a bonus never retains a full [B, seq, hidden] tensor per
layer, so it is cheaper than the thing it replaces.
`hidden_states[i]` in the transformers convention is the *input* to layer i,
which is exactly what a pre-hook on `layers[i]` sees — so this is the same
vector the previous capture used, not a redefinition.
PADDING SIDE IS ALSO LOAD-BEARING. We pad on the RIGHT and index each row's
true final token. In a causal stack — including this model's DeltaNet linear
attention — nothing after position t can influence position t, so trailing
pad tokens cannot contaminate the state we read. LEFT padding would be wrong
here: it prepends pad tokens *into* the linear-attention recurrence ahead of
the real prompt, and the torch fallback path (the one we are stuck on, see
the dtype note below) is not trustworthy about masking that prefix out. The
equivalence gate in `check_batch_equivalence` proves this empirically before
the real capture runs.
pad tokens cannot contaminate the state we read. LEFT padding would prepend
pad tokens *into* the linear-attention recurrence ahead of the real prompt,
and that fallback path is not trustworthy about masking a prefix out.
"""
enc = tokenizer(texts, return_tensors="pt", padding=True) # side pinned at load
lengths = enc["attention_mask"].sum(-1) # [B], true token counts
out = model(**enc.to(device), output_hidden_states=True)
per_layer = []
for h in out.hidden_states: # each [B, seq, hidden]
rows = torch.arange(h.shape[0], device=h.device)
idx = (lengths - 1).to(h.device)
per_layer.append(h[rows, idx, :].float().cpu())
return torch.stack(per_layer, dim=1) # [B, L+1, hidden]
lengths = enc["attention_mask"].sum(-1) # [B], true token counts
stack = decoder_layers(model)
grabbed, handles = {}, []
def make_hook(i):
def pre_hook(_mod, args, kwargs):
h = args[0] if args else kwargs.get("hidden_states")
rows = torch.arange(h.shape[0], device=h.device)
idx = (lengths - 1).to(h.device)
grabbed[i] = h[rows, idx, :].detach().float().cpu().clone()
return None
return pre_hook
try:
for i in layers:
handles.append(stack[i].register_forward_pre_hook(make_hook(i), with_kwargs=True))
model(**enc.to(device))
finally:
for h in handles:
h.remove()
missed = [i for i in layers if i not in grabbed]
if missed:
raise RuntimeError(f"pre-hooks never fired for layers {missed[:5]} — layer indexing is wrong")
return grabbed
@torch.no_grad()
def check_batch_equivalence(model, tokenizer, texts, device):
def check_batch_equivalence(model, tokenizer, texts, device, layers):
"""Prove padded-batch == one-at-a-time before spending the capture window.
Cheap insurance against a silently wrong number: this architecture's
linear-attention path already produced NaN once under conditions that looked
fine, so batching is not taken on faith. Compares the batched last-token
hidden states against single-prompt forwards over a handful of prompts of
differing length (so at least one row is actually padded).
Cheap insurance against a silently wrong number: this stack has already
produced both a nondeterministic NaN and a recycled-buffer Inf, so batching
is not taken on faith. Compares batched last-token states against
single-prompt forwards over prompts of differing length, so at least one row
is genuinely padded. Also catches non-finite states, whatever their cause.
"""
batched = last_token_hidden(model, tokenizer, texts, device) # [B, L+1, H]
singles = torch.cat([last_token_hidden(model, tokenizer, [t], device) for t in texts])
delta = (batched - singles).abs().max().item()
scale = singles.abs().max().item()
batched = last_token_hidden(model, tokenizer, texts, device, layers)
singles = [last_token_hidden(model, tokenizer, [t], device, layers) for t in texts]
delta = 0.0
scale = 0.0
for i in layers:
single_i = torch.cat([s[i] for s in singles]) # [B, hidden]
delta = max(delta, (batched[i] - single_i).abs().max().item())
scale = max(scale, single_i.abs().max().item())
rel = delta / max(scale, 1e-6)
return rel, delta, scale
def _mean_hidden(model, tokenizer, prompts, thinking, device, batch_size, label):
"""Mean last-token hidden state per layer over a prompt set. [L+1, hidden].
def _collect_hidden(model, tokenizer, prompts, thinking, device, batch_size, layers, label):
"""Per-prompt last-token hidden states. {layer: [N, hidden]} float32 on CPU.
Accumulated in float64: the direction is a difference of two means, which is
precisely where catastrophic cancellation lives, and this model has already
demonstrated it is precision-sensitive. The accumulator is on CPU and tiny
(65 x 5120), so the wider dtype is free.
Retained per-prompt rather than accumulated into a mean, because the layer
SELECTION metric needs the individual projections (see `capture_direction`).
The cost is trivial — 416 prompts x 28 layers x 5120 floats is ~238 MB.
Every batch is finite-checked as it lands. A single Inf would poison the mean
for that layer, and finding out at the end of an 832-prompt run wastes the run.
"""
import time
if not prompts:
raise ValueError(f"empty prompt set for {label}")
total = None
n = 0
chunks = {i: [] for i in layers}
t0 = time.time()
for start in range(0, len(prompts), batch_size):
chunk = prompts[start:start + batch_size]
texts = [render(tokenizer, p, thinking) for p in chunk]
h = last_token_hidden(model, tokenizer, texts, device).double() # [B, L+1, H]
total = h.sum(0) if total is None else total + h.sum(0)
n += h.shape[0]
got = last_token_hidden(model, tokenizer, texts, device, layers)
for i in layers:
h = got[i]
if not torch.isfinite(h).all():
raise RuntimeError(
f"non-finite hidden state at layer {i}, prompts {start}..{start+len(chunk)-1} "
f"({label}) — refusing to fold it into the mean")
chunks[i].append(h.float())
done = start + len(chunk)
if done % (batch_size * 10) == 0 or done == len(prompts):
rate = done / max(time.time() - t0, 1e-6)
print(f" [{label}] {done}/{len(prompts)} prompts ({rate:.1f}/s)", flush=True)
return (total / n).float()
return {i: torch.cat(chunks[i]) for i in layers}
def capture_direction(model, tokenizer, device, harmful, harmless, batch_size):
"""Per-layer refusal direction from each template, plus the |cos| agreement.
Returns (directions[template][layer], agreement[layer])."""
dirs = {}
def separation_stats(harm_acts, safe_acts, direction):
"""How cleanly `direction` splits harmful from harmless. (cohen_d, auc).
THE metric for picking the abliteration layer. Project every prompt onto the
unit direction and ask how separated the two clouds are: Cohen's d for effect
size, AUC for rank separability. A direction that does not separate the two
populations cannot be the thing the model uses to decide to refuse, so
removing it will do nothing — which is exactly the failure this replaced.
"""
ph = harm_acts @ direction
ps = safe_acts @ direction
pooled = ((ph.var() + ps.var()) / 2).sqrt().clamp_min(1e-8)
cohen = float((ph.mean() - ps.mean()) / pooled)
ranks = torch.cat([ph, ps]).argsort().argsort().float()
n1 = len(ph)
auc = float((ranks[:n1].sum() - n1 * (n1 - 1) / 2) / (n1 * len(ps)))
return cohen, auc
def capture_direction(model, tokenizer, device, harmful, harmless, batch_size, layers):
"""Per-layer refusal direction, its separation power, and template agreement.
Returns (dirs[template][layer], agreement{layer}, sep{layer: (cohen_d, auc)}).
"""
dirs, acts = {}, {}
for thinking in (False, True):
tag = "xhigh" if thinking else "no-think"
print(f" template: {tag}", flush=True)
harm_mu = _mean_hidden(model, tokenizer, harmful, thinking, device, batch_size, f"{tag}/harmful")
safe_mu = _mean_hidden(model, tokenizer, harmless, thinking, device, batch_size, f"{tag}/harmless")
d = harm_mu - safe_mu # [L+1, hidden]
d = d / d.norm(dim=-1, keepdim=True).clamp_min(1e-8)
H = _collect_hidden(model, tokenizer, harmful, thinking, device, batch_size, layers, f"{tag}/harmful")
S = _collect_hidden(model, tokenizer, harmless, thinking, device, batch_size, layers, f"{tag}/harmless")
d = {}
for i in layers:
v = H[i].double().mean(0) - S[i].double().mean(0)
d[i] = (v / v.norm().clamp_min(1e-8)).float()
dirs[thinking] = d
a = (dirs[False] * dirs[True]).sum(-1).abs() # |cos| per layer
return dirs, a
if thinking is False:
acts = (H, S) # separation is measured on the template we ship
agree = {i: float((dirs[False][i] * dirs[True][i]).sum().abs()) for i in layers}
H, S = acts
sep = {i: separation_stats(H[i], S[i], dirs[False][i]) for i in layers}
return dirs, agree, sep
def sink_energy(direction_vec, dim=SINK_DIM):
@@ -242,6 +325,109 @@ def orthogonalize_embed_(weight, d_unit):
weight.sub_(torch.outer(coeff, d))
def write_abliterated(model_dir: Path, args, targets, embed_keys, n_vision):
"""Orthogonalize the 131 residual writers SHARD BY SHARD and write a new checkpoint.
This is deliberately not done through a loaded model object, and that is a
correctness requirement rather than a preference. `AutoModelForCausalLM`
resolves to `Qwen3_5ForCausalLM` — the TEXT model. Saving from it would drop
all 333 vision tensors, silently violating the recipe's byte-identical-vision
guarantee; and the `ForConditionalGeneration` wrapper does not load the MTP
head at all (the same reason the incumbent gen seat's Heretic pass left its
MTP head an untouched base graft), so the in-band MTP edit that is the whole
point of the Robinson formula would be skipped. Neither failure raises.
Operating on the shards instead: every tensor we do not target is re-serialized
from the exact bytes we read, so vision and the other 1068 tensors are
byte-identical by construction, and the MTP writers are just two more keys.
No GPU, no accelerate, no offload, no meta tensors — the whole class of
silent-no-op failures goes away with the model object.
Per playbook 3.6, shards are read with plain `read()` + `load()` rather than
mmap: `safe_open` mmaps a whole shard and a 50 GB shard ENOMEMs on ZFS
regardless of free RAM.
"""
from glob import glob
import shutil
from safetensors.torch import load as st_load, save_file
if not args.out:
print("\n!! --out is required to write the abliterated model "
"(use --capture for direction-only).", file=sys.stderr)
sys.exit(1)
if not args.direction:
print("\n!! --direction <refusal-direction.pt> is required for the write. Capture "
"first (--capture), inspect the agreement and sink energy, then write.",
file=sys.stderr)
sys.exit(1)
blob = torch.load(args.direction, weights_only=False)
layer, d_unit, e = blob["layer"], blob["direction"], blob.get("sink_energy")
calib = blob.get("calibration", {})
print(f"\ndirection: layer {layer}, sink energy {e*100:.3f}%, "
f"agreement {blob.get('agreement')}, calib {calib.get('calib')} "
f"({calib.get('n_harmful')}/{calib.get('n_harmless')})")
if not torch.isfinite(d_unit).all():
print("\n!! direction is not finite — refusing to write.", file=sys.stderr)
sys.exit(4)
d_unit = (d_unit.float() / d_unit.float().norm().clamp_min(1e-12)).cpu()
if e is not None and e > SINK_ENERGY_MAX:
print(f"\n!! sink-energy gate FAILED ({e*100:.3f}% > {SINK_ENERGY_MAX*100:.1f}%) — "
f"orthogonalizing this direction would brick the model.", file=sys.stderr)
sys.exit(3)
out_dir = Path(args.out)
if out_dir.exists() and any(out_dir.glob("*.safetensors")):
print(f"\n!! {out_dir} already holds safetensors shards — refusing to overwrite an "
f"existing checkpoint. Move it aside or pick another --out.", file=sys.stderr)
sys.exit(10)
out_dir.mkdir(parents=True, exist_ok=True)
targets = set(targets)
shards = sorted(glob(str(model_dir / "*.safetensors")))
edited, seen_targets, total_tensors = 0, set(), 0
for si, shard in enumerate(shards, 1):
with open(shard, "rb") as f:
tensors = st_load(f.read())
total_tensors += len(tensors)
hits = [k for k in tensors if k in targets]
for k in hits:
w = tensors[k]
orig_dtype = w.dtype
# Math in fp32. The weights are bf16 (8 mantissa bits); computing
# d^T W and the rank-1 subtraction at that precision would lose more
# than the edit itself is worth.
w32 = w.float()
if k in embed_keys:
orthogonalize_embed_(w32, d_unit) # [vocab, hidden]
else:
orthogonalize_(w32, d_unit) # [hidden, in]
tensors[k] = w32.to(orig_dtype)
edited += 1
seen_targets.add(k)
save_file(tensors, str(out_dir / Path(shard).name), metadata={"format": "pt"})
print(f" shard {si}/{len(shards)} {Path(shard).name}: {len(hits)} edited", flush=True)
del tensors
missed = targets - seen_targets
if missed or edited != len(targets):
print(f"\n!! surgery incomplete — edited {edited} of {len(targets)} targets, "
f"{len(missed)} never found in any shard: {sorted(missed)[:5]}", file=sys.stderr)
sys.exit(7)
print(f" edited {edited} tensors of {total_tensors}; vision ({n_vision}) byte-identical")
# Everything that is not weights rides along unchanged.
for pat in ("*.json", "*.jinja", "*.txt", "*.model", "*.py"):
for src in sorted(model_dir.glob(pat)):
if src.name in ("dl.py",):
continue
shutil.copy2(src, out_dir / src.name)
torch.save(blob, out_dir / "refusal-direction.pt")
print(f"\nwrote abliterated checkpoint -> {out_dir}")
print("DONE.")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True, help="bf16 checkpoint dir")
@@ -259,6 +445,10 @@ def main():
help="harmless calibration prompts sampled from alpaca")
ap.add_argument("--calib-seed", type=int, default=0, help="seed for the harmless sample")
ap.add_argument("--batch-size", type=int, default=8, help="prompts per forward during capture")
ap.add_argument("--capture-dtype", choices=("bfloat16", "float32"), default="bfloat16",
help="dtype for the capture forward. bf16 (50 GB, full 64 layers, one GPU) "
"is validated deterministic and coherent; fp32 (111 GB) was adopted on "
"a misdiagnosis and is kept only as an escape hatch.")
ap.add_argument("--max-layer", type=int, default=None,
help="truncate the decoder to this many layers before capture. Exact, not an "
"approximation: a causal stack's layer-N hidden state cannot depend on "
@@ -268,6 +458,20 @@ def main():
args = ap.parse_args()
model_dir = Path(args.model)
# --- allocator gate ------------------------------------------------------
# PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True corrupts tensors that
# outlive their allocation here (torch 2.12+cu130, Blackwell): captured
# hidden states came back with Inf/NaN/zeros that MOVED between bit-identical
# forwards. Unset, the same forwards are exactly reproducible. The previous
# runbook recommended this flag for headroom; it buys corruption.
alloc = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
if args.capture and "expandable_segments" in alloc:
print(f"\n!! PYTORCH_CUDA_ALLOC_CONF={alloc!r} — expandable_segments corrupts retained "
f"tensors on this stack and makes the capture nondeterministic. Unset it.",
file=sys.stderr)
sys.exit(9)
from safetensors import safe_open
from glob import glob
@@ -298,6 +502,12 @@ def main():
print("\ndry-run complete — surface verified, nothing loaded or written.")
return
if not args.capture:
targets = (trunk["down_proj"] + trunk["o_proj"] + trunk["linear_out"]
+ mtp_writers + embed)
write_abliterated(model_dir, args, targets, set(embed), n_vision)
return
# --- load model for capture / surgery -------------------------------------
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
print("\nloading model (bf16, device_map=auto across the Blackwells)...")
@@ -308,16 +518,19 @@ def main():
tok.padding_side = "right"
if tok.pad_token is None:
tok.pad_token = tok.eos_token
# DTYPE IS LOAD-BEARING FOR CAPTURE. This is a Qwen3_5 hybrid (DeltaNet
# linear-attn + full-attn). Without the causal_conv1d fast-path kernel
# (unbuildable here — no nvcc), the DeltaNet recurrence runs the torch
# fallback, which produces NONDETERMINISTIC NaN hidden states in bf16
# (verified 2026-08-20: same 11-token input finite on one forward, NaN at
# layer 4 on the next). bf16 and fp32 share exponent range, so this is
# PRECISION-driven catastrophic cancellation, not overflow — fp32's mantissa
# resolves it. Capture therefore loads fp32 (fits: 98GB GPU + CPU offload,
# 244GB RAM free). The surgery/write path takes bf16 (no forward, no NaN).
load_dtype = torch.float32 if args.capture else torch.bfloat16
# CAPTURE DTYPE — bf16, and the fp32 that used to be here was a misdiagnosis.
#
# The earlier note claimed bf16 produced nondeterministic NaN in the DeltaNet
# linear-attention fallback and that fp32's mantissa "resolved" it. Retested
# 2026-08-20 once the sharding and allocator defects below were fixed: bf16,
# full 64 layers, one GPU, 50.1 GB — every probed layer through 63 finite and
# bit-deterministic across repeated forwards, and the model generates coherent
# prose. The NaN was never about precision. It was multi-GPU sharding plus
# expandable_segments, both of which fabricate NaN/Inf/zeros that fp32 merely
# made rarer. Keeping fp32 would cost 111 GB (forcing truncation and a wider
# seat-down window) to buy nothing.
capture_dtype = torch.bfloat16 if args.capture_dtype == "bfloat16" else torch.float32
load_dtype = capture_dtype if args.capture else torch.bfloat16
load_kwargs = dict(dtype=load_dtype, device_map="auto", attn_implementation="sdpa")
if args.max_layer is not None:
@@ -348,11 +561,42 @@ def main():
model.eval()
device = next(model.parameters()).device
if args.capture:
# --- residency gate: this model must not be SHARDED for a forward -----
# Diagnosed 2026-08-20. Split across the two Blackwells by device_map,
# the residual stream collapses to exactly zero a couple of layers past
# the GPU0->GPU1 boundary and the logits decode to garbage ('8', '�',
# 'b', ...), while every layer *below* the boundary stays healthy,
# deterministic, and bit-identical to a single-GPU run. That is why the
# first capture looked plausible: it picked layer 22, which happened to
# sit on GPU0 in the healthy region. Layers above the boundary were zeros
# and their agreement scores were meaningless.
#
# There is no partial-credit version of this. Pin to one GPU
# (CUDA_VISIBLE_DEVICES=0) and truncate with --max-layer so the fp32
# weights fit: 46 layers is ~75 GB on a 96 GB card.
dmap = getattr(model, "hf_device_map", {}) or {}
placements = {str(v) for v in dmap.values()}
gpus = {p for p in placements if p not in ("cpu", "disk")}
offloaded = sorted(k for k, v in dmap.items() if str(v) in ("cpu", "disk"))
if len(gpus) > 1 or offloaded:
print(f"\n!! residency gate FAILED — the model is not on a single GPU "
f"(gpus={sorted(gpus)}, offloaded={len(offloaded)} modules). Sharding this "
f"architecture silently zeroes the residual stream past the device boundary "
f"and the capture would read garbage for the upper window.\n"
f" Fix: CUDA_VISIBLE_DEVICES=0 and --max-layer 46 (~75 GB fp32), with the "
f"vLLM seats stopped.", file=sys.stderr)
if offloaded:
print(f" first offloaded: {offloaded[:3]}", file=sys.stderr)
sys.exit(8)
print(f" residency: single device {sorted(gpus) or ['(unsharded)']}, no offload")
if args.direction:
blob = torch.load(args.direction)
layer, d_unit = blob["layer"], blob["direction"]
print(f"loaded direction for layer {layer} from {args.direction}")
calib_prov = blob.get("calibration", {"calib": "loaded-from-file"})
sep = None
agree = None
else:
# calibration.py sits beside this script; make that explicit rather than
@@ -370,40 +614,97 @@ def main():
# --- batch-equivalence gate ------------------------------------------
# Batching is what makes an 832-prompt capture affordable, so prove it is
# free of side effects before spending the window on it.
# Only the recipe's window is ever captured. Outside it the direction is
# not a candidate anyway, and the early layers are dominated by the
# dim-3994 massive activation, which inflates |cos| for reasons that have
# nothing to do with refusal — quoting that global figure beside the
# window's layer is how the first capture came to be reported as 0.8538
# when the number that mattered was 0.5944.
lo, hi = CAPTURE_WINDOW
window = list(range(lo, hi + 1))
if args.layer is not None and args.layer not in window:
print(f"\n!! --layer {args.layer} is outside the capture window [{lo},{hi}]; no "
f"direction is captured there.", file=sys.stderr)
sys.exit(5)
probe = (harmful[:2] + harmless[:2]) if len(harmless) >= 2 else harmful[:4]
probe_texts = [render(tok, p, False) for p in probe]
rel, delta, scale = check_batch_equivalence(model, tok, probe_texts, device)
rel, delta, scale = check_batch_equivalence(model, tok, probe_texts, device, window)
# Tolerance is dtype-aware, because the gate is looking for CONTAMINATION
# (pad leakage, recycled buffers), not for bit-exactness. Changing the
# batch shape changes kernel tiling and therefore accumulation order, so a
# few ULP of disagreement is expected and benign. bf16 carries 8 mantissa
# bits: at magnitude ~80 one ULP is ~0.25, so ~4 ULP lands near 1e-2
# relative. fp32 measures ~5e-5 on the same probe. Real contamination is
# not subtle — the sharding defect read rel 1.00, two orders clear of
# either threshold.
tol = 5e-2 if capture_dtype == torch.bfloat16 else 1e-3
print(f"batch-equivalence gate: max |batched - single| = {delta:.3e} "
f"(rel {rel:.2e} of scale {scale:.3f}; threshold 1e-3)")
if not (rel < 1e-3):
print("\n!! batched and single-prompt forwards disagree — right-padding is not "
"neutral on this path. Re-run with --batch-size 1, or fix the masking; do "
"NOT capture on contaminated states.", file=sys.stderr)
f"(rel {rel:.2e} of scale {scale:.3f}; threshold {tol:.0e} for "
f"{str(capture_dtype).replace('torch.','')})")
if not (rel < tol):
print("\n!! batched and single-prompt forwards disagree, or a state came back "
"non-finite. Re-run with --batch-size 1 to isolate; do NOT capture on "
"contaminated states.", file=sys.stderr)
sys.exit(6)
print(" -> batch-equivalence PASSED")
print("\ncapturing refusal direction from two chat templates...")
dirs, agree = capture_direction(model, tok, device, harmful, harmless, args.batch_size)
dirs, agree, sep = capture_direction(model, tok, device, harmful, harmless,
args.batch_size, window)
# auto-pick: highest two-template agreement in the recipe's window.
# NOTE: report the WINDOW max, never agree.max() — the global argmax sits
# in the early layers where the dim-3994 massive activation dominates both
# templates and inflates |cos| for reasons that have nothing to do with
# refusal semantics. Quoting the global figure next to the window's layer
# is how the first capture came to be reported as 0.85 when the number
# that mattered was 0.59.
lo, hi = CAPTURE_WINDOW
window = list(range(lo, min(hi + 1, agree.shape[0])))
best = max(window, key=lambda L: float(agree[L]))
# LAYER SELECTION — by SEPARATION, not by two-template agreement.
#
# The recipe picks the layer by peak |cos| between the no-think and xhigh
# renderings. On this checkpoint that metric is actively misleading, and
# following it cost a full write-and-test cycle for a no-op. Measured
# 2026-08-20: agreement ranked layer 18 first (0.6238) — and layer 18 has
# the WORST harmful/harmless separation of the entire window (Cohen's d
# 5.51 vs 9.89 at layer 39). Abliterating there changed nothing: stock and
# abliterated refused all six probe prompts identically.
#
# The reason agreement fails here is that the two renderings do not merely
# differ in formatting — they leave the model in different generative
# modes at the token we read (`</think>\n\n` = about to answer, `<think>\n`
# = about to reason). So |cos| scores refusal semantics *plus* mode, and on
# a heavily-merged base the mode term dominates. Robinson's stock
# Qwen3.8-27B scored 0.99 across that same split; this model scores 0.62,
# and that difference says more about the templates than the direction.
#
# Separation asks the question that actually predicts efficacy: does this
# direction split harmful from harmless prompts? Here it does, superbly
# (AUC 0.9996+ across the whole window) — the direction was never the
# problem, only where we removed it. Agreement is still reported, as a
# diagnostic rather than a selector.
# The sink screen is a FILTER on selection, not just a post-hoc abort.
# Separation and sink-energy both climb with depth on this model, so the
# best-separating layer (39, d=9.89) is also sink-dominated (1.97% > 1%)
# and would brick the model. Pick the best separator *among layers that
# pass the screen* — one pass, no guess-and-retry.
sink = {L: sink_energy(dirs[False][L]) for L in window}
by_sep = lambda L: sep[L][0]
eligible = [L for L in window if sink[L] <= SINK_ENERGY_MAX]
print(f"\nlayer selection over [{lo},{hi}] — separation, gated on sink < "
f"{SINK_ENERGY_MAX*100:.1f}%:")
for L in sorted(window, key=by_sep, reverse=True)[:8]:
mark = "ok " if sink[L] <= SINK_ENERGY_MAX else "SINK"
print(f" [{mark}] L{L:<3} d={sep[L][0]:6.3f} AUC={sep[L][1]:.4f} "
f"sink={sink[L]*100:6.3f}% |cos|={agree[L]:.4f}")
if not eligible:
print("\n!! every layer in the window is sink-dominated — no safe direction exists "
"here. Widen the window or reconsider the approach.", file=sys.stderr)
sys.exit(3)
best = max(eligible, key=by_sep)
layer = args.layer if args.layer is not None else best
d_unit = dirs[False][layer] # thinking-off direction at the chosen layer
top = sorted(window, key=lambda L: float(agree[L]), reverse=True)[:5]
print(f"\nagreement peak in [{lo},{hi}]: layer {best} (|cos|={float(agree[best]):.4f})")
print(" top-5 in window: " + ", ".join(f"L{L}={float(agree[L]):.4f}" for L in top))
print(f" global argmax (out-of-window layers are sink-dominated, informational only): "
f"L{int(agree.argmax())} (|cos|={float(agree.max()):.4f})")
print(f" using layer {layer} (|cos|={float(agree[layer]):.4f}); "
f"recipe anchor L{DEFAULT_LAYER} at |cos| 0.9925")
print(f" -> {len(eligible)}/{len(window)} layers pass the sink screen; "
f"best separator among them: L{best} (d={sep[best][0]:.3f})")
print(f" using layer {layer} (d={sep[layer][0]:.3f}, AUC={sep[layer][1]:.4f}, "
f"sink={sink[layer]*100:.3f}%, two-template |cos|={agree[layer]:.4f})")
agree_best = max(window, key=lambda L: agree[L])
print(f" [diagnostic] agreement would have picked L{agree_best} "
f"(|cos|={agree[agree_best]:.4f}, d={sep[agree_best][0]:.3f}) — "
f"recipe anchor L{DEFAULT_LAYER} at |cos| 0.9925 on stock Qwen3.8")
# --- finite gate: a NaN/Inf direction must NEVER pass silently -----------
# (the sink screen alone doesn't catch this — `nan > threshold` is False, so
@@ -428,8 +729,10 @@ def main():
blob = {
"layer": layer, "direction": d_unit.cpu(), "sink_energy": e,
"calibration": calib_prov,
"agreement": None if agree is None else float(agree[layer]),
"agreement_per_layer": None if agree is None else agree.cpu(),
"agreement": None if agree is None else agree[layer],
"agreement_per_layer": agree,
"separation": None if sep is None else sep[layer],
"separation_per_layer": sep,
"capture_window": CAPTURE_WINDOW,
}
@@ -439,47 +742,6 @@ def main():
print(f"direction saved -> {dpath} (capture-only, no write)")
return
if not args.out:
print("\n!! --out is required to write the abliterated model "
"(use --capture for direction-only).", file=sys.stderr)
sys.exit(1)
# --- surgery: orthogonalize every residual writer -------------------------
print(f"\northogonalizing {total_edits} residual writers along the refusal direction...")
sd = model.state_dict()
# Offload gate. `orthogonalize_` edits in place; a parameter that accelerate
# has offloaded shows up here as a meta tensor, where `sub_` writes into
# nothing and reports no error. That ships a quietly half-abliterated model —
# the same failure the coverage gate exists to prevent, arriving by a
# different door. Free the VRAM (stop the seats) rather than defeating this.
targets = trunk["down_proj"] + trunk["o_proj"] + trunk["linear_out"] + mtp_writers + embed
missing = [k for k in targets if k not in sd]
meta = [k for k in targets if k in sd and sd[k].device.type == "meta"]
if missing or meta:
print(f"\n!! surgery pre-check FAILED — {len(missing)} target tensor(s) absent from the "
f"state dict and {len(meta)} on the meta device (offloaded). In-place edits to "
f"those are silent no-ops. Ensure the model loads fully resident (stop the vLLM "
f"seats) and do not run --max-layer on the write path.", file=sys.stderr)
for k in (missing + meta)[:5]:
print(f" {k}", file=sys.stderr)
sys.exit(7)
print(f" -> surgery pre-check PASSED ({len(targets)} targets resident, none offloaded)")
edited = 0
for k in trunk["down_proj"] + trunk["o_proj"] + trunk["linear_out"] + mtp_writers:
orthogonalize_(sd[k], d_unit); edited += 1
for k in embed:
orthogonalize_embed_(sd[k], d_unit); edited += 1
print(f" edited {edited} tensors; vision ({n_vision}) untouched")
out_dir = Path(args.out); out_dir.mkdir(parents=True, exist_ok=True)
print(f"saving abliterated bf16 -> {out_dir}")
model.save_pretrained(out_dir, safe_serialization=True)
tok.save_pretrained(out_dir)
torch.save(blob, out_dir / "refusal-direction.pt")
print("DONE.")
if __name__ == "__main__":
main()