Spaces:
Sleeping
Sleeping
Commit History
-
Speed up primary element stats and cut peak memory
Replace the per-head np.histogram + scipy stats with one fused pass: O(n) uniform-bin histogram, shared-central-moment mean/std/skew/kurtosis (bit-identical to scipy), and a histogram plug-in differential_entropy that drops scipy's O(n log n) Vasicek sort. ~18x faster per-head element stats. Build W_QK one head at a time instead of the full (n_heads, d_model, d_model) tensor, removing an ~838 MB transient and a pointless device round-trip (W_QK is only histogrammed on CPU). This is what relieves swap thrashing for the large models. Also add a gated svd_via_gram option (off by default) deriving W_Q/W_K U,S from the d_head Gram eigendecomposition instead of a wide MPS->CPU-fallback SVD. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>fa476b7
-
Fix transform/primary path mismatch when revision is null
reprocess_metrics looked for {model} but primary writes {model}_main, so revision-less runs silently skipped (the Snakemake sentinel was still touched, hiding the breakage). Use the same "main" fallback string in both places. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>2dcdcdf -
Fix two OLMo 2 pipeline failures
mlflow.log_input expects an mlflow.data.Dataset, not a raw DatasetInput entity — use MlflowClient().log_inputs() which takes the lower-level type directly. torch.linalg.svd uses LAPACK gesdd which fails to converge on ill-conditioned weight matrices seen in OLMo 2. Add _svd() helper that catches LinAlgError and retries via numpy.linalg.svd (gesvd driver, always converges). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>7b59fb4
-
Add d_head metadata, top-k-svd-d-head option, and SVD truncation experiment
- Log d_head explicitly in metadata.json (correct for GQA models where d_head != d_model/n_heads); strip redundant head_dim alias from serialized output - Add top_k_svd_d_head param to process_model: enables low-rank SVD with N=d_head without needing to know d_head at call time - Expose --low-rank-svd, --top-k-svd, --top-k-svd-d-head in run_primary.py; log all three to MLflow - _get_d_head in cross_model.py reads d_head directly from metadata when available - Add scripts/experiment_svd_truncation.py: runs d_head-truncated pipeline, compares compute time and per-head singular value accuracy against nominal results on disk Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>e9f64da -
Namespace all outputs under experiment_dir, add MLflow run metadata
All pipeline outputs and sentinels now live under {output_dir}/{mlflow_experiment}/ instead of scattered across the project root and output_dir. Prevents clobbering across experiments and mirrors the MLflow hierarchy. Changes: - Snakefile: derive config["experiment_dir"], pre-create it, override eval.out dynamically; update all sentinel expand() paths - rules/*: sentinel touch paths and inter-rule deps use experiment_dir; correlations/figures dirs moved under experiment_dir - config.yaml: remove hardcoded eval.out (now derived in Snakefile) - scripts/run_*.py: add mlflow.note.content description, model tag, and DatasetInput (local source → output path) to every run Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>0bf8fb1
-
Wire QK decomposition weight types into primary pipeline
W_Q_gram, W_K_gram, QK_alignment are already computed in head_analyzer.py. This commit exposes weight_types through config.yaml → run_primary.py → process_model() → primary.smk so they can be overridden at runtime. Tests updated to cover all 6 weight types; histogram test restricted to raw weight types since gram entries exceed standard bin range. compute_qk_decomposition.py removed (data now lives in primary dataset). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>a07e79a -
Rename src modules for head/pair symmetry; add pair_figures.smk; split eval_metrics
src/transformer_analysis/: weight_analysis → head_pipeline attn_head_analysis → head_analyzer metrics → head_metrics correlation_analysis → pair_pipeline head_correlations → pair_analyzer eval_metrics.py (new — implementation extracted from compute_perplexity.py) scripts/: plot_correlations.py → plot_pair_correlations.py plot_corr_figures.py → plot_pair_correlations_batch.py compute_perplexity.py → deleted (impl in eval_metrics; thin wrapper run_eval.py) run_eval.py (new thin wrapper) rules/: correlations.smk — switched to sentinel output (consistent with other rules) pair_figures.smk (new — regenerate pair figures from .npz without re-running correlations) eval.smk — updated to call run_eval.py Snakefile: added correlations_all, pair_figures_all convenience targets Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>a81fd74 -
feat: add Snakemake pipeline, MLflow tracking, metrics refactor, tests
Pipeline / workflow: - Snakefile + config.yaml orchestrate primary → transform → merge stages - rules/primary.smk, transform.smk, correlations.smk, merge.smk - Sentinel-file pattern (done/*.done) avoids directory-output conflicts - Incremental cross-model merge: appends new models without clobbering; --refresh flag for explicit replacement after re-transform - Per-model checkpoint merge rule for Pythia step-collapse CLI wrappers (scripts/): - run_primary.py — process_model() + MLflow logging - run_transform.py — reprocess_metrics() + MLflow logging (git sha tagged) - run_correlations.py — thin wrapper for run_multi_circuit_analysis() - run_merge.py — cross-model incremental merge + checkpoint merge Metrics refactor: - src/transformer_analysis/metrics.py: canonical home for all stat functions, bin configs, normality_metrics, singular_value_metrics, and adaptive binning helpers (make_weight_bins, make_sv_bins) - histogram_utils.py: backward-compat re-export shim; no logic c…24e1932