Buckets:
| #!/usr/bin/env python3 | |
| """Dynamic per-token top-K SAE cache, stored SPARSE as GPU-loadable torch tensors. | |
| Per token, per layer: keep that token's own top-K (=150) features (JumpReLU magnitude) from | |
| the full 16,384 dictionary, identity preserved by position. Columns = union of features that | |
| appear in any training token's top-K (per layer), capped to the CAP most-frequent per layer | |
| (kept fraction ~100%). Stored as CSR-by-token (col indices + values + per-token pointers) so | |
| the whole thing is a few GB and lives resident on the GPU during training -- no per-batch | |
| CPU<->GPU transfer, no network reads. | |
| """ | |
| 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_sparse") | |
| 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; CAP_PER_LAYER = 6000; 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()] | |
| print(f"records={len(records)} tokens={total} K={K} 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] | |
| 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] | |
| 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) | |
| 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() | |
| pre = (act - b_dec) @ W + b_enc | |
| feat = pre * (pre > thr) | |
| vals, idx = torch.topk(feat, K, dim=-1) | |
| zero = vals <= 0 | |
| idx = idx.clone(); idx[zero] = -1 | |
| topk_idx[lo:hi, li, :] = idx.cpu().numpy() | |
| topk_val[lo:hi, li, :] = vals.to(torch.float16).cpu().numpy() | |
| if 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) | |
| # union columns per layer (cap to most-frequent), global column ids | |
| widths, luts, selected = [], [], {} | |
| for li, L in enumerate(LAYERS): | |
| present = np.where(freq[li] > 0)[0] | |
| keep = present if len(present) <= CAP_PER_LAYER else present[np.argsort(freq[li][present])[::-1][:CAP_PER_LAYER]] | |
| keep = np.sort(keep) | |
| lut = np.full(d_sae, -1, dtype=np.int64); lut[keep] = np.arange(len(keep)) | |
| luts.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"widths={widths} TOTAL_WIDTH={W_total}", flush=True) | |
| # map each (token,layer,k) to a global union column; build CSR-by-token | |
| gcol = np.full((total, len(LAYERS), K), -1, dtype=np.int64) | |
| kept = 0; fired = 0 | |
| for li in range(len(LAYERS)): | |
| idx = topk_idx[:, li, :]; valid = idx >= 0; fired += int(valid.sum()) | |
| mapped = np.full(idx.shape, -1, dtype=np.int64) | |
| mapped[valid] = luts[li][idx[valid]] | |
| ok = mapped >= 0; kept += int(ok.sum()) | |
| gcol[:, li, :][ok] = layer_off[li] + mapped[ok] | |
| print(f"kept {100*kept/fired:.2f}% of firing feats across all layers", flush=True) | |
| gcol_flat = gcol.reshape(total, -1) # (total, 10*K) | |
| val_flat = topk_val.reshape(total, -1) | |
| mask = gcol_flat >= 0 | |
| nnz_per_tok = mask.sum(axis=1).astype(np.int64) | |
| tok_ptr = np.concatenate(([0], np.cumsum(nnz_per_tok))) | |
| col = gcol_flat[mask].astype(np.int32) # token-major (C order) | |
| val = val_flat[mask].astype(np.float16) | |
| print(f"nnz={col.size} mean nnz/token={col.size/total:.0f}", flush=True) | |
| sp = np.load(SRC / "split_indices.npz") | |
| payload = { | |
| "col": torch.from_numpy(col), | |
| "val": torch.from_numpy(val), | |
| "tok_ptr": torch.from_numpy(tok_ptr), | |
| "labels": torch.from_numpy(np.load(SRC / "labels.npy").astype(np.int64)), | |
| "roles": torch.from_numpy(np.load(SRC / "role_ids.npy").astype(np.int64)), | |
| "lengths": torch.from_numpy(lengths.astype(np.int64)), | |
| "width": W_total, | |
| "train": torch.from_numpy(sp["train"].astype(np.int64)), | |
| "validation": torch.from_numpy(sp["validation"].astype(np.int64)), | |
| "test": torch.from_numpy(sp["test"].astype(np.int64)), | |
| } | |
| torch.save(payload, OUT / "sparse.pt") | |
| np.savez(OUT / "selected_features.npz", **selected) | |
| (OUT / "meta.json").write_text(json.dumps( | |
| {"width": W_total, "widths": widths, "K": K, "cap": CAP_PER_LAYER, | |
| "kept_fraction": kept / fired, "nnz": int(col.size), "tokens": total}, indent=2) + "\n") | |
| import os | |
| print(f"wrote {OUT}/sparse.pt ({os.path.getsize(OUT/'sparse.pt')/1e9:.2f} GB) width={W_total}", flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 6.98 kB
- Xet hash:
- 9d1f10713d546dff5a5c7d33c90398c2e468c4c63ce427776bc5a9edc0444338
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.