Buckets:
| #!/usr/bin/env python3 | |
| """Build a 150-features-per-layer SAE cache (sae1500) from the existing sae500 pipeline. | |
| Reproduces the exact conventions of the sae500_2k_clean_prompt_response cache but with | |
| PER_LAYER_TOPK=150 instead of 50. Labels, role_ids, split_indices, and lengths are | |
| feature-count-independent and are copied unchanged from the sae500 cache. Only features.npy | |
| (now 1500-wide) and selected_features.npz (150 indices/layer) are regenerated. | |
| Conventions (validated against the sae500 cache, corr>0.999 / 0.3% median rel-error): | |
| * SAE release gemma-scope-2-1b-it-res-all, layer_{N}_width_16k_l0_small, layers 16-25 | |
| * JumpReLU encode: pre=(act-b_dec)@W_enc+b_enc ; feat = pre*(pre>threshold) | |
| * hidden_states[layer+1] (index 0 is embeddings), attn=sdpa, bf16 model, features float16 | |
| * top-K features per layer by summed activation over TRAIN-record tokens only | |
| """ | |
| from __future__ import annotations | |
| import json, re, sys | |
| from pathlib import Path | |
| import numpy as np, torch | |
| from sae_lens import SAE | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| sys.path.insert(0, "/workspace/scope/modeling") | |
| from models import TOPICS | |
| SRC = Path("/workspace/scope/cache/sae500_2k_clean_prompt_response") | |
| OUT = Path("/workspace/scope/cache/sae1500_2k_clean_prompt_response") | |
| DATASET = Path("/workspace/scope/runs/2k_run_5.6/dataset.cleaned.jsonl") | |
| GEMMA = "google/gemma-3-1b-it" | |
| REL = "gemma-scope-2-1b-it-res-all" | |
| LAYERS = list(range(16, 26)) | |
| TOPK = 150 | |
| CHUNK = 8192 | |
| TAG_RE = re.compile(r"</?(?:" + "|".join(re.escape(t) for t in TOPICS) + r")>") | |
| strip = lambda s: TAG_RE.sub("", s) | |
| DEVICE = "cuda" | |
| def ids_for(rec, tok): | |
| x = tok.apply_chat_template( | |
| [{"role": "user", "content": strip(rec["prompt"])}, | |
| {"role": "assistant", "content": strip(rec["response"])}], | |
| tokenize=True, add_generation_prompt=False) | |
| if x and not isinstance(x[0], int): | |
| x = x[0].ids | |
| return list(x) | |
| def main(): | |
| torch.set_grad_enabled(False) | |
| OUT.mkdir(parents=True, exist_ok=True) | |
| lengths = np.load(SRC / "lengths.npy") | |
| offsets = np.concatenate(([0], np.cumsum(lengths, dtype=np.int64))) | |
| total_tokens = int(offsets[-1]) | |
| train_idx = np.load(SRC / "split_indices.npz")["train"] | |
| stored_sel = np.load(SRC / "selected_features.npz") | |
| ref_feats = np.load(SRC / "features.npy", mmap_mode="r") # (T,500) for validation | |
| records = [json.loads(l)["record"] for l in DATASET.open()] | |
| assert len(records) == len(lengths), f"{len(records)} records vs {len(lengths)} lengths" | |
| print(f"records={len(records)} total_tokens={total_tokens} train_records={len(train_idx)}", flush=True) | |
| tok = AutoTokenizer.from_pretrained(GEMMA) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| GEMMA, torch_dtype=torch.bfloat16, attn_implementation="sdpa").to(DEVICE).eval() | |
| # --- Pass 1: prefill every record, store all 10 layers' hidden states (float16) --- | |
| hidden = np.empty((total_tokens, len(LAYERS), 1152), dtype=np.float16) | |
| with torch.inference_mode(): | |
| for i, rec in enumerate(records): | |
| input_ids = ids_for(rec, tok) | |
| lo, hi = int(offsets[i]), int(offsets[i + 1]) | |
| if len(input_ids) != hi - lo: | |
| raise RuntimeError(f"record {i}: tokenizer={len(input_ids)} cache={hi-lo}") | |
| out = model(input_ids=torch.tensor([input_ids], device=DEVICE), | |
| use_cache=False, output_hidden_states=True, return_dict=True) | |
| for lo_off, L in enumerate(LAYERS): | |
| hidden[lo:hi, lo_off, :] = out.hidden_states[L + 1][0].to(torch.float16).cpu().numpy() | |
| if (i + 1) % 250 == 0: | |
| print(f" prefilled {i+1}/{len(records)}", flush=True) | |
| del model | |
| torch.cuda.empty_cache() | |
| print("prefill complete", flush=True) | |
| # train-token row indices (all tokens of train records) | |
| train_rows = np.concatenate([np.arange(offsets[i], offsets[i + 1]) for i in train_idx]) | |
| print(f"train tokens={len(train_rows)}", flush=True) | |
| # --- Per-layer: select top-150 over train tokens, then encode all tokens --- | |
| columns = [] | |
| selected = {} | |
| sel_overlap = [] | |
| feat_corr = [] | |
| for lo_off, L in enumerate(LAYERS): | |
| sae = SAE.from_pretrained(release=REL, sae_id=f"layer_{L}_width_16k_l0_small", device=DEVICE) | |
| sae.eval() | |
| W = sae.W_enc.float(); b_enc = sae.b_enc.float(); b_dec = sae.b_dec.float() | |
| thr = sae.threshold.float() | |
| # selection: sum jumprelu activations over train tokens (full 16k) | |
| sums = torch.zeros(W.shape[1], device=DEVICE) | |
| layer_hidden = hidden[:, lo_off, :] | |
| for s in range(0, len(train_rows), CHUNK): | |
| rows = train_rows[s:s + CHUNK] | |
| act = torch.from_numpy(np.asarray(layer_hidden[rows], dtype=np.float32)).to(DEVICE) | |
| pre = (act - b_dec) @ W + b_enc | |
| sums += (pre * (pre > thr)).sum(0) | |
| idx = torch.topk(sums, TOPK).indices | |
| idx = torch.sort(idx).values | |
| selected[f"layer_{L}"] = idx.cpu().numpy().astype(np.int64) | |
| # validation: does our top-50 reproduce the stored 50? | |
| top50 = set(torch.topk(sums, 50).indices.cpu().numpy().tolist()) | |
| stored50 = set(stored_sel[f"layer_{L}"].tolist()) | |
| sel_overlap.append(len(top50 & stored50)) | |
| assert set(stored50).issubset(set(selected[f"layer_{L}"].tolist())) or True # info only | |
| # encode all tokens with selected 150 indices | |
| idx_dev = idx | |
| col = np.empty((total_tokens, TOPK), dtype=np.float16) | |
| for s in range(0, total_tokens, CHUNK): | |
| act = torch.from_numpy(np.asarray(layer_hidden[s:s + CHUNK], dtype=np.float32)).to(DEVICE) | |
| pre = (act - b_dec) @ W[:, idx_dev] + b_enc[idx_dev] | |
| col[s:s + CHUNK] = (pre * (pre > thr[idx_dev])).to(torch.float16).cpu().numpy() | |
| columns.append(col) | |
| # validation: re-encode the stored-50 indices and correlate with the reference cache | |
| stored_idx = torch.as_tensor(stored_sel[f"layer_{L}"], device=DEVICE) | |
| ref_col = np.asarray(ref_feats[:, lo_off * 50:(lo_off + 1) * 50], dtype=np.float32) | |
| chk = np.empty((total_tokens, 50), dtype=np.float32) | |
| for s in range(0, total_tokens, CHUNK): | |
| act = torch.from_numpy(np.asarray(layer_hidden[s:s + CHUNK], dtype=np.float32)).to(DEVICE) | |
| pre = (act - b_dec) @ W[:, stored_idx] + b_enc[stored_idx] | |
| chk[s:s + CHUNK] = (pre * (pre > thr[stored_idx])).to(torch.float16).float().cpu().numpy() | |
| feat_corr.append(float(np.corrcoef(chk.ravel(), ref_col.ravel())[0, 1])) | |
| del sae | |
| torch.cuda.empty_cache() | |
| print(f" layer {L}: top50∩stored={sel_overlap[-1]}/50 feat_corr(stored-50)={feat_corr[-1]:.5f}", flush=True) | |
| features = np.concatenate(columns, axis=1) # (T, 1500) | |
| assert features.shape == (total_tokens, TOPK * len(LAYERS)) | |
| np.save(OUT / "features.npy", features) | |
| np.savez(OUT / "selected_features.npz", **selected) | |
| # copy feature-count-independent files unchanged | |
| for name in ["labels.npy", "lengths.npy", "role_ids.npy"]: | |
| np.save(OUT / name, np.load(SRC / name)) | |
| sp = np.load(SRC / "split_indices.npz") | |
| np.savez(OUT / "split_indices.npz", **{k: sp[k] for k in sp.files}) | |
| meta = json.loads((SRC / "metadata.json").read_text()) | |
| meta["selected_features_per_layer"] = TOPK | |
| meta["total_features_per_token"] = TOPK * len(LAYERS) | |
| meta["encode_convention"] = "jumprelu threshold, attn=sdpa (rebuilt from sae500 pipeline)" | |
| meta["built_from"] = str(SRC) | |
| meta["selection_overlap_top50_vs_stored"] = sel_overlap | |
| meta["feature_corr_stored50_vs_ref"] = feat_corr | |
| (OUT / "metadata.json").write_text(json.dumps(meta, indent=2) + "\n") | |
| print(f"\nwrote {OUT}", flush=True) | |
| print(f"features.npy shape={features.shape} dtype={features.dtype}", flush=True) | |
| print(f"selection overlap (top50 vs stored, per layer): {sel_overlap} mean={np.mean(sel_overlap):.1f}/50", flush=True) | |
| print(f"feature corr (stored-50 re-encode vs ref cache): min={min(feat_corr):.5f} mean={np.mean(feat_corr):.5f}", flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 8.19 kB
- Xet hash:
- 5ec12710a272028537baefb034787bf833e277c9621950ba2390bd265f990a95
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.