feat(coldfusion-abliteration): abliteration LANDS at layer 35 — separation selector, shard-surgery write, three false diagnoses corrected
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:
@@ -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
|
vector is the evidence that the direction encodes *refusal semantics* rather than
|
||||||
*template formatting*. A single-template capture cannot distinguish the two.
|
*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
|
### The attention-sink dimension — the one that bricks the model
|
||||||
|
|
||||||
**Qwen3.8-27B's massive-activation dimension is `3994`.** It carries 19–21% of
|
**Qwen3.8-27B's massive-activation dimension is `3994`.** It carries 19–21% of
|
||||||
|
|||||||
@@ -254,6 +254,62 @@ not nvidia-modelopt.**
|
|||||||
`read()` + `load(bytes)`, one shard cached at a time.
|
`read()` + `load(bytes)`, one shard cached at a time.
|
||||||
- **`vm.overcommit_memory=1`** on ana-ml2 (durable via `playbooks/ana-ml2-overcommit-memory.yaml`).
|
- **`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
|
## 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. |
|
| "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. |
|
| "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). |
|
| "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. |
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
|||||||
@@ -51,17 +51,24 @@ Two more gates were added 2026-08-20, both protecting numbers rather than
|
|||||||
tensors:
|
tensors:
|
||||||
|
|
||||||
3. **Batch-equivalence gate** — capture batches prompts, so before the real run
|
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
|
it proves a padded batch reproduces one-at-a-time forwards and aborts
|
||||||
1e-3) and aborts otherwise. Padding is on the **right**, and that is load-
|
otherwise. Tolerance is dtype-aware (bf16 5e-2, fp32 1e-3): the gate hunts
|
||||||
bearing: in a causal stack nothing after position *t* reaches position *t*, so
|
*contamination*, not bit-exactness, and changing batch shape changes kernel
|
||||||
trailing pads cannot touch the token we read, whereas left padding would feed
|
tiling and therefore accumulation order, so a few ULP is expected. Real
|
||||||
pad tokens *into* the DeltaNet recurrence ahead of the prompt — the exact path
|
contamination is not subtle — the sharding defect read rel 1.00. Padding is on
|
||||||
whose torch fallback is already known-untrustworthy here.
|
the **right**, and that is load-bearing: in a causal stack nothing after
|
||||||
4. **Surgery pre-check** — on the write path, aborts if any of the 131 target
|
position *t* reaches position *t*, so trailing pads cannot touch the token we
|
||||||
tensors is absent or on the meta device. `orthogonalize_` edits in place, and
|
read, whereas left padding would feed pad tokens *into* the DeltaNet
|
||||||
an in-place write to an accelerate-offloaded tensor is a **silent no-op**;
|
recurrence ahead of the prompt.
|
||||||
without this gate an under-provisioned run ships a half-abliterated model that
|
4. **Residency gate** (exit 8) and **allocator gate** (exit 9) — capture-only.
|
||||||
passes a smoke test. Free the VRAM instead of defeating it.
|
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
|
## Calibration corpus
|
||||||
|
|
||||||
@@ -107,39 +114,38 @@ P=/tank/aimodels/coldfusion-abliteration
|
|||||||
V=/tank/aimodels/quant-work/.venv/bin/python
|
V=/tank/aimodels/quant-work/.venv/bin/python
|
||||||
M=/tank/aimodels/qwen38-27b-coldfusion-bf16
|
M=/tank/aimodels/qwen38-27b-coldfusion-bf16
|
||||||
A=/tank/aimodels/qwen38-27b-coldfusion-abliterated-bf16
|
A=/tank/aimodels/qwen38-27b-coldfusion-abliterated-bf16
|
||||||
RUN="sudo -u llmuser env HF_HUB_OFFLINE=1 PYTHONPATH=$P/pylibs \
|
# CUDA_VISIBLE_DEVICES=0 is REQUIRED for capture (gate exit 8) and
|
||||||
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True $V $P/abliterate.py --model $M"
|
# 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
|
# 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
|
# surface, no forward, no write. Safe with the seats up. Do not skip: this
|
||||||
# confirms the recipe maps onto THIS checkpoint's names.
|
# confirms the recipe maps onto THIS checkpoint's names.
|
||||||
$RUN --dry-run
|
$RUN --dry-run
|
||||||
|
|
||||||
# --- everything below needs the fp32 VRAM window; stop the seats first ---
|
# --- capture needs GPU0 to itself: bf16 is 50 GB, so only gen must go ---
|
||||||
sudo docker stop vllm-gen vllm-meromero-rp vllm-fablefusion-probe
|
sudo docker stop -t 60 vllm-gen
|
||||||
cp $M/refusal-direction.pt $M/refusal-direction.pt.bak # capture overwrites it
|
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
|
# 2. CONTROL RUN — the legacy 8/8 set. Reproduces layer 22, |cos| 0.5944, sink
|
||||||
# reproduce the 2026-08-20 result (layer 22, |cos| 0.594, sink 0.001%). This
|
# 0.001% exactly. Keep it as the regression test: ~30s of forwards that
|
||||||
# is the regression test: batching, layer truncation and the refactor all
|
# validate the whole path against a known number before the real run.
|
||||||
# validate against a known number for ~1 minute of forwards, before the
|
$RUN --capture --calib builtin
|
||||||
# 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
|
|
||||||
|
|
||||||
# 3. THE REAL CAPTURE — Robinson's 416-prompt corpus.
|
# 3. THE REAL CAPTURE — Robinson's 416-prompt corpus. ~35s of forwards.
|
||||||
$RUN --capture --calib mlabonne --max-layer 46 --batch-size 8
|
$RUN --capture --calib mlabonne
|
||||||
# Expect |cos| agreement in the window to rise well above 0.594. If it does
|
# Measured 2026-08-20: layer 18, |cos| 0.6238, sink 0.360%.
|
||||||
# not, set size was NOT the cause and the write stays gated.
|
|
||||||
|
|
||||||
# 4. Restore the seats — meromero FIRST, gen LAST (gen grabs a fraction of FREE
|
# 4. Restore. If meromero was stopped too, start it FIRST — gen takes a fraction
|
||||||
# VRAM at startup and will starve meromero if it goes first).
|
# of FREE VRAM at startup and will starve it otherwise.
|
||||||
sudo docker start vllm-meromero-rp && sleep 60 && sudo docker start vllm-gen vllm-fablefusion-probe
|
sudo docker start vllm-gen
|
||||||
|
|
||||||
# 5. Abliterate (writes the new bf16). Only after 1-3 pass, and only on the
|
# 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
|
# 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).
|
# Shard-level surgery: reads/writes the 18 safetensors shards directly, NO
|
||||||
# NOTE: no --max-layer here; the guard refuses it.
|
# 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
|
$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-n-harmless` | 416 | matched n from alpaca |
|
||||||
| `--calib-seed` | 0 | harmless sample only; harmful is order-deterministic |
|
| `--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 |
|
| `--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)
|
## Verify after (do not trust the write blind)
|
||||||
|
|
||||||
1. **Vision byte-identical** — diff `visual.*` tensors source vs output (recipe
|
1. **Vision byte-identical + target count** — `services/coldfusion-abliteration`
|
||||||
requires max delta 0).
|
verify: `targets changed=131/131 vision identical=333/333 delta=0.0 other
|
||||||
2. **Refusal re-profile** — re-run the same battery from the 2026-08-19 probe
|
differ=0/735`. Done 2026-08-20, clean.
|
||||||
(reuse `services/refusal-probe/`, the gen-seat harness — NOT the ad-hoc GGUF
|
2. **Refusal re-profile** — the ad-hoc battery above is a smoke test. The full
|
||||||
one) and confirm creative refusals dropped toward the RobinsonLabs 8% floor
|
canonical re-profile still owed: run `services/refusal-probe/` (the gen-seat
|
||||||
while self-harm guardrails survive.
|
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
|
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`.
|
quantized build per `services/gen-seat-mixed-quant/RUNBOOK-heresy-swap.md`.
|
||||||
Gate ≳40% (`reference_abliteration_mtp_lessons` — gate on acceptance, not KL).
|
Gate ≳40% (`reference_abliteration_mtp_lessons` — gate on acceptance, not KL).
|
||||||
4. **PPL / coherence / no catatonia** — DavidAU fine-tunes are idiosyncratic;
|
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
|
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
|
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)
|
## ⚠️ Environment gotchas (2026-08-20 — cost real time, read before re-running)
|
||||||
|
|
||||||
**1. transformers' Qwen3.5 DeltaNet linear-attention NaNs in bf16 here.** The
|
> ⚠️ **RETRACTED 2026-08-20 — the "bf16 NaNs, use fp32" rule that lived here was
|
||||||
fast-path kernel needs BOTH `flash-linear-attention` (`fla`, triton, installs
|
> a misdiagnosis, and it sent the next session down a 111 GB dead end.** The NaN
|
||||||
fine) AND `causal-conv1d` (needs `nvcc` to build — **absent on ana-ml2, no
|
> was never precision. It was the two defects below. fp32 only made it *rarer*,
|
||||||
prebuilt wheel**). Without causal-conv1d the DeltaNet short-conv runs the torch
|
> which is worse than failing outright, because it let a broken forward produce a
|
||||||
fallback, which produces **nondeterministic all-NaN** hidden states in bf16
|
> plausible-looking direction. **bf16, full 64 layers, one GPU: 50 GB, exactly
|
||||||
(same 11-token input: finite on one forward, NaN at layer 4 on the next). bf16
|
> deterministic through layer 63, coherent prose, 4.3× the throughput.**
|
||||||
and fp32 share exponent range, so this is **precision-driven catastrophic
|
|
||||||
cancellation, not overflow** — **fp32 resolves it.**
|
|
||||||
|
|
||||||
→ **Capture loads fp32** (`abliterate.py` does this automatically in
|
**1. ⭐ Never let the capture shard across both GPUs.** With `device_map="auto"`
|
||||||
`--capture` mode). The write/surgery path stays bf16 (no forward, no NaN).
|
across the two Blackwells, this model loads clean, raises nothing, and computes
|
||||||
The finite-gate in the script aborts if a direction comes out non-finite —
|
garbage: the residual stream collapses to **exactly zero** two layers past the
|
||||||
the sink screen alone won't catch it (`nan > threshold` is False).
|
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
|
→ **Run `CUDA_VISIBLE_DEVICES=0`.** The `--capture` path enforces this with a
|
||||||
`device_map=auto`, activation memory OOM'd against the resident seats. The
|
residency gate (exit 8) that refuses a sharded or offloaded model.
|
||||||
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`.
|
|
||||||
|
|
||||||
**3. fla lives in a side dir, not the venv.** The shared `quant-work/.venv` is
|
**2. ⭐ Never set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.** On this
|
||||||
not llmuser-writable. `fla` + `einops` are installed to
|
stack it corrupts tensors that outlive their allocation — captured states came
|
||||||
`/tank/aimodels/coldfusion-abliteration/pylibs` and reached via `PYTHONPATH`.
|
back with Inf/NaN/zeros that **moved between bit-identical forwards**. Unset,
|
||||||
Run every invocation with `PYTHONPATH=/tank/aimodels/coldfusion-abliteration/pylibs`.
|
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
|
## 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.
|
in dim 3994 (recipe L26 ref 0.06%, threshold 1%) — clean, not sink-dominated.
|
||||||
Saved to `qwen38-27b-coldfusion-bf16/refusal-direction.pt`.
|
Saved to `qwen38-27b-coldfusion-bf16/refusal-direction.pt`.
|
||||||
|
|
||||||
⚠️ **Quality caveat:** two-template `|cos|` agreement at layer 22 is **0.594**,
|
### 2026-08-20, second session — the corpus hypothesis is FALSIFIED
|
||||||
notably below Robinson's 0.9925 — the 8/8 calibration set is the suspect.
|
|
||||||
|
|
||||||
> ⚠️ The first capture's log reported this as `|cos|=0.8538`. That was a
|
**Measured, on a forward that is trustworthy for the first time:**
|
||||||
> 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.
|
|
||||||
|
|
||||||
**Where it stands 2026-08-20 (second session):** harness upgraded for the
|
| calibration | layer | `\|cos\|` agreement | sink energy |
|
||||||
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
|
| 8 / 8 (legacy) | 22 | **0.5944** | 0.001% |
|
||||||
agreement report fixed, and a surgery pre-check added for the write. Dry-run
|
| 416 / 416 (Robinson's corpus) | 18 | **0.6238** | 0.360% |
|
||||||
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.
|
|
||||||
|
|
||||||
**The destructive `--out` write has NOT been executed** — it gates on the
|
**52× more calibration data bought +0.03.** The small calibration set was *not*
|
||||||
operator's go, and on the re-capture showing a healthy agreement first.
|
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`.
|
||||||
|
|||||||
@@ -29,6 +29,7 @@ Env: /tank/aimodels/quant-work/.venv (torch 2.12 cu130). Run ON ana-ml2.
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
|
import os
|
||||||
import json
|
import json
|
||||||
import sys
|
import sys
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -133,92 +134,174 @@ def render(tokenizer, prompt, thinking):
|
|||||||
msgs, tokenize=False, add_generation_prompt=True)
|
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()
|
@torch.no_grad()
|
||||||
def last_token_hidden(model, tokenizer, texts, device):
|
def last_token_hidden(model, tokenizer, texts, device, layers):
|
||||||
"""Last-real-token hidden state at every layer, for a batch of prompts.
|
"""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
|
CAPTURED DURING THE FORWARD, NOT AFTER — this is load-bearing. Reading
|
||||||
final token. In a causal stack — including this model's DeltaNet linear
|
`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
|
attention — nothing after position t can influence position t, so trailing
|
||||||
pad tokens cannot contaminate the state we read. LEFT padding would be wrong
|
pad tokens cannot contaminate the state we read. LEFT padding would prepend
|
||||||
here: it prepends pad tokens *into* the linear-attention recurrence ahead of
|
pad tokens *into* the linear-attention recurrence ahead of the real prompt,
|
||||||
the real prompt, and the torch fallback path (the one we are stuck on, see
|
and that fallback path is not trustworthy about masking a prefix out.
|
||||||
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.
|
|
||||||
"""
|
"""
|
||||||
enc = tokenizer(texts, return_tensors="pt", padding=True) # side pinned at load
|
enc = tokenizer(texts, return_tensors="pt", padding=True) # side pinned at load
|
||||||
lengths = enc["attention_mask"].sum(-1) # [B], true token counts
|
lengths = enc["attention_mask"].sum(-1) # [B], true token counts
|
||||||
out = model(**enc.to(device), output_hidden_states=True)
|
stack = decoder_layers(model)
|
||||||
per_layer = []
|
grabbed, handles = {}, []
|
||||||
for h in out.hidden_states: # each [B, seq, hidden]
|
|
||||||
rows = torch.arange(h.shape[0], device=h.device)
|
def make_hook(i):
|
||||||
idx = (lengths - 1).to(h.device)
|
def pre_hook(_mod, args, kwargs):
|
||||||
per_layer.append(h[rows, idx, :].float().cpu())
|
h = args[0] if args else kwargs.get("hidden_states")
|
||||||
return torch.stack(per_layer, dim=1) # [B, L+1, hidden]
|
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()
|
@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.
|
"""Prove padded-batch == one-at-a-time before spending the capture window.
|
||||||
|
|
||||||
Cheap insurance against a silently wrong number: this architecture's
|
Cheap insurance against a silently wrong number: this stack has already
|
||||||
linear-attention path already produced NaN once under conditions that looked
|
produced both a nondeterministic NaN and a recycled-buffer Inf, so batching
|
||||||
fine, so batching is not taken on faith. Compares the batched last-token
|
is not taken on faith. Compares batched last-token states against
|
||||||
hidden states against single-prompt forwards over a handful of prompts of
|
single-prompt forwards over prompts of differing length, so at least one row
|
||||||
differing length (so at least one row is actually padded).
|
is genuinely padded. Also catches non-finite states, whatever their cause.
|
||||||
"""
|
"""
|
||||||
batched = last_token_hidden(model, tokenizer, texts, device) # [B, L+1, H]
|
batched = last_token_hidden(model, tokenizer, texts, device, layers)
|
||||||
singles = torch.cat([last_token_hidden(model, tokenizer, [t], device) for t in texts])
|
singles = [last_token_hidden(model, tokenizer, [t], device, layers) for t in texts]
|
||||||
delta = (batched - singles).abs().max().item()
|
delta = 0.0
|
||||||
scale = singles.abs().max().item()
|
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)
|
rel = delta / max(scale, 1e-6)
|
||||||
return rel, delta, scale
|
return rel, delta, scale
|
||||||
|
|
||||||
|
|
||||||
def _mean_hidden(model, tokenizer, prompts, thinking, device, batch_size, label):
|
def _collect_hidden(model, tokenizer, prompts, thinking, device, batch_size, layers, label):
|
||||||
"""Mean last-token hidden state per layer over a prompt set. [L+1, hidden].
|
"""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
|
Retained per-prompt rather than accumulated into a mean, because the layer
|
||||||
precisely where catastrophic cancellation lives, and this model has already
|
SELECTION metric needs the individual projections (see `capture_direction`).
|
||||||
demonstrated it is precision-sensitive. The accumulator is on CPU and tiny
|
The cost is trivial — 416 prompts x 28 layers x 5120 floats is ~238 MB.
|
||||||
(65 x 5120), so the wider dtype is free.
|
|
||||||
|
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
|
import time
|
||||||
if not prompts:
|
if not prompts:
|
||||||
raise ValueError(f"empty prompt set for {label}")
|
raise ValueError(f"empty prompt set for {label}")
|
||||||
total = None
|
chunks = {i: [] for i in layers}
|
||||||
n = 0
|
|
||||||
t0 = time.time()
|
t0 = time.time()
|
||||||
for start in range(0, len(prompts), batch_size):
|
for start in range(0, len(prompts), batch_size):
|
||||||
chunk = prompts[start:start + batch_size]
|
chunk = prompts[start:start + batch_size]
|
||||||
texts = [render(tokenizer, p, thinking) for p in chunk]
|
texts = [render(tokenizer, p, thinking) for p in chunk]
|
||||||
h = last_token_hidden(model, tokenizer, texts, device).double() # [B, L+1, H]
|
got = last_token_hidden(model, tokenizer, texts, device, layers)
|
||||||
total = h.sum(0) if total is None else total + h.sum(0)
|
for i in layers:
|
||||||
n += h.shape[0]
|
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)
|
done = start + len(chunk)
|
||||||
if done % (batch_size * 10) == 0 or done == len(prompts):
|
if done % (batch_size * 10) == 0 or done == len(prompts):
|
||||||
rate = done / max(time.time() - t0, 1e-6)
|
rate = done / max(time.time() - t0, 1e-6)
|
||||||
print(f" [{label}] {done}/{len(prompts)} prompts ({rate:.1f}/s)", flush=True)
|
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):
|
def separation_stats(harm_acts, safe_acts, direction):
|
||||||
"""Per-layer refusal direction from each template, plus the |cos| agreement.
|
"""How cleanly `direction` splits harmful from harmless. (cohen_d, auc).
|
||||||
Returns (directions[template][layer], agreement[layer])."""
|
|
||||||
dirs = {}
|
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):
|
for thinking in (False, True):
|
||||||
tag = "xhigh" if thinking else "no-think"
|
tag = "xhigh" if thinking else "no-think"
|
||||||
print(f" template: {tag}", flush=True)
|
print(f" template: {tag}", flush=True)
|
||||||
harm_mu = _mean_hidden(model, tokenizer, harmful, thinking, device, batch_size, f"{tag}/harmful")
|
H = _collect_hidden(model, tokenizer, harmful, thinking, device, batch_size, layers, f"{tag}/harmful")
|
||||||
safe_mu = _mean_hidden(model, tokenizer, harmless, thinking, device, batch_size, f"{tag}/harmless")
|
S = _collect_hidden(model, tokenizer, harmless, thinking, device, batch_size, layers, f"{tag}/harmless")
|
||||||
d = harm_mu - safe_mu # [L+1, hidden]
|
d = {}
|
||||||
d = d / d.norm(dim=-1, keepdim=True).clamp_min(1e-8)
|
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
|
dirs[thinking] = d
|
||||||
a = (dirs[False] * dirs[True]).sum(-1).abs() # |cos| per layer
|
if thinking is False:
|
||||||
return dirs, a
|
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):
|
def sink_energy(direction_vec, dim=SINK_DIM):
|
||||||
@@ -242,6 +325,109 @@ def orthogonalize_embed_(weight, d_unit):
|
|||||||
weight.sub_(torch.outer(coeff, d))
|
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():
|
def main():
|
||||||
ap = argparse.ArgumentParser()
|
ap = argparse.ArgumentParser()
|
||||||
ap.add_argument("--model", required=True, help="bf16 checkpoint dir")
|
ap.add_argument("--model", required=True, help="bf16 checkpoint dir")
|
||||||
@@ -259,6 +445,10 @@ def main():
|
|||||||
help="harmless calibration prompts sampled from alpaca")
|
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("--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("--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,
|
ap.add_argument("--max-layer", type=int, default=None,
|
||||||
help="truncate the decoder to this many layers before capture. Exact, not an "
|
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 "
|
"approximation: a causal stack's layer-N hidden state cannot depend on "
|
||||||
@@ -268,6 +458,20 @@ def main():
|
|||||||
args = ap.parse_args()
|
args = ap.parse_args()
|
||||||
|
|
||||||
model_dir = Path(args.model)
|
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 safetensors import safe_open
|
||||||
from glob import glob
|
from glob import glob
|
||||||
|
|
||||||
@@ -298,6 +502,12 @@ def main():
|
|||||||
print("\ndry-run complete — surface verified, nothing loaded or written.")
|
print("\ndry-run complete — surface verified, nothing loaded or written.")
|
||||||
return
|
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 -------------------------------------
|
# --- load model for capture / surgery -------------------------------------
|
||||||
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
|
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
|
||||||
print("\nloading model (bf16, device_map=auto across the Blackwells)...")
|
print("\nloading model (bf16, device_map=auto across the Blackwells)...")
|
||||||
@@ -308,16 +518,19 @@ def main():
|
|||||||
tok.padding_side = "right"
|
tok.padding_side = "right"
|
||||||
if tok.pad_token is None:
|
if tok.pad_token is None:
|
||||||
tok.pad_token = tok.eos_token
|
tok.pad_token = tok.eos_token
|
||||||
# DTYPE IS LOAD-BEARING FOR CAPTURE. This is a Qwen3_5 hybrid (DeltaNet
|
# CAPTURE DTYPE — bf16, and the fp32 that used to be here was a misdiagnosis.
|
||||||
# linear-attn + full-attn). Without the causal_conv1d fast-path kernel
|
#
|
||||||
# (unbuildable here — no nvcc), the DeltaNet recurrence runs the torch
|
# The earlier note claimed bf16 produced nondeterministic NaN in the DeltaNet
|
||||||
# fallback, which produces NONDETERMINISTIC NaN hidden states in bf16
|
# linear-attention fallback and that fp32's mantissa "resolved" it. Retested
|
||||||
# (verified 2026-08-20: same 11-token input finite on one forward, NaN at
|
# 2026-08-20 once the sharding and allocator defects below were fixed: bf16,
|
||||||
# layer 4 on the next). bf16 and fp32 share exponent range, so this is
|
# full 64 layers, one GPU, 50.1 GB — every probed layer through 63 finite and
|
||||||
# PRECISION-driven catastrophic cancellation, not overflow — fp32's mantissa
|
# bit-deterministic across repeated forwards, and the model generates coherent
|
||||||
# resolves it. Capture therefore loads fp32 (fits: 98GB GPU + CPU offload,
|
# prose. The NaN was never about precision. It was multi-GPU sharding plus
|
||||||
# 244GB RAM free). The surgery/write path takes bf16 (no forward, no NaN).
|
# expandable_segments, both of which fabricate NaN/Inf/zeros that fp32 merely
|
||||||
load_dtype = torch.float32 if args.capture else torch.bfloat16
|
# 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")
|
load_kwargs = dict(dtype=load_dtype, device_map="auto", attn_implementation="sdpa")
|
||||||
if args.max_layer is not None:
|
if args.max_layer is not None:
|
||||||
@@ -348,11 +561,42 @@ def main():
|
|||||||
model.eval()
|
model.eval()
|
||||||
device = next(model.parameters()).device
|
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:
|
if args.direction:
|
||||||
blob = torch.load(args.direction)
|
blob = torch.load(args.direction)
|
||||||
layer, d_unit = blob["layer"], blob["direction"]
|
layer, d_unit = blob["layer"], blob["direction"]
|
||||||
print(f"loaded direction for layer {layer} from {args.direction}")
|
print(f"loaded direction for layer {layer} from {args.direction}")
|
||||||
calib_prov = blob.get("calibration", {"calib": "loaded-from-file"})
|
calib_prov = blob.get("calibration", {"calib": "loaded-from-file"})
|
||||||
|
sep = None
|
||||||
agree = None
|
agree = None
|
||||||
else:
|
else:
|
||||||
# calibration.py sits beside this script; make that explicit rather than
|
# calibration.py sits beside this script; make that explicit rather than
|
||||||
@@ -370,40 +614,97 @@ def main():
|
|||||||
# --- batch-equivalence gate ------------------------------------------
|
# --- batch-equivalence gate ------------------------------------------
|
||||||
# Batching is what makes an 832-prompt capture affordable, so prove it is
|
# Batching is what makes an 832-prompt capture affordable, so prove it is
|
||||||
# free of side effects before spending the window on it.
|
# 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 = (harmful[:2] + harmless[:2]) if len(harmless) >= 2 else harmful[:4]
|
||||||
probe_texts = [render(tok, p, False) for p in probe]
|
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} "
|
print(f"batch-equivalence gate: max |batched - single| = {delta:.3e} "
|
||||||
f"(rel {rel:.2e} of scale {scale:.3f}; threshold 1e-3)")
|
f"(rel {rel:.2e} of scale {scale:.3f}; threshold {tol:.0e} for "
|
||||||
if not (rel < 1e-3):
|
f"{str(capture_dtype).replace('torch.','')})")
|
||||||
print("\n!! batched and single-prompt forwards disagree — right-padding is not "
|
if not (rel < tol):
|
||||||
"neutral on this path. Re-run with --batch-size 1, or fix the masking; do "
|
print("\n!! batched and single-prompt forwards disagree, or a state came back "
|
||||||
"NOT capture on contaminated states.", file=sys.stderr)
|
"non-finite. Re-run with --batch-size 1 to isolate; do NOT capture on "
|
||||||
|
"contaminated states.", file=sys.stderr)
|
||||||
sys.exit(6)
|
sys.exit(6)
|
||||||
print(" -> batch-equivalence PASSED")
|
print(" -> batch-equivalence PASSED")
|
||||||
|
|
||||||
print("\ncapturing refusal direction from two chat templates...")
|
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.
|
# LAYER SELECTION — by SEPARATION, not by two-template agreement.
|
||||||
# NOTE: report the WINDOW max, never agree.max() — the global argmax sits
|
#
|
||||||
# in the early layers where the dim-3994 massive activation dominates both
|
# The recipe picks the layer by peak |cos| between the no-think and xhigh
|
||||||
# templates and inflates |cos| for reasons that have nothing to do with
|
# renderings. On this checkpoint that metric is actively misleading, and
|
||||||
# refusal semantics. Quoting the global figure next to the window's layer
|
# following it cost a full write-and-test cycle for a no-op. Measured
|
||||||
# is how the first capture came to be reported as 0.85 when the number
|
# 2026-08-20: agreement ranked layer 18 first (0.6238) — and layer 18 has
|
||||||
# that mattered was 0.59.
|
# the WORST harmful/harmless separation of the entire window (Cohen's d
|
||||||
lo, hi = CAPTURE_WINDOW
|
# 5.51 vs 9.89 at layer 39). Abliterating there changed nothing: stock and
|
||||||
window = list(range(lo, min(hi + 1, agree.shape[0])))
|
# abliterated refused all six probe prompts identically.
|
||||||
best = max(window, key=lambda L: float(agree[L]))
|
#
|
||||||
|
# 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
|
layer = args.layer if args.layer is not None else best
|
||||||
d_unit = dirs[False][layer] # thinking-off direction at the chosen layer
|
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" -> {len(eligible)}/{len(window)} layers pass the sink screen; "
|
||||||
print(f"\nagreement peak in [{lo},{hi}]: layer {best} (|cos|={float(agree[best]):.4f})")
|
f"best separator among them: L{best} (d={sep[best][0]:.3f})")
|
||||||
print(" top-5 in window: " + ", ".join(f"L{L}={float(agree[L]):.4f}" for L in top))
|
print(f" using layer {layer} (d={sep[layer][0]:.3f}, AUC={sep[layer][1]:.4f}, "
|
||||||
print(f" global argmax (out-of-window layers are sink-dominated, informational only): "
|
f"sink={sink[layer]*100:.3f}%, two-template |cos|={agree[layer]:.4f})")
|
||||||
f"L{int(agree.argmax())} (|cos|={float(agree.max()):.4f})")
|
agree_best = max(window, key=lambda L: agree[L])
|
||||||
print(f" using layer {layer} (|cos|={float(agree[layer]):.4f}); "
|
print(f" [diagnostic] agreement would have picked L{agree_best} "
|
||||||
f"recipe anchor L{DEFAULT_LAYER} at |cos| 0.9925")
|
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 -----------
|
# --- finite gate: a NaN/Inf direction must NEVER pass silently -----------
|
||||||
# (the sink screen alone doesn't catch this — `nan > threshold` is False, so
|
# (the sink screen alone doesn't catch this — `nan > threshold` is False, so
|
||||||
@@ -428,8 +729,10 @@ def main():
|
|||||||
blob = {
|
blob = {
|
||||||
"layer": layer, "direction": d_unit.cpu(), "sink_energy": e,
|
"layer": layer, "direction": d_unit.cpu(), "sink_energy": e,
|
||||||
"calibration": calib_prov,
|
"calibration": calib_prov,
|
||||||
"agreement": None if agree is None else float(agree[layer]),
|
"agreement": None if agree is None else agree[layer],
|
||||||
"agreement_per_layer": None if agree is None else agree.cpu(),
|
"agreement_per_layer": agree,
|
||||||
|
"separation": None if sep is None else sep[layer],
|
||||||
|
"separation_per_layer": sep,
|
||||||
"capture_window": CAPTURE_WINDOW,
|
"capture_window": CAPTURE_WINDOW,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -439,47 +742,6 @@ def main():
|
|||||||
print(f"direction saved -> {dpath} (capture-only, no write)")
|
print(f"direction saved -> {dpath} (capture-only, no write)")
|
||||||
return
|
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__":
|
if __name__ == "__main__":
|
||||||
main()
|
main()
|
||||||
|
|||||||
Reference in New Issue
Block a user