#!/usr/bin/env python3
"""Robinson-formula abliteration of Qwen3.8-27B (MTP-aware, vision-preserving).

Implements the recipe documented in
`docs/pfi/abliteration-recipe-qwen38.md` (captured from RobinsonLabs). Single
refusal direction, orthogonalized out of every residual-stream WRITER, including
the MTP head in-band and preserving the vision tower byte-identical.

This is deliberately conservative and gated. It REFUSES to write a byte unless:
  1. the residual-writer coverage identity holds
     (o_proj + linear_out == num_hidden_layers), and
  2. the chosen layer's direction is not dominated by the attention-sink
     dimension (Qwen3.8-27B: dim 3994 — orthogonalizing it out bricks the model).

Both gates come straight from the recipe; both are failure modes that otherwise
ship a model that loads and runs but is half-abliterated or emits garbage.

Modes:
  --dry-run     enumerate + classify tensors, run BOTH gates on the static
                surface (no model forward, no capture, no write). Run this first
                against the real checkpoint to confirm the tensor map.
  --capture     load the model, derive the refusal direction from two chat
                templates, screen dim 3994, save the direction + report. No write.
  (default)     capture (or --direction <file>) then orthogonalize and save the
                abliterated bf16 to --out.

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

import torch

# --- architecture constants (Qwen3.8-27B, verified against the recipe) --------
NUM_LAYERS = 64
HIDDEN = None  # read from config at runtime
FULL_ATTN_INTERVAL = 4            # self_attn.o_proj on every 4th layer -> 16
EXPECT_O_PROJ = 16
EXPECT_LINEAR_OUT = 48
EXPECT_DOWN_PROJ = 64
EXPECT_MTP_WRITERS = 2            # mtp.layers.0: o_proj + down_proj
SINK_DIM = 3994                   # massive-activation dim; must NOT be orthogonalized out
SINK_ENERGY_MAX = 0.01           # >1% of direction energy in dim 3994 at chosen layer -> abort
DEFAULT_LAYER = 26               # recipe peak-agreement layer (|cos| 0.9925); auto-picked, this is the sanity anchor

# Residual-writer suffixes. A tensor writing INTO the residual stream has output
# dim == hidden_size; orthogonalizing removes its ability to write along the
# refusal direction. Classified by suffix so this survives minor name drift; the
# coverage gate below catches any misclassification.
WRITER_SUFFIXES = {
    "mlp.down_proj.weight": "down_proj",
    "self_attn.o_proj.weight": "o_proj",
    "linear_attn.out_proj.weight": "linear_out",
}
MTP_WRITER_SUFFIXES = ("o_proj.weight", "down_proj.weight")  # within an mtp block
EMBED_SUFFIX = "embed_tokens.weight"
VISION_MARKERS = (".visual.", "visual.")  # never touched


def is_vision(name: str) -> bool:
    return any(m in name for m in VISION_MARKERS)


def is_mtp(name: str) -> bool:
    return ".mtp." in name or name.startswith("mtp.") or ".nextn." in name


def classify_writers(keys):
    """Map every residual-writer tensor to its class. Vision is excluded up front."""
    trunk = {"down_proj": [], "o_proj": [], "linear_out": []}
    mtp_writers, embed = [], []
    for k in keys:
        if is_vision(k):
            continue
        if is_mtp(k):
            if k.endswith(MTP_WRITER_SUFFIXES) and ("o_proj" in k or "down_proj" in k):
                mtp_writers.append(k)
            continue
        for suf, cls in WRITER_SUFFIXES.items():
            if k.endswith(suf):
                trunk[cls].append(k)
                break
        if k.endswith(EMBED_SUFFIX):
            embed.append(k)
    return trunk, mtp_writers, embed


def coverage_gate(trunk, mtp_writers, embed):
    """Hard gate — the recipe's o_proj(16) + linear_out(48) == 64 identity,
    plus down_proj==64, MTP writers==2, exactly one embed. Returns (ok, report)."""
    n_o, n_lin, n_dn = len(trunk["o_proj"]), len(trunk["linear_out"]), len(trunk["down_proj"])
    checks = {
        "o_proj + linear_out == NUM_LAYERS": (n_o + n_lin == NUM_LAYERS, f"{n_o}+{n_lin}={n_o+n_lin} vs {NUM_LAYERS}"),
        "o_proj == 16": (n_o == EXPECT_O_PROJ, f"{n_o} vs {EXPECT_O_PROJ}"),
        "linear_out == 48": (n_lin == EXPECT_LINEAR_OUT, f"{n_lin} vs {EXPECT_LINEAR_OUT}"),
        "down_proj == 64": (n_dn == EXPECT_DOWN_PROJ, f"{n_dn} vs {EXPECT_DOWN_PROJ}"),
        "mtp writers == 2": (len(mtp_writers) == EXPECT_MTP_WRITERS, f"{len(mtp_writers)} vs {EXPECT_MTP_WRITERS}"),
        "exactly one embed_tokens": (len(embed) == 1, f"{len(embed)}"),
    }
    ok = all(v[0] for v in checks.values())
    return ok, checks


def load_config(model_dir: Path):
    cfg = json.loads((model_dir / "config.json").read_text())
    tc = cfg.get("text_config", cfg)
    return cfg, tc


# --- direction capture (Arditi-style, two-template agreement) ------------------
#
# Calibration corpora live in calibration.py. The capture window (the layer range
# the direction may be picked from) comes straight from the recipe.
CAPTURE_WINDOW = (18, 45)         # inclusive; recipe reports |cos| 0.96-0.99 here


def render(tokenizer, prompt, thinking):
    msgs = [{"role": "user", "content": prompt}]
    kw = {}
    # Qwen chat templates gate thinking via enable_thinking; xhigh path injects
    # an extra system block, shifting positions — exactly the two renderings the
    # recipe captured to prove the direction is refusal-semantic, not template.
    try:
        return tokenizer.apply_chat_template(
            msgs, tokenize=False, add_generation_prompt=True,
            enable_thinking=thinking, **kw)
    except TypeError:
        return tokenizer.apply_chat_template(
            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, layers):
    """Last-real-token hidden state at the requested layers, for a batch.

    Returns {layer: [B, hidden]} on CPU in float32.

    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 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
    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, layers):
    """Prove padded-batch == one-at-a-time before spending the capture window.

    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, 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 _collect_hidden(model, tokenizer, prompts, thinking, device, batch_size, layers, label):
    """Per-prompt last-token hidden states. {layer: [N, hidden]} float32 on CPU.

    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}")
    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]
        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 {i: torch.cat(chunks[i]) for i in layers}


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)
        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
        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):
    e = (direction_vec[dim] ** 2) / (direction_vec ** 2).sum().clamp_min(1e-12)
    return float(e)


def orthogonalize_(weight, d_unit):
    """Remove the refusal component from a residual-WRITE matrix in place.
    weight: [out=hidden, in]; W <- (I - d d^T) W = W - d (d^T W)."""
    d = d_unit.to(weight.dtype).to(weight.device)
    coeff = d @ weight                # [in]
    weight.sub_(torch.outer(d, coeff))


def orthogonalize_embed_(weight, d_unit):
    """embed_tokens [vocab, hidden]: strip the refusal component from each row.
    E <- E - (E d) d^T."""
    d = d_unit.to(weight.dtype).to(weight.device)
    coeff = weight @ d                # [vocab]
    weight.sub_(torch.outer(coeff, d))


def write_abliterated(model_dir: Path, args, targets, embed_keys, n_vision):
    """Orthogonalize the 131 residual writers SHARD BY SHARD and write a new checkpoint.

    This is deliberately not done through a loaded model object, and that is a
    correctness requirement rather than a preference. `AutoModelForCausalLM`
    resolves to `Qwen3_5ForCausalLM` — the TEXT model. Saving from it would drop
    all 333 vision tensors, silently violating the recipe's byte-identical-vision
    guarantee; and the `ForConditionalGeneration` wrapper does not load the MTP
    head at all (the same reason the incumbent gen seat's Heretic pass left its
    MTP head an untouched base graft), so the in-band MTP edit that is the whole
    point of the Robinson formula would be skipped. Neither failure raises.

    Operating on the shards instead: every tensor we do not target is re-serialized
    from the exact bytes we read, so vision and the other 1068 tensors are
    byte-identical by construction, and the MTP writers are just two more keys.
    No GPU, no accelerate, no offload, no meta tensors — the whole class of
    silent-no-op failures goes away with the model object.

    Per playbook 3.6, shards are read with plain `read()` + `load()` rather than
    mmap: `safe_open` mmaps a whole shard and a 50 GB shard ENOMEMs on ZFS
    regardless of free RAM.
    """
    from glob import glob
    import shutil
    from safetensors.torch import load as st_load, save_file

    if not args.out:
        print("\n!! --out is required to write the abliterated model "
              "(use --capture for direction-only).", file=sys.stderr)
        sys.exit(1)
    if not args.direction:
        print("\n!! --direction <refusal-direction.pt> is required for the write. Capture "
              "first (--capture), inspect the agreement and sink energy, then write.",
              file=sys.stderr)
        sys.exit(1)

    blob = torch.load(args.direction, weights_only=False)
    layer, d_unit, e = blob["layer"], blob["direction"], blob.get("sink_energy")
    calib = blob.get("calibration", {})
    print(f"\ndirection: layer {layer}, sink energy {e*100:.3f}%, "
          f"agreement {blob.get('agreement')}, calib {calib.get('calib')} "
          f"({calib.get('n_harmful')}/{calib.get('n_harmless')})")

    if not torch.isfinite(d_unit).all():
        print("\n!! direction is not finite — refusing to write.", file=sys.stderr)
        sys.exit(4)
    d_unit = (d_unit.float() / d_unit.float().norm().clamp_min(1e-12)).cpu()
    if e is not None and e > SINK_ENERGY_MAX:
        print(f"\n!! sink-energy gate FAILED ({e*100:.3f}% > {SINK_ENERGY_MAX*100:.1f}%) — "
              f"orthogonalizing this direction would brick the model.", file=sys.stderr)
        sys.exit(3)

    out_dir = Path(args.out)
    if out_dir.exists() and any(out_dir.glob("*.safetensors")):
        print(f"\n!! {out_dir} already holds safetensors shards — refusing to overwrite an "
              f"existing checkpoint. Move it aside or pick another --out.", file=sys.stderr)
        sys.exit(10)
    out_dir.mkdir(parents=True, exist_ok=True)

    targets = set(targets)
    shards = sorted(glob(str(model_dir / "*.safetensors")))
    edited, seen_targets, total_tensors = 0, set(), 0
    for si, shard in enumerate(shards, 1):
        with open(shard, "rb") as f:
            tensors = st_load(f.read())
        total_tensors += len(tensors)
        hits = [k for k in tensors if k in targets]
        for k in hits:
            w = tensors[k]
            orig_dtype = w.dtype
            # Math in fp32. The weights are bf16 (8 mantissa bits); computing
            # d^T W and the rank-1 subtraction at that precision would lose more
            # than the edit itself is worth.
            w32 = w.float()
            if k in embed_keys:
                orthogonalize_embed_(w32, d_unit)      # [vocab, hidden]
            else:
                orthogonalize_(w32, d_unit)            # [hidden, in]
            tensors[k] = w32.to(orig_dtype)
            edited += 1
            seen_targets.add(k)
        save_file(tensors, str(out_dir / Path(shard).name), metadata={"format": "pt"})
        print(f"  shard {si}/{len(shards)} {Path(shard).name}: {len(hits)} edited", flush=True)
        del tensors

    missed = targets - seen_targets
    if missed or edited != len(targets):
        print(f"\n!! surgery incomplete — edited {edited} of {len(targets)} targets, "
              f"{len(missed)} never found in any shard: {sorted(missed)[:5]}", file=sys.stderr)
        sys.exit(7)
    print(f"  edited {edited} tensors of {total_tensors}; vision ({n_vision}) byte-identical")

    # Everything that is not weights rides along unchanged.
    for pat in ("*.json", "*.jinja", "*.txt", "*.model", "*.py"):
        for src in sorted(model_dir.glob(pat)):
            if src.name in ("dl.py",):
                continue
            shutil.copy2(src, out_dir / src.name)
    torch.save(blob, out_dir / "refusal-direction.pt")
    print(f"\nwrote abliterated checkpoint -> {out_dir}")
    print("DONE.")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True, help="bf16 checkpoint dir")
    ap.add_argument("--out", help="output dir for the abliterated bf16")
    ap.add_argument("--layer", type=int, default=None, help="override the auto-picked direction layer")
    ap.add_argument("--direction", help="load a saved direction .pt instead of capturing")
    ap.add_argument("--dry-run", action="store_true", help="enumerate + gate only; no forward, no write")
    ap.add_argument("--capture", action="store_true", help="capture direction + screen sink; no write")
    ap.add_argument("--calib", choices=("mlabonne", "builtin"), default="mlabonne",
                    help="calibration corpus: 'mlabonne' = the recipe's 416-prompt AdvBench "
                         "train split + alpaca (default); 'builtin' = the legacy inline 8/8")
    ap.add_argument("--calib-n-harmful", type=int, default=416,
                    help="harmful calibration prompts (416 = the full train split, as the recipe used)")
    ap.add_argument("--calib-n-harmless", type=int, default=416,
                    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 "
                         "layers above N, so any value > the capture window's top (45) leaves the "
                         "chosen direction bit-identical while cutting fp32 weight residency and "
                         "forward cost by the dropped fraction. Off by default.")
    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

    # --- static surface: enumerate every tensor key from the shards -----------
    keys = []
    for shard in sorted(glob(str(model_dir / "*.safetensors"))):
        with safe_open(shard, framework="pt") as f:
            keys.extend(f.keys())
    trunk, mtp_writers, embed = classify_writers(keys)
    ok, checks = coverage_gate(trunk, mtp_writers, embed)

    n_vision = sum(1 for k in keys if is_vision(k))
    print(f"tensors: {len(keys)} total | vision preserved: {n_vision}")
    print(f"writers: down_proj={len(trunk['down_proj'])} o_proj={len(trunk['o_proj'])} "
          f"linear_out={len(trunk['linear_out'])} mtp={len(mtp_writers)} embed={len(embed)}")
    print("COVERAGE GATE:")
    for name, (passed, detail) in checks.items():
        print(f"  [{'PASS' if passed else 'FAIL'}] {name}  ({detail})")
    if not ok:
        print("\n!! coverage gate FAILED — tensor names do not match the recipe. "
              "Inspect the checkpoint; do NOT abliterate blind.", file=sys.stderr)
        sys.exit(2)
    print("  -> coverage gate PASSED")
    total_edits = len(trunk["down_proj"]) + len(trunk["o_proj"]) + len(trunk["linear_out"]) + len(mtp_writers) + len(embed)
    print(f"  -> {total_edits} tensors would be orthogonalized (recipe expects 131)")

    if args.dry_run:
        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)...")
    tok = AutoTokenizer.from_pretrained(model_dir)
    # Right padding + a real pad id, both required by the batched capture. See
    # the padding-side note in last_token_hidden(): right is the correct side
    # here, not merely a preference.
    tok.padding_side = "right"
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    # 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:
        lo, hi = CAPTURE_WINDOW
        if not args.capture:
            # A truncated model would save_pretrained as a truncated CHECKPOINT.
            # Capture-only, no exceptions.
            print("\n!! --max-layer is a capture-only optimization; on the write path it would "
                  "emit a decoder missing its upper layers. Drop it, or add --capture.",
                  file=sys.stderr)
            sys.exit(5)
        if args.max_layer <= hi:
            print(f"\n!! --max-layer {args.max_layer} would truncate inside the capture window "
                  f"[{lo},{hi}] — the direction layer must exist. Use > {hi}.", file=sys.stderr)
            sys.exit(5)
        cfg_obj = AutoConfig.from_pretrained(model_dir)
        tcfg = getattr(cfg_obj, "text_config", cfg_obj)
        tcfg.num_hidden_layers = args.max_layer
        # layer_types is per-layer (linear_attention / full_attention every 4th);
        # it has to be truncated in step or the built stack disagrees with itself.
        if getattr(tcfg, "layer_types", None):
            tcfg.layer_types = list(tcfg.layer_types)[:args.max_layer]
        load_kwargs["config"] = cfg_obj
        print(f"  truncating decoder to {args.max_layer}/{NUM_LAYERS} layers for capture "
              f"(exact for any layer <= {args.max_layer}; window top is {hi})")

    model = AutoModelForCausalLM.from_pretrained(model_dir, **load_kwargs)
    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
        # relying on the caller's cwd landing in the right place.
        sys.path.insert(0, str(Path(__file__).resolve().parent))
        from calibration import load_calibration
        harmful, harmless, calib_prov = load_calibration(
            args.calib, args.calib_n_harmful, args.calib_n_harmless, args.calib_seed)
        print(f"\ncalibration corpus: {calib_prov['source']}")
        print(f"  harmful={calib_prov['n_harmful']}  harmless={calib_prov['n_harmless']}  "
              f"seed={calib_prov['seed']}"
              + (f"  held-out reserved={calib_prov['heldout_reserved']}"
                 if "heldout_reserved" in calib_prov else ""))

        # --- 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, 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 {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, sep = capture_direction(model, tok, device, harmful, harmless,
                                             args.batch_size, window)

        # LAYER SELECTION — by SEPARATION, not by two-template agreement.
        #
        # The recipe picks the layer by peak |cos| between the no-think and xhigh
        # renderings. On this checkpoint that metric is actively misleading, and
        # following it cost a full write-and-test cycle for a no-op. Measured
        # 2026-08-20: agreement ranked layer 18 first (0.6238) — and layer 18 has
        # the WORST harmful/harmless separation of the entire window (Cohen's d
        # 5.51 vs 9.89 at layer 39). Abliterating there changed nothing: stock and
        # abliterated refused all six probe prompts identically.
        #
        # The reason agreement fails here is that the two renderings do not merely
        # differ in formatting — they leave the model in different generative
        # modes at the token we read (`</think>\n\n` = about to answer, `<think>\n`
        # = about to reason). So |cos| scores refusal semantics *plus* mode, and on
        # a heavily-merged base the mode term dominates. Robinson's stock
        # Qwen3.8-27B scored 0.99 across that same split; this model scores 0.62,
        # and that difference says more about the templates than the direction.
        #
        # Separation asks the question that actually predicts efficacy: does this
        # direction split harmful from harmless prompts? Here it does, superbly
        # (AUC 0.9996+ across the whole window) — the direction was never the
        # problem, only where we removed it. Agreement is still reported, as a
        # diagnostic rather than a selector.
        # The sink screen is a FILTER on selection, not just a post-hoc abort.
        # Separation and sink-energy both climb with depth on this model, so the
        # best-separating layer (39, d=9.89) is also sink-dominated (1.97% > 1%)
        # and would brick the model. Pick the best separator *among layers that
        # pass the screen* — one pass, no guess-and-retry.
        sink = {L: sink_energy(dirs[False][L]) for L in window}
        by_sep = lambda L: sep[L][0]
        eligible = [L for L in window if sink[L] <= SINK_ENERGY_MAX]
        print(f"\nlayer selection over [{lo},{hi}] — separation, gated on sink < "
              f"{SINK_ENERGY_MAX*100:.1f}%:")
        for L in sorted(window, key=by_sep, reverse=True)[:8]:
            mark = "ok " if sink[L] <= SINK_ENERGY_MAX else "SINK"
            print(f"    [{mark}] L{L:<3} d={sep[L][0]:6.3f}  AUC={sep[L][1]:.4f}  "
                  f"sink={sink[L]*100:6.3f}%  |cos|={agree[L]:.4f}")
        if not eligible:
            print("\n!! every layer in the window is sink-dominated — no safe direction exists "
                  "here. Widen the window or reconsider the approach.", file=sys.stderr)
            sys.exit(3)
        best = max(eligible, key=by_sep)
        layer = args.layer if args.layer is not None else best
        d_unit = dirs[False][layer]  # thinking-off direction at the chosen layer
        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
    #  a NaN direction would "pass" the sink gate. This is the real guard.)
    if not torch.isfinite(d_unit).all():
        frac = float(torch.isfinite(d_unit).float().mean())
        print(f"\n!! captured direction is NOT finite (finite frac {frac:.3f}) — "
              "the forward pass produced NaN/Inf. Check attn_implementation and the "
              "fla/linear-attn path; do NOT abliterate on this direction.", file=sys.stderr)
        sys.exit(4)

    # --- attention-sink screen (the brick-the-model gate) ---------------------
    e = sink_energy(d_unit)
    print(f"attention-sink screen: dim {SINK_DIM} carries {e*100:.3f}% of layer-{layer} direction energy "
          f"(recipe L26 ref: 0.06%; abort threshold {SINK_ENERGY_MAX*100:.1f}%)")
    if e > SINK_ENERGY_MAX:
        print("\n!! sink-energy gate FAILED — orthogonalizing this direction would brick the model. "
              "Pick a different layer.", file=sys.stderr)
        sys.exit(3)
    print("  -> sink screen PASSED")

    blob = {
        "layer": layer, "direction": d_unit.cpu(), "sink_energy": e,
        "calibration": calib_prov,
        "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,
    }

    if args.capture:
        dpath = model_dir / "refusal-direction.pt"
        torch.save(blob, dpath)
        print(f"direction saved -> {dpath}  (capture-only, no write)")
        return


if __name__ == "__main__":
    main()