sae: loader + build_sae_data + eval_sae (cross-uplift metric); weights_only=False for torch 2.13

scripts/build_sae_data.py

@@ -0,0 +1,88 @@
+"""Build an SAE-feature post-train dataset in the exact vec-bank format pretrain.py consumes.
+
+Per feature: direction = unit(W_enc[:,f]) (encoder column, probe side — the same side eval_sae.py
+rewards); targets = the feature's top --targets corpus windows (by per-window peak act) decoded to
+text. Features are split (seeded, --eval-frac) into train/eval and ONLY train features enter
+vecs.f32/records.jsonl — the eval half is reserved for cross-uplift (eval_sae.py --split
+<out>/split.json). build_stats.json carries n_examples so scripts/rl.py --data-dir <out> also
+works on this bank unchanged.
+
+    python scripts/build_sae_data.py --out-dir data/sae --targets 3
+"""
+import argparse
+import json
+import os
+import random
+
+import numpy as np
+import torch
+from transformers import AutoTokenizer
+
+from mxf.config import D_MODEL, MODEL
+from mxf.sae import load_max_acts, load_sae
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--sae-path", default=None, help="ae.pt path; default: hf_hub_download")
+    ap.add_argument("--maxacts-path", default=None, help="max-acts .pt path; default: hf_hub_download")
+    ap.add_argument("--out-dir", default="data/sae")
+    ap.add_argument("--n-features", type=int, default=0, help="0 = ALL alive features; else seeded sample")
+    ap.add_argument("--targets", type=int, default=3, help="top corpus windows per feature")
+    ap.add_argument("--min-act", type=float, default=0.0, help="drop features with corpus peak <= this")
+    ap.add_argument("--eval-frac", type=float, default=0.5)
+    ap.add_argument("--seed", type=int, default=0)
+    a = ap.parse_args()
+    os.makedirs(a.out_dir, exist_ok=True)
+    rng = random.Random(a.seed)
+
+    tok = AutoTokenizer.from_pretrained(MODEL)
+    sae = load_sae(a.sae_path)                                # cpu fp32
+    data = load_max_acts(a.maxacts_path)
+    tokens, acts = data["max_tokens"], data["max_acts"]       # [F, N, L]
+    assert acts.shape[0] == sae.d_sae, f"max-acts F={acts.shape[0]} != SAE F={sae.d_sae}"
+    dataset_max = acts.amax(dim=(1, 2))                       # [F] corpus peak per feature
+
+    alive = (dataset_max > a.min_act).nonzero(as_tuple=True)[0].tolist()
+    if a.n_features and a.n_features < len(alive):
+        alive = rng.sample(alive, a.n_features)
+    rng.shuffle(alive)
+    n_eval = int(len(alive) * a.eval_frac)
+    ev, train = sorted(alive[:n_eval]), sorted(alive[n_eval:])
+    json.dump({"seed": a.seed, "eval_frac": a.eval_frac, "min_act": a.min_act,
+               "train": train, "eval": ev}, open(f"{a.out_dir}/split.json", "w"))
+    print(f"{len(alive)} alive features (min_act {a.min_act}) -> {len(train)} train / {len(ev)} eval",
+          flush=True)
+
+    dirs = torch.nn.functional.normalize(sae.W_enc, dim=0)    # [d, F] unit encoder columns
+    rows, skipped = [], 0
+    for j, f in enumerate(train):
+        peak = acts[f].amax(dim=-1)                           # [N] per-window peaks
+        for w in peak.argsort(descending=True)[: a.targets].tolist():
+            if peak[w] <= 0:
+                break                                         # desc order: rest are dead too
+            text = tok.decode(tokens[f, w].tolist(), skip_special_tokens=True).strip()
+            if len(text) < 3:
+                skipped += 1
+                continue
+            rows.append((f, text, peak[w].item()))
+        if (j + 1) % 5000 == 0:
+            print(f"{j + 1}/{len(train)} features, {len(rows)} records", flush=True)
+
+    assert rows, "no records minted — check --min-act / max-acts file"
+    vecs = np.memmap(f"{a.out_dir}/vecs.f32", dtype=np.float32, mode="w+",
+                     shape=(len(rows), D_MODEL))
+    with open(f"{a.out_dir}/records.jsonl", "w") as recs:
+        for n, (f, text, act) in enumerate(rows):             # one row per (feature, target)
+            vecs[n] = dirs[:, f].numpy()
+            recs.write(json.dumps({"vec_idx": n, "target_text": text, "feature": f,
+                                   "act": round(act, 3)}) + "\n")
+    vecs.flush()
+    stats = {"n_examples": len(rows), "n_train_features": len(train), "n_eval_features": len(ev),
+             "targets_per_feature": a.targets, "skipped_short": skipped, "seed": a.seed}
+    json.dump(stats, open(f"{a.out_dir}/build_stats.json", "w"), indent=1)
+    print(f"BUILD_SAE_DATA_DONE {stats}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

scripts/eval_sae.py

@@ -0,0 +1,168 @@
+"""Stage 7: CROSS-UPLIFT eval — does cluster-direction training transfer to unseen SAE features?
+
+Per held-out feature: condition on unit(W_enc[:,f]) (inject@INJECT_LAYER at the marker), generate
+greedy + best-of-N at T via HF generate, then score every text STANDALONE through the CLEAN base
+(adapter disabled, no injection) at READ_LAYER: SAE-encode feature f, max over kept positions
+(pos-0 attention-sink + norm-outlier guards — the old SAE line's exact protocol, so numbers are
+directly comparable). dataset_max = the feature's corpus peak from the max-acts dump.
+
+    python scripts/eval_sae.py --adapter checkpoints/pretrain/final --split data/sae/split.json \
+        --best-of 16 --out results/xuplift_pretrain.json
+    python scripts/eval_sae.py --adapter none --n-features 256 --out results/xuplift_base.json
+"""
+import argparse
+import contextlib
+import json
+import os
+import random
+
+import torch
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from mxf.config import INJECT_LAYER, MODEL, READ_LAYER, STEER_COEFF
+from mxf.inject import get_layer, hooked, make_inject_hook, read_resid
+from mxf.prompts import build_prompt_ids
+from mxf.sae import load_max_acts, load_sae
+
+NORM_FILTER_MULT = 10.0  # drop positions with resid norm > mult * batch median (sink guard)
+
+
+@torch.no_grad()
+def generate(model, tok, prompt_ids, marker, dirs, a, device, greedy):
+    """One batched HF generate() with per-row direction injection. All prompts are the identical
+    token sequence, so `marker` is prefill-absolute; decode steps (len-1 forwards under KV cache)
+    skip the hook — see make_inject_hook."""
+    B = dirs.shape[0]
+    ids = torch.tensor([prompt_ids] * B, dtype=torch.long, device=device)
+    attn = torch.ones_like(ids, dtype=torch.bool)
+    hook = make_inject_hook([dirs[i : i + 1] for i in range(B)], [[marker]] * B,
+                            STEER_COEFF, device, torch.bfloat16)
+    kw = dict(max_new_tokens=a.max_new_tokens, min_new_tokens=a.min_new_tokens,
+              pad_token_id=tok.pad_token_id, do_sample=not greedy, use_cache=True)
+    if not greedy:
+        kw.update(temperature=a.temperature, top_p=1.0, top_k=0)
+    with hooked(get_layer(model, INJECT_LAYER), hook):
+        out = model.generate(input_ids=ids, attention_mask=attn, **kw)
+    stop = {tok.eos_token_id, tok.pad_token_id}  # <|im_end|> or <|endoftext|> (also pad)
+    texts = []
+    for row in out:
+        comp = row[len(prompt_ids):].tolist()
+        stops = [j for j, t in enumerate(comp) if t in stop]
+        texts.append(tok.decode(comp[: stops[0]] if stops else comp,
+                                skip_special_tokens=True).strip())
+    return texts
+
+
+@torch.no_grad()
+def score(texts, feats, model, clean, tok, sae, device, a):
+    """enc_act[i] = max_t relu((x_t-b_dec)·W_enc[:,f_i]+b_enc[f_i]) at READ_LAYER — standalone
+    re-tokenization (no chat template), clean base. Empty/all-filtered rows score 0."""
+    r = torch.zeros(len(texts))
+    valid = [i for i, t in enumerate(texts) if t.strip()]
+    prev = tok.padding_side
+    tok.padding_side = "right"  # position 0 must be the first real token
+    try:
+        for s in range(0, len(valid), a.score_batch):
+            idxs = valid[s : s + a.score_batch]
+            enc = tok([texts[i] for i in idxs], return_tensors="pt", padding=True, truncation=True,
+                      max_length=a.max_new_tokens + 32, add_special_tokens=False).to(device)
+            if enc["input_ids"].shape[1] == 0:
+                continue
+            with clean():
+                h, mask = read_resid(model, READ_LAYER, dict(enc), pool="all")  # fp32 [b,T,d],[b,T]
+            norms = h.norm(dim=-1)
+            keep = mask & (norms <= NORM_FILTER_MULT * norms[mask].median())
+            keep[:, 0] = False
+            b = torch.arange(len(idxs), device=device)
+            per = sae.encode_features(h, [feats[i] for i in idxs])[b, :, b]  # row i, feature i: [b,T]
+            best = per.masked_fill(~keep, 0.0).max(1).values  # relu>=0: 0-fill == old no-kept guard
+            for row, i in enumerate(idxs):
+                r[i] = best[row].item()
+    finally:
+        tok.padding_side = prev
+    return r
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--adapter", required=True, help="maxact-fast checkpoint dir, or 'none' for base")
+    ap.add_argument("--sae-path", default=None, help="ae.pt path; default: hf_hub_download")
+    ap.add_argument("--maxacts-path", default=None, help="max-acts .pt path; default: hf_hub_download")
+    ap.add_argument("--split", default=None, help="split.json from build_sae_data; uses its 'eval' half")
+    ap.add_argument("--n-features", type=int, default=0,
+                    help="cap the split's eval list; without --split, seeded sample of alive features")
+    ap.add_argument("--best-of", type=int, default=16)
+    ap.add_argument("--temperature", type=float, default=1.0)
+    ap.add_argument("--batch-features", type=int, default=8, help="sampled gen batch = this * best-of")
+    ap.add_argument("--max-new-tokens", type=int, default=96)
+    ap.add_argument("--min-new-tokens", type=int, default=16)
+    ap.add_argument("--score-batch", type=int, default=128)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--out", required=True)
+    a = ap.parse_args()
+    assert a.split or a.n_features > 0, "need --split or --n-features"
+    torch.manual_seed(a.seed)
+    device = "cuda:0"
+
+    tok = AutoTokenizer.from_pretrained(MODEL)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    prompt_ids, mpos = build_prompt_ids(tok)
+    marker = mpos[0]
+
+    dataset_max = load_max_acts(a.maxacts_path)["max_acts"].amax(dim=(1, 2))  # [F] corpus peaks
+    if a.split:
+        feats = json.load(open(a.split))["eval"]
+        if a.n_features:
+            feats = feats[: a.n_features]
+    else:
+        alive = (dataset_max > 0).nonzero(as_tuple=True)[0].tolist()
+        feats = sorted(random.Random(a.seed).sample(alive, min(a.n_features, len(alive))))
+
+    sae = load_sae(a.sae_path, device)
+    model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16,
+                                                 attn_implementation="sdpa", device_map={"": device})
+    if a.adapter != "none":
+        model = PeftModel.from_pretrained(model, a.adapter)
+    model.eval()
+    # scoring must see the clean base; with adapter='none' there is nothing to disable
+    clean = model.disable_adapter if a.adapter != "none" else contextlib.nullcontext
+    print(f"{len(feats)} eval features | adapter {a.adapter} | marker @{marker}", flush=True)
+
+    results = []
+    for s in range(0, len(feats), a.batch_features):
+        fb = feats[s : s + a.batch_features]
+        g_texts = generate(model, tok, prompt_ids, marker, sae.enc_dirs(fb), a, device, greedy=True)
+        flat = [f for f in fb for _ in range(a.best_of)]
+        s_texts = generate(model, tok, prompt_ids, marker, sae.enc_dirs(flat), a, device, greedy=False)
+        g_act = score(g_texts, fb, model, clean, tok, sae, device, a)
+        s_act = score(s_texts, flat, model, clean, tok, sae, device, a).view(len(fb), a.best_of)
+        for i, f in enumerate(fb):
+            bi = int(s_act[i].argmax())
+            results.append({"feature": int(f), "dataset_max": dataset_max[f].item(),
+                            "greedy_text": g_texts[i], "greedy_act": g_act[i].item(),
+                            "best_text": s_texts[i * a.best_of + bi], "best_act": s_act[i, bi].item()})
+        print(f"{len(results)}/{len(feats)} features", flush=True)
+
+    dmax = torch.tensor([max(r["dataset_max"], 1e-6) for r in results])
+    g = torch.tensor([r["greedy_act"] for r in results])
+    b = torch.tensor([r["best_act"] for r in results])
+    summary = {
+        "adapter": a.adapter, "n_features": len(results), "best_of": a.best_of,
+        "temperature": a.temperature,
+        "greedy/normalized_act": (g / dmax).mean().item(),
+        "greedy/normalized_act_median": (g / dmax).median().item(),
+        "greedy/beat_frac": (g > dmax).float().mean().item(),
+        f"best_of_{a.best_of}/normalized_act": (b / dmax).mean().item(),
+        f"best_of_{a.best_of}/normalized_act_median": ((b / dmax).median().item()),
+        f"best_of_{a.best_of}/beat_frac": (b > dmax).float().mean().item(),
+    }
+    print(json.dumps(summary, indent=2))
+    os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True)
+    json.dump({"summary": summary, "results": results}, open(a.out, "w"))
+    print("EVAL_SAE_DONE", flush=True)
+
+
+if __name__ == "__main__":
+    main()

src/mxf/sae.py

@@ -0,0 +1,61 @@
+"""BatchTopK SAE (adamkarvonen/qwen3-8b-saes, dictionary_learning format) — minimal loader.
+
+The conditioning/reward direction for feature f is the UNIT ENCODER COLUMN unit(W_enc[:,f])
+(probe side — same space as the cluster-probe directions, drops into inject@1/read@L as-is).
+Scoring uses raw per-sample encoder activations, never the batch-topk gate: we want a smooth
+reward signal below the firing threshold.
+"""
+import torch
+from huggingface_hub import hf_hub_download
+
+from mxf.config import D_MODEL, READ_LAYER
+
+SAE_REPO = "adamkarvonen/qwen3-8b-saes"
+SAE_FILENAME = f"saes_Qwen_Qwen3-8B_batch_top_k/resid_post_layer_{READ_LAYER}/trainer_2/ae.pt"
+MAX_ACTS_REPO = "adamkarvonen/sae_max_acts"          # dataset repo
+MAX_ACTS_FILENAME = (f"acts_Qwen_Qwen3-8B_layer_{READ_LAYER}_trainer_2_"
+                     "layer_percent_75_context_length_32.pt")
+
+
+class BatchTopKSAE:
+    """W_enc [d,F], W_dec [F,d], b_enc [F], b_dec [d]."""
+
+    def __init__(self, W_enc, W_dec, b_enc, b_dec):
+        self.W_enc, self.W_dec, self.b_enc, self.b_dec = W_enc, W_dec, b_enc, b_dec
+        self.d_in, self.d_sae = W_enc.shape
+
+    def encode_features(self, acts_BLD, feature_ids):
+        """Pre-topk post-ReLU encoder activations relu((x-b_dec)@W_enc[:,f]+b_enc[f]).
+        acts_BLD [B,L,d] -> [B,L,len(feature_ids)]."""
+        idx = torch.as_tensor(feature_ids, device=self.W_enc.device)
+        return torch.relu((acts_BLD - self.b_dec) @ self.W_enc[:, idx] + self.b_enc[idx])
+
+    def enc_dirs(self, feature_ids):
+        """Unit encoder columns [len(ids), d] — the conditioning/reward directions."""
+        idx = torch.as_tensor(feature_ids, device=self.W_enc.device)
+        return torch.nn.functional.normalize(self.W_enc[:, idx].T, dim=-1)
+
+
+def load_sae(path=None, device="cpu", dtype=torch.float32):
+    """path=None resolves via hf_hub_download (uses HF_HOME cache on the box)."""
+    path = path or hf_hub_download(repo_id=SAE_REPO, filename=SAE_FILENAME)
+    params = torch.load(path, map_location="cpu", weights_only=False)
+    key_map = {"encoder.weight": "W_enc", "decoder.weight": "W_dec", "encoder.bias": "b_enc",
+               "bias": "b_dec", "b_dec": "b_dec"}   # dictionary_learning aliases for b_dec
+    t = {key_map[k]: v.to(dtype) for k, v in params.items() if k in key_map}
+    sae = BatchTopKSAE(t["W_enc"].T.contiguous().to(device),  # nn.Linear stores [out, in]
+                       t["W_dec"].T.contiguous().to(device),
+                       t["b_enc"].to(device), t["b_dec"].to(device))
+    assert sae.d_in == D_MODEL, f"SAE d_in {sae.d_in} != D_MODEL {D_MODEL}"
+    nrm = sae.W_dec.norm(dim=1)
+    assert torch.allclose(nrm, torch.ones_like(nrm), atol=1e-2), "decoder rows must be unit norm"
+    return sae
+
+
+def load_max_acts(path=None):
+    """{"max_tokens": [F,N,L] long, "max_acts": [F,N,L] float} on cpu. path=None resolves via HF."""
+    path = path or hf_hub_download(repo_id=MAX_ACTS_REPO, filename=MAX_ACTS_FILENAME,
+                                   repo_type="dataset")
+    data = torch.load(path, map_location="cpu", weights_only=False)
+    assert "max_tokens" in data and "max_acts" in data, f"unexpected keys: {list(data.keys())}"
+    return data