#!/usr/bin/env python3 """Capture the refusal direction for Qwen3.8-Flash-Next (qwen4_exp), 177.4B bf16. WHY THIS IS NOT capture_qwen.py ------------------------------- qwen4_exp carries `hc_count`=4 parallel hyper-connection streams: Qwen4ExpTextModel.forward does `hidden_states = hidden_states.repeat(1, 1, hc_count)` (modeling_qwen4_exp.py:1417), so the trunk tensor is (B, S, 4*2560 = 10240). Two consequences kill the 27B script: 1. `output_hidden_states=True` returns a RAGGED tuple. `_can_record_outputs["hidden_states"]` is the DecoderLayer class (modeling:1320) and `OutputRecorder.capture_initial_hidden_state` defaults True (output_capturing.py:57,112-113), so entries 0..47 are the (B,S,10240) trunk -- but `@capture_outputs` with tie_last_hidden_states=True then OVERWRITES entry 48 with `last_hidden_state`, which is (B,S,2560) (output_capturing.py:270-277). Feeding that to compute_direction.py yields 10240-dim garbage or a stack() crash, and the raw trunk after the last decoder layer is destroyed outright. 2. The 2560-dim vectors that ARE exposed (`last_hidden_state`, and the `mixed_input` returned by every `*_hyper_connection`) live in an hc_norm-WARPED basis: Qwen4ExpTextRMSNorm applies a learned per-coordinate `output * (1.0 + self.weight)` (modeling:177) over the 10240 vector, and `mixed_input` is additionally scaled by a DATA-DEPENDENT sigmoid gate `input_mix_weight` (modeling:959-965). A direction found there is diag(1+w)-distorted relative to the space the 151 residual writers actually write into. So we capture the RAW 10240 trunk via forward PRE-hooks and fold it ourselves: v = trunk[0, -1, :].float().view(4, 2560).mean(0) That is the correct linear readout of the writer basis. Every writer's output `out` is injected as `hidden_states = hyper_input + (out.unsqueeze(-2) * injection_weights.unsqueeze(-1)).flatten(-2)` (modeling:1236-1237, 1242-1243) with `injection_weights = 2*sigmoid(...)` in (0,2) -- a strictly POSITIVE per-stream SCALAR. Hence mean_k(stream_k) = embed + sum_l mean_k(w_{l,k}) * out_l, an exact positive-weighted accumulation of every writer's output in the raw 2560 basis, with no norm and no gate in between. All four streams start life identical (`repeat`), so they share one basis. CAPTURE POINTS: 49 forward pre-hooks -- `layers[i]` for i in 0..47 (args[0] = trunk entering layer i = residual AFTER i layers) plus `hyper_connection_mixer` (args[0] = raw trunk after layer 47 = residual after 48 layers). All 49 are (B,S,10240); index i means exactly what apply_ablation_flashnext.py's ABLIT_META documents: "index i = residual stream AFTER i decoder layers (index 0 = embedding output)". This also RECOVERS index 48, which output_hidden_states throws away. Writes work_/{harmful,harmless}/layer_XX.pt, each [n_prompts, 2560] -- the same format scripts/compute_direction.py expects -- and then writes work_/refusal_direction.pt itself with the exact 10-key payload apply_ablation_flashnext.py reads for --dir-source. """ import os, sys, json, argparse, time sys.path.insert(0, "/workspace/datasets/dsv4/qwen35") sys.path.insert(0, "/workspace/datasets/dsv4/scripts") import torch from pathlib import Path from transformers import AutoTokenizer from prompts import HARMFUL, HARMLESS HIDDEN = 2560 HC = 4 NGRAM_PARAM = "ple.ple_embedding.ngram_embedding.weight" def load_model(path, device, threads, reserve, per_gpu): """CPU or multi-GPU. On GPU the ~95 GiB n-gram table is kept on CPU: transformers declares it in `_no_placement_params` (modeling:1261-1264) and NGramEmbedding.forward does an explicit device round-trip (modeling:1113-1114), so offloading it is mathematically EXACT.""" import transformers kw = dict(dtype=torch.bfloat16, low_cpu_mem_usage=True, attn_implementation="sdpa") if device == "cpu": torch.set_num_threads(threads) print(f"[cpu] torch threads={torch.get_num_threads()}", flush=True) kw["device_map"] = {"": "cpu"} else: import gpuguard gpuguard.assert_supported() if per_gpu: mm = {i: f"{per_gpu}GiB" for i in range(torch.cuda.device_count())} mm["cpu"] = "220GiB" else: mm = gpuguard.auto_max_memory(reserve, cpu="220GiB") print(f"[gpu] max_memory={mm}", flush=True) kw["device_map"] = "auto" kw["max_memory"] = mm last = None # AutoModelForCausalLM -> Qwen4ExpForCausalLM drops model.visual.* (333 tensors) AND mtp.* # cleanly; the ImageTextToText class keeps the vision tower we never use. Try text-only first. for cls_name in ("AutoModelForCausalLM", "AutoModelForImageTextToText"): try: cls = getattr(transformers, cls_name) m = cls.from_pretrained(path, **kw) print(f"loaded via {cls_name}", flush=True) return m except Exception as e: last = e print(f" {cls_name} failed: {type(e).__name__}: {str(e)[:300]}", flush=True) raise last def text_model(model): """Qwen4ExpForCausalLM -> model.model ; Qwen4ExpForConditionalGeneration -> model.model.language_model.""" m = model.model return getattr(m, "language_model", m) def unit(x): x = x.float() return x / x.norm().clamp_min(1e-8) def compute_payload(harm_dir, harmless_dir, n_points): """Same math as /workspace/datasets/dsv4/scripts/compute_direction.py, inlined so the capture is self-contained and the payload schema is guaranteed to match apply_ablation_flashnext.py.""" per_layer, separation, n_h, n_hl, layers, centered = {}, {}, {}, {}, [], [] for lid in range(n_points): fh, fl = harm_dir / f"layer_{lid:02d}.pt", harmless_dir / f"layer_{lid:02d}.pt" if not (fh.exists() and fl.exists()): continue H = torch.load(fh, map_location="cpu").float() L = torch.load(fl, map_location="cpu").float() assert H.shape[-1] == HIDDEN, f"layer_{lid:02d} harmful dim {H.shape[-1]} != {HIDDEN}" assert L.shape[-1] == HIDDEN, f"layer_{lid:02d} harmless dim {L.shape[-1]} != {HIDDEN}" d = H.mean(0) - L.mean(0) separation[lid] = float(d.norm()) # raw ||diff||, NOT normalized (grows with depth) per_layer[str(lid)] = unit(d) # str keys: apply script does per_layer[str(i)] n_h[str(lid)], n_hl[str(lid)] = int(H.shape[0]), int(L.shape[0]) layers.append(lid) centered.append(H - L.mean(0, keepdim=True)) assert layers, "no layer files found" layer_units = torch.stack([per_layer[str(l)] for l in layers], 0) # [n_layers, 2560] broad = unit(layer_units.mean(0)) deep = per_layer[str(max(layers))] M = torch.cat([torch.cat(centered, 0), layer_units], 0) _, S, Vh = torch.linalg.svd(M, full_matrices=False) k = min(4, Vh.shape[0]) return { "broad": broad, # (2560,) "deep": deep, # (2560,) "directions": Vh[:k].contiguous().float(), # (k, 2560) "singular_values": [float(x) for x in S[:k]], "per_layer": per_layer, # {str(lid): (2560,)} "separation": separation, # {int lid: float} "n_harmful": n_h, "n_harmless": n_hl, "layers": layers, "n_directions": k, } def verify_payload(p): """apply_ablation_flashnext.py:126-138 reads exactly these. Getting it wrong wastes the run.""" for k in ("broad", "deep", "directions", "per_layer", "separation"): assert k in p, f"payload missing required key {k!r}" assert p["broad"].shape == (HIDDEN,), p["broad"].shape assert p["deep"].shape == (HIDDEN,), p["deep"].shape assert p["directions"].ndim == 2 and p["directions"].shape[-1] == HIDDEN, p["directions"].shape best = max(p["separation"], key=lambda k: p["separation"][k]) assert str(best) in p["per_layer"], f"separation key {best!r} has no per_layer[{str(best)!r}]" for lid, v in p["per_layer"].items(): assert isinstance(lid, str), f"per_layer key {lid!r} must be str" assert v.shape == (HIDDEN,), f"per_layer[{lid}] {tuple(v.shape)} != ({HIDDEN},)" print(f" payload OK: {len(p['per_layer'])} layers, dirs={tuple(p['directions'].shape)}, " f"--dir-source best -> layer:{best}", flush=True) def main(): ap = argparse.ArgumentParser() base = Path("/workspace/datasets/dsv4/qwen4") ap.add_argument("--model", default=str(base / "flashnext")) ap.add_argument("--work-root", default=str(base)) ap.add_argument("--device", choices=["cpu", "gpu"], default="gpu") ap.add_argument("--threads", type=int, default=48) ap.add_argument("--reserve", type=float, default=8.0, help="GiB left free per GPU (shared host)") # 80GiB verified on meta: ngram -> cpu (95.4GiB), layers -> cuda:1/2/3 at 78.5/77.3/78.5 GiB, # nothing else spilled, ~61GiB headroom per card. Above ~100GiB the no_placement escape stops # firing and the 95GiB ngram table lands on a GPU instead. ap.add_argument("--per-gpu", default="80", help="fixed GiB/GPU cap; empty -> gpuguard.auto_max_memory") ap.add_argument("--modes", default="nothink") ap.add_argument("--reasoning-effort", default="medium", help="think mode only; 'xhigh' (the template default) injects an extra system turn") ap.add_argument("--limit", type=int, default=0, help="use only the first N prompts of each set") ap.add_argument("--dry-run", action="store_true", help="tiny randomly-initialised CPU model, no checkpoint load, full pipeline") ap.add_argument("--smoke", action="store_true", help="short greedy generate before capture") ap.add_argument("--force", action="store_true", help="ignore checkpoints, recapture everything") args = ap.parse_args() if args.dry_run: return dry_run(args) t0 = time.time() tok = AutoTokenizer.from_pretrained(args.model) model = load_model(args.model, args.device, args.threads, args.reserve, args.per_gpu) model.eval() lm = text_model(model) n_layers = len(lm.layers) n_points = n_layers + 1 # 0..47 pre-layer + 48 pre-mixer dev = lm.embed_tokens.weight.device if dev.type == "meta": dev = torch.device("cuda:0") print(f"loaded in {time.time()-t0:.0f}s | layers={n_layers} points={n_points} embed_dev={dev}", flush=True) if args.device == "gpu": ng = next((p for n, p in lm.named_parameters() if n.endswith(NGRAM_PARAM)), None) if ng is not None: print(f" ngram table device = {ng.device} (want cpu; {ng.numel()*2/2**30:.1f} GiB)", flush=True) for i in range(torch.cuda.device_count()): free, tot = torch.cuda.mem_get_info(i) print(f" cuda:{i} used={(tot-free)/2**30:.1f}GiB free={free/2**30:.1f}GiB", flush=True) run(args, tok, model, lm, dev, n_points) def run(args, tok, model, lm, dev, n_points): buf = {} def mk(i): def pre_hook(_mod, hook_args): buf[i] = hook_args[0] return pre_hook handles = [lm.layers[i].register_forward_pre_hook(mk(i)) for i in range(len(lm.layers))] handles.append(lm.hyper_connection_mixer.register_forward_pre_hook(mk(len(lm.layers)))) print(f"registered {len(handles)} forward pre-hooks " f"(layers 0..{len(lm.layers)-1} + hyper_connection_mixer)", flush=True) @torch.inference_mode() def capture(prompts, outdir, thinking): outdir = Path(outdir) outdir.mkdir(parents=True, exist_ok=True) tag = f"{outdir.parent.name}/{outdir.name}" done = outdir / "_partial.pt" acc, start = None, 0 if done.exists() and not args.force: ck = torch.load(done, map_location="cpu", weights_only=False) if ck["n_prompts"] == len(prompts) and ck["thinking"] == thinking: acc, start = ck["acc"], ck["done"] print(f" [{tag}] resuming from prompt {start}/{len(prompts)}", flush=True) if acc is None: acc = {lid: [] for lid in range(n_points)} # running sum of the UNFOLDED 10240 vector, for the per-stream sanity check below full = {lid: torch.zeros(HC * HIDDEN, dtype=torch.float64) for lid in range(n_points)} if start >= len(prompts): print(f" [{tag}] already complete", flush=True) t0 = time.time() for i in range(start, len(prompts)): kw = dict(tokenize=False, add_generation_prompt=True, enable_thinking=thinking) if thinking: # the template's default reasoning_effort='xhigh' silently prepends a whole # system turn ("Reasoning effort is set to xhigh..."); 'medium' emits none. kw["reasoning_effort"] = args.reasoning_effort txt = tok.apply_chat_template([{"role": "user", "content": prompts[i]}], **kw) # batch=1, no padding -> create_recurrent_attention_mask returns None, so the PLE # eos-substitution (modeling:1409-1413) never fires and the DeltaNet conv sees no pad. enc = tok(txt, return_tensors="pt", add_special_tokens=False).to(dev) buf.clear() # ple_input_ids is derived from input_ids automatically (modeling:1360-1362). model(input_ids=enc["input_ids"], use_cache=False) assert len(buf) == n_points, f"got {len(buf)} hook fires, want {n_points}" for lid in range(n_points): h = buf[lid] assert h.shape[-1] == HC * HIDDEN, \ f"point {lid} dim {h.shape[-1]} != {HC*HIDDEN}; trunk layout changed" raw = h[0, -1, :].float().cpu() full[lid] += raw.double() v = raw.view(HC, HIDDEN).mean(0) assert v.shape == (HIDDEN,), v.shape acc[lid].append(v) buf.clear() if i == start: print(f" [{tag}] points={n_points} trunk_dim={HC*HIDDEN} -> folded {HIDDEN} " f"| tokens={enc['input_ids'].shape[1]} first_forward={time.time()-t0:.1f}s", flush=True) if (i + 1) % 4 == 0 or i + 1 == len(prompts): torch.save({"acc": acc, "done": i + 1, "n_prompts": len(prompts), "thinking": thinking}, done) print(f" [{tag}] {i+1}/{len(prompts)} " f"({(time.time()-t0)/max(i+1-start,1):.1f}s/prompt)", flush=True) nonfinite = 0 for lid, vecs in acc.items(): t = torch.stack(vecs, 0) if not torch.isfinite(t).all(): nonfinite += 1 assert t.shape[-1] == HIDDEN, t.shape torch.save(t, outdir / f"layer_{lid:02d}.pt") # only a full pass gives an unbiased per-stream mean; a resumed run keeps the old file if start == 0 and len(prompts) > 0: torch.save({lid: (full[lid] / len(prompts)).float() for lid in full}, outdir / "stream_means.pt") print(f" saved [{len(prompts)}x{n_points}] -> {outdir} nonfinite_layers={nonfinite} " f"({time.time()-t0:.0f}s)", flush=True) harmful = HARMFUL[: args.limit] if args.limit else HARMFUL harmless = HARMLESS[: args.limit] if args.limit else HARMLESS print(f"prompt sets: harmful={len(harmful)} harmless={len(harmless)}", flush=True) for mode in args.modes.split(","): thinking = (mode == "think") root = Path(args.work_root) / f"work_{mode}" print(f"\n### mode={mode} enable_thinking={thinking} ###", flush=True) capture(harmful, root / "harmful", thinking) capture(harmless, root / "harmless", thinking) p = compute_payload(root / "harmful", root / "harmless", n_points) verify_payload(p) sanity(root, p, n_points) out = root / "refusal_direction.pt" torch.save(p, out) print(f" wrote {out}", flush=True) top = sorted(p["separation"].items(), key=lambda kv: -kv[1])[:6] print(f" separation top6: {[(l, round(s,2)) for l,s in top]}", flush=True) print("CAPTURE DONE", flush=True) def sanity(root, p, n_points): """The shared-basis assumption is what makes the 4-stream MEAN legal. Check it on the real weights: the 4 PER-STREAM diff-of-means should be near-parallel, since they differ only by accumulated POSITIVE scalars (injection_weights = 2*sigmoid(.) in (0,2), modeling:968). If any pair drops below ~0.9 the assumption has broken -- stop and re-derive before surgery.""" fh = root / "harmful" / "stream_means.pt" fl = root / "harmless" / "stream_means.pt" print(f" [sanity] cos(broad,deep)={float(p['broad'] @ p['deep']):.3f}", flush=True) if not (fh.exists() and fl.exists()): print(" [sanity] stream_means.pt missing (resumed run) -- per-stream check SKIPPED", flush=True) return mh, ml = torch.load(fh, map_location="cpu"), torch.load(fl, map_location="cpu") worst, worst_lid = 1.0, -1 for lid in sorted(set(mh) & set(ml)): d = (mh[lid] - ml[lid]).float().view(HC, HIDDEN) if d.norm() < 1e-6: continue # layer 0 is degenerate (same last token) u = d / d.norm(dim=1, keepdim=True).clamp_min(1e-8) c = (u @ u.T) lo = float(c[~torch.eye(HC, dtype=torch.bool)].min()) if lo < worst: worst, worst_lid = lo, lid print(f" [sanity] worst pairwise per-stream cosine = {worst:.4f} at point {worst_lid} " f"({'OK' if worst >= 0.9 else 'WARNING -- stream mean may be cancelling signal'})", flush=True) def dry_run(args): """Tiny randomly-initialised qwen4_exp on CPU: exercises hooks, fold, payload and schema verification end to end in seconds, without touching the 336 GiB checkpoint or any GPU.""" import warnings warnings.filterwarnings("ignore") from transformers import AutoConfig from transformers.models.qwen4_exp.modeling_qwen4_exp import Qwen4ExpForCausalLM torch.manual_seed(0) c = AutoConfig.from_pretrained(args.model).text_config c.hidden_size = HIDDEN # keep the real hidden size so the 2560 asserts are real c.num_hidden_layers = 3 c.head_dim = 32; c.num_attention_heads = 4; c.num_key_value_heads = 2 c.vocab_size = 1024; c.num_experts = 4; c.num_experts_per_tok = 2 c.moe_intermediate_size = 32; c.shared_expert_intermediate_size = 32; c.intermediate_size = 64 c.hc_lowrank = 16 c.layer_types = ["linear_attention", "linear_attention", "full_attention"] c.ple_layer_ids = [2]; c.split_ngram_parts = 2 c.linear_num_value_heads = 4; c.linear_num_key_heads = 2 c.linear_key_head_dim = 32; c.linear_value_head_dim = 32 c.ngram_vocab_size_base = 2000; c.ple_embed_dim = HIDDEN print(f"[dry-run] tiny qwen4_exp hidden={c.hidden_size} layers={c.num_hidden_layers} " f"hc_count={c.hc_count}", flush=True) model = Qwen4ExpForCausalLM(c).eval() lm = text_model(model) tok = AutoTokenizer.from_pretrained(args.model) dev = torch.device("cpu") class TinyTok: """Real template, but ids clamped into the tiny vocab.""" def apply_chat_template(self, *a, **k): return tok.apply_chat_template(*a, **k) def __call__(self, *a, **k): e = tok(*a, **k) e["input_ids"] = e["input_ids"] % c.vocab_size return e args.limit = args.limit or 4 args.work_root = args.work_root if args.work_root.startswith("/tmp") else "/tmp/fn_dryrun" print(f"[dry-run] work-root={args.work_root} limit={args.limit}", flush=True) run(args, TinyTok(), model, lm, dev, c.num_hidden_layers + 1) print("[dry-run] OK -- hooks, 4-stream fold, 2560 asserts and payload schema all pass.\n" "[dry-run] NOTE: a low per-stream cosine WARNING here is expected and meaningless -- the\n" "[dry-run] weights are random and hc_norm.weight is zero-initialised. Only the\n" "[dry-run] warning from a REAL capture run is informative.", flush=True) if __name__ == "__main__": main()