Document cost-aware atlas workflow

README.md

@@ -183,91 +183,113 @@ python app.py \
 ```
 
 
-### Consolidated full-feature command
+### Cost-aware full atlas workflow
 
-The recommended single pass for a finished atlas uses mean-pooling, stores raw
-per-token activations, runs all optional analyses, and writes everything into the
-same atlas directory:
+For large GPU runs, do **not** start with every optional feature turned on. The
+proven workflow is to build two component halves, then combine them at the atlas
+layer. This keeps failures recoverable and avoids redoing expensive census work.
+
+Run MLP-side components first:
 
 ```bash
 python app.py \
     --model <hf-model-id> \
     --corpus /home/user/app/prompts/prompts_balanced.jsonl \
-    --outdir outputs/<model-slug>-census \
-    --atlas atlas/<model-slug> \
+    --outdir outputs/<model-slug>-mlp \
+    --atlas atlas/<model-slug>-mlp \
     --layers all \
     --batch-size 8 \
     --max-length 128 \
-    --components mlp,gate,up,attn,heads,q,k,v \
+    --components mlp,gate,up \
     --pooling mean \
-    --store-per-token \
-    --per-token-analysis \
-    --logit-lens \
-    --sub-zero \
-    --sub-zero-corpora /home/user/app/prompts \
     --positive /home/user/app/prompts/authentic.jsonl \
     --negative /home/user/app/prompts/corporate.jsonl \
     --timing-every 25
 ```
 
-This produces:
+Then run the attention-side components:
 
-- `outputs/<model-slug>-census/l*_census_raw.npz` - last-token, mean-token and
-  optional ragged per-token activation arrays for every requested component.
-- `outputs/<model-slug>-census/analysis/` - feature separation scores, feature
-  fingerprints and per-token attribution summaries.
-- `outputs/<model-slug>-census/logit_lens_scores.json` - top-k unembedding
-  projections for `mlp`, `gate`, `up` and `attn` features.
-- `outputs/<model-slug>-census/subzero_report.json` - Sub-Zero DAS rotational
-  fingerprints for the requested corpora.
-- `atlas/<model-slug>/` - the merged SQLite atlas with all of the above indexed.
+```bash
+python app.py \
+    --model <hf-model-id> \
+    --corpus /home/user/app/prompts/prompts_balanced.jsonl \
+    --outdir outputs/<model-slug>-attn \
+    --atlas atlas/<model-slug>-attn \
+    --layers all \
+    --batch-size 8 \
+    --max-length 128 \
+    --components attn,heads,q,k,v \
+    --pooling mean \
+    --positive /home/user/app/prompts/authentic.jsonl \
+    --negative /home/user/app/prompts/corporate.jsonl \
+    --timing-every 10
+```
+
+Expected full feature count for Llama-3.1-8B-class dimensions is:
+
+```text
+mlp/gate/up: 32 * (14336 + 14336 + 14336) = 1,376,256
+attn side:   32 * (4096 + 4096 + 4096 + 1024 + 1024) = 458,752
+total:       1,835,008
+```
+
+If a run is interrupted, rerun with `--skip-census` and/or
+`--skip-existing-analysis` only after confirming the expected `.npz` files and
+analysis outputs exist.
 
 ### Combine outputs into the atlas
 
-The `app.py --sub-zero` command above already combines everything. Use these
-manual merge commands only if you ran pieces separately or need to rebuild the
-atlas from existing outputs.
+Use these commands to build a final combined atlas from existing MLP and
+attention halves. This is the preferred merge path; do not raw-merge large
+`.npz` files.
 
 ```bash
 python -m qwip_atlas.build_atlas \
-    --atlas atlas/<model-slug> \
+    --atlas atlas/<model-slug>-combined \
     init \
     --model-id <hf-model-id> \
-    --census outputs/<model-slug>-census/l0_census_raw.npz
+    --census outputs/<model-slug>-mlp/l0_census_raw.npz
 
 python -m qwip_atlas.build_atlas \
-    --atlas atlas/<model-slug> \
+    --atlas atlas/<model-slug>-combined \
     merge-all-layers \
-    --census-dir outputs/<model-slug>-census \
-    --analysis-dir outputs/<model-slug>-census/analysis \
+    --census-dir outputs/<model-slug>-mlp \
+    --analysis-dir outputs/<model-slug>-mlp/analysis \
     --no-census-copy
 
 python -m qwip_atlas.build_atlas \
-    --atlas atlas/<model-slug> \
+    --atlas atlas/<model-slug>-combined \
+    merge-all-layers \
+    --census-dir outputs/<model-slug>-attn \
+    --analysis-dir outputs/<model-slug>-attn/analysis \
+    --no-census-copy
+
+python -m qwip_atlas.build_atlas \
+    --atlas atlas/<model-slug>-combined \
     merge-ov \
-    --report outputs/<model-slug>-census/ov_circuit_scores.json
+    --report outputs/<model-slug>-attn/ov_circuit_scores.json
 
 python -m qwip_atlas.build_atlas \
-    --atlas atlas/<model-slug> \
+    --atlas atlas/<model-slug>-combined \
     merge-compliance-behaviour \
-    --report outputs/<model-slug>-census/compliance_behaviour_scores.json
+    --report outputs/<model-slug>-attn/compliance_behaviour_scores.json
 
 python -m qwip_atlas.build_atlas \
-    --atlas atlas/<model-slug> \
+    --atlas atlas/<model-slug>-combined \
     merge-subzero \
-    --report outputs/<model-slug>-census/subzero_report.json
+    --report outputs/<model-slug>-attn/subzero_report.json
 
 python -m qwip_atlas.build_atlas \
-    --atlas atlas/<model-slug> \
+    --atlas atlas/<model-slug>-combined \
     merge-logit-lens \
-    --report outputs/<model-slug>-census/logit_lens_scores.json
+    --report outputs/<model-slug>-combined-logit_lens_scores.json
 
 python -m qwip_atlas.build_atlas \
-    --atlas atlas/<model-slug> \
+    --atlas atlas/<model-slug>-combined \
     index
 
 python -m qwip_atlas.build_atlas \
-    --atlas atlas/<model-slug> \
+    --atlas atlas/<model-slug>-combined \
     status
 ```
 
@@ -349,6 +371,49 @@ python app.py \
     --skip-existing-analysis
 ```
 
+## Validated Llama-3.1-8B-Instruct run
+
+The RunPod GPU workflow for `meta-llama/Llama-3.1-8B-Instruct` was validated
+with separate output roots:
+
+- `outputs/llama8b-mlp` -> `mlp,gate,up`
+- `outputs/llama8b-attn` -> `attn,heads,q,k,v`
+- `atlas/llama8b-combined` -> final combined atlas
+
+Expected SQLite sanity checks after `index`:
+
+```text
+features:
+  attn  131072
+  heads 131072
+  q     131072
+  k      32768
+  v      32768
+  mlp   458752
+  gate  458752
+  up    458752
+
+per_head: 2560
+ov_circuits: 1024
+compliance_behaviour_features: 458752
+compliance_behaviour_per_head: 2560
+subzero_layer: 32
+subzero_svs: populated for down_proj/gate_proj/up_proj
+```
+
+The attention census `.npz` files must include all of:
+
+```text
+attn_last, attn_mean,
+attn_heads_last, attn_heads_mean,
+q_heads_last, q_heads_mean,
+k_heads_last, k_heads_mean,
+v_heads_last, v_heads_mean
+```
+
+If `heads/q/k/v` are missing, pull a version including the fixed fast-path
+census writer before rerunning the attention half.
+
 ## Memory knobs
 
 Defaults are safe for ~8 GB:
@@ -402,32 +467,52 @@ This writes `analysis/per_token/l<N>_<component>_per_token.json` with, per top f
 
 ### Logit-lens projection
 
-Project top feature directions onto the model's unembedding matrix to see which tokens each feature promotes or suppresses. This loads the full model weights, so it is off by default:
+Project top feature directions onto the model's unembedding matrix to see which
+tokens each feature promotes or suppresses. This loads the full model weights
+and can be CPU-heavy, so it is off by default.
 
-```bash
-python app.py \
-    --model <hf-model-id> \
-    --corpus /home/user/app/prompts/prompts_balanced.jsonl \
-    --outdir outputs/<model>-census \
-    --atlas atlas/<model> \
-    --logit-lens \
-    --components mlp,gate,up
-```
+Current support:
+
+- `attn`: direct d_model basis projection.
+- `mlp`, `gate`, `up`: projected through the layer `down_proj`.
+- `heads`: projected through the layer `o_proj`.
+- `q`, `k`, `v`: intentionally skipped. These do not have a context-free
+  unembedding direction because they affect attention dynamically.
+
+The default `--top-k 64` is per layer per component. For a combined atlas with
+`mlp,gate,up,attn,heads`, that is at most `32 * 5 * 64 = 10240` feature
+projections. Do not run every feature by default on a paid pod; full coverage
+for supported components can exceed one million rows and produces a large JSON
+with promoted/suppressed token payloads.
 
-The output is merged into the atlas as `logit_lens`. You can also run it manually:
+Recommended manual run for a final combined atlas:
 
 ```bash
-python analyze_logit_lens.py \
+nohup python analyze_logit_lens.py \
     --model <hf-model-id> \
-    --atlas atlas/<model> \
-    --output outputs/<model>-census/logit_lens_scores.json \
+    --atlas atlas/<model-slug>-combined \
+    --output outputs/<model-slug>-combined-logit_lens_scores.json \
     --top-k 64 \
-    --components mlp,gate,up
+    --tokens-per-feature 8 \
+    --components mlp,gate,up,attn,heads \
+    > outputs/<model-slug>-combined-logit_lens.nohup.log 2>&1 &
 
+tail -f outputs/<model-slug>-combined-logit_lens.nohup.log
+```
+
+`analyze_logit_lens.py` writes the JSON only at the end. If SSH disconnects or
+you interrupt it before completion, there may be no new output file to merge.
+After it completes, merge and re-index:
+
+```bash
 python -m qwip_atlas.build_atlas \
-    --atlas atlas/<model> \
+    --atlas atlas/<model-slug>-combined \
     merge-logit-lens \
-    --report outputs/<model>-census/logit_lens_scores.json
+    --report outputs/<model-slug>-combined-logit_lens_scores.json
+
+python -m qwip_atlas.build_atlas \
+    --atlas atlas/<model-slug>-combined \
+    index
 ```
 
 ### Advanced census options
@@ -490,11 +575,12 @@ Use `/home/user/app/prompts/prompts_balanced.jsonl` in the `--corpus` examples b
 
 The prompt axis was renamed from `pos.jsonl` / `neg.jsonl` to `authentic.jsonl` / `corporate.jsonl` to make the semantic role obvious. The defaults in `app.py` and `make_compliance_corpora.py` have been updated to match. If you keep your own curated corpora, just point `--positive` and `--negative` at them.
 
-### Sub-Zero corpora path (legacy)
+### Sub-Zero corpora path and scope
 
-Sub-Zero is the older DAS rotational probe. It is still wired in, but the
-recommended resolution features are now mean-pooling, per-token attribution and
-logit-lens projection.
+Sub-Zero is the DAS rotational probe. In the current default configuration it
+probes MLP-side projection weights (`gate_proj`, `up_proj`, `down_proj`) across
+the selected layers. It does not mean every atlas component has its own
+Sub-Zero projection entry.
 
 When you do run it, keep the Sub-Zero corpora in the same prompts directory used
 for the rest of the pipeline. In the Space that directory is
@@ -509,6 +595,18 @@ python app.py \
     --sub-zero-corpora /home/user/app/prompts
 ```
 
+If your prompt files are named `neutral_stems.jsonl` and `red_team_stems.jsonl`,
+create symlinks before running Sub-Zero:
+
+```bash
+ln -sf neutral_stems.jsonl /home/user/app/prompts/neutral.jsonl
+ln -sf red_team_stems.jsonl /home/user/app/prompts/red_team.jsonl
+```
+
+The SVD stage is checkpointed under `outputs/<model>/sub_zero_ckpt`. A SIGKILL
+during SVD usually means the pod was memory-killed; rerun after pulling the
+latest code so unused left-singular vectors are not retained in the checkpoint.
+
 ### OV-circuit compliance score placeholder
 
 `analyze_ov_circuits.py` currently writes `compliance_score` and `layer_comp_strength` as `None` for every head. The merge path is wired, but the analyzer does not yet consume a compliance-behaviour report to fill those columns. Until that patch lands, treat those two fields as placeholders.