diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 0000000..8ae9917 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,19 @@ +## Summary + + + +## Trace provenance (when applicable) + +- Model and revision: +- GPU and dtype: +- Layer and estimator: +- Language / concept: +- Demo or preview: + +## Validation + +- [ ] `python experiments/validate_trace_gallery.py` +- [ ] `pytest -q` +- [ ] `cd apps/web && npm ci && npm audit --audit-level=high && npm run build` +- [ ] No model weights, caches, secrets, personal data, or restricted prompts +- [ ] Frontend displays baked rigor values and does not recompute sharedness diff --git a/.github/workflows/validate.yml b/.github/workflows/validate.yml new file mode 100644 index 0000000..f214e6d --- /dev/null +++ b/.github/workflows/validate.yml @@ -0,0 +1,44 @@ +name: Validate code and traces + +on: + push: + branches: [main] + pull_request: + workflow_dispatch: + +permissions: + contents: read + +jobs: + validate: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + cache-dependency-path: pyproject.toml + + - name: Install Python package and test runner + run: python -m pip install -e . pytest + + - name: Validate committed trace gallery + run: python experiments/validate_trace_gallery.py + + - name: Run Python tests + run: pytest -q + + - uses: actions/setup-node@v4 + with: + node-version: "20" + cache: npm + cache-dependency-path: apps/web/package-lock.json + + - name: Build web application + working-directory: apps/web + run: | + npm ci + npm audit --audit-level=high + npm run build diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..a0f7b83 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,73 @@ +# Contributing to J7Scope + +感谢你帮助扩展跨语言 J-space trace。最有价值的贡献是：可复现、明确标注局限、不会把 +preview 当作研究结论的真实模型采集。 + +## 快速路径：Colab 采集 + +1. 在 Google Colab 打开 `notebooks/capture_colab.ipynb`。 +2. 选择 **Runtime → Change runtime type → GPU**。 +3. 按顺序运行所有 cell，填写 prompt、语言、概念和唯一 `trace_id`。 +4. 下载 notebook 生成的 ZIP，在本地解压。 + +默认 Qwen2.5-1.5B + 少量随机探针用于验证贡献流程，产物会标记 +`preview: true`，不是可引用的 M1 证据。GPU 型号、dtype、模型 revision、Jacobian +estimator 和张量 SHA-1 会写进 manifest。 + +也可以在任何 CUDA 机器直接运行： + +```bash +pip install -e . +python experiments/capture_trace.py \ + --trace-id community-deception-en \ + --language en \ + --concept deception \ + --prompt "In one sentence, explain why deception can be tempting." +``` + +先用 `--dry-run` 检查配置，不会下载模型或占用 GPU。 + +## Trace 命名与内容 + +- `trace_id`：1–80 个小写字母、数字、点、下划线或连字符；建议 + `community---`。 +- prompt 不得包含个人信息、机密内容、API key 或未获许可的数据。 +- 模型必须是贡献者有权使用和重新发布派生小型 artifact 的权重。 +- 不提交模型权重、Jacobian `.pt`、Hugging Face cache 或其他大文件。 +- 不手工修改 `tokens.jsonl` 中的 `rigor.sharedness`。严谨层只能由 + `j7scope.rigor` 生成。 +- 社区 trace 默认保留 `preview: true`；研究级标记由维护者在复核协议、语料和 + estimator 收敛后处理。 + +## 加入 Gallery + +把完整 trace 目录复制到 `results/traces//`，然后从仓库根目录运行： + +```bash +python -c "from j7scope.trace import rebuild_trace_index; rebuild_trace_index('results/traces')" +python experiments/validate_trace_gallery.py +pytest -q +``` + +提交 `manifest.json`、`tokens.jsonl`、`metrics.json`，以及平行 trace 才需要的 +`align.json`。同时提交重建后的 `results/traces/index.json`。 + +## Pull request 检查清单 + +- [ ] Trace 在本地通过 `experiments/validate_trace_gallery.py` +- [ ] `trace_id` 与目录名一致且 index 已重建 +- [ ] `is_demo` / `preview` 标记真实准确 +- [ ] manifest 含 model revision、dtype、device、estimator 和 Jacobian SHA-1 +- [ ] prompt 和 artifact 不含个人、机密或受限信息 +- [ ] 未提交模型权重、cache 或大型二进制 +- [ ] PR 描述说明模型、GPU、语言、概念和采集目的 + +## Code contributions + +Python 变更应补测试；前端不得重算 sharedness/null 指标，只显示 trace 中已烤入的值。 +运行： + +```bash +pytest -q +cd apps/web && npm ci && npm audit --audit-level=high && npm run build +``` diff --git a/README.md b/README.md index a641725..353df04 100644 --- a/README.md +++ b/README.md @@ -219,7 +219,69 @@ h_en = jlens.collect_residual(pairs["deception-01"]["en"]["text"]) print(jlens.readout(h_zh), jlens.readout(h_en)) # 两种语言的读出一致吗？ ``` -### 8.1 J-Space Explorer 前端（离线 artifact） +在投入 7B GPU 采集前，可先在随机初始化的 tiny-Qwen2 上复跑 Jacobian 数学一致性检查； +它不下载模型，也不产生研究结论： + +```bash +python experiments/validate_jacobian.py --n-probes 4096 +``` + +报告中的 `exact_vs_full_graph` 比较 `torch.func.jacrev` 尾部重放与独立的全模型逐维 +VJP，`paper_batched_vs_scalar_vjp` 检查上游 replicated-batch position reduction 与 +逐标量实现，两者 +都应该接近数值零误差；`stochastic_vs_exact` 则量化当前随机探针估计的有限样本误差。 +当前 capture 路径拟合的是 cloze 的同位置 `∂h_L[p]/∂h_l[p]` 变体； +`estimate_jacobian_paper` 才会汇总未来 target positions。正式 P3 运行前必须先锁定并记录 +所用口径。 + +上游预训练 lens 的云 GPU 复现入口已经固定到发布时使用的 Qwen3.5-4B 配置： + +```bash +pip install -e '.[upstream]' +python experiments/reproduce_upstream.py --preflight-only +python experiments/reproduce_upstream.py +``` + +正式 paper estimator 的拟合入口支持显存预检、运行计时、峰值显存记录以及逐 prompt +原子 checkpoint。默认仓库语料仅用于 smoke test，不是研究证据；正式运行流式抽取 +WikiText 1000 条： + +```bash +pip install -e '.[fit]' +python experiments/fit_paper_jacobian.py --dry-run +python experiments/fit_paper_jacobian.py --preflight-only +python experiments/fit_paper_jacobian.py \ + --dataset Salesforce/wikitext \ + --dataset-config wikitext-103-raw-v1 \ + --max-prompts 1000 --min-chars 200 +``` + +中断后原命令重跑会从 `results/*.checkpoint.pt` 续传，并校验模型、层、estimator 参数和 +完整有序语料 SHA-1。最终 `*.pt` 是可直接传给 sidecar `--jacobian-file` 的 float32 +矩阵，同名 JSON 保存模型 revision、硬件、耗时、峰值显存和 tensor SHA-1。 + +在线 GPU 选型、Runpod 操作、Colab 边界和学术 credits 申请步骤见 +[`docs/gpu-cloud-runbook.md`](docs/gpu-cloud-runbook.md)；英文申请底稿见 +[`docs/gpu-credit-proposal.md`](docs/gpu-credit-proposal.md)。 + +### 8.1 社区 GPU / Colab 采集 + +[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/arthurpanhku/j7scope/blob/main/notebooks/capture_colab.ipynb) + +Colab notebook 用 Qwen2.5-1.5B-Instruct 采一条自定义真实模型 trace，自动适配 T4 的 +`float16` 与新 GPU 的 `bfloat16`，完成 schema/provenance/rigor 校验后下载 ZIP。默认 +产物明确标为 `preview`，不作为研究证据。任何 CUDA 机器也可直接运行： + +```bash +python experiments/capture_trace.py \ + --trace-id community-deception-en \ + --language en --concept deception \ + --prompt "In one sentence, explain why deception can be tempting." +``` + +提交步骤和数据边界见 [`CONTRIBUTING.md`](CONTRIBUTING.md)。 + +### 8.2 J-Space Explorer 前端（离线 artifact） `apps/web` 是离线浏览已跑完实验 run 的前端，读取 `results/runs//` 下一组稳定 artifact（`manifest.json` / `readouts.jsonl` / `patches.jsonl` / `projections.json` / `layer_scan.json` / `metrics.json`）。本地开发可先生成一个明确标记为 demo 的假数据 run： @@ -239,8 +301,13 @@ cd apps/web && npm install && npm run dev -- --port 5173 - [x] P1 Trace schema v1 + `--record` + Replay 回放模式（mock 可验证，零 GPU） - [x] P2 静态 Gallery 站（Gallery/Replay/Compare + 深链 + 导出 SVG/PNG/JSON/BibTeX + Pages 部署工作流） - [ ] P3 首批真实 trace + Zenodo DOI（与 M1 首批数据合并一次租卡） -- [ ] P4 Colab 采集笔记本 + 社区提交流程 +- [x] P4（实现）Colab 自定义采集 + T4 dtype 适配 + 贡献规范/模板 + trace CI +- [ ] P4（验收）首个外部贡献者产出并合入通过 CI 的 preview trace - [x] P5（切片）跨会话 cross-trace 指标 + 并排对齐高亮 + null 带 + 方法学页（可展开审计） +- [x] P3 前置（切片）精确 position-local Jacobian + upstream paper reduction + tiny-Qwen2 数学校验 +- [x] P3 前置（切片）固定上游 Qwen3.5-4B/lens revision 的云 GPU 复现脚本与显存预检 +- [x] P3 前置（实现）replicated-batch 正式 fitter + 原子 checkpoint/resume + WikiText 入口 +- [ ] P3 前置（剩余）在云 GPU 执行上游已知案例并据实锁定正式拟合口径 **研究线（M1–M3）** diff --git a/apps/serve/README.md b/apps/serve/README.md index adfe22e..f8a5924 100644 --- a/apps/serve/README.md +++ b/apps/serve/README.md @@ -52,9 +52,18 @@ generated token's residual is read out through it. ```bash python -m j7scope_serve --backend hf \ --model Qwen/Qwen2.5-7B-Instruct --layer 18 \ - --jacobian-cache ~/.cache/j7scope + --model-revision main --dtype auto \ + --jacobian-file ../../results/jacobian-qwen2.5-7b-l18.pt ``` +`--jacobian-file` 直接加载 +[`experiments/fit_paper_jacobian.py`](../../experiments/fit_paper_jacobian.py) +生成的正式矩阵；省略它时才会用小语料和随机 position-local estimator 临时拟合，可通过 +`--jacobian-cache` 缓存，但这种快速路径只适合 wiring/preview，不应当作为研究结果。 +`--dtype auto` 在支持 BF16 的 CUDA GPU 上选择 `bfloat16`，在 T4 等旧卡上选择 +`float16`，CPU 则选择 `float32`。录制 manifest 会保存解析后的模型 revision、实际 +device/dtype、estimator 和 Jacobian SHA-1。 + ## Record & replay (platform P1) Live sessions are ephemeral; **traces** are the citable, deep-linkable artifact diff --git a/apps/serve/j7scope_serve/__main__.py b/apps/serve/j7scope_serve/__main__.py index 3a43511..dd31dc6 100644 --- a/apps/serve/j7scope_serve/__main__.py +++ b/apps/serve/j7scope_serve/__main__.py @@ -36,9 +36,17 @@ def main() -> None: p.add_argument("--token-delay", type=float, default=0.04, help="seconds between streamed tokens (mock pacing / demo feel)") p.add_argument("--device", default=None, help="hf: torch device override") + p.add_argument("--dtype", choices=["auto", "bfloat16", "float16", "float32"], + default="auto", help="hf: model dtype (auto handles T4 vs newer GPUs)") + p.add_argument("--model-revision", default="main", + help="hf: Hugging Face model branch, tag, or commit") p.add_argument("--max-new-tokens", type=int, default=256, help="hf: generation cap") p.add_argument("--jacobian-cache", default=None, help="hf: directory to cache the fitted Jacobian") + p.add_argument("--jacobian-file", default=None, + help="hf: precomputed Jacobian tensor from the paper fitter") + p.add_argument("--jacobian-metadata", default=None, + help="hf: fitter metadata JSON (default: same stem as tensor)") p.add_argument("--record", default=None, metavar="DIR", help="record each session as a Trace v1 under DIR//") p.add_argument("--traces", default=None, metavar="DIR", @@ -50,7 +58,10 @@ def main() -> None: kw["model_name"] = args.model if args.backend == "hf": kw.update(device=args.device, max_new_tokens=args.max_new_tokens, - cache_dir=args.jacobian_cache) + cache_dir=args.jacobian_cache, dtype=args.dtype, + model_revision=args.model_revision, + jacobian_path=args.jacobian_file, + jacobian_metadata_path=args.jacobian_metadata) kw.setdefault("model_name", "Qwen/Qwen2.5-7B-Instruct") backend = make_backend(args.backend, **kw) diff --git a/apps/serve/j7scope_serve/backends.py b/apps/serve/j7scope_serve/backends.py index 31cb888..6993343 100644 --- a/apps/serve/j7scope_serve/backends.py +++ b/apps/serve/j7scope_serve/backends.py @@ -21,6 +21,8 @@ import hashlib import itertools +import json +from pathlib import Path from typing import Iterator, List, NamedTuple, Sequence, Tuple @@ -34,6 +36,7 @@ class Backend: model_name: str layer: int is_demo: bool = False + jacobian_estimator: str = "not_applicable" def generate(self, messages: Sequence[dict], **params) -> Iterator[Step]: raise NotImplementedError @@ -86,6 +89,7 @@ def _jitter(seed: str) -> float: class MockBackend(Backend): is_demo = True + jacobian_estimator = "synthetic" def __init__(self, model_name: str = "j7scope-mock", layer: int = 18, topk: int = 24, per_lang: int = 8): @@ -152,14 +156,31 @@ class HFBackend(Backend): def __init__(self, model_name: str = "Qwen/Qwen2.5-7B-Instruct", layer: int = 18, topk: int = 24, max_new_tokens: int = 256, device: str = None, jacobian_corpus: Sequence[str] = None, n_probes: int = 16, - cache_dir: str = None): + cache_dir: str = None, dtype: str = "auto", + model_revision: str = "main", jacobian_path: str = None, + jacobian_metadata_path: str = None): self.model_name = model_name self.layer = layer self.topk = topk self.max_new_tokens = max_new_tokens self.device = device + self.dtype = dtype + self.model_revision = model_revision + self.model_revision_resolved = None self.jacobian_corpus = list(jacobian_corpus or _DEFAULT_JACOBIAN_CORPUS) self.n_probes = n_probes + self.jacobian_path = jacobian_path + self.jacobian_metadata_path = jacobian_metadata_path + self.jacobian_estimator = ( + "paper_replicated_batch_vjp" + if jacobian_path + else "position_local_gaussian_vjp" + ) + self.jacobian_corpus_id = "generic-v1" + self.jacobian_n_prompts = len(self.jacobian_corpus) + self.jacobian_position = -1 + self.jacobian_seed = 0 + self.jacobian_sha1 = None self.cache_dir = cache_dir self._model = None self._tok = None @@ -169,18 +190,46 @@ def __init__(self, model_name: str = "Qwen/Qwen2.5-7B-Instruct", layer: int = 18 # ---- warm-up ---------------------------------------------------------- def load(self) -> None: - import torch # noqa: F401 (import-time check + used below) + import torch from j7scope.fitting import JLens, load_model, _Capture - model, tok = load_model(self.model_name, device=self.device) + dtype = self._resolve_dtype(torch) + model, tok = load_model( + self.model_name, + device=self.device, + dtype=dtype, + revision=self.model_revision, + ) + self.device = str(model.device) + self.dtype = str(dtype).removeprefix("torch.") + self.model_revision_resolved = ( + getattr(model.config, "_commit_hash", None) or self.model_revision + ) jlens = JLens(model, tok, layer=self.layer) - J = self._load_cached_jacobian() + J = self._load_precomputed_jacobian() + if J is None: + J = self._load_cached_jacobian() if J is not None: + expected_shape = ( + model.config.hidden_size, + model.config.hidden_size, + ) + if not hasattr(J, "shape") or tuple(J.shape) != expected_shape: + raise ValueError( + f"Jacobian shape must be {expected_shape}, " + f"got {getattr(J, 'shape', None)}" + ) jlens.J = J else: - jlens.estimate_jacobian(self.jacobian_corpus, n_probes=self.n_probes) + jlens.estimate_jacobian( + self.jacobian_corpus, + n_probes=self.n_probes, + position=self.jacobian_position, + seed=self.jacobian_seed, + ) self._save_cached_jacobian(jlens.J) + self.jacobian_sha1 = self._tensor_sha1(jlens.J) # Persistent hook: capture the residual at layer l on every forward. cap = _Capture() @@ -188,17 +237,103 @@ def load(self) -> None: self._model, self._tok, self._jlens, self._capture = model, tok, jlens, cap + def _resolve_dtype(self, torch): + allowed = {"auto", "bfloat16", "float16", "float32"} + if self.dtype not in allowed: + raise ValueError( + "dtype must be auto, bfloat16, float16, or float32" + ) + if self.dtype != "auto": + return getattr(torch, self.dtype) + device = self.device or ("cuda" if torch.cuda.is_available() else "cpu") + if str(device).startswith("cuda"): + return ( + torch.bfloat16 + if torch.cuda.is_bf16_supported() + else torch.float16 + ) + if str(device).startswith("mps"): + return torch.float16 + return torch.float32 + + @staticmethod + def _tensor_sha1(tensor) -> str: + return hashlib.sha1( + tensor.detach().float().cpu().contiguous().numpy().tobytes() + ).hexdigest() + def _jacobian_cache_path(self): if not self.cache_dir: return None import os + revision = self.model_revision_resolved or self.model_revision key = hashlib.sha1( - (self.model_name + f"|L{self.layer}|" + "\n".join(self.jacobian_corpus) - + f"|p{self.n_probes}").encode("utf-8") + (self.model_name + f"|revision{revision}" + + f"|dtype{self.dtype}|L{self.layer}|" + "\n".join(self.jacobian_corpus) + + f"|{self.jacobian_estimator}|position{self.jacobian_position}" + + f"|p{self.n_probes}|seed{self.jacobian_seed}").encode("utf-8") ).hexdigest()[:16] os.makedirs(self.cache_dir, exist_ok=True) return os.path.join(self.cache_dir, f"jacobian-{key}.pt") + def _load_precomputed_jacobian(self): + if not self.jacobian_path: + return None + import torch + jacobian = torch.load( + self.jacobian_path, map_location="cpu", weights_only=True + ) + metadata_path = Path( + self.jacobian_metadata_path + or Path(self.jacobian_path).with_suffix(".json") + ) + if not metadata_path.exists(): + raise FileNotFoundError( + f"Jacobian metadata not found: {metadata_path}. " + "Keep the fitter's .pt and .json outputs together, or pass " + "--jacobian-metadata." + ) + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + configuration = metadata.get("configuration", {}) + provenance = metadata.get("provenance", {}) + result = metadata.get("result", {}) + mismatches = [] + expected_configuration = { + "model": self.model_name, + "layer": self.layer, + } + for key, expected in expected_configuration.items(): + if configuration.get(key) != expected: + mismatches.append( + f"{key}={configuration.get(key)!r} (expected {expected!r})" + ) + resolved_revision = provenance.get("model_revision_resolved") + if ( + resolved_revision + and self.model_revision_resolved + and resolved_revision != self.model_revision_resolved + ): + mismatches.append( + f"model revision={resolved_revision!r} " + f"(loaded {self.model_revision_resolved!r})" + ) + actual_sha1 = self._tensor_sha1(jacobian) + if result.get("tensor_sha1") != actual_sha1: + mismatches.append("tensor SHA-1") + if provenance.get("estimator") != "paper_replicated_batch_vjp": + mismatches.append("estimator") + if mismatches: + raise ValueError( + "precomputed Jacobian metadata mismatch: " + + ", ".join(mismatches) + ) + self.jacobian_estimator = provenance["estimator"] + self.jacobian_corpus_id = provenance.get("corpus_sha1") + self.jacobian_n_prompts = provenance.get("prompts_used") + self.n_probes = None + self.jacobian_position = None + return jacobian + def _load_cached_jacobian(self): path = self._jacobian_cache_path() if not path: @@ -264,5 +399,7 @@ def make_backend(kind: str, **kw) -> Backend: if kind == "hf": return HFBackend(**{k: v for k, v in kw.items() if k in ("model_name", "layer", "topk", "max_new_tokens", - "device", "jacobian_corpus", "n_probes", "cache_dir")}) + "device", "jacobian_corpus", "n_probes", "cache_dir", + "dtype", "model_revision", "jacobian_path", + "jacobian_metadata_path")}) raise ValueError(f"unknown backend: {kind!r} (use 'mock' or 'hf')") diff --git a/apps/serve/j7scope_serve/recorder.py b/apps/serve/j7scope_serve/recorder.py index 707ae6a..c546b07 100644 --- a/apps/serve/j7scope_serve/recorder.py +++ b/apps/serve/j7scope_serve/recorder.py @@ -56,6 +56,8 @@ def record_trace( lexicon: Dict[str, str], trace_id: Optional[str] = None, concept: Optional[str] = None, + language: Optional[str] = None, + capture_tool: str = "j7scope_serve --record", seed: int = 0, ) -> Optional[Path]: """Write one Trace v1 from buffered token records. Returns the trace dir. @@ -67,7 +69,7 @@ def record_trace( return None trace_id = trace_id or ("rec-" + uuid.uuid4().hex[:8]) - language = _dominant_language(buffered) + language = language or _dominant_language(buffered) tokens = rigor.compute_trace_rigor(buffered, lexicon, seed=seed) @@ -76,17 +78,36 @@ def record_trace( "trace_id": trace_id, "kind": "single", "model": backend.model_name, + "revision": getattr(backend, "model_revision_resolved", None), "layer": backend.layer, "language": language, "concept": concept, "prompt": prompt, "jacobian": { - "corpus_id": "mock-synthetic" if is_demo else "generic-v1", - "sha1": None, + "corpus_id": ( + "mock-synthetic" + if is_demo + else getattr(backend, "jacobian_corpus_id", "generic-v1") + ), + "estimator": getattr( + backend, "jacobian_estimator", + "synthetic" if is_demo else "unknown", + ), + "position": getattr(backend, "jacobian_position", None), + "n_prompts": getattr( + backend, + "jacobian_n_prompts", + len(getattr(backend, "jacobian_corpus", [])), + ), + "n_probes": getattr(backend, "n_probes", None), + "seed": getattr(backend, "jacobian_seed", None), + "sha1": getattr(backend, "jacobian_sha1", None), }, "capture": { - "tool": "j7scope_serve --record", + "tool": capture_tool, "backend": type(backend).__name__, + "device": getattr(backend, "device", None), + "dtype": getattr(backend, "dtype", None), "created_at": datetime.now(timezone.utc).isoformat(), }, "is_demo": is_demo, diff --git a/apps/web/package-lock.json b/apps/web/package-lock.json index 9657d41..dc3bd04 100644 --- a/apps/web/package-lock.json +++ b/apps/web/package-lock.json @@ -3480,9 +3480,9 @@ "peer": true }, "node_modules/nanoid": { - "version": "3.3.15", - "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.15.tgz", - "integrity": "sha512-y7Wygv/7mEOvxTuEQDB8StXdMRBWf1kR/tlhAzBRUFkB2jfcLOAxO/SHmOO2zgz1pVgK29/kyupn059/bCHdjA==", + "version": "3.3.16", + "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.16.tgz", + "integrity": "sha512-bzlKTyNJ7+LdGIIwy8ijFpIqEQIvafahV7eYykJ8Cvh42EdJeODoJ6gUJXpQJvej1BddH8OqTXZNE/KfbWAu8Q==", "dev": true, "funding": [ { @@ -3751,9 +3751,9 @@ "peer": true }, "node_modules/postcss": { - "version": "8.5.16", - "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.16.tgz", - "integrity": "sha512-vuwillviilfKZsg0VGj5R/YwwcHx4SLsIOI/7K6mQkWx+l5cUHTjj5g0AasTBcyXsbfTgrwsUNmVUb5xVwyPwg==", + "version": "8.5.23", + "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.23.tgz", + "integrity": "sha512-g50586zr4bZmwFiTlflMu8E0bDTb5I5gertgwAKmsdUlTQIhZtunzUlD1WSzwcVWPoAVpsrA6vlfCD7oXvRwgg==", "dev": true, "funding": [ { @@ -3771,7 +3771,7 @@ ], "license": "MIT", "dependencies": { - "nanoid": "^3.3.12", + "nanoid": "^3.3.16", "picocolors": "^1.1.1", "source-map-js": "^1.2.1" }, diff --git a/data/jacobian_fit_smoke.jsonl b/data/jacobian_fit_smoke.jsonl new file mode 100644 index 0000000..b341bc2 --- /dev/null +++ b/data/jacobian_fit_smoke.jsonl @@ -0,0 +1,12 @@ +{"id":"smoke-en-01","text":"Scientific progress often depends on careful measurement, open criticism, repeated experiments, and a willingness to revise an explanation when new evidence contradicts it."} +{"id":"smoke-en-02","text":"During a difficult negotiation, both groups listed their priorities, identified the points they could compromise on, and gradually constructed an agreement that neither side considered perfect."} +{"id":"smoke-en-03","text":"A historian comparing several independent records must consider who created each document, what information the author could access, and which incentives may have shaped the account."} +{"id":"smoke-en-04","text":"When people learn a complex skill, short cycles of practice and feedback usually reveal misconceptions faster than passive review, especially when the examples vary in structure."} +{"id":"smoke-en-05","text":"The engineer traced the unexpected failure through several connected systems, tested each assumption separately, and discovered that a harmless configuration change had altered the timing."} +{"id":"smoke-en-06","text":"A good public explanation distinguishes observed facts from interpretation, acknowledges meaningful uncertainty, and gives readers enough context to evaluate whether the conclusion follows from the evidence."} +{"id":"smoke-zh-01","text":"科学研究的进展往往依赖精确测量、公开批评和可重复实验，也依赖研究者在新证据与旧解释冲突时主动修正自己的判断。"} +{"id":"smoke-zh-02","text":"在一次艰难的谈判中，双方先列出各自最重要的目标，再寻找可以让步的部分，最后逐步形成了一个并不完美但可以执行的协议。"} +{"id":"smoke-zh-03","text":"历史学家比较多份独立记录时，需要判断文件由谁撰写、作者能够接触哪些信息，以及当时的利益关系可能怎样影响叙述。"} +{"id":"smoke-zh-04","text":"学习复杂技能时，短周期的练习和反馈通常比被动复习更快暴露误解，尤其是在例题的表面形式不断变化而核心结构保持一致的时候。"} +{"id":"smoke-zh-05","text":"工程师沿着多个相互连接的系统追查故障，分别检验每一个假设，最终发现一项看似无害的配置修改改变了关键操作的执行顺序。"} +{"id":"smoke-zh-06","text":"高质量的公共说明应当区分观察事实与主观解释，明确承认重要的不确定性，并提供足够背景让读者自行判断结论是否真正得到证据支持。"} diff --git a/docs/gpu-cloud-runbook.md b/docs/gpu-cloud-runbook.md new file mode 100644 index 0000000..7bf5598 --- /dev/null +++ b/docs/gpu-cloud-runbook.md @@ -0,0 +1,161 @@ +# J7Scope 云 GPU 运行手册 + +> 更新时间：2026-07-27。价格与库存会变化，启动实例前以供应商控制台为准。 + +## 1. 结论与推荐配置 + +不需要购买实体 GPU。J7Scope 把昂贵计算集中在拟合和 trace 采集阶段，最适合租用 +按秒或按小时计费的云 GPU，任务结束即释放。 + +| 任务 | 最低建议 | 稳妥选择 | 说明 | +|---|---:|---:|---| +| 上游 Qwen3.5-4B 已知案例复现 | 16 GB | 24 GB L4 / 3090 | 只加载已拟合 lens，先验证读出质量 | +| Qwen2.5-7B position-local 采集 | 24 GB | 40–48 GB | bf16 模型约 14 GB，反向与激活还需空间 | +| 7B paper estimator 正式拟合 | 40 GB | 48–80 GB | 可降低 `dim_batch` 换显存，但运行会更慢 | + +当前首选是 **Runpod Pod + 48 GB A40/A6000**：交互式 shell、磁盘持久化和按小时 +计费都适合一次性研究批次。2026-07-27 官方页面列出的 Pod 参考价为 A40 +48 GB `$0.44/h`、RTX A6000 48 GB `$0.53/h`、L40S 48 GB `$0.99/h`； +24 GB L4 为 `$0.39/h`。这只是 GPU 实例标价，磁盘等费用另计。 + +官方链接： + +- [Runpod GPU 定价](https://www.runpod.io/pricing) +- [Runpod Pod 管理文档](https://docs.runpod.io/pods/manage-pods) +- [Google Colab 资源限制说明](https://research.google.com/colaboratory/faq.html) + +## 2. 直接租用：Runpod + +这条路径无需提交研究资助申请，注册、充值后即可创建实例。 + +1. 在 [Runpod](https://www.runpod.io/) 注册账户并在 Billing 添加支付方式/余额。 +2. 进入 **Pods → Deploy**，选择官方 PyTorch 模板。 +3. 上游 4B 复现选择 24 GB；正式 7B 拟合优先选择 48 GB。磁盘建议至少 80 GB。 +4. 启动后打开 Web Terminal 或 SSH，执行： + +```bash +git clone https://github.com/arthurpanhku/j7scope.git +cd j7scope +python -m venv .venv +source .venv/bin/activate +pip install -U pip +pip install -e '.[upstream,fit]' + +export HF_HOME=/workspace/hf_cache +python experiments/reproduce_upstream.py --preflight-only +python experiments/reproduce_upstream.py +``` + +成功后会生成 `results/upstream-reproduction.json`，其中记录实际 GPU、模型 commit、 +lens snapshot、各层 J-lens/logit-lens top token 和 `euro` 概念命中情况。 + +## 3. 正式拟合与断点恢复 + +先用 48 GB 卡执行显存预检，再用两条语料测峰值显存和单条耗时。benchmark 使用独立 +checkpoint，避免与正式 1000 条语料的 corpus SHA-1 冲突： + +```bash +python experiments/fit_paper_jacobian.py --preflight-only + +python experiments/fit_paper_jacobian.py \ + --dataset Salesforce/wikitext \ + --dataset-config wikitext-103-raw-v1 \ + --max-prompts 2 --min-chars 200 \ + --checkpoint results/jacobian-benchmark.checkpoint.pt \ + --output results/jacobian-benchmark.pt +``` + +`results/jacobian-benchmark.json` 会记录实际耗时和 `peak_cuda_memory_gb`。确认显存余量后 +运行正式 1000 条拟合；若 Pod/SSH 中断，使用完全相同的命令即可从下一个 prompt +继续： + +```bash +python experiments/fit_paper_jacobian.py \ + --dataset Salesforce/wikitext \ + --dataset-config wikitext-103-raw-v1 \ + --dataset-revision main \ + --max-prompts 1000 --min-chars 200 +``` + +正式归档时应把 `--model-revision` 和 `--dataset-revision` 从 `main` 换成运行前确认的 +commit。即使远端数据发生变化，checkpoint 的完整有序语料 SHA-1 也会拒绝混合续跑。 +`dim_batch=8` 显存不足时可以降低到 4 或 2；这不改变估计定义，但会增加运行时间。 + +默认的 `data/jacobian_fit_smoke.jsonl` 只有 12 条中英语句，只用于验证代码、输出格式 +和恢复机制，不足以形成研究结论。正式结果包括： + +- `results/jacobian-qwen2.5-7b-l18.checkpoint.pt`：可恢复的累积和与进度； +- `results/jacobian-qwen2.5-7b-l18.pt`：sidecar 可直接加载的最终 float32 矩阵； +- `results/jacobian-qwen2.5-7b-l18.json`：revision、语料/矩阵 SHA-1、硬件和性能元数据。 + +加载正式矩阵： + +```bash +cd apps/serve +python -m j7scope_serve --backend hf \ + --model Qwen/Qwen2.5-7B-Instruct --layer 18 \ + --jacobian-file ../../results/jacobian-qwen2.5-7b-l18.pt +``` + +开始大下载前先运行 `--preflight-only`；显存不足时脚本会直接退出。复制回结果和需要 +保留的 cache 后，在控制台 **Terminate Pod**。只停止实例仍可能产生磁盘费用，操作前 +查看控制台的实时费用摘要。 + +公共 Qwen 模型与公开 lens 不需要 Hugging Face token。若以后使用受限模型，把 token +放入平台 Secret/环境变量，禁止写入仓库、notebook 或 trace。 + +## 4. Colab 能否使用 + +可以，但建议只用于 P4 社区采集笔记本和小规模预演： + +- 免费层 GPU 类型和配额不保证，昂贵资源受限； +- 免费 notebook 最长通常不超过 12 小时，实际还受可用性和使用模式影响； +- 即使是付费方案，具体 GPU 仍取决于当时库存。 + +如果 Colab 分配到 24 GB L4，可以运行上游 4B 复现；若只分配到约 16 GB T4，不建议 +强行跑正式批次。P3 需要固定硬件、不中断和可核算成本，因此使用 Pod 更合适。 + +## 5. 申请免费的学术云额度 + +直接租用可以马上开始；学术额度适合覆盖后续 M1–M3 批量实验。 + +### Google Cloud Research Credits（优先） + +1. 用学校/研究机构身份创建 Google Cloud 账户和 Billing Account。 +2. 用官方 Pricing Calculator 估算 GPU、CPU、磁盘和存储费用。 +3. 通过 [GCP Research Credits 申请指南](https://support.google.com/google-cloud-higher-ed/answer/10724468) + 的在线表单提交研究计划和费用估算。 +4. 官方说明通常需要 6–8 周审核；申请全年开放。符合条件的博士生可申请每年 + `$1,000` GCP credits。 + +建议申请 `$1,000`，项目类型写成“finite proof-of-concept + repeatable open research +tool”，强调 trace、代码、实验配置和 Zenodo 数据都会公开。 + +### AWS Cloud Credit for Research + +1. 先建立 AWS 账户并取得 12 位 account ID。 +2. 使用学校/机构邮箱填写 + [AWS Cloud Credit for Research](https://aws.amazon.com/government-education/research-and-technical-computing/cloud-credit-for-research/) + 申请。 +3. 研究生申请上限为 `$5,000`，教师/全职研究人员没有该上限。 + +AWS 官方当前说明无法处理 **greater China region** 的申请；如果申请人/机构地址属于 +该范围，应优先走 GCP 或直接租用 Runpod，并在提交前向 AWS 确认香港地址是否可受理。 + +### NVIDIA Academic Grant + +[NVIDIA Academic Grant Program](https://www.nvidia.com/en-us/industries/higher-education-research/academic-grant-program/) +按主题和申请周期征集 proposal，适合中长期研究合作或硬件支持，不适合作为眼前 P3 +批次的唯一依赖。 + +## 6. 申请材料清单 + +- 项目名称与 150–250 字英文摘要； +- 研究问题：跨语言 J-space 是否共享； +- 有限、可验证的交付物：上游复现、M1 指标、真实 trace、Zenodo DOI； +- 模型、GPU 类型、预计 GPU 小时和费用计算； +- 数据治理：只使用公开模型、公开/合成语料，不处理个人敏感数据； +- 开放科学计划：Apache-2.0 代码、版本化 schema、公开 artifact 和可复现实验配置； +- 时间线：复现 → 小批次 → 全概念 → DOI 发布。 + +可直接使用 [`gpu-credit-proposal.md`](gpu-credit-proposal.md) 作为英文申请底稿。 diff --git a/docs/gpu-credit-proposal.md b/docs/gpu-credit-proposal.md new file mode 100644 index 0000000..b43b50e --- /dev/null +++ b/docs/gpu-credit-proposal.md @@ -0,0 +1,71 @@ +# J7Scope GPU Cloud Credit Proposal + +## Project title + +J7Scope: Reproducible Cross-Lingual Validation of the Jacobian-Lens Global Workspace + +## Abstract + +J7Scope is an open-source interpretability project testing whether the +verbalizable “global workspace” identified by Jacobian-lens methods is shared +across Chinese and English. We will reproduce the published pretrained +Jacobian lens on a known Qwen case, then fit and evaluate a bilingual lens on +open-weight Qwen models. The project separates GPU-dependent capture from a +zero-GPU static replay platform: every experiment produces a versioned trace +with model revision, estimator definition, Jacobian provenance, shuffled-pair +null baselines, confidence intervals, and deep links to individual tokens. + +Cloud credits will fund a finite validation and capture campaign rather than +an ongoing service. Deliverables will include (1) an upstream reproduction +report, (2) cross-lingual overlap/CKA/SVCCA results with null and same-language +controls, (3) a curated bilingual trace dataset, and (4) a public Zenodo +release with DOI. Code is Apache-2.0 and all small artifacts, configurations, +and methodology documentation will be publicly available. + +## Research question + +Do semantically matched Chinese and English prompts converge on a shared +Jacobian-lens subspace beyond shuffled-pair and same-language baselines? + +## Work plan + +1. Reproduce the published Qwen3.5-4B pretrained-lens walkthrough and archive + resolved model/lens revisions and layer-level readouts. +2. Lock the estimator definition by comparing position-local and published + future-target reductions. +3. Fit Qwen2.5-7B on a generic, language-balanced corpus. +4. Run M1 on deception and manipulation probes, then expand to all registered + concept categories with same-language controls. +5. Capture bilingual traces and publish code, metrics, data cards, and a + versioned dataset with DOI. + +## Requested resources + +- Requested credit: USD 1,000. +- GPU: NVIDIA A100 40/80 GB, L40S 48 GB, or equivalent. +- Planned usage: up to 200 GPU-hours, adjusted after the upstream reproduction + and a measured ten-prompt fitting benchmark. +- Storage: 200 GB temporary model/cache storage and under 20 GB durable + experiment artifacts. +- Compute is batchable and will be stopped between experiments. + +## Reproducibility and public benefit + +Every run records the exact model revision, upstream code commit, lens +snapshot, estimator, corpus identifier, layer, seed, dtype, and GPU. Results +will be released through the public J7Scope repository and Zenodo. The static +viewer requires no ongoing GPU resources, so the credit creates a durable +public research artifact rather than a service with recurring compute cost. + +## Data and safety + +The project uses public open-weight models and public, authored, or synthetic +probe text. It does not process personal, medical, financial, or otherwise +sensitive user data. Model licenses and upstream Apache-2.0 attribution will +be preserved. + +## Timeline + +- Weeks 1–2: upstream reproduction and estimator lock. +- Weeks 3–5: first bilingual M1 batch and robustness controls. +- Weeks 6–8: expanded concepts, trace capture, data card, and Zenodo release. diff --git a/docs/platform-plan.md b/docs/platform-plan.md index 7281604..cb83511 100644 --- a/docs/platform-plan.md +++ b/docs/platform-plan.md @@ -86,6 +86,8 @@ traces// "prompt": "…", "jacobian": { // J_l 溯源：可复现的前提 "corpus_id": "generic-v1", + "estimator": "position_local_gaussian_vjp", + "position": -1, "n_prompts": 8, "n_probes": 16, "seed": 0, "sha1": "" }, @@ -178,16 +180,42 @@ traces// > 的碎 token（`WaitForSeconds`、`.ToDecimal`、`闹`…）。原因是：(1) 随机 Jacobian 估计 > 严重欠采样（512 个样本估 896×896 矩阵，需要数千以上），(2) 0.5B 的 workspace 本就弱。 > **含义：小模型 + CPU 可行预算下没有信号，这不是平台问题，正是 P3 需要 GPU + 更大模型的实证理由。** -> **P3 前置动作**：先按 `estimate_jacobian` 的 TODO，用 `torch.func.jacrev` 精确 Jacobian -> 和上游 `anthropics/jacobian-lens` 在一个已知案例上校验拟合数学，确认 7B 跑出来的是信号而非噪声，再租卡。 +> **P3 前置动作**： +> +> - [x] `j7scope.fitting.exact_jacobian_for_prompt` 用 `torch.func.jacrev` 实现 +> position-local 精确 Jacobian；已同时通过解析线性尾部与随机 tiny-Qwen2 的独立 +> full-graph VJP 校验。`experiments/validate_jacobian.py` 可复跑精确/随机估计对比， +> 全程不下载模型。 +> - [x] `paper_jacobian_for_prompt` 独立实现上游 paper reduction：对每个 source +> position 汇总当前及未来 target positions，再跨 source positions 平均；执行路径 +> 与上游一致，复制 prompt 后做普通 VJP，并已在 tiny-Qwen2 上与逐标量 VJP 零误差 +> 交叉校验。J7Scope 当前 capture 路径仍是 +> cloze 用的同位置 `∂h_L[p]/∂h_l[p]` 变体，两者不会静默混用。 +> - [x] `estimate_jacobian_paper` 已加入逐 prompt 原子 checkpoint/resume、完整语料 +> SHA-1 和模型/层/参数一致性校验；`experiments/fit_paper_jacobian.py` 提供 48 GB +> GPU preflight、WikiText streaming、benchmark、峰值显存/耗时与最终矩阵元数据。 +> - [x] `experiments/reproduce_upstream.py` 已固定上游发布配置 +> `Qwen/Qwen3.5-4B` + `neuronpedia/jacobian-lens@qwen-n1000`，记录解析后的 model/lens +> revision、GPU 和逐层 J-lens/logit-lens 读出；支持无下载 `--dry-run` 与下载前 +> `--preflight-only`。云 GPU 操作与 credits 申请见 `docs/gpu-cloud-runbook.md`。 +> - [ ] 在云 GPU 执行该已知案例，并据真实读出锁定 P3 正式 estimator；本地随机 tiny +> 模型只能校验数学实现，不能校验 lens 质量。 +> - [ ] 上述校验通过并锁定正式 estimator 后再租卡，避免把 GPU 预算花在数学口径错误 +> 或尚未收敛的矩阵上。 ### P4 · Colab 采集笔记本 + 社区提交流程 - **目标**：把 GPU 成本外包给社区，形成 trace 供给。 - **交付物** - - `notebooks/capture_colab.ipynb`：T4 免费层跑 Qwen2.5-1.5B-Instruct，10 分钟采一条自定义 trace。 - - `CONTRIBUTING.md` + trace 提交模板 + CI 校验（schema、`is_demo` 标注、`sharedness` 未在前端重算）。 -- **验收**：外部贡献者能在 Colab 免费层产出通过 CI 的 trace 并合入 gallery。 + - [x] `notebooks/capture_colab.ipynb`：免费 GPU 跑 Qwen2.5-1.5B-Instruct， + `dtype=auto` 在 T4 使用 fp16、新卡使用 bf16；采一条自定义 preview trace 并下载 ZIP。 + - [x] `experiments/capture_trace.py`：非交互单 trace 入口；记录 model revision、GPU、 + dtype、estimator 与 Jacobian SHA-1，输出后自动重建并验证 gallery。 + - [x] `CONTRIBUTING.md` + PR 模板 + GitHub Actions 校验（schema、索引、provenance、 + `is_demo`/`preview`、sharedness definition、前端未重算严谨层）。 +- **验收**： + - [x] 本地无 GPU dry-run、demo gallery 端到端生成/校验、CI 命令通过。 + - [ ] 首个外部贡献者在 Colab 免费层产出通过 CI 的 trace 并合入 gallery。 ### P5 · 严谨层完全体（并排对齐 + null 带 + 可展开审计） diff --git a/experiments/build_demo_trace.py b/experiments/build_demo_trace.py index 11612c0..19c657e 100644 --- a/experiments/build_demo_trace.py +++ b/experiments/build_demo_trace.py @@ -134,7 +134,15 @@ def build_trace(narrative, language: str, lexicon, *, trace_id: str, "language": language, "concept": None, "prompt": "(synthetic demo narrative)", - "jacobian": {"corpus_id": "demo-synthetic", "sha1": None}, + "jacobian": { + "corpus_id": "demo-synthetic", + "estimator": "synthetic", + "position": None, + "n_prompts": 0, + "n_probes": None, + "seed": seed, + "sha1": None, + }, "capture": {"tool": "build_demo_trace.py", "created_at": datetime.now(timezone.utc).isoformat()}, "is_demo": True, diff --git a/experiments/capture_small.py b/experiments/capture_small.py index 1ec7e2f..de3f0e7 100644 --- a/experiments/capture_small.py +++ b/experiments/capture_small.py @@ -20,7 +20,6 @@ import json import sys import time -from datetime import datetime, timezone from pathlib import Path ROOT = Path(__file__).resolve().parents[1] @@ -59,6 +58,7 @@ def main() -> None: from j7scope_serve.backends import HFBackend from j7scope_serve.recorder import build_lexicon, record_trace from j7scope_serve.protocol import bucket_readout, script_of + from j7scope.trace import rebuild_trace_index print(f"loading {args.model} on {args.device} (first run downloads ~1GB)…") t0 = time.time() @@ -88,18 +88,22 @@ def main() -> None: print(f" (no tokens produced for {trace_id})") continue path = record_trace(args.out, backend=be, prompt=prompt, buffered=buffered, - lexicon=lexicon, trace_id=trace_id) + lexicon=lexicon, trace_id=trace_id, language=lang, + concept="deception", + capture_tool="experiments/capture_small.py") _mark_preview(path, model=args.model, layer=args.layer, language=lang) print(f" wrote {path} ({len(buffered)} tokens, {time.time() - t1:.0f}s)") written.append(trace_id) - _refresh_index(args.out) + rebuild_trace_index(args.out) print(f"done. real small-model traces: {written}") print("view: python -m j7scope_serve --backend mock --traces results/traces -> http://127.0.0.1:8799/") def _mark_preview(trace_dir: Path, *, model: str, layer: int, language: str) -> None: """Flag the manifest as a small-model preview (real read-outs, not research).""" + from j7scope.artifacts import write_json + mpath = trace_dir / "manifest.json" m = json.loads(mpath.read_text(encoding="utf-8")) m["preview"] = True @@ -107,30 +111,7 @@ def _mark_preview(trace_dir: Path, *, model: str, layer: int, language: str) -> m["note"] = ("Real J-lens read-out from a small model on CPU. Pipeline preview, " "not a research result; cross-lingual numbers are not meaningful at " "this scale.") - mpath.write_text(json.dumps(m, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8") - - -def _refresh_index(out: Path) -> None: - """Rebuild traces/index.json from whatever trace dirs exist.""" - rows = [] - for d in sorted(out.glob("*/")): - mf = d / "manifest.json" - if not mf.exists(): - continue - m = json.loads(mf.read_text(encoding="utf-8")) - rows.append({ - "trace_id": m["trace_id"], - "label": m.get("label", m["trace_id"]), - "language": m.get("language", "?"), - "is_demo": bool(m.get("is_demo", False)), - "preview": bool(m.get("preview", False)), - "n_tokens": sum(1 for _ in (d / "tokens.jsonl").open(encoding="utf-8")), - "parallel_group": m.get("parallel_group"), - }) - (out / "index.json").write_text( - json.dumps({"schema_version": 1, "traces": rows}, ensure_ascii=False, - indent=2, sort_keys=True) + "\n", encoding="utf-8") + write_json(mpath, m) if __name__ == "__main__": diff --git a/experiments/capture_trace.py b/experiments/capture_trace.py new file mode 100644 index 0000000..5d1c919 --- /dev/null +++ b/experiments/capture_trace.py @@ -0,0 +1,192 @@ +"""Capture one custom, real-model Trace v1 artifact. + +This is the non-interactive entry point used by the P4 Colab notebook and cloud +GPU contributors. Community captures are marked ``preview: true`` until they +have been reviewed against the project research protocol. + +Example: + python experiments/capture_trace.py \ + --trace-id community-deception-en \ + --language en \ + --concept deception \ + --prompt "In one sentence, explain why deception can be tempting." +""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +import time +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] +for import_path in (ROOT, ROOT / "apps" / "serve"): + if str(import_path) not in sys.path: + sys.path.insert(0, str(import_path)) + +from j7scope.artifacts import write_json # noqa: E402 +from j7scope.trace import rebuild_trace_index, validate_trace_gallery # noqa: E402 + + +TRACE_ID_RE = re.compile(r"^[a-z0-9][a-z0-9._-]{0,79}$") +DEFAULT_JACOBIAN_CORPUS = [ + "The weather today is quite", + "In the history of science, the most important", + "She opened the door and saw", + "The best way to learn a new skill is to", + "他走进房间，看到桌上放着一封", + "关于这个问题，我认为最关键的是", +] + + +def _configuration(args) -> dict: + return { + "trace_id": args.trace_id, + "language": args.language, + "concept": args.concept, + "model": args.model, + "model_revision": args.model_revision, + "layer": args.layer, + "device": args.device, + "dtype": args.dtype, + "n_probes": args.n_probes, + "max_new_tokens": args.max_new_tokens, + "output": str(args.out), + } + + +def capture(args) -> Path: + from j7scope_serve.backends import HFBackend + from j7scope_serve.protocol import bucket_readout, script_of + from j7scope_serve.recorder import build_lexicon, record_trace + + backend = HFBackend( + model_name=args.model, + model_revision=args.model_revision, + layer=args.layer, + topk=args.topk, + max_new_tokens=args.max_new_tokens, + device=args.device, + dtype=args.dtype, + jacobian_corpus=DEFAULT_JACOBIAN_CORPUS, + n_probes=args.n_probes, + cache_dir=str(args.cache_dir), + ) + print( + f"Loading {args.model}@{args.model_revision} on {args.device} " + f"(dtype={args.dtype})…", + flush=True, + ) + started = time.monotonic() + backend.load() + print( + f"Model and Jacobian ready in {time.monotonic() - started:.1f}s " + f"({backend.device}, {backend.dtype})", + flush=True, + ) + + buffered = [] + generation_started = time.monotonic() + for step in backend.generate([{"role": "user", "content": args.prompt}]): + buffered.append( + { + "seq": len(buffered), + "ts_rel": round(time.monotonic() - generation_started, 3), + "token": step.token, + "token_script": script_of(step.token), + "readout": bucket_readout(step.topk, per_lang=args.per_lang), + } + ) + if not buffered: + raise RuntimeError("model produced no tokens; no trace was written") + + trace_dir = record_trace( + args.out, + backend=backend, + prompt=args.prompt, + buffered=buffered, + lexicon=build_lexicon(ROOT, backend), + trace_id=args.trace_id, + concept=args.concept, + language=args.language, + capture_tool="experiments/capture_trace.py", + ) + manifest_path = trace_dir / "manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest.update( + { + "label": args.label or f"Community preview · {args.trace_id}", + "preview": True, + "note": ( + "Real model capture contributed through the P4 capture tool. " + "Preview only until estimator convergence and research review." + ), + } + ) + manifest["capture"]["community_submission"] = True + write_json(manifest_path, manifest) + rebuild_trace_index(args.out) + + problems = validate_trace_gallery(args.out) + if problems: + joined = "\n".join(f"- {problem}" for problem in problems) + raise RuntimeError(f"captured trace failed gallery validation:\n{joined}") + print(f"Wrote {len(buffered)} tokens to {trace_dir}", flush=True) + return trace_dir + + +def parse_args(argv=None): + parser = argparse.ArgumentParser() + parser.add_argument("--trace-id", required=True) + parser.add_argument("--prompt", required=True) + parser.add_argument("--language", choices=["en", "zh", "other"], required=True) + parser.add_argument("--concept", default=None) + parser.add_argument("--label", default=None) + parser.add_argument("--model", default="Qwen/Qwen2.5-1.5B-Instruct") + parser.add_argument("--model-revision", default="main") + parser.add_argument("--layer", type=int, default=14) + parser.add_argument("--device", default="cuda") + parser.add_argument( + "--dtype", + choices=["auto", "bfloat16", "float16", "float32"], + default="auto", + ) + parser.add_argument("--n-probes", type=int, default=8) + parser.add_argument("--max-new-tokens", type=int, default=48) + parser.add_argument("--topk", type=int, default=24) + parser.add_argument("--per-lang", type=int, default=8) + parser.add_argument( + "--out", + type=Path, + default=ROOT / "results" / "community-traces", + ) + parser.add_argument( + "--cache-dir", + type=Path, + default=ROOT / ".cache" / "jacobian", + ) + parser.add_argument("--dry-run", action="store_true") + return parser.parse_args(argv) + + +def main(argv=None) -> None: + args = parse_args(argv) + if not TRACE_ID_RE.fullmatch(args.trace_id): + raise SystemExit( + "--trace-id must be 1-80 lowercase letters, digits, dots, " + "underscores, or hyphens, and cannot start with punctuation" + ) + for name in ("n_probes", "max_new_tokens", "topk", "per_lang"): + if getattr(args, name) < 1: + raise SystemExit(f"--{name.replace('_', '-')} must be >= 1") + if args.dry_run: + print(json.dumps(_configuration(args), indent=2, sort_keys=True)) + return + capture(args) + + +if __name__ == "__main__": + main() diff --git a/experiments/fit_paper_jacobian.py b/experiments/fit_paper_jacobian.py new file mode 100644 index 0000000..9b8a27f --- /dev/null +++ b/experiments/fit_paper_jacobian.py @@ -0,0 +1,390 @@ +"""Fit a production-style paper-reduction Jacobian on a CUDA GPU. + +The local JSONL corpus is intentionally small and is only a smoke test. A +research run should use a pretraining-like corpus such as WikiText and retain +the checkpoint, final tensor, and metadata JSON together. + +Smoke test configuration: + python experiments/fit_paper_jacobian.py --dry-run + +Cloud GPU: + pip install -e '.[fit]' + python experiments/fit_paper_jacobian.py --preflight-only + python experiments/fit_paper_jacobian.py \ + --dataset Salesforce/wikitext \ + --dataset-config wikitext-103-raw-v1 \ + --max-prompts 1000 --min-chars 200 +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from j7scope.artifacts import write_json # noqa: E402 +from j7scope.fitting import ( # noqa: E402 + JLens, + _atomic_torch_save, + _prompt_corpus_sha1, + load_model, +) + + +DEFAULT_MODEL = "Qwen/Qwen2.5-7B-Instruct" +DEFAULT_CORPUS = ROOT / "data" / "jacobian_fit_smoke.jsonl" +DEFAULT_OUTPUT = ROOT / "results" / "jacobian-qwen2.5-7b-l18.pt" +DEFAULT_CHECKPOINT = ( + ROOT / "results" / "jacobian-qwen2.5-7b-l18.checkpoint.pt" +) + + +def _load_local_prompts( + path: Path, *, text_key: str, max_prompts: int, min_chars: int +) -> list[str]: + if not path.exists(): + raise FileNotFoundError(f"corpus not found: {path}") + prompts = [] + with path.open(encoding="utf-8") as corpus_file: + for line_number, line in enumerate(corpus_file, start=1): + line = line.strip() + if not line: + continue + if path.suffix.casefold() == ".jsonl": + try: + row = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError( + f"{path}:{line_number}: invalid JSON" + ) from exc + if isinstance(row, str): + text = row + elif isinstance(row, dict) and isinstance( + row.get(text_key), str + ): + text = row[text_key] + else: + raise ValueError( + f"{path}:{line_number}: expected a string or " + f"object with string field {text_key!r}" + ) + else: + text = line + text = text.strip() + if len(text) < min_chars: + continue + prompts.append(text) + if len(prompts) >= max_prompts: + break + if not prompts: + raise ValueError("corpus contains no prompts after filtering") + return prompts + + +def _load_dataset_prompts(args) -> list[str]: + try: + from datasets import load_dataset + except ImportError as exc: + raise RuntimeError( + "Hugging Face dataset support is missing; " + "run `pip install -e '.[fit]'`" + ) from exc + + dataset = load_dataset( + args.dataset, + args.dataset_config, + split=args.dataset_split, + revision=args.dataset_revision, + streaming=True, + ) + prompts = [] + for row in dataset: + text = row.get(args.text_key) + if not isinstance(text, str): + raise ValueError( + f"dataset row has no string field {args.text_key!r}" + ) + text = text.strip() + if len(text) < args.min_chars: + continue + prompts.append(text) + if len(prompts) >= args.max_prompts: + break + if not prompts: + raise ValueError("dataset contains no prompts after filtering") + if len(prompts) < args.max_prompts: + raise ValueError( + f"dataset ended after {len(prompts)} usable prompts; " + f"requested {args.max_prompts}" + ) + return prompts + + +def _load_prompts(args) -> list[str]: + if args.dataset: + return _load_dataset_prompts(args) + return _load_local_prompts( + args.corpus, + text_key=args.text_key, + max_prompts=args.max_prompts, + min_chars=args.min_chars, + ) + + +def _resolve_dtype(torch, requested: str, device: str): + if requested != "auto": + return getattr(torch, requested) + if device.startswith("cuda"): + return ( + torch.bfloat16 + if torch.cuda.is_bf16_supported() + else torch.float16 + ) + return torch.float32 + + +def _gpu_preflight(torch, *, device: str, min_vram_gb: float, + allow_cpu: bool) -> dict: + if not device.startswith("cuda"): + if not allow_cpu: + raise RuntimeError( + "paper Jacobian fitting is intended for CUDA; pass " + "--allow-cpu only for a tiny diagnostic run" + ) + return { + "device": device, + "cuda_available": bool(torch.cuda.is_available()), + "warning": "CPU mode is diagnostic only.", + } + if not torch.cuda.is_available(): + raise RuntimeError( + "CUDA GPU not found. Select a cloud GPU runtime and rerun " + "--preflight-only." + ) + index = torch.device(device).index or 0 + properties = torch.cuda.get_device_properties(index) + vram_gb = properties.total_memory / 1024**3 + if vram_gb < min_vram_gb: + raise RuntimeError( + f"GPU {properties.name!r} has {vram_gb:.1f} GiB VRAM; " + f"this run requires at least {min_vram_gb:.1f} GiB. " + "Lower --dim-batch and explicitly lower --min-vram-gb only " + "after a measured smoke run." + ) + return { + "device": str(torch.device(device)), + "cuda_available": True, + "gpu_name": properties.name, + "vram_gb": round(vram_gb, 2), + "bf16_supported": bool(torch.cuda.is_bf16_supported()), + "cuda_version": torch.version.cuda, + } + + +def _tensor_sha1(tensor) -> str: + return hashlib.sha1( + tensor.detach().float().cpu().contiguous().numpy().tobytes() + ).hexdigest() + + +def _configuration(args) -> dict: + source = ( + { + "kind": "huggingface_dataset", + "dataset": args.dataset, + "config": args.dataset_config, + "split": args.dataset_split, + "revision_requested": args.dataset_revision, + "text_key": args.text_key, + } + if args.dataset + else { + "kind": "local", + "path": str(args.corpus), + "text_key": args.text_key, + "smoke_only": args.corpus.resolve() == DEFAULT_CORPUS.resolve(), + } + ) + return { + "model": args.model, + "model_revision_requested": args.model_revision, + "layer": args.layer, + "device": args.device, + "dtype_requested": args.dtype, + "skip_first": args.skip_first, + "dim_batch": args.dim_batch, + "max_prompts": args.max_prompts, + "min_chars": args.min_chars, + "checkpoint_every": args.checkpoint_every, + "resume": not args.no_resume, + "corpus": source, + "checkpoint": str(args.checkpoint), + "output": str(args.output), + "metadata": str(args.metadata), + "min_vram_gb": args.min_vram_gb, + } + + +def fit(args) -> dict: + import torch + + hardware = _gpu_preflight( + torch, + device=args.device, + min_vram_gb=args.min_vram_gb, + allow_cpu=args.allow_cpu, + ) + if args.preflight_only: + return { + "status": "preflight_passed", + "configuration": _configuration(args), + "hardware": hardware, + "torch_version": torch.__version__, + } + + prompts = _load_prompts(args) + dtype = _resolve_dtype(torch, args.dtype, args.device) + model, tokenizer = load_model( + args.model, + device=args.device, + dtype=dtype, + revision=args.model_revision, + ) + lens = JLens(model, tokenizer, layer=args.layer) + + if args.device.startswith("cuda"): + torch.cuda.reset_peak_memory_stats(torch.device(args.device)) + started = time.monotonic() + jacobian = lens.estimate_jacobian_paper( + prompts, + skip_first=args.skip_first, + dim_batch=args.dim_batch, + checkpoint_path=args.checkpoint, + checkpoint_every=args.checkpoint_every, + resume=not args.no_resume, + ) + elapsed_seconds = time.monotonic() - started + _atomic_torch_save(jacobian, args.output) + + checkpoint = torch.load( + args.checkpoint, map_location="cpu", weights_only=True + ) + report = { + "status": "completed", + "created_at": datetime.now(timezone.utc).isoformat(), + "scope": ( + "Smoke/benchmark only; not research evidence." + if not args.dataset and args.corpus.resolve() == DEFAULT_CORPUS.resolve() + else "Paper-reduction Jacobian fit." + ), + "configuration": _configuration(args), + "hardware": hardware, + "software": { + "python": sys.version.split()[0], + "torch": torch.__version__, + }, + "provenance": { + "model_revision_resolved": ( + getattr(model.config, "_commit_hash", None) + or args.model_revision + ), + "corpus_sha1": _prompt_corpus_sha1(prompts), + "prompts_loaded": len(prompts), + "prompts_used": int(checkpoint["n_used"]), + "prompts_skipped": checkpoint["skipped"], + "estimator": "paper_replicated_batch_vjp", + }, + "result": { + "shape": list(jacobian.shape), + "dtype": str(jacobian.dtype).removeprefix("torch."), + "tensor_sha1": _tensor_sha1(jacobian), + "elapsed_seconds": round(elapsed_seconds, 3), + "peak_cuda_memory_gb": ( + round( + torch.cuda.max_memory_allocated( + torch.device(args.device) + ) / 1024**3, + 3, + ) + if args.device.startswith("cuda") + else None + ), + "output": str(args.output), + }, + } + write_json(args.metadata, report) + return report + + +def parse_args(argv=None): + parser = argparse.ArgumentParser() + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--model-revision", default="main") + parser.add_argument("--layer", type=int, default=18) + parser.add_argument("--device", default="cuda:0") + parser.add_argument( + "--dtype", + choices=["auto", "bfloat16", "float16", "float32"], + default="auto", + ) + parser.add_argument("--skip-first", type=int, default=16) + parser.add_argument("--dim-batch", type=int, default=8) + parser.add_argument("--max-prompts", type=int, default=1000) + parser.add_argument("--min-chars", type=int, default=0) + parser.add_argument("--corpus", type=Path, default=DEFAULT_CORPUS) + parser.add_argument("--dataset", default=None) + parser.add_argument("--dataset-config", default=None) + parser.add_argument("--dataset-split", default="train") + parser.add_argument("--dataset-revision", default="main") + parser.add_argument("--text-key", default="text") + parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT) + parser.add_argument("--checkpoint-every", type=int, default=1) + parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT) + parser.add_argument("--metadata", type=Path, default=None) + parser.add_argument("--min-vram-gb", type=float, default=40.0) + parser.add_argument("--no-resume", action="store_true") + parser.add_argument("--allow-cpu", action="store_true") + parser.add_argument("--preflight-only", action="store_true") + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args(argv) + if args.metadata is None: + args.metadata = args.output.with_suffix(".json") + return args + + +def main(argv=None) -> None: + args = parse_args(argv) + integer_bounds = { + "layer": 0, + "skip_first": 0, + "dim_batch": 1, + "max_prompts": 1, + "min_chars": 0, + "checkpoint_every": 1, + } + for name, minimum in integer_bounds.items(): + if getattr(args, name) < minimum: + raise SystemExit( + f"--{name.replace('_', '-')} must be >= {minimum}" + ) + if args.min_vram_gb <= 0: + raise SystemExit("--min-vram-gb must be > 0") + report = ( + {"status": "dry_run", "configuration": _configuration(args)} + if args.dry_run + else fit(args) + ) + print(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/experiments/reproduce_upstream.py b/experiments/reproduce_upstream.py new file mode 100644 index 0000000..65cb342 --- /dev/null +++ b/experiments/reproduce_upstream.py @@ -0,0 +1,280 @@ +"""Reproduce the upstream pretrained Jacobian-lens walkthrough on a cloud GPU. + +Defaults are pinned to Anthropic's published Qwen3.5-4B walkthrough and its +Neuronpedia-hosted n=1000 lens. The resulting JSON records resolved model/lens +revisions, GPU details, and J-lens vs logit-lens top tokens at the known +two-hop prompt. It is the final P3 math/quality gate before fitting our own 7B +lens. + +Install: + pip install -e '.[upstream]' + +Cloud-GPU usage: + python experiments/reproduce_upstream.py --preflight-only + python experiments/reproduce_upstream.py + +Configuration-only check (no GPU or downloads): + python experiments/reproduce_upstream.py --dry-run +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import sys +from datetime import datetime, timezone +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from j7scope.artifacts import write_json # noqa: E402 + + +UPSTREAM_COMMIT = "581d398613e5602a5af361e1c34d3a92ea82ba8e" +DEFAULT_MODEL = "Qwen/Qwen3.5-4B" +DEFAULT_LENS_REPO = "neuronpedia/jacobian-lens" +DEFAULT_LENS_REVISION = "qwen-n1000" +DEFAULT_LENS_FILE = ( + "qwen3.5-4b/jlens/Salesforce-wikitext/" + "Qwen3.5-4B_jacobian_lens_n1000.pt" +) +DEFAULT_PROMPT = "Fact: The currency used in the country shaped like a boot is" + + +def _configuration(args) -> dict: + return { + "model": args.model, + "model_revision_requested": args.model_revision, + "lens_repo": args.lens_repo, + "lens_revision_requested": args.lens_revision, + "lens_file": args.lens_file, + "prompt": args.prompt, + "position": -2, + "top_k": args.top_k, + "min_vram_gb": args.min_vram_gb, + "requires_cuda": not args.allow_cpu, + "upstream_commit": UPSTREAM_COMMIT, + "upstream_dependency_installed": importlib.util.find_spec("jlens") is not None, + } + + +def _gpu_preflight(torch, *, min_vram_gb: float, allow_cpu: bool) -> dict: + if not torch.cuda.is_available(): + if allow_cpu: + return { + "device": "cpu", + "cuda_available": False, + "warning": "CPU mode is diagnostic only and may be very slow.", + } + raise RuntimeError( + "CUDA GPU not found. Select a cloud GPU runtime, then rerun " + "--preflight-only before downloading model weights." + ) + + props = torch.cuda.get_device_properties(0) + vram_gb = props.total_memory / 1024**3 + if vram_gb < min_vram_gb: + raise RuntimeError( + f"GPU {props.name!r} has {vram_gb:.1f} GiB VRAM; " + f"this run requires at least {min_vram_gb:.1f} GiB." + ) + return { + "device": "cuda:0", + "cuda_available": True, + "gpu_name": props.name, + "vram_gb": round(vram_gb, 2), + "bf16_supported": bool(torch.cuda.is_bf16_supported()), + "cuda_version": torch.version.cuda, + } + + +def _snapshot_revision(path: str) -> str | None: + parts = Path(path).parts + try: + return parts[parts.index("snapshots") + 1] + except (ValueError, IndexError): + return None + + +def _top_tokens(logits, tokenizer, k: int) -> list[dict]: + top = logits.float().topk(k) + return [ + { + "token_id": int(token_id), + "token": tokenizer.decode([int(token_id)]), + "logit": round(float(logit), 5), + } + for token_id, logit in zip( + top.indices.detach().cpu(), top.values.detach().cpu(), strict=True + ) + ] + + +def _contains_euro(rows: list[dict]) -> bool: + return any("euro" in row["token"].casefold() for row in rows) + + +def run(args) -> dict: + try: + import jlens + import torch + import transformers + from huggingface_hub import hf_hub_download + except ImportError as exc: + raise RuntimeError( + "upstream validation dependencies are missing; run " + "`pip install -e '.[upstream]'` first" + ) from exc + + hardware = _gpu_preflight( + torch, min_vram_gb=args.min_vram_gb, allow_cpu=args.allow_cpu + ) + if args.preflight_only: + return { + "status": "preflight_passed", + "configuration": _configuration(args), + "hardware": hardware, + "torch_version": torch.__version__, + "transformers_version": transformers.__version__, + } + + device = torch.device(hardware["device"]) + dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 + hf_model = transformers.AutoModelForCausalLM.from_pretrained( + args.model, + revision=args.model_revision, + dtype=dtype, + ).to(device) + tokenizer = transformers.AutoTokenizer.from_pretrained( + args.model, + revision=args.model_revision, + ) + model = jlens.from_hf(hf_model, tokenizer) + + lens_path = hf_hub_download( + args.lens_repo, + filename=args.lens_file, + revision=args.lens_revision, + ) + lens = jlens.JacobianLens.load(lens_path) + layers = [ + model.n_layers // 4, + model.n_layers // 2, + model.n_layers // 4 * 3, + model.n_layers - 2, + ] + lens_logits, model_logits, _ = lens.apply( + model, + args.prompt, + layers=layers, + positions=[-2], + ) + logit_logits, _, _ = lens.apply( + model, + args.prompt, + layers=layers, + positions=[-2], + use_jacobian=False, + ) + + layer_rows = [] + for layer in layers: + jacobian_tokens = _top_tokens( + lens_logits[layer][0], tokenizer, args.top_k + ) + logit_tokens = _top_tokens( + logit_logits[layer][0], tokenizer, args.top_k + ) + layer_rows.append( + { + "layer": layer, + "jacobian_lens": jacobian_tokens, + "logit_lens": logit_tokens, + "jacobian_contains_euro": _contains_euro(jacobian_tokens), + "logit_contains_euro": _contains_euro(logit_tokens), + } + ) + + report = { + "status": "completed", + "created_at": datetime.now(timezone.utc).isoformat(), + "scope": "Upstream pretrained-lens reproduction; not a J7Scope M1 result.", + "configuration": _configuration(args), + "hardware": hardware, + "software": { + "python": sys.version.split()[0], + "torch": torch.__version__, + "transformers": transformers.__version__, + }, + "provenance": { + "model_revision_resolved": getattr(hf_model.config, "_commit_hash", None), + "lens_revision_resolved": _snapshot_revision(lens_path), + "lens_n_prompts": lens.n_prompts, + "lens_d_model": lens.d_model, + }, + "layers": layer_rows, + "model_top_tokens": _top_tokens(model_logits[0], tokenizer, args.top_k), + "quality_gate": { + "expected_concept": "euro", + "jacobian_layers_matching": [ + row["layer"] for row in layer_rows + if row["jacobian_contains_euro"] + ], + "logit_layers_matching": [ + row["layer"] for row in layer_rows + if row["logit_contains_euro"] + ], + }, + } + write_json(args.output, report) + return report + + +def parse_args(argv=None): + parser = argparse.ArgumentParser() + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--model-revision", default="main") + parser.add_argument("--lens-repo", default=DEFAULT_LENS_REPO) + parser.add_argument("--lens-revision", default=DEFAULT_LENS_REVISION) + parser.add_argument("--lens-file", default=DEFAULT_LENS_FILE) + parser.add_argument("--prompt", default=DEFAULT_PROMPT) + parser.add_argument("--top-k", type=int, default=10) + parser.add_argument("--min-vram-gb", type=float, default=16.0) + parser.add_argument( + "--output", + type=Path, + default=ROOT / "results" / "upstream-reproduction.json", + ) + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--preflight-only", action="store_true") + parser.add_argument( + "--allow-cpu", + action="store_true", + help="Diagnostic only; the pretrained reproduction is intended for CUDA.", + ) + return parser.parse_args(argv) + + +def main(argv=None) -> None: + args = parse_args(argv) + if args.top_k < 1: + raise SystemExit("--top-k must be >= 1") + if args.min_vram_gb <= 0: + raise SystemExit("--min-vram-gb must be > 0") + report = ( + { + "status": "dry_run", + "configuration": _configuration(args), + } + if args.dry_run + else run(args) + ) + print(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/experiments/validate_jacobian.py b/experiments/validate_jacobian.py new file mode 100644 index 0000000..ec4b0a5 --- /dev/null +++ b/experiments/validate_jacobian.py @@ -0,0 +1,180 @@ +"""Validate J7Scope's position-local Jacobian math without model downloads. + +The script builds a randomly initialized, tiny Qwen2 decoder and compares: + +1. the tail-replay ``torch.func.jacrev`` result; +2. an independent full-model, one-output-dimension-at-a-time VJP; +3. the stochastic Gaussian-probe estimator used by the current capture path. + +This checks implementation consistency, not lens quality or the paper's +corpus/position reduction. No result from this synthetic model is research +evidence. + +Usage: + python experiments/validate_jacobian.py --n-probes 4096 +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import torch +from transformers import Qwen2Config, Qwen2ForCausalLM + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from j7scope.fitting import ( # noqa: E402 + _Capture, + _decoder_layers, + JLens, + exact_jacobian_for_prompt, + jacobian_error, + paper_jacobian_for_prompt, +) + + +class _Batch(dict): + def to(self, device): + return _Batch({key: value.to(device) for key, value in self.items()}) + + +class _FixedTokenizer: + def __call__(self, prompt, return_tensors): + if return_tensors != "pt": + raise ValueError("this validation tokenizer only returns PyTorch tensors") + return _Batch( + input_ids=torch.tensor([[1, 2, 3, 4]], dtype=torch.long), + attention_mask=torch.ones(1, 4, dtype=torch.long), + ) + + def decode(self, token_ids): + return str(token_ids[0]) + + +def _full_graph_vjp(model, tokenizer, prompt: str, layer: int) -> torch.Tensor: + """Independent exact Jacobian using one scalar VJP per output dimension.""" + source_capture, target_capture = _Capture(), _Capture() + layers = _decoder_layers(model) + handles = [ + layers[layer].register_forward_hook(source_capture), + layers[-1].register_forward_hook(target_capture), + ] + try: + inputs = tokenizer(prompt, return_tensors="pt").to(model.device) + with torch.enable_grad(): + model(**inputs, use_cache=False) + rows = [] + for output_index in range(model.config.hidden_size): + (gradient,) = torch.autograd.grad( + target_capture.value[0, -1, output_index].float(), + source_capture.value, + retain_graph=output_index < model.config.hidden_size - 1, + ) + rows.append(gradient[0, -1, :]) + return torch.stack(rows).detach().float().cpu() + finally: + for handle in handles: + handle.remove() + + +def _paper_scalar_vjp(model, tokenizer, prompt: str, layer: int, + skip_first: int) -> torch.Tensor: + """Slow scalar-VJP reference for the paper estimator's position reduction.""" + source_capture, target_capture = _Capture(), _Capture() + layers = _decoder_layers(model) + handles = [ + layers[layer].register_forward_hook(source_capture), + layers[-1].register_forward_hook(target_capture), + ] + try: + inputs = tokenizer(prompt, return_tensors="pt").to(model.device) + with torch.enable_grad(): + model(**inputs, use_cache=False) + valid_positions = torch.arange( + skip_first, + target_capture.value.shape[1] - 1, + device=model.device, + ) + rows = [] + for output_index in range(model.config.hidden_size): + objective = target_capture.value[ + 0, valid_positions, output_index + ].sum() + (gradient,) = torch.autograd.grad( + objective, + source_capture.value, + retain_graph=output_index < model.config.hidden_size - 1, + ) + rows.append(gradient[0, valid_positions, :].mean(dim=0)) + return torch.stack(rows).detach().float().cpu() + finally: + for handle in handles: + handle.remove() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--n-probes", type=int, default=4096) + parser.add_argument("--seed", type=int, default=0) + args = parser.parse_args() + + torch.manual_seed(args.seed) + config = Qwen2Config( + vocab_size=16, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=3, + num_attention_heads=2, + num_key_value_heads=1, + max_position_embeddings=32, + use_cache=False, + ) + model = Qwen2ForCausalLM(config).eval().requires_grad_(False) + tokenizer = _FixedTokenizer() + prompt = "synthetic fixed-token prompt" + layer = 0 + + exact = exact_jacobian_for_prompt( + model, tokenizer, prompt, layer, position=-1, chunk_size=1 + ) + full_graph = _full_graph_vjp(model, tokenizer, prompt, layer) + paper = paper_jacobian_for_prompt( + model, tokenizer, prompt, layer, skip_first=0, dim_batch=3 + ) + paper_scalar = _paper_scalar_vjp( + model, tokenizer, prompt, layer, skip_first=0 + ) + lens = JLens(model, tokenizer, layer) + stochastic = lens.estimate_jacobian( + [prompt], + n_probes=args.n_probes, + position=-1, + seed=args.seed, + show_progress=False, + ) + + report = { + "model": "random-tiny-qwen2", + "d_model": config.hidden_size, + "layer": layer, + "position": -1, + "n_probes": args.n_probes, + "seed": args.seed, + "exact_vs_full_graph": jacobian_error(exact, full_graph), + "paper_batched_vs_scalar_vjp": jacobian_error(paper, paper_scalar), + "stochastic_vs_exact": jacobian_error(stochastic, exact), + "scope": ( + "Implementation check only: synthetic position-local and paper-" + "reduction Jacobians, not a research result." + ), + } + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/experiments/validate_trace_gallery.py b/experiments/validate_trace_gallery.py new file mode 100644 index 0000000..b8d5abc --- /dev/null +++ b/experiments/validate_trace_gallery.py @@ -0,0 +1,66 @@ +"""Validate a Trace v1 gallery and frontend rigor-layer boundaries.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from j7scope.trace import validate_trace_gallery # noqa: E402 + + +FORBIDDEN_FRONTEND_FORMULAS = ( + "cross_lang_overlap -", + "same_lang_baseline -", + "(obs - null", + "(obs-null", +) + + +def validate_frontend_boundary(repo_root: Path) -> list[str]: + """Ensure executable frontend code displays, but does not derive, sharedness.""" + problems = [] + source_roots = (repo_root / "apps" / "site", repo_root / "apps" / "web" / "src") + for source_root in source_roots: + if not source_root.exists(): + continue + for path in source_root.rglob("*"): + if path.suffix not in {".js", ".ts", ".tsx"}: + continue + text = path.read_text(encoding="utf-8") + for formula in FORBIDDEN_FRONTEND_FORMULAS: + if formula in text: + problems.append( + f"{path.relative_to(repo_root)} recomputes rigor formula " + f"({formula!r}); frontend must display trace values only" + ) + return problems + + +def main(argv=None) -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "trace_root", + nargs="?", + type=Path, + default=ROOT / "results" / "traces", + ) + parser.add_argument("--repo-root", type=Path, default=ROOT) + args = parser.parse_args(argv) + + problems = validate_trace_gallery(args.trace_root) + problems.extend(validate_frontend_boundary(args.repo_root)) + if problems: + for problem in problems: + print(f"ERROR: {problem}", file=sys.stderr) + raise SystemExit(1) + print(f"Validated trace gallery: {args.trace_root}") + + +if __name__ == "__main__": + main() diff --git a/j7scope/fitting.py b/j7scope/fitting.py index 9819d59..fdd7783 100644 --- a/j7scope/fitting.py +++ b/j7scope/fitting.py @@ -14,22 +14,42 @@ Adapted from the approach of `jlens.fitting` in anthropics/jacobian-lens (Apache-2.0, see NOTICE), reimplemented against HuggingFace Qwen2-style dense decoders (model.model.layers / model.model.norm / lm_head). + +Estimator scope: J7Scope currently fits the position-local cloze variant +``∂h_L[p] / ∂h_l[p]``. The upstream paper implementation instead sums +cotangents over all valid current-and-future target positions and averages over +source positions. Both are Jacobian lenses, but their fitted matrices must not +be presented as numerically interchangeable. """ from __future__ import annotations +import hashlib +import json +import os +import tempfile +from pathlib import Path + import torch from tqdm import tqdm +def _hidden_states(output): + """Return a decoder layer's hidden-state tensor across HF API versions.""" + return output[0] if isinstance(output, tuple) else output + + def load_model(name: str = "Qwen/Qwen2.5-7B-Instruct", device: str | None = None, - dtype: torch.dtype = torch.bfloat16): + dtype: torch.dtype = torch.bfloat16, + revision: str | None = None): """Load a HF causal LM + tokenizer, frozen and in eval mode.""" from transformers import AutoModelForCausalLM, AutoTokenizer device = device or ("cuda" if torch.cuda.is_available() else "cpu") - tokenizer = AutoTokenizer.from_pretrained(name) - model = AutoModelForCausalLM.from_pretrained(name, torch_dtype=dtype).to(device) + tokenizer = AutoTokenizer.from_pretrained(name, revision=revision) + model = AutoModelForCausalLM.from_pretrained( + name, torch_dtype=dtype, revision=revision + ).to(device) # Params stay frozen: autograd graphs are built from the captured residual # leaf only (see _Capture), which keeps Jacobian passes cheap. model.eval().requires_grad_(False) @@ -57,7 +77,7 @@ def __call__(self, module, args, output): # both, and preserve the shape of what we return so the next layer gets # the same type it expected. is_tuple = isinstance(output, tuple) - h = output[0] if is_tuple else output + h = _hidden_states(output) if torch.is_grad_enabled() and not h.requires_grad: h = h.detach().requires_grad_(True) self.value = h @@ -73,12 +93,249 @@ def capture_residual(model, tokenizer, prompt: str, layer: int) -> torch.Tensor: try: inputs = tokenizer(prompt, return_tensors="pt").to(model.device) with torch.no_grad(): - model(**inputs) + model(**inputs, use_cache=False) finally: handle.remove() return cap.value +class _LayerInputCapture: + """Capture one decoder invocation's args/kwargs for tail replay.""" + + def __init__(self): + self.args = None + self.kwargs = None + + def __call__(self, module, args, kwargs): + self.args = args + self.kwargs = kwargs + + +def _replace_hidden_states(args, kwargs, hidden_states): + """Replace a decoder call's hidden-state argument without mutating inputs.""" + if args: + return (hidden_states, *args[1:]), kwargs + if "hidden_states" in kwargs: + replaced = dict(kwargs) + replaced["hidden_states"] = hidden_states + return args, replaced + raise RuntimeError("decoder layer call has no hidden_states argument") + + +def _resolve_position(position: int, seq_len: int) -> int: + resolved = position + seq_len if position < 0 else position + if not 0 <= resolved < seq_len: + raise IndexError( + f"position {position} is out of range for sequence length {seq_len}" + ) + return resolved + + +def exact_jacobian_for_prompt(model, tokenizer, prompt: str, layer: int, + position: int = -1, + chunk_size: int | None = 1) -> torch.Tensor: + """Compute the exact position-local Jacobian for one prompt with ``jacrev``. + + The returned matrix is ``∂h_L[position] / ∂h_layer[position]`` with shape + ``(d_out, d_in)``. A no-grad forward first captures the residual and the + constant arguments (causal mask, rotary embeddings, positions) supplied to + every downstream decoder layer. ``torch.func.jacrev`` then differentiates a + pure replay of that tail, replacing only the selected source-position + vector. This avoids differentiating the unused prefix of the network. + + ``chunk_size=1`` is deliberately conservative: exact Jacobians are a + validation tool and a full vmap over thousands of output dimensions can + exhaust memory. Raise it on small models when memory permits. + """ + layers = _decoder_layers(model) + if not 0 <= layer < len(layers) - 1: + raise ValueError( + f"layer must be in [0, {len(layers) - 2}] so the lens has a tail to map through" + ) + if chunk_size is not None and chunk_size < 1: + raise ValueError("chunk_size must be >= 1 or None") + + source_capture = _Capture() + tail_captures = [_LayerInputCapture() for _ in layers[layer + 1:]] + handles = [layers[layer].register_forward_hook(source_capture)] + handles.extend( + tail_layer.register_forward_pre_hook(capture, with_kwargs=True) + for tail_layer, capture in zip( + layers[layer + 1:], tail_captures, strict=True + ) + ) + try: + inputs = tokenizer(prompt, return_tensors="pt").to(model.device) + with torch.no_grad(): + model(**inputs, use_cache=False) + finally: + for handle in handles: + handle.remove() + + if source_capture.value is None or any(c.args is None for c in tail_captures): + raise RuntimeError("failed to capture decoder activations for exact Jacobian") + + source = source_capture.value.detach() + position = _resolve_position(position, source.shape[1]) + position_index = torch.tensor([position], device=source.device) + + def replay_tail(source_vector): + hidden = source.index_copy( + 1, position_index, source_vector.reshape(1, 1, -1) + ) + for tail_layer, capture in zip( + layers[layer + 1:], tail_captures, strict=True + ): + args, kwargs = _replace_hidden_states( + capture.args, capture.kwargs, hidden + ) + hidden = _hidden_states(tail_layer(*args, **kwargs)) + return hidden[0, position, :].float() + + source_vector = source[0, position, :] + jacobian = torch.func.jacrev( + replay_tail, chunk_size=chunk_size + )(source_vector) + return jacobian.detach().float().cpu() + + +def paper_jacobian_for_prompt(model, tokenizer, prompt: str, layer: int, + skip_first: int = 16, + dim_batch: int = 8) -> torch.Tensor: + """Compute the upstream paper's exact position-reduced estimator. + + For each output dimension, a cotangent is placed at every valid target + position. Causal attention makes the gradient at source position ``p`` the + sum of contributions from valid targets at or after ``p``; those source + gradients are then averaged. This matches the reduction documented by + ``anthropics/jacobian-lens`` and intentionally differs from J7Scope's + position-local cloze estimator. + + ``skip_first`` excludes attention-sink positions and the final position is + always excluded because it has no next-token target. ``dim_batch`` controls + how many output rows are evaluated together. As in the upstream fitter, + the prompt is replicated across that batch and differentiated with an + ordinary VJP; this avoids relying on vmap-compatible backward kernels. + """ + layers = _decoder_layers(model) + if not 0 <= layer < len(layers) - 1: + raise ValueError( + f"layer must be in [0, {len(layers) - 2}] so the lens has a tail to map through" + ) + if skip_first < 0: + raise ValueError("skip_first must be >= 0") + if dim_batch < 1: + raise ValueError("dim_batch must be >= 1") + + source_capture, target_capture = _Capture(), _Capture() + handles = [ + layers[layer].register_forward_hook(source_capture), + layers[-1].register_forward_hook(target_capture), + ] + try: + inputs = tokenizer(prompt, return_tensors="pt").to(model.device) + replicated_inputs = { + name: ( + value.expand(dim_batch, *value.shape[1:]) + if torch.is_tensor(value) + and value.ndim > 0 + and value.shape[0] == 1 + else value + ) + for name, value in inputs.items() + } + with torch.enable_grad(): + model(**replicated_inputs, use_cache=False) + source = source_capture.value + target = target_capture.value + _, seq_len, d_model = target.shape + valid_positions = torch.arange( + skip_first, seq_len - 1, device=target.device + ) + if valid_positions.numel() == 0: + raise ValueError( + f"prompt too short: seq_len={seq_len}, " + f"need > {skip_first + 1} tokens" + ) + + jacobian = torch.zeros(d_model, d_model, dtype=torch.float32) + for dim_start in range(0, d_model, dim_batch): + n_dims = min(dim_batch, d_model - dim_start) + batch_indices = torch.arange(n_dims, device=target.device) + cotangent = torch.zeros_like(target) + cotangent[ + batch_indices[:, None], + valid_positions[None, :], + dim_start + batch_indices[:, None], + ] = 1 + (gradient,) = torch.autograd.grad( + target, + source, + grad_outputs=cotangent, + retain_graph=dim_start + n_dims < d_model, + ) + source_positions = valid_positions.to(gradient.device) + rows = gradient[ + :n_dims, source_positions, : + ].float().mean(dim=1) + jacobian[dim_start:dim_start + n_dims] = rows.cpu() + finally: + for handle in handles: + handle.remove() + + return jacobian + + +_PAPER_CHECKPOINT_FORMAT = "j7scope.paper_jacobian_checkpoint.v1" + + +def _prompt_corpus_sha1(prompts: list[str]) -> str: + serialized = json.dumps( + prompts, ensure_ascii=False, separators=(",", ":") + ).encode("utf-8") + return hashlib.sha1(serialized).hexdigest() + + +def _atomic_torch_save(payload, path: str | Path) -> Path: + """Atomically replace a torch artifact, keeping partial writes invisible.""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + file_descriptor, temporary_name = tempfile.mkstemp( + prefix=f".{path.name}.", suffix=".tmp", dir=path.parent + ) + os.close(file_descriptor) + temporary_path = Path(temporary_name) + try: + torch.save(payload, temporary_path) + os.replace(temporary_path, path) + finally: + temporary_path.unlink(missing_ok=True) + return path + + +def jacobian_error(candidate: torch.Tensor, reference: torch.Tensor) -> dict: + """Return scale-aware diagnostics comparing two Jacobian matrices.""" + if candidate.shape != reference.shape: + raise ValueError( + f"Jacobian shapes differ: {tuple(candidate.shape)} != {tuple(reference.shape)}" + ) + candidate = candidate.detach().float().cpu() + reference = reference.detach().float().cpu() + delta = candidate - reference + reference_norm = reference.norm() + denominator = reference_norm.clamp_min(torch.finfo(torch.float32).eps) + cosine_denominator = ( + candidate.norm() * reference_norm + ).clamp_min(torch.finfo(torch.float32).eps) + return { + "relative_frobenius": float(delta.norm() / denominator), + "cosine_similarity": float( + (candidate.flatten() @ reference.flatten()) / cosine_denominator + ), + "max_absolute": float(delta.abs().max()), + } + + class JLens: """J-lens readout at a fixed layer. @@ -111,9 +368,11 @@ def estimate_jacobian(self, prompts, n_probes: int = 64, position: int = -1, The estimate is position-local (readout position onto itself), the variant the probe corpus is built for (cloze at the final token). - TODO: cross-check against exact per-position Jacobians - (torch.func.jacrev) on a subsample, and against upstream jlens.fitting. + Use :meth:`estimate_jacobian_exact` on a small prompt/model to measure + estimator error before committing an expensive capture run. """ + if n_probes < 1: + raise ValueError("n_probes must be >= 1") d = self.d_model gen = torch.Generator().manual_seed(seed) J = torch.zeros(d, d, dtype=torch.float32) @@ -128,7 +387,7 @@ def estimate_jacobian(self, prompts, n_probes: int = 64, position: int = -1, for prompt in iterator: inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device) with torch.enable_grad(): - self.model(**inputs) + self.model(**inputs, use_cache=False) h_l = cap_l.value out = cap_L.value[0, position, :].float() for _ in range(n_probes): @@ -141,10 +400,187 @@ def estimate_jacobian(self, prompts, n_probes: int = 64, position: int = -1, for h in handles: h.remove() + if n_terms == 0: + raise ValueError("prompts must contain at least one item") self.J = J / n_terms self._J_dev = None return self.J + def estimate_jacobian_exact(self, prompts, position: int = -1, + chunk_size: int | None = 1, + show_progress: bool = True) -> torch.Tensor: + """Exact mean position-local Jacobian over a small validation sample. + + This computes every Jacobian row and is therefore much more expensive + than :meth:`estimate_jacobian`. It is intended to validate estimator + direction and convergence on tiny models or a very small prompt sample, + not to fit a production 7B lens. + """ + jacobian_sum = torch.zeros( + self.d_model, self.d_model, dtype=torch.float32 + ) + n_prompts = 0 + iterator = ( + tqdm(prompts, desc=f"exact J @ layer {self.layer}") + if show_progress else prompts + ) + for prompt in iterator: + jacobian_sum += exact_jacobian_for_prompt( + self.model, + self.tokenizer, + prompt, + self.layer, + position=position, + chunk_size=chunk_size, + ) + n_prompts += 1 + if n_prompts == 0: + raise ValueError("prompts must contain at least one item") + self.J = jacobian_sum / n_prompts + self._J_dev = None + return self.J + + def estimate_jacobian_paper(self, prompts, skip_first: int = 16, + dim_batch: int = 8, + show_progress: bool = True, + checkpoint_path: str | Path | None = None, + checkpoint_every: int = 1, + resume: bool = True) -> torch.Tensor: + """Exact mean Jacobian using the upstream paper's position reduction. + + This is exposed separately so experiments cannot silently mix it with + the position-local capture estimator. It is exact but expensive; use a + pretraining-like corpus with sequences long enough to leave positions + after ``skip_first``. + + When ``checkpoint_path`` is set, the running float32 sum is atomically + saved every ``checkpoint_every`` visited prompts and once at completion. + Resuming validates the full ordered prompt corpus and all estimator + parameters before continuing, so a stale partial fit cannot be mixed + into a new run. Too-short prompts are recorded and skipped. + """ + if checkpoint_every < 1: + raise ValueError("checkpoint_every must be >= 1") + prompts = list(prompts) + if not prompts: + raise ValueError("prompts must contain at least one item") + if not all(isinstance(prompt, str) for prompt in prompts): + raise TypeError("every prompt must be a string") + + corpus_sha1 = _prompt_corpus_sha1(prompts) + checkpoint_path = ( + Path(checkpoint_path) if checkpoint_path is not None else None + ) + parameters = { + "layer": self.layer, + "d_model": self.d_model, + "skip_first": skip_first, + "dim_batch": dim_batch, + } + model_identity = { + "name": getattr(self.model.config, "_name_or_path", None), + "revision": getattr(self.model.config, "_commit_hash", None), + } + state = { + "format": _PAPER_CHECKPOINT_FORMAT, + "completed": False, + "parameters": parameters, + "model": model_identity, + "corpus_sha1": corpus_sha1, + "prompts_total": len(prompts), + "next_index": 0, + "n_used": 0, + "skipped": [], + "jacobian_sum": torch.zeros( + self.d_model, self.d_model, dtype=torch.float32 + ), + } + + if checkpoint_path is not None and resume and checkpoint_path.exists(): + loaded = torch.load( + checkpoint_path, map_location="cpu", weights_only=True + ) + if not isinstance(loaded, dict): + raise ValueError("invalid paper Jacobian checkpoint payload") + expected = { + "format": _PAPER_CHECKPOINT_FORMAT, + "parameters": parameters, + "model": model_identity, + "corpus_sha1": corpus_sha1, + "prompts_total": len(prompts), + } + mismatches = [ + key for key, value in expected.items() + if loaded.get(key) != value + ] + if mismatches: + raise ValueError( + "paper Jacobian checkpoint does not match this run: " + + ", ".join(mismatches) + ) + jacobian_sum = loaded.get("jacobian_sum") + if ( + not torch.is_tensor(jacobian_sum) + or tuple(jacobian_sum.shape) != (self.d_model, self.d_model) + ): + raise ValueError( + "paper Jacobian checkpoint has an invalid jacobian_sum" + ) + state.update(loaded) + state["jacobian_sum"] = jacobian_sum.float().cpu() + + start_index = int(state["next_index"]) + if not 0 <= start_index <= len(prompts): + raise ValueError("paper Jacobian checkpoint has an invalid next_index") + iterator = range(start_index, len(prompts)) + if show_progress: + iterator = tqdm( + iterator, + total=len(prompts), + initial=start_index, + desc=f"paper J @ layer {self.layer}", + ) + + for prompt_index in iterator: + prompt = prompts[prompt_index] + try: + prompt_jacobian = paper_jacobian_for_prompt( + self.model, + self.tokenizer, + prompt, + self.layer, + skip_first=skip_first, + dim_batch=dim_batch, + ) + except ValueError as exc: + if not str(exc).startswith("prompt too short:"): + raise + state["skipped"].append( + {"index": prompt_index, "reason": str(exc)} + ) + else: + state["jacobian_sum"] += prompt_jacobian + state["n_used"] += 1 + state["next_index"] = prompt_index + 1 + if ( + checkpoint_path is not None + and state["next_index"] % checkpoint_every == 0 + ): + _atomic_torch_save(state, checkpoint_path) + + if state["n_used"] == 0: + if checkpoint_path is not None: + _atomic_torch_save(state, checkpoint_path) + raise ValueError( + "no prompts were long enough for the paper Jacobian estimator" + ) + state["completed"] = True + if checkpoint_path is not None: + _atomic_torch_save(state, checkpoint_path) + self.J = state["jacobian_sum"] / state["n_used"] + self._J_dev = None + return self.J + # ---- readout ----------------------------------------------------------- def collect_residual(self, prompt: str, position: int = -1) -> torch.Tensor: diff --git a/j7scope/trace.py b/j7scope/trace.py index 8bbdfc9..b47e463 100644 --- a/j7scope/trace.py +++ b/j7scope/trace.py @@ -16,6 +16,7 @@ from __future__ import annotations +import json from pathlib import Path from typing import Iterable, List, Mapping, Optional, Sequence, Union @@ -68,7 +69,6 @@ def write_trace( def read_trace(trace_dir: PathLike) -> dict: """Read a Trace v1 into memory. ``align`` is None when absent.""" - import json paths = trace_paths(trace_dir) manifest = json.loads(Path(paths["manifest"]).read_text(encoding="utf-8")) tokens = read_jsonl(paths["tokens"]) @@ -123,3 +123,157 @@ def validate_trace(trace: Mapping, *, require_rigor: bool = True) -> List[str]: def validate_trace_dir(trace_dir: PathLike, *, require_rigor: bool = True) -> List[str]: return validate_trace(read_trace(trace_dir), require_rigor=require_rigor) + + +def trace_index_row(trace_dir: PathLike) -> dict: + """Build one gallery-index row from a trace directory.""" + trace = read_trace(trace_dir) + manifest, tokens = trace["manifest"], trace["tokens"] + return { + "trace_id": manifest["trace_id"], + "label": manifest.get("label", manifest["trace_id"]), + "model": manifest.get("model"), + "layer": manifest.get("layer"), + "language": manifest.get("language", "?"), + "concept": manifest.get("concept"), + "is_demo": bool(manifest.get("is_demo", False)), + "preview": bool(manifest.get("preview", False)), + "doi": manifest.get("doi"), + "n_tokens": len(tokens), + "parallel_group": manifest.get("parallel_group"), + } + + +def rebuild_trace_index(trace_root: PathLike) -> Path: + """Rebuild ``index.json`` deterministically from all child traces.""" + trace_root = Path(trace_root) + rows = [ + trace_index_row(trace_dir) + for trace_dir in sorted(trace_root.iterdir()) + if trace_dir.is_dir() and (trace_dir / TRACE_FILES["manifest"]).exists() + ] + return write_json( + trace_root / "index.json", + {"schema_version": TRACE_SCHEMA_VERSION, "traces": rows}, + ) + + +def validate_trace_gallery(trace_root: PathLike) -> List[str]: + """Validate gallery indexing, trace schema, rigor, and provenance.""" + from .rigor import SHAREDNESS_DEFINITION + + trace_root = Path(trace_root) + index_path = trace_root / "index.json" + if not index_path.exists(): + return [f"gallery missing index: {index_path}"] + try: + index = json.loads(index_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + return [f"gallery index is unreadable: {exc}"] + + problems: List[str] = [] + if index.get("schema_version") != TRACE_SCHEMA_VERSION: + problems.append( + f"gallery index schema_version must be {TRACE_SCHEMA_VERSION}" + ) + rows = index.get("traces") + if not isinstance(rows, list): + return problems + ["gallery index.traces must be a list"] + + indexed_ids = [ + row.get("trace_id") for row in rows + if isinstance(row, Mapping) and isinstance(row.get("trace_id"), str) + ] + if len(indexed_ids) != len(rows): + problems.append("every gallery index row must be an object with trace_id") + duplicates = sorted( + trace_id for trace_id in set(indexed_ids) + if indexed_ids.count(trace_id) > 1 + ) + if duplicates: + problems.append(f"gallery index has duplicate trace_ids: {duplicates}") + + discovered_ids = sorted( + path.name for path in trace_root.iterdir() + if path.is_dir() and (path / TRACE_FILES["manifest"]).exists() + ) + missing = sorted(set(discovered_ids) - set(indexed_ids)) + extra = sorted(set(indexed_ids) - set(discovered_ids)) + if missing: + problems.append(f"traces missing from gallery index: {missing}") + if extra: + problems.append(f"gallery index references missing traces: {extra}") + + rows_by_id = { + row["trace_id"]: row for row in rows + if isinstance(row, Mapping) and row.get("trace_id") in discovered_ids + } + for trace_id in discovered_ids: + trace_dir = trace_root / trace_id + try: + trace = read_trace(trace_dir) + except (OSError, json.JSONDecodeError) as exc: + problems.append(f"{trace_id}: unreadable trace: {exc}") + continue + problems.extend( + f"{trace_id}: {problem}" for problem in validate_trace(trace) + ) + manifest = trace["manifest"] + tokens = trace["tokens"] + metrics = trace["metrics"] + if manifest.get("trace_id") != trace_id: + problems.append( + f"{trace_id}: manifest.trace_id does not match directory name" + ) + if metrics.get("trace_id") != trace_id: + problems.append(f"{trace_id}: metrics.trace_id does not match") + if metrics.get("n_tokens") != len(tokens): + problems.append(f"{trace_id}: metrics.n_tokens does not match tokens.jsonl") + seqs = [token.get("seq") for token in tokens] + if seqs != list(range(len(tokens))): + problems.append(f"{trace_id}: token seq values must be contiguous from 0") + + row = rows_by_id.get(trace_id, {}) + expected_row = trace_index_row(trace_dir) + for key in ("language", "is_demo", "n_tokens", "parallel_group"): + if row.get(key) != expected_row[key]: + problems.append(f"{trace_id}: index.{key} does not match trace") + + for token_index, token in enumerate(tokens): + sharedness = token.get("rigor", {}).get("sharedness", {}) + if sharedness.get("definition") != SHAREDNESS_DEFINITION: + problems.append( + f"{trace_id}: token[{token_index}] sharedness definition drifted" + ) + + if not manifest.get("is_demo", False): + jacobian = manifest.get("jacobian", {}) + capture = manifest.get("capture", {}) + for key in ("revision",): + if not manifest.get(key): + problems.append(f"{trace_id}: real trace missing manifest.{key}") + for key in ("estimator", "sha1"): + if not jacobian.get(key): + problems.append(f"{trace_id}: real trace missing jacobian.{key}") + for key in ("device", "dtype"): + if not capture.get(key): + problems.append(f"{trace_id}: real trace missing capture.{key}") + + align = trace.get("align") + if align is not None: + members = set(align.get("members", {}).values()) + unknown_members = sorted(members - set(discovered_ids)) + if unknown_members: + problems.append( + f"{trace_id}: align references missing members {unknown_members}" + ) + for pair_index, pair in enumerate(align.get("position_map", [])): + if ( + not isinstance(pair, list) + or len(pair) != 2 + or not all(isinstance(value, int) and value >= 0 for value in pair) + ): + problems.append( + f"{trace_id}: align.position_map[{pair_index}] is invalid" + ) + return problems diff --git a/notebooks/capture_colab.ipynb b/notebooks/capture_colab.ipynb new file mode 100644 index 0000000..e583795 --- /dev/null +++ b/notebooks/capture_colab.ipynb @@ -0,0 +1,182 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "intro", + "metadata": {}, + "source": [ + "# J7Scope community trace capture (Colab GPU)\n", + "\n", + "Capture one custom Trace Schema v1 artifact with a real open-weight model, then download a validated ZIP for a gallery pull request.\n", + "\n", + "> **Preview boundary:** the default Qwen2.5-1.5B model and small stochastic Jacobian budget prove the capture pipeline; they do not establish a research result. The output is marked `preview: true`. Runtime and GPU type depend on Colab availability." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "clone-repository", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import os\n", + "import subprocess\n", + "import sys\n", + "\n", + "REPO = Path('/content/j7scope')\n", + "if not REPO.exists():\n", + " subprocess.run([\n", + " 'git', 'clone', '--depth', '1',\n", + " 'https://github.com/arthurpanhku/j7scope.git', str(REPO)\n", + " ], check=True)\n", + "os.chdir(REPO)\n", + "print(f'Working in {REPO}')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "install-package", + "metadata": {}, + "outputs": [], + "source": [ + "%pip install -q -e ." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "gpu-preflight", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "\n", + "if not torch.cuda.is_available():\n", + " raise RuntimeError('No CUDA GPU. Choose Runtime > Change runtime type > GPU, then reconnect.')\n", + "props = torch.cuda.get_device_properties(0)\n", + "print({\n", + " 'gpu': props.name,\n", + " 'vram_gib': round(props.total_memory / 1024**3, 2),\n", + " 'bf16_supported': torch.cuda.is_bf16_supported(),\n", + " 'selected_dtype': 'bfloat16' if torch.cuda.is_bf16_supported() else 'float16',\n", + "})" + ] + }, + { + "cell_type": "markdown", + "id": "configure-heading", + "metadata": {}, + "source": [ + "## Configure one trace\n", + "\n", + "Use a unique lowercase `trace_id`. Do not put personal, confidential, copyrighted, or credential-bearing text in the prompt." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "capture-parameters", + "metadata": {}, + "outputs": [], + "source": [ + "TRACE_ID = 'community-deception-en' # @param {type:'string'}\n", + "LANGUAGE = 'en' # @param ['en', 'zh', 'other']\n", + "CONCEPT = 'deception' # @param {type:'string'}\n", + "PROMPT = 'In one sentence, explain why deception can be tempting.' # @param {type:'string'}\n", + "MODEL = 'Qwen/Qwen2.5-1.5B-Instruct' # @param {type:'string'}\n", + "MODEL_REVISION = 'main' # @param {type:'string'}\n", + "LAYER = 14 # @param {type:'integer'}\n", + "N_PROBES = 8 # @param {type:'integer'}\n", + "MAX_NEW_TOKENS = 48 # @param {type:'integer'}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "run-capture", + "metadata": {}, + "outputs": [], + "source": [ + "OUTPUT = Path('/content/j7scope-output/traces')\n", + "CACHE = Path('/content/j7scope-cache/jacobian')\n", + "command = [\n", + " sys.executable, 'experiments/capture_trace.py',\n", + " '--trace-id', TRACE_ID,\n", + " '--language', LANGUAGE,\n", + " '--concept', CONCEPT,\n", + " '--prompt', PROMPT,\n", + " '--model', MODEL,\n", + " '--model-revision', MODEL_REVISION,\n", + " '--layer', str(LAYER),\n", + " '--n-probes', str(N_PROBES),\n", + " '--max-new-tokens', str(MAX_NEW_TOKENS),\n", + " '--device', 'cuda',\n", + " '--dtype', 'auto',\n", + " '--out', str(OUTPUT),\n", + " '--cache-dir', str(CACHE),\n", + "]\n", + "print('Starting capture. The first run downloads model weights and fits J_l...')\n", + "subprocess.run(command, check=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "validate-archive", + "metadata": {}, + "outputs": [], + "source": [ + "subprocess.run([\n", + " sys.executable, 'experiments/validate_trace_gallery.py', str(OUTPUT)\n", + "], check=True)\n", + "\n", + "import shutil\n", + "archive = shutil.make_archive(\n", + " f'/content/{TRACE_ID}', 'zip', root_dir=OUTPUT.parent, base_dir=OUTPUT.name\n", + ")\n", + "print(f'Validated archive: {archive}')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "download-archive", + "metadata": {}, + "outputs": [], + "source": [ + "from google.colab import files\n", + "files.download(archive)" + ] + }, + { + "cell_type": "markdown", + "id": "submit-guidance", + "metadata": {}, + "source": [ + "## Submit\n", + "\n", + "Follow [`CONTRIBUTING.md`](https://github.com/arthurpanhku/j7scope/blob/main/CONTRIBUTING.md): copy the trace directory into `results/traces/`, rebuild `index.json`, run validation, and open a pull request. Keep `preview: true` unless a maintainer has completed research review." + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "name": "J7Scope capture_colab.ipynb", + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml index a01e12a..d2b86f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,10 @@ dependencies = [ [project.optional-dependencies] dev = ["pytest", "jupyter", "matplotlib"] +fit = ["datasets>=3.0"] +upstream = [ + "jlens @ git+https://github.com/anthropics/jacobian-lens.git@581d398613e5602a5af361e1c34d3a92ea82ba8e", +] [tool.setuptools] packages = ["j7scope"] diff --git a/tests/test_fit_paper_jacobian_cli.py b/tests/test_fit_paper_jacobian_cli.py new file mode 100644 index 0000000..66c640c --- /dev/null +++ b/tests/test_fit_paper_jacobian_cli.py @@ -0,0 +1,45 @@ +import json +import sys +from pathlib import Path + +import pytest + + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from experiments.fit_paper_jacobian import ( + DEFAULT_CORPUS, + _configuration, + _load_local_prompts, + parse_args, +) + + +def test_default_smoke_corpus_is_valid(): + prompts = _load_local_prompts( + DEFAULT_CORPUS, text_key="text", max_prompts=1000, min_chars=0 + ) + + assert len(prompts) == 12 + assert all(len(prompt) > 40 for prompt in prompts) + + +def test_local_corpus_reports_invalid_jsonl(tmp_path): + corpus = tmp_path / "broken.jsonl" + corpus.write_text('{"text": "ok"}\nnot-json\n', encoding="utf-8") + + with pytest.raises(ValueError, match="invalid JSON"): + _load_local_prompts( + corpus, text_key="text", max_prompts=10, min_chars=0 + ) + + +def test_dry_run_configuration_marks_default_as_smoke(): + args = parse_args(["--dry-run"]) + configuration = _configuration(args) + + assert configuration["corpus"]["smoke_only"] is True + assert configuration["resume"] is True + json.dumps(configuration) diff --git a/tests/test_fitting.py b/tests/test_fitting.py new file mode 100644 index 0000000..e424189 --- /dev/null +++ b/tests/test_fitting.py @@ -0,0 +1,354 @@ +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +import j7scope.fitting as fitting +from j7scope.fitting import ( + _Capture, + _decoder_layers, + JLens, + exact_jacobian_for_prompt, + jacobian_error, + paper_jacobian_for_prompt, +) + + +class _Batch(dict): + def to(self, device): + return _Batch({key: value.to(device) for key, value in self.items()}) + + +class _Tokenizer: + def __call__(self, prompt, return_tensors): + assert return_tensors == "pt" + token_ids = [int(token) for token in prompt.split()] + return _Batch(input_ids=torch.tensor([token_ids], dtype=torch.long)) + + def decode(self, token_ids): + return str(token_ids[0]) + + +class _LinearDecoderLayer(nn.Module): + def __init__(self, matrix, *, returns_tuple=False): + super().__init__() + self.register_buffer("matrix", torch.tensor(matrix, dtype=torch.float32)) + self.returns_tuple = returns_tuple + + def forward(self, hidden_states, *, gain): + output = gain * (hidden_states @ self.matrix.T) + return (output,) if self.returns_tuple else output + + +class _Core(nn.Module): + def __init__(self, matrices): + super().__init__() + self.embed_tokens = nn.Embedding(8, len(matrices[0])) + self.layers = nn.ModuleList( + [ + _LinearDecoderLayer(matrix, returns_tuple=(index == 1)) + for index, matrix in enumerate(matrices) + ] + ) + self.norm = nn.Identity() + + +class _TinyCausalLM(nn.Module): + def __init__(self, matrices): + super().__init__() + width = len(matrices[0]) + self.config = SimpleNamespace(hidden_size=width) + self.model = _Core(matrices) + self.lm_head = nn.Linear(width, 8, bias=False) + self.gain = 1.0 + self.last_batch_size = None + self.eval().requires_grad_(False) + + @property + def device(self): + return self.model.embed_tokens.weight.device + + @property + def dtype(self): + return self.model.embed_tokens.weight.dtype + + def forward(self, input_ids, *, use_cache): + assert use_cache is False + self.last_batch_size = input_ids.shape[0] + hidden = self.model.embed_tokens(input_ids) + for layer in self.model.layers: + output = layer(hidden, gain=self.gain) + hidden = output[0] if isinstance(output, tuple) else output + return SimpleNamespace(logits=self.lm_head(hidden)) + + +@pytest.fixture +def linear_lens(): + matrices = [ + [[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, -1.0]], + [[1.0, 2.0, 0.0], [0.0, -1.0, 1.0], [0.5, 0.0, 1.0]], + [[2.0, 0.0, 1.0], [-1.0, 1.0, 0.0], [0.0, 0.5, 1.0]], + ] + model = _TinyCausalLM(matrices) + lens = JLens(model, _Tokenizer(), layer=0) + expected = torch.tensor(matrices[2]) @ torch.tensor(matrices[1]) + return lens, expected + + +def test_exact_jacobian_matches_analytic_tail(linear_lens): + lens, expected = linear_lens + + actual = exact_jacobian_for_prompt( + lens.model, lens.tokenizer, "1 2 3", layer=0, position=-1 + ) + + torch.testing.assert_close(actual, expected) + assert jacobian_error(actual, expected) == { + "relative_frobenius": 0.0, + "cosine_similarity": pytest.approx(1.0), + "max_absolute": 0.0, + } + + +def test_exact_estimator_sets_mean_jacobian(linear_lens): + lens, expected = linear_lens + + actual = lens.estimate_jacobian_exact( + ["1 2", "3 4 5"], show_progress=False + ) + + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(lens.J, expected) + + +def test_stochastic_estimator_converges_to_exact(linear_lens): + lens, expected = linear_lens + + estimate = lens.estimate_jacobian( + ["1 2 3"], n_probes=5000, seed=7, show_progress=False + ) + diagnostics = jacobian_error(estimate, expected) + + assert diagnostics["relative_frobenius"] < 0.06 + assert diagnostics["cosine_similarity"] > 0.998 + + +def test_exact_jacobian_matches_full_graph_vjp_on_tiny_qwen(): + from transformers import Qwen2Config, Qwen2ForCausalLM + + config = Qwen2Config( + vocab_size=16, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=3, + num_attention_heads=2, + num_key_value_heads=1, + max_position_embeddings=32, + use_cache=False, + ) + model = Qwen2ForCausalLM(config).eval().requires_grad_(False) + tokenizer = _Tokenizer() + prompt = "1 2 3 4" + + exact = exact_jacobian_for_prompt( + model, tokenizer, prompt, layer=0, position=-1, chunk_size=1 + ) + + source_capture, target_capture = _Capture(), _Capture() + layers = _decoder_layers(model) + handles = [ + layers[0].register_forward_hook(source_capture), + layers[-1].register_forward_hook(target_capture), + ] + try: + inputs = tokenizer(prompt, return_tensors="pt").to(model.device) + with torch.enable_grad(): + model(**inputs, use_cache=False) + rows = [] + for output_index in range(config.hidden_size): + (gradient,) = torch.autograd.grad( + target_capture.value[0, -1, output_index].float(), + source_capture.value, + retain_graph=output_index < config.hidden_size - 1, + ) + rows.append(gradient[0, -1, :]) + full_graph_vjp = torch.stack(rows).float().cpu() + finally: + for handle in handles: + handle.remove() + + torch.testing.assert_close(exact, full_graph_vjp) + + +def test_paper_estimator_matches_scalar_vjps_on_tiny_qwen(): + from transformers import Qwen2Config, Qwen2ForCausalLM + + config = Qwen2Config( + vocab_size=16, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=3, + num_attention_heads=2, + num_key_value_heads=1, + max_position_embeddings=32, + use_cache=False, + ) + model = Qwen2ForCausalLM(config).eval().requires_grad_(False) + tokenizer = _Tokenizer() + prompt = "1 2 3 4" + valid_positions = torch.tensor([0, 1, 2]) + + paper = paper_jacobian_for_prompt( + model, tokenizer, prompt, layer=0, skip_first=0, dim_batch=3 + ) + + source_capture, target_capture = _Capture(), _Capture() + layers = _decoder_layers(model) + handles = [ + layers[0].register_forward_hook(source_capture), + layers[-1].register_forward_hook(target_capture), + ] + try: + inputs = tokenizer(prompt, return_tensors="pt").to(model.device) + with torch.enable_grad(): + model(**inputs, use_cache=False) + rows = [] + for output_index in range(config.hidden_size): + objective = target_capture.value[ + 0, valid_positions, output_index + ].sum() + (gradient,) = torch.autograd.grad( + objective, + source_capture.value, + retain_graph=output_index < config.hidden_size - 1, + ) + rows.append(gradient[0, valid_positions, :].mean(dim=0)) + scalar_vjps = torch.stack(rows).float().cpu() + finally: + for handle in handles: + handle.remove() + + torch.testing.assert_close(paper, scalar_vjps) + + +def test_paper_estimator_replicates_prompt_batch(linear_lens): + lens, expected = linear_lens + + actual = paper_jacobian_for_prompt( + lens.model, + lens.tokenizer, + "1 2 3", + layer=0, + skip_first=0, + dim_batch=2, + ) + + assert lens.model.last_batch_size == 2 + torch.testing.assert_close(actual, expected) + + +def test_paper_estimator_resumes_atomic_checkpoint( + linear_lens, monkeypatch, tmp_path +): + lens, expected = linear_lens + checkpoint = tmp_path / "paper-fit.checkpoint.pt" + real_estimator = fitting.paper_jacobian_for_prompt + calls = [] + + def fail_second_prompt(*args, **kwargs): + calls.append(args[2]) + if len(calls) == 2: + raise RuntimeError("simulated interruption") + return real_estimator(*args, **kwargs) + + monkeypatch.setattr( + fitting, "paper_jacobian_for_prompt", fail_second_prompt + ) + prompts = ["1 2 3", "3 4 5"] + with pytest.raises(RuntimeError, match="simulated interruption"): + lens.estimate_jacobian_paper( + prompts, + skip_first=0, + dim_batch=2, + show_progress=False, + checkpoint_path=checkpoint, + ) + + partial = torch.load(checkpoint, map_location="cpu", weights_only=True) + assert partial["completed"] is False + assert partial["next_index"] == 1 + assert partial["n_used"] == 1 + + actual = lens.estimate_jacobian_paper( + prompts, + skip_first=0, + dim_batch=2, + show_progress=False, + checkpoint_path=checkpoint, + ) + completed = torch.load(checkpoint, map_location="cpu", weights_only=True) + + assert calls == ["1 2 3", "3 4 5", "3 4 5"] + assert completed["completed"] is True + assert completed["next_index"] == 2 + assert completed["n_used"] == 2 + torch.testing.assert_close(actual, expected) + + +def test_paper_checkpoint_rejects_different_corpus(linear_lens, tmp_path): + lens, _ = linear_lens + checkpoint = tmp_path / "paper-fit.checkpoint.pt" + lens.estimate_jacobian_paper( + ["1 2 3"], + skip_first=0, + dim_batch=2, + show_progress=False, + checkpoint_path=checkpoint, + ) + + with pytest.raises(ValueError, match="corpus_sha1"): + lens.estimate_jacobian_paper( + ["1 2 4"], + skip_first=0, + dim_batch=2, + show_progress=False, + checkpoint_path=checkpoint, + ) + + +def test_exact_jacobian_rejects_invalid_position_and_chunk(linear_lens): + lens, _ = linear_lens + + with pytest.raises(IndexError, match="out of range"): + exact_jacobian_for_prompt( + lens.model, lens.tokenizer, "1 2", layer=0, position=2 + ) + with pytest.raises(ValueError, match="chunk_size"): + exact_jacobian_for_prompt( + lens.model, lens.tokenizer, "1 2", layer=0, chunk_size=0 + ) + + +def test_paper_jacobian_rejects_too_short_prompt(linear_lens): + lens, _ = linear_lens + + with pytest.raises(ValueError, match="prompt too short"): + paper_jacobian_for_prompt( + lens.model, lens.tokenizer, "1 2", layer=0, skip_first=1 + ) + + +@pytest.mark.parametrize( + "method", + [ + "estimate_jacobian", + "estimate_jacobian_exact", + "estimate_jacobian_paper", + ], +) +def test_estimators_reject_empty_prompts(linear_lens, method): + lens, _ = linear_lens + + with pytest.raises(ValueError, match="at least one"): + getattr(lens, method)([], show_progress=False) diff --git a/tests/test_recorder.py b/tests/test_recorder.py new file mode 100644 index 0000000..465cd0c --- /dev/null +++ b/tests/test_recorder.py @@ -0,0 +1,117 @@ +import json +import sys +from pathlib import Path + +import pytest +import torch + + +ROOT = Path(__file__).resolve().parents[1] +SERVE_APP = ROOT / "apps" / "serve" +if str(SERVE_APP) not in sys.path: + sys.path.insert(0, str(SERVE_APP)) + +from j7scope_serve.backends import HFBackend, MockBackend # noqa: E402 +from j7scope_serve.protocol import bucket_readout, script_of # noqa: E402 +from j7scope_serve.recorder import build_lexicon, record_trace # noqa: E402 + + +def test_recorded_trace_includes_jacobian_estimator_provenance(tmp_path): + backend = MockBackend() + buffered = [] + for seq, step in enumerate(backend.generate([])): + buffered.append( + { + "seq": seq, + "ts_rel": seq / 10, + "token": step.token, + "token_script": script_of(step.token), + "readout": bucket_readout(step.topk, per_lang=8), + } + ) + if len(buffered) == 3: + break + + trace_dir = record_trace( + tmp_path, + backend=backend, + prompt="test", + buffered=buffered, + lexicon=build_lexicon(ROOT, backend), + trace_id="provenance-test", + ) + manifest = json.loads( + (trace_dir / "manifest.json").read_text(encoding="utf-8") + ) + + assert manifest["jacobian"] == { + "corpus_id": "mock-synthetic", + "estimator": "synthetic", + "n_probes": None, + "n_prompts": 0, + "position": None, + "seed": None, + "sha1": None, + } + + +def test_hf_backend_auto_dtype_handles_t4_and_cpu(): + class _Cuda: + @staticmethod + def is_available(): + return True + + @staticmethod + def is_bf16_supported(): + return False + + class _Torch: + cuda = _Cuda() + bfloat16 = "bf16" + float16 = "fp16" + float32 = "fp32" + + backend = HFBackend(device="cuda", dtype="auto") + assert backend._resolve_dtype(_Torch) == "fp16" + + backend.device = "cpu" + assert backend._resolve_dtype(_Torch) == "fp32" + + +def test_hf_backend_loads_precomputed_paper_jacobian(tmp_path): + path = tmp_path / "jacobian.pt" + expected = torch.eye(3) + torch.save(expected, path) + backend = HFBackend(jacobian_path=str(path)) + metadata = { + "configuration": { + "model": backend.model_name, + "layer": backend.layer, + }, + "provenance": { + "model_revision_resolved": "main", + "corpus_sha1": "corpus-hash", + "prompts_used": 1000, + "estimator": "paper_replicated_batch_vjp", + }, + "result": {"tensor_sha1": backend._tensor_sha1(expected)}, + } + path.with_suffix(".json").write_text( + json.dumps(metadata), encoding="utf-8" + ) + + actual = backend._load_precomputed_jacobian() + + assert backend.jacobian_estimator == "paper_replicated_batch_vjp" + assert backend.jacobian_corpus_id == "corpus-hash" + assert backend.jacobian_n_prompts == 1000 + torch.testing.assert_close(actual, expected) + + +def test_hf_backend_rejects_precomputed_jacobian_without_metadata(tmp_path): + path = tmp_path / "jacobian.pt" + torch.save(torch.eye(3), path) + backend = HFBackend(jacobian_path=str(path)) + + with pytest.raises(FileNotFoundError, match="metadata"): + backend._load_precomputed_jacobian() diff --git a/tests/test_reproduce_upstream.py b/tests/test_reproduce_upstream.py new file mode 100644 index 0000000..0f36a99 --- /dev/null +++ b/tests/test_reproduce_upstream.py @@ -0,0 +1,28 @@ +import json +import subprocess +import sys +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] +SCRIPT = ROOT / "experiments" / "reproduce_upstream.py" + + +def test_upstream_reproduction_dry_run_needs_no_gpu_or_downloads(): + completed = subprocess.run( + [sys.executable, str(SCRIPT), "--dry-run"], + check=True, + capture_output=True, + text=True, + cwd=ROOT, + ) + report = json.loads(completed.stdout) + + assert report["status"] == "dry_run" + config = report["configuration"] + assert config["model"] == "Qwen/Qwen3.5-4B" + assert config["lens_repo"] == "neuronpedia/jacobian-lens" + assert config["lens_revision_requested"] == "qwen-n1000" + assert config["position"] == -2 + assert config["requires_cuda"] is True + assert len(config["upstream_commit"]) == 40 diff --git a/tests/test_trace_gallery.py b/tests/test_trace_gallery.py new file mode 100644 index 0000000..97d1e28 --- /dev/null +++ b/tests/test_trace_gallery.py @@ -0,0 +1,76 @@ +import json +import subprocess +import sys +from pathlib import Path + +from j7scope.trace import validate_trace_gallery + + +ROOT = Path(__file__).resolve().parents[1] + + +def test_demo_builder_produces_valid_gallery(tmp_path): + subprocess.run( + [ + sys.executable, + str(ROOT / "experiments" / "build_demo_trace.py"), + "--out", + str(tmp_path), + ], + check=True, + capture_output=True, + text=True, + cwd=ROOT, + ) + + assert validate_trace_gallery(tmp_path) == [] + + +def test_gallery_validator_detects_index_drift(tmp_path): + subprocess.run( + [ + sys.executable, + str(ROOT / "experiments" / "build_demo_trace.py"), + "--out", + str(tmp_path), + ], + check=True, + capture_output=True, + text=True, + cwd=ROOT, + ) + index_path = tmp_path / "index.json" + index = json.loads(index_path.read_text(encoding="utf-8")) + index["traces"][0]["n_tokens"] += 1 + index_path.write_text(json.dumps(index), encoding="utf-8") + + problems = validate_trace_gallery(tmp_path) + + assert any("index.n_tokens does not match" in problem for problem in problems) + + +def test_capture_cli_dry_run_needs_no_gpu(): + completed = subprocess.run( + [ + sys.executable, + str(ROOT / "experiments" / "capture_trace.py"), + "--trace-id", + "community-deception-en", + "--language", + "en", + "--concept", + "deception", + "--prompt", + "Explain deception briefly.", + "--dry-run", + ], + check=True, + capture_output=True, + text=True, + cwd=ROOT, + ) + config = json.loads(completed.stdout) + + assert config["model"] == "Qwen/Qwen2.5-1.5B-Instruct" + assert config["dtype"] == "auto" + assert config["trace_id"] == "community-deception-en"