Buckets:
| #!/usr/bin/env python3 | |
| """Dynamic per-token top-K SAE cache (identity preserved by position). | |
| Contrast with the global-pooled fixed selection (sae500/sae1500): there we picked ONE | |
| set of K indices per layer (top-K by summed activation over all training tokens) and froze | |
| it for every token. Here, for EVERY token we keep that token's OWN top-K features (by | |
| JumpReLU magnitude) within each layer's full 16,384-feature dictionary, placed at their true | |
| feature positions. Columns = the union of features that ever appear in any training token's | |
| top-K (per layer), so identity is preserved and no ever-firing feature is dropped. | |
| Same base pipeline as build_sae1500_cache.py (Gemma-3-1b-it bf16/sdpa, gemma-scope JumpReLU, | |
| hidden_states[layer+1], float16). Labels/role_ids/splits/lengths copied unchanged. | |
| """ | |
| 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/dyn150_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)) | |
| K = 150 # per-token, per-layer | |
| CAP_PER_LAYER = 2000 # bound union width (quota-safe ~19GB dense); keep most-frequent features | |
| CHUNK = 4096 | |
| DEVICE = "cuda" | |
| TAG_RE = re.compile(r"</?(?:" + "|".join(re.escape(t) for t in TOPICS) + r")>") | |
| strip = lambda s: TAG_RE.sub("", s) | |
| 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 = int(offsets[-1]) | |
| train_idx = set(int(i) for i in np.load(SRC / "split_indices.npz")["train"]) | |
| records = [json.loads(l)["record"] for l in DATASET.open()] | |
| assert len(records) == len(lengths) | |
| print(f"records={len(records)} total_tokens={total} K={K}/layer/token cap={CAP_PER_LAYER}", flush=True) | |
| tok = AutoTokenizer.from_pretrained(GEMMA) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| GEMMA, torch_dtype=torch.bfloat16, attn_implementation="sdpa").to(DEVICE).eval() | |
| saes = [] | |
| for L in LAYERS: | |
| s = SAE.from_pretrained(release=REL, sae_id=f"layer_{L}_width_16k_l0_small", device=DEVICE); s.eval() | |
| saes.append((L, s.W_enc.float(), s.b_enc.float(), s.b_dec.float(), s.threshold.float())) | |
| d_sae = saes[0][1].shape[1] | |
| # per-token top-K store (index/value) and per-layer appearance frequency over TRAIN tokens | |
| topk_idx = np.full((total, len(LAYERS), K), -1, dtype=np.int32) | |
| topk_val = np.zeros((total, len(LAYERS), K), dtype=np.float16) | |
| freq = [np.zeros(d_sae, dtype=np.int64) for _ in LAYERS] | |
| l0_sum = np.zeros(len(LAYERS), dtype=np.int64) | |
| 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}: tok={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) | |
| is_train = i in train_idx | |
| for li, (L, W, b_enc, b_dec, thr) in enumerate(saes): | |
| act = out.hidden_states[L + 1][0].float() # (T,1152) | |
| pre = (act - b_dec) @ W + b_enc # (T,16384) | |
| feat = pre * (pre > thr) # JumpReLU | |
| l0_sum[li] += int((feat > 0).sum()) | |
| vals, idx = torch.topk(feat, K, dim=-1) # (T,K) | |
| vals = vals.to(torch.float16) | |
| zero = vals <= 0 | |
| idx = idx.clone(); idx[zero] = -1 # drop non-firing padding | |
| topk_idx[lo:hi, li, :] = idx.cpu().numpy() | |
| topk_val[lo:hi, li, :] = vals.cpu().numpy() | |
| if is_train: | |
| flat = idx[~zero].reshape(-1).cpu().numpy() | |
| if flat.size: | |
| freq[li] += np.bincount(flat, minlength=d_sae) | |
| if (i + 1) % 250 == 0: | |
| print(f" prefilled {i+1}/{len(records)}", flush=True) | |
| del model, saes | |
| torch.cuda.empty_cache() | |
| print("prefill complete", flush=True) | |
| for li, L in enumerate(LAYERS): | |
| u = int((freq[li] > 0).sum()) | |
| print(f" layer {L}: mean L0={l0_sum[li]/total:6.1f} union(train top-{K})={u}", flush=True) | |
| # choose columns per layer = union, capped to CAP_PER_LAYER most frequent | |
| colmap = [] # per layer: lookup array d_sae -> local col (or -1) | |
| widths = [] | |
| selected = {} | |
| for li, L in enumerate(LAYERS): | |
| present = np.where(freq[li] > 0)[0] | |
| if len(present) > CAP_PER_LAYER: | |
| keep = present[np.argsort(freq[li][present])[::-1][:CAP_PER_LAYER]] | |
| else: | |
| keep = present | |
| keep = np.sort(keep) | |
| lut = np.full(d_sae, -1, dtype=np.int32) | |
| lut[keep] = np.arange(len(keep), dtype=np.int32) | |
| colmap.append(lut); widths.append(len(keep)); selected[f"layer_{L}"] = keep.astype(np.int64) | |
| layer_off = np.concatenate(([0], np.cumsum(widths))) | |
| W_total = int(layer_off[-1]) | |
| print(f"union widths per layer: {widths} TOTAL WIDTH={W_total}", flush=True) | |
| # scatter per-token top-K values into dense union-width features.npy (float16, memmap) | |
| feats = np.lib.format.open_memmap(OUT / "features.npy", mode="w+", dtype=np.float16, shape=(total, W_total)) | |
| feats[:] = 0 | |
| kept_frac = np.zeros(len(LAYERS)) | |
| for li, L in enumerate(LAYERS): | |
| lut = colmap[li]; off = int(layer_off[li]) | |
| idx_l = topk_idx[:, li, :] # (total,K) | |
| val_l = topk_val[:, li, :] | |
| valid = idx_l >= 0 | |
| local = np.full(idx_l.shape, -1, dtype=np.int64) | |
| local[valid] = lut[idx_l[valid]] # map to local col (or -1 if capped out) | |
| ok = local >= 0 | |
| kept_frac[li] = ok.sum() / max(1, valid.sum()) | |
| rows = np.repeat(np.arange(total)[:, None], K, axis=1)[ok] | |
| cols = off + local[ok] | |
| feats[rows, cols] = val_l[ok] | |
| print(f" scattered layer {L}: width={widths[li]} kept {kept_frac[li]*100:.1f}% of firing feats", flush=True) | |
| feats.flush(); del feats | |
| 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}) | |
| np.savez(OUT / "selected_features.npz", **selected) | |
| meta = json.loads((SRC / "metadata.json").read_text()) | |
| meta.update({"selection": f"dynamic per-token top-{K}/layer, identity-preserved (union columns)", | |
| "per_token_k": K, "cap_per_layer": CAP_PER_LAYER, "layer_widths": widths, | |
| "total_width": W_total, "mean_l0_per_layer": [float(x/total) for x in l0_sum], | |
| "kept_fraction_per_layer": [float(x) for x in kept_frac]}) | |
| (OUT / "metadata.json").write_text(json.dumps(meta, indent=2) + "\n") | |
| print(f"\nwrote {OUT} features shape=({total},{W_total})", flush=True) | |
| print(f"TOTAL_WIDTH={W_total}", flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 7.9 kB
- Xet hash:
- 67e46cf20d780dbf3dfe37ab7f9627337a2b25445bc5b2cd6c0ef2ea0661ac77
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.