diff --git a/docs/pfi/abliteration-recipe-qwen38.md b/docs/pfi/abliteration-recipe-qwen38.md index deb9944..7d00825 100644 --- a/docs/pfi/abliteration-recipe-qwen38.md +++ b/docs/pfi/abliteration-recipe-qwen38.md @@ -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 (`\n\n` = about to answer vs `\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 diff --git a/docs/pfi/model-quantization-playbook.md b/docs/pfi/model-quantization-playbook.md index b6bcb7c..f74024e 100644 --- a/docs/pfi/model-quantization-playbook.md +++ b/docs/pfi/model-quantization-playbook.md @@ -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. | --- diff --git a/services/coldfusion-abliteration/README.md b/services/coldfusion-abliteration/README.md index ab97bee..c209cbd 100644 --- a/services/coldfusion-abliteration/README.md +++ b/services/coldfusion-abliteration/README.md @@ -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 `\n\n` vs `\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 `…\n\n\n\n` → about to write **the answer** +- `xhigh` ends `…\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`. diff --git a/services/coldfusion-abliteration/abliterate.py b/services/coldfusion-abliteration/abliterate.py index fd5e971..f11e4e6 100644 --- a/services/coldfusion-abliteration/abliterate.py +++ b/services/coldfusion-abliteration/abliterate.py @@ -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 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 (`\n\n` = about to answer, `\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()