Speed up experiment on Apple Silicon/MPS (verified result-preserving)

[rohanpoudel2]

Summary

Speeds up the experiment pipeline on Apple Silicon/MPS (and generally) by removing redundant compute, tokenization, and device↔host synchronizations from the hot paths. Every change is deliberately result-preserving: outputs are identical up to floating-point reduction order, and the reasoning for why each optimization does not alter results is captured inline as comments at each site.

No changes to the scientific contract, protocol, or any statistic's definition — only how they are computed.

Changes

Model forward passes — skip work that gets thrown away

• extraction/task_extractor.py: load the base AutoModel instead of AutoModelForCausalLM (no LM head), and drop transformer blocks above the highest requested layer, neutralizing the final norm so retained blocks' captured hidden states stay bit-identical. Also gathers all requested layers in a single device→host transfer instead of one per layer, and switches to torch.inference_mode().
• cli/run_llm_judge_baselines.py: when the model supports it and padding is left-side, request only the final logit row (logits_to_keep=1) so the LM-head projection skips the whole padded sequence. Falls back to the previous masked-argmax path otherwise.

Tokenization — stop tokenizing twice

• cli/extract_text_embeddings.py: when nothing in a batch exceeds max_length, pad the already-tokenized ids instead of re-tokenizing with truncation=True.
• cli/run_llm_judge_baselines.py: derive per-example lengths from the attention mask of a single padded tokenization instead of a separate unpadded pass.
• cli/generate_task_rollouts.py: build the chat-template prompt encoding once per scenario rather than once per replicate (sampling randomness is still seeded per generate() call, so replicates remain independent).

Statistics — vectorize the bootstrap hot loops

• evaluation/hierarchical_statistics.py: collapse rows to a dense seed×group cell-mean matrix once and resample via integer index draws + np.ix_, replacing ~n_boot * n_seeds * n_groups dict lookups with two array means. Random stream is consumed identically, so reported statistics stay bit-identical.
• evaluation/metrics.py: precompute each group's row indices once instead of rescanning groups every bootstrap iteration.
• data/generation_confidence.py: batch the per-token log-prob/entropy/margin/top-1 summaries over the step-score matrix in chunks, collapsing per-token Python/kernel dispatch that dominates wall-clock on MPS.

Memory & caching

• cli/run_task_sweep.py: clear the per-layer bundle cache between layers so host/unified memory is bounded by a single layer's bundles.
• probes/sae_probe.py: cache loaded SAEs by pinned identity so the sweep doesn't re-download/re-instantiate the immutable encoder on every run.

I/O

• cli/run_task_sweep.py: flush every summary row (so a resumed process sees completed runs) but only fsync periodically, with a final durable sync once all runs are recorded. Prediction files are still written atomically-and-fsynced before their summary row, and any lost summary tail is recomputed on resume.

Verification

Changes are constructed to preserve results exactly; correctness rationale is documented at each call site. Recommend confirming on a representative run that summary outputs and reported statistics match the pre-change baseline.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.