bergson 2.2.0__tar.gz → 2.2.2__tar.gz

{bergson-2.2.0 → bergson-2.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bergson
-Version: 2.2.0
+Version: 2.2.2
 Summary: Tracing the memory of neural nets with data attribution
 License: MIT License
 Keywords: interpretability,explainable-ai
@@ -60,7 +60,7 @@ Requires-Dist: matplotlib; extra == "plot"
 Dynamic: license-file
 
 # Bergson
-Bergson is a python library which provides scalable, state-of-the-art data attribution methods for large language models. Data attribution methods estimate the effect on a behavior of interest of removing data points from a model's training corpus. We support [EK-FAC](https://arxiv.org/abs/2308.03296) (2023), [TrackStar](https://arxiv.org/abs/2410.17413v3) (2024), [SOURCE](https://arxiv.org/abs/2405.12186) (2024), [MAGIC](https://arxiv.org/abs/2504.16430) (2025), gradient cosine similarity, and more.
+Bergson is a python library which provides scalable, state-of-the-art data attribution methods for large language models. Data attribution methods estimate the effect on a behavior of interest of removing data points from a model's training corpus. We support [EK-FAC](https://arxiv.org/abs/2308.03296) (2023), [TrackStar](https://arxiv.org/abs/2410.17413v3) (2024), [SOURCE](https://arxiv.org/abs/2405.12186) (2024), [MAGIC](https://arxiv.org/abs/2504.16430) (2025), [ASTRA](https://arxiv.org/abs/2507.14740) (2025), gradient cosine similarity, and more.
 
 We try to make research and application straightforward. You can reproduce any of your runs with a single line, use a few CLI flags to save and re-use useful intermediate artifacts, do everything using just the CLI or just code, train models here or use existing ones, tune and evaluate your methods, and scale them up to multi-node runs with 70B+ parameters. More information is available in the [Bergson](https://bergson.readthedocs.io) docs.
 
@@ -86,23 +86,23 @@ pip install -e .
 
 Indicative performance of data attribution methods in the finetuning regime - see the [leaderboard](https://bergson.readthedocs.io/en/latest/leaderboard.html) for more methods.
 
-[Linear datamodeling score](https://arxiv.org/abs/2303.14186) (LDS) is the accuracy of a method for producing global data rankings by influence. The query loss difference (QLD) shows how much model loss for a held-out query can be increased by retraining without the most highly ranked data by influence (here the top 1%), compared to a random removal baseline.
+[Linear datamodeling score](https://arxiv.org/abs/2303.14186) (LDS) is the accuracy of a method for producing global data rankings by influence. The query loss difference (QLD) shows how much model loss for a held-out query can be increased by retraining without the most highly ranked data by influence (here the top 1%), compared to a random removal baseline. The per-token QLD removes the loss terms of the top 1% of training tokens, ranked by the forward-mode equivalent of each method's scores.
 
-| Method | Proponent QLD [95% CI] | LDS [95% CI] |
-|:---|:---:|:---:|
-| MAGIC | 0.100 [0.090, 0.112] | 0.931 [0.925, 0.936] |
-| EK-FAC + ASTRA | 0.074 [0.063, 0.087] | 0.643 [0.624, 0.660] |
-| EK-FAC | 0.070 [0.058, 0.082] | 0.454 [0.426, 0.479] |
-| KFAC | 0.067 [0.056, 0.080] | 0.420 [0.391, 0.446] |
-| BM25 | 0.062 [0.048, 0.076] | 0.220 [0.185, 0.252] |
-| [Qwen3-Embedding-8B](https://huggingface.co/spaces/mteb/leaderboard) semantic search | 0.049 [0.038, 0.061] | 0.132 [0.093, 0.169] |
-| TrackStar (no optimizer correction, projection 64) | 0.045 [0.036, 0.055] | 0.270 [0.240, 0.295] |
-| TRAK (8-model ensemble) | 0.032 [0.024, 0.040] | 0.138 [0.111, 0.165] |
-| SOURCE (Adam) | 0.024 [0.018, 0.030] | 0.154 [0.126, 0.181] |
-| Gradient cosine similarity | 0.021 [0.016, 0.027] | 0.156 [0.131, 0.181] |
-| Activation similarity | 0.000 [-0.000, 0.001] | 0.110 [0.070, 0.149] |
+| Method | QLD | LDS | Per-token QLD |
+|:---|:---:|:---:|:---:|
+| MAGIC | 0.100 [0.090, 0.112] | 0.931 [0.925, 0.936] | 1.375 [1.348, 1.403] |
+| EK-FAC + ASTRA | 0.074 [0.063, 0.087] | 0.643 [0.624, 0.660] | 1.082 [1.047, 1.116] |
+| Eigenvalue-corrected Shampoo | 0.071 [0.060, 0.082] | 0.517 [0.491, 0.539] | 0.922 [0.887, 0.957] |
+| SOURCE (Adam, EK-FAC) | 0.071 [0.060, 0.084] | 0.473 [0.446, 0.498] | 0.908 [0.877, 0.942] |
+| EK-FAC | 0.070 [0.058, 0.082] | 0.454 [0.426, 0.479] | 0.863 [0.833, 0.894] |
+| BM25 | 0.062 [0.048, 0.076] | 0.220 [0.185, 0.252] | 0.677 [0.650, 0.704] |
+| Semantic search ([Qwen3 8B](https://huggingface.co/Qwen/Qwen3-Embedding-8B)) | 0.049 [0.038, 0.061] | 0.132 [0.093, 0.169] | 0.008 [0.006, 0.010] |
+| TrackStar (no optimizer correction, projection 64) | 0.045 [0.036, 0.055] | 0.270 [0.240, 0.295] | 0.489 [0.465, 0.515] |
+| TRAK (8-model ensemble) | 0.032 [0.024, 0.040] | 0.138 [0.111, 0.165] | 0.119 [0.112, 0.125] |
+| Gradient cosine similarity | 0.021 [0.016, 0.027] | 0.156 [0.131, 0.181] | 0.217 [0.206, 0.229] |
+| Activation similarity | 0.000 [-0.000, 0.001] | 0.110 [0.070, 0.149] | 0.028 [0.018, 0.039] |
 
-Results for GPT-2 finetuned on 4 epochs of the WikiText corpus. Held-out loss dropped from 3.545 to 3.111 over training.
+Results with 95% confidence intervals for GPT-2 finetuned on 4 epochs of the WikiText corpus. Held-out loss dropped from 3.545 to 3.111 over training.
 
 ## Functionality
 
@@ -110,7 +110,7 @@ Results for GPT-2 finetuned on 4 epochs of the WikiText corpus. Held-out loss dr
 
 `bergson magic` runs a powerful attribution method that backpropagates through the training process to compute the gradient of a model behavior loss with respect to a weighting placed on each training item. It is powered by our twice-differentiable trainer, which can also be called directly using `bergson train`.
 
-**Note: unrolled differentiation efficacy is proportional to [metasmoothness](https://bergson.readthedocs.io/en/latest/magic.html#metasmoothness). Untuned training hyperparameters can result in disappointing linear datamodeling scores. Check your run's estimated metasmoothness with `bergson metasmoothness`.**
+**Note: unrolled differentiation efficacy is proportional to [metasmoothness](https://bergson.readthedocs.io/en/latest/magic.html#metasmoothness), which is low in some settings, including early pretraining steps. Check your run's estimated metasmoothness with `bergson metasmoothness`.**
 
 `bergson approxunrolling` is an approximation of `bergson magic` that uses a handful of training checkpoints to run the multi-step SOURCE attribution pipeline. This is roughly equivalent to an influence function averaged over several checkpoints.
 
@@ -120,45 +120,39 @@ To build a train‑time gradient store, use our HF Trainer callback. This will i
 
 Bergson supports on-disk gradient stores and on-the-fly queries, and per-token and per-sequence attribution.
 
-`bergson trackstar`, `bergson ekfac`, and `bergson trak` all orchestrate multi-step attribution recipes over a model checkpoint. `bergson trak` also supports [ensembling](https://arxiv.org/abs/2303.14186) over independently trained models.
+The CLI commands `bergson trackstar`, `bergson ekfac`, and `bergson trak` all orchestrate multi-step attribution recipes over a model checkpoint. `bergson trak` also supports [ensembling](https://arxiv.org/abs/2303.14186) over independently trained models.
 
 At a lower level, you can build your own gradient store for efficient serial queries using `bergson build`. Collection-time gradient compression makes the store space-efficient, and a FAISS integration enables fast KNN search over large stores - see `bergson query`, or `Attributor` in the programmatic interface. For small queries and methods that don't use gradient compression (e.g., EK-FAC), you can score a dataset in a single pass using an in-memory query index of precomputed gradients. Dataset items may be scored using max, mean, and individual scoring strategies, enabling [LESS](https://arxiv.org/pdf/2402.04333)-style data filtering. See `bergson score` and `bergson build`.
 
 Per-module and per-attention head gradients can be extracted from the store.
 
-**Note: influence functions can be sensitive to both Hessian approximation inversion hyperparameters (in tiny models) and metasmoothness. Untuned hyperparameters can result in a disappointing linear datamodeling score.**
+**Note: influence functions are sensitive to the Hessian approximation damping hyperparameter in tiny models. Untuned hyperparameters can result in a suboptimal linear datamodeling score.**
 
 ### Evaluate
 
 Use `bergson validate` and `bergson recall` to compute LDS and recall@k metrics respectively.
 
-# Examples
+# Getting Started
 
-There are many example YAMLs in the `examples` directory. use `bergson <yaml_path>` to try them. For example, to MAGIC-attribute a GPT-2 WikiText fine-tune:
+There are many example YAMLs in the `examples` directory, including various paper experiment replications. Use `bergson <yaml_path>` to run them. For example, to MAGIC-attribute a GPT-2 WikiText fine-tune:
 
 ```bash
 bergson examples/magic/gpt2_wikitext_tiny.yaml
 ```
 
-Or check out a notebook for programmatic usage:
-
-| | Notebook | Description |
-|---|----------|-------------|
-| [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/EleutherAI/bergson/blob/main/notebooks/poison_detection.ipynb) | **Poison Detection** | Detect poisoned training examples with gradient attribution (T4, ~5 min) |
-| [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/EleutherAI/bergson/blob/main/notebooks/style_ablation.ipynb) | **Style Ablation** | Suppress style to recover semantic matching (A100, ~20 min) |
-
-Construct and query an on-disk index of randomly projected gradients:
+You can use the same fields used to specify experiments in the YAMLs to run experiments directly in the Bergson CLI. For example, to construct and query an on-disk index of randomly projected gradients from the CLI:
 
 ```bash
 bergson build runs/index --model EleutherAI/pythia-14m --dataset NeelNanda/pile-10k --truncation --token_batch_size 4096 --projection_dim 16
 bergson query --index runs/index --unit_norm
 ```
 
-Collect TrackStar attribution scores for an I.I.D sample query:
+Or check out a notebook for programmatic usage:
 
-```bash
-bergson trackstar runs/trackstar --model EleutherAI/pythia-14m --query.dataset NeelNanda/pile-10k --data.dataset NeelNanda/pile-10k --data.truncation --token_batch_size 4096 --query.truncation --query.split "train[:20]"
-```
+| | Notebook | Description |
+|---|----------|-------------|
+| [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/EleutherAI/bergson/blob/main/notebooks/poison_detection.ipynb) | **Poison Detection** | Detect poisoned training examples with gradient attribution (T4, ~5 min) |
+| [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/EleutherAI/bergson/blob/main/notebooks/style_ablation.ipynb) | **Style Ablation** | Suppress style to recover semantic matching (A100, ~20 min) |
 
 # Development
 

{bergson-2.2.0 → bergson-2.2.2}/README.md

@@ -1,5 +1,5 @@
 # Bergson
-Bergson is a python library which provides scalable, state-of-the-art data attribution methods for large language models. Data attribution methods estimate the effect on a behavior of interest of removing data points from a model's training corpus. We support [EK-FAC](https://arxiv.org/abs/2308.03296) (2023), [TrackStar](https://arxiv.org/abs/2410.17413v3) (2024), [SOURCE](https://arxiv.org/abs/2405.12186) (2024), [MAGIC](https://arxiv.org/abs/2504.16430) (2025), gradient cosine similarity, and more.
+Bergson is a python library which provides scalable, state-of-the-art data attribution methods for large language models. Data attribution methods estimate the effect on a behavior of interest of removing data points from a model's training corpus. We support [EK-FAC](https://arxiv.org/abs/2308.03296) (2023), [TrackStar](https://arxiv.org/abs/2410.17413v3) (2024), [SOURCE](https://arxiv.org/abs/2405.12186) (2024), [MAGIC](https://arxiv.org/abs/2504.16430) (2025), [ASTRA](https://arxiv.org/abs/2507.14740) (2025), gradient cosine similarity, and more.
 
 We try to make research and application straightforward. You can reproduce any of your runs with a single line, use a few CLI flags to save and re-use useful intermediate artifacts, do everything using just the CLI or just code, train models here or use existing ones, tune and evaluate your methods, and scale them up to multi-node runs with 70B+ parameters. More information is available in the [Bergson](https://bergson.readthedocs.io) docs.
 
@@ -25,23 +25,23 @@ pip install -e .
 
 Indicative performance of data attribution methods in the finetuning regime - see the [leaderboard](https://bergson.readthedocs.io/en/latest/leaderboard.html) for more methods.
 
-[Linear datamodeling score](https://arxiv.org/abs/2303.14186) (LDS) is the accuracy of a method for producing global data rankings by influence. The query loss difference (QLD) shows how much model loss for a held-out query can be increased by retraining without the most highly ranked data by influence (here the top 1%), compared to a random removal baseline.
+[Linear datamodeling score](https://arxiv.org/abs/2303.14186) (LDS) is the accuracy of a method for producing global data rankings by influence. The query loss difference (QLD) shows how much model loss for a held-out query can be increased by retraining without the most highly ranked data by influence (here the top 1%), compared to a random removal baseline. The per-token QLD removes the loss terms of the top 1% of training tokens, ranked by the forward-mode equivalent of each method's scores.
 
-| Method | Proponent QLD [95% CI] | LDS [95% CI] |
-|:---|:---:|:---:|
-| MAGIC | 0.100 [0.090, 0.112] | 0.931 [0.925, 0.936] |
-| EK-FAC + ASTRA | 0.074 [0.063, 0.087] | 0.643 [0.624, 0.660] |
-| EK-FAC | 0.070 [0.058, 0.082] | 0.454 [0.426, 0.479] |
-| KFAC | 0.067 [0.056, 0.080] | 0.420 [0.391, 0.446] |
-| BM25 | 0.062 [0.048, 0.076] | 0.220 [0.185, 0.252] |
-| [Qwen3-Embedding-8B](https://huggingface.co/spaces/mteb/leaderboard) semantic search | 0.049 [0.038, 0.061] | 0.132 [0.093, 0.169] |
-| TrackStar (no optimizer correction, projection 64) | 0.045 [0.036, 0.055] | 0.270 [0.240, 0.295] |
-| TRAK (8-model ensemble) | 0.032 [0.024, 0.040] | 0.138 [0.111, 0.165] |
-| SOURCE (Adam) | 0.024 [0.018, 0.030] | 0.154 [0.126, 0.181] |
-| Gradient cosine similarity | 0.021 [0.016, 0.027] | 0.156 [0.131, 0.181] |
-| Activation similarity | 0.000 [-0.000, 0.001] | 0.110 [0.070, 0.149] |
+| Method | QLD | LDS | Per-token QLD |
+|:---|:---:|:---:|:---:|
+| MAGIC | 0.100 [0.090, 0.112] | 0.931 [0.925, 0.936] | 1.375 [1.348, 1.403] |
+| EK-FAC + ASTRA | 0.074 [0.063, 0.087] | 0.643 [0.624, 0.660] | 1.082 [1.047, 1.116] |
+| Eigenvalue-corrected Shampoo | 0.071 [0.060, 0.082] | 0.517 [0.491, 0.539] | 0.922 [0.887, 0.957] |
+| SOURCE (Adam, EK-FAC) | 0.071 [0.060, 0.084] | 0.473 [0.446, 0.498] | 0.908 [0.877, 0.942] |
+| EK-FAC | 0.070 [0.058, 0.082] | 0.454 [0.426, 0.479] | 0.863 [0.833, 0.894] |
+| BM25 | 0.062 [0.048, 0.076] | 0.220 [0.185, 0.252] | 0.677 [0.650, 0.704] |
+| Semantic search ([Qwen3 8B](https://huggingface.co/Qwen/Qwen3-Embedding-8B)) | 0.049 [0.038, 0.061] | 0.132 [0.093, 0.169] | 0.008 [0.006, 0.010] |
+| TrackStar (no optimizer correction, projection 64) | 0.045 [0.036, 0.055] | 0.270 [0.240, 0.295] | 0.489 [0.465, 0.515] |
+| TRAK (8-model ensemble) | 0.032 [0.024, 0.040] | 0.138 [0.111, 0.165] | 0.119 [0.112, 0.125] |
+| Gradient cosine similarity | 0.021 [0.016, 0.027] | 0.156 [0.131, 0.181] | 0.217 [0.206, 0.229] |
+| Activation similarity | 0.000 [-0.000, 0.001] | 0.110 [0.070, 0.149] | 0.028 [0.018, 0.039] |
 
-Results for GPT-2 finetuned on 4 epochs of the WikiText corpus. Held-out loss dropped from 3.545 to 3.111 over training.
+Results with 95% confidence intervals for GPT-2 finetuned on 4 epochs of the WikiText corpus. Held-out loss dropped from 3.545 to 3.111 over training.
 
 ## Functionality
 
@@ -49,7 +49,7 @@ Results for GPT-2 finetuned on 4 epochs of the WikiText corpus. Held-out loss dr
 
 `bergson magic` runs a powerful attribution method that backpropagates through the training process to compute the gradient of a model behavior loss with respect to a weighting placed on each training item. It is powered by our twice-differentiable trainer, which can also be called directly using `bergson train`.
 
-**Note: unrolled differentiation efficacy is proportional to [metasmoothness](https://bergson.readthedocs.io/en/latest/magic.html#metasmoothness). Untuned training hyperparameters can result in disappointing linear datamodeling scores. Check your run's estimated metasmoothness with `bergson metasmoothness`.**
+**Note: unrolled differentiation efficacy is proportional to [metasmoothness](https://bergson.readthedocs.io/en/latest/magic.html#metasmoothness), which is low in some settings, including early pretraining steps. Check your run's estimated metasmoothness with `bergson metasmoothness`.**
 
 `bergson approxunrolling` is an approximation of `bergson magic` that uses a handful of training checkpoints to run the multi-step SOURCE attribution pipeline. This is roughly equivalent to an influence function averaged over several checkpoints.
 
@@ -59,45 +59,39 @@ To build a train‑time gradient store, use our HF Trainer callback. This will i
 
 Bergson supports on-disk gradient stores and on-the-fly queries, and per-token and per-sequence attribution.
 
-`bergson trackstar`, `bergson ekfac`, and `bergson trak` all orchestrate multi-step attribution recipes over a model checkpoint. `bergson trak` also supports [ensembling](https://arxiv.org/abs/2303.14186) over independently trained models.
+The CLI commands `bergson trackstar`, `bergson ekfac`, and `bergson trak` all orchestrate multi-step attribution recipes over a model checkpoint. `bergson trak` also supports [ensembling](https://arxiv.org/abs/2303.14186) over independently trained models.
 
 At a lower level, you can build your own gradient store for efficient serial queries using `bergson build`. Collection-time gradient compression makes the store space-efficient, and a FAISS integration enables fast KNN search over large stores - see `bergson query`, or `Attributor` in the programmatic interface. For small queries and methods that don't use gradient compression (e.g., EK-FAC), you can score a dataset in a single pass using an in-memory query index of precomputed gradients. Dataset items may be scored using max, mean, and individual scoring strategies, enabling [LESS](https://arxiv.org/pdf/2402.04333)-style data filtering. See `bergson score` and `bergson build`.
 
 Per-module and per-attention head gradients can be extracted from the store.
 
-**Note: influence functions can be sensitive to both Hessian approximation inversion hyperparameters (in tiny models) and metasmoothness. Untuned hyperparameters can result in a disappointing linear datamodeling score.**
+**Note: influence functions are sensitive to the Hessian approximation damping hyperparameter in tiny models. Untuned hyperparameters can result in a suboptimal linear datamodeling score.**
 
 ### Evaluate
 
 Use `bergson validate` and `bergson recall` to compute LDS and recall@k metrics respectively.
 
-# Examples
+# Getting Started
 
-There are many example YAMLs in the `examples` directory. use `bergson <yaml_path>` to try them. For example, to MAGIC-attribute a GPT-2 WikiText fine-tune:
+There are many example YAMLs in the `examples` directory, including various paper experiment replications. Use `bergson <yaml_path>` to run them. For example, to MAGIC-attribute a GPT-2 WikiText fine-tune:
 
 ```bash
 bergson examples/magic/gpt2_wikitext_tiny.yaml
 ```
 
-Or check out a notebook for programmatic usage:
-
-| | Notebook | Description |
-|---|----------|-------------|
-| [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/EleutherAI/bergson/blob/main/notebooks/poison_detection.ipynb) | **Poison Detection** | Detect poisoned training examples with gradient attribution (T4, ~5 min) |
-| [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/EleutherAI/bergson/blob/main/notebooks/style_ablation.ipynb) | **Style Ablation** | Suppress style to recover semantic matching (A100, ~20 min) |
-
-Construct and query an on-disk index of randomly projected gradients:
+You can use the same fields used to specify experiments in the YAMLs to run experiments directly in the Bergson CLI. For example, to construct and query an on-disk index of randomly projected gradients from the CLI:
 
 ```bash
 bergson build runs/index --model EleutherAI/pythia-14m --dataset NeelNanda/pile-10k --truncation --token_batch_size 4096 --projection_dim 16
 bergson query --index runs/index --unit_norm
 ```
 
-Collect TrackStar attribution scores for an I.I.D sample query:
+Or check out a notebook for programmatic usage:
 
-```bash
-bergson trackstar runs/trackstar --model EleutherAI/pythia-14m --query.dataset NeelNanda/pile-10k --data.dataset NeelNanda/pile-10k --data.truncation --token_batch_size 4096 --query.truncation --query.split "train[:20]"
-```
+| | Notebook | Description |
+|---|----------|-------------|
+| [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/EleutherAI/bergson/blob/main/notebooks/poison_detection.ipynb) | **Poison Detection** | Detect poisoned training examples with gradient attribution (T4, ~5 min) |
+| [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/EleutherAI/bergson/blob/main/notebooks/style_ablation.ipynb) | **Style Ablation** | Suppress style to recover semantic matching (A100, ~20 min) |
 
 # Development
 

{bergson-2.2.0 → bergson-2.2.2}/bergson/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.2.0"
+__version__ = "2.2.2"
 
 import logging
 

{bergson-2.2.0 → bergson-2.2.2}/bergson/config/config.py

@@ -885,6 +885,16 @@ class ScoreConfig(Serializable):
     capability (e.g. in influence functions). False for unrolled
     differentiation."""
 
+    token_influence: Literal["gradient", "output"] = "gradient"
+    """What each row scores with ``attribute_tokens``. ``gradient`` scores row
+    ``t`` by the per-token gradient at position ``t``, which is position ``t``'s
+    effect on the loss of every later token. ``output`` scores row ``t`` by the
+    loss on token ``t + 1`` alone, the output token influence of Grosse et al.
+    (2023). Both kinds of row sum to the per-example score, so without
+    ``attribute_tokens``, ``output`` only changes how that score is computed.
+    ``output`` costs one forward-mode pass per query column, so aggregate the
+    query when you can, and needs an unprojected query and dot-product scoring."""
+
     candidates: CandidateConfig = field(default_factory=CandidateConfig)
     """Score only the rows an earlier run ranked highest. The store then has one
     row per candidate; ``candidates.npy`` beside it holds their training rows."""

{bergson-2.2.0 → bergson-2.2.2}/bergson/hessians/eigenvectors.py

@@ -85,8 +85,9 @@ class LambdaCollector(HookCollectorBase):
     then computes outer products for diagonal correction terms.
 
     Distributed, each module belongs to the rank that owned its covariances,
-    which loads its full eigenvectors from every rank's shard and receives every
-    rank's positions for it; teardown saves the usual row shards.
+    which takes its full eigenvectors from ``eigenvectors`` or joins them from
+    every rank's shard, and receives every rank's positions for it; teardown
+    saves the usual row shards.
     """
 
     path: str
@@ -104,6 +105,11 @@ class LambdaCollector(HookCollectorBase):
     num_documents: int = 1
     """Divides the summed corrections into a mean over documents."""
 
+    eigenvectors: tuple[dict[str, Tensor], dict[str, Tensor]] | None = None
+    """The activation and gradient eigenvectors of the modules this rank owns,
+    in full, as returned by :func:`eigendecompose_owned`. Read from the shard
+    files when not given."""
+
     def setup(self) -> None:
         """Load eigenvectors and initialize storage."""
         self.shard_computer = ShardedMul()
@@ -121,6 +127,13 @@ class LambdaCollector(HookCollectorBase):
             self._setup_owned(eigen_src)
             return
 
+        if self.eigenvectors is not None:
+            self.eigen_a, self.eigen_g = (
+                {k: v.to(self.device, self.dtype) for k, v in vectors.items()}
+                for vectors in self.eigenvectors
+            )
+            return
+
         # Load precomputed eigenvectors
         self.eigen_a = load_file(
             os.path.join(
@@ -141,7 +154,7 @@ class LambdaCollector(HookCollectorBase):
 
     def _setup_owned(self, eigen_src: str) -> None:
         """Load the full eigenvectors of the modules this rank owns, joining the
-        row shards every rank wrote."""
+        row shards every rank wrote unless they were passed in."""
         self.owners = assign_module_owners(self.target_info, self.world_size)
         owned = [name for name, owner in self.owners.items() if owner == self.rank]
 
@@ -160,8 +173,14 @@ class LambdaCollector(HookCollectorBase):
                 for name, rows in full.items()
             }
 
-        self.eigen_a = load("eigen_activation_sharded")
-        self.eigen_g = load("eigen_gradient_sharded")
+        if self.eigenvectors is not None:
+            self.eigen_a, self.eigen_g = (
+                {name: vectors[name].to(self.device, self.dtype) for name in owned}
+                for vectors in self.eigenvectors
+            )
+        else:
+            self.eigen_a = load("eigen_activation_sharded")
+            self.eigen_g = load("eigen_gradient_sharded")
         self.correction_shapes = {
             name: (out_dim, in_dim + collect_bias)
             for name, (_, (out_dim, in_dim), collect_bias) in self.target_info.items()
@@ -322,6 +341,86 @@ def _compute_full_matrix(
     return full_matrix
 
 
+def _eigh(
+    key: str, matrix: Tensor, total_processed: Tensor, dtype: torch.dtype
+) -> tuple[Tensor, Tensor]:
+    """Eigenvalues and eigenvectors of ``matrix / total_processed``, computed in
+    fp64 on ``total_processed``'s device and returned there in ``dtype``."""
+    matrix_normalized = matrix.to(total_processed.device, torch.float64)
+    # Free the caller's copy early when it passed the only reference.
+    del matrix
+    matrix_normalized = matrix_normalized + matrix_normalized.T
+    matrix_normalized.div_(2 * total_processed)
+
+    if not torch.isfinite(matrix_normalized).all():
+        raise ValueError(
+            f"Covariance matrix for {key} contains NaNs or Infs. "
+            "Consider using fp32."
+        )
+
+    try:
+        eigenvalues, eigenvectors = torch.linalg.eigh(matrix_normalized)
+    except Exception as e:
+        raise RuntimeError(f"Eigendecomposition failed for {key}") from e
+
+    del matrix_normalized
+    return eigenvalues.to(dtype), eigenvectors.to(dtype).contiguous()
+
+
+def eigendecompose_owned(
+    covariances: dict[str, Tensor],
+    shapes: dict[str, tuple[int, ...]],
+    owners: dict[str, int] | None,
+    total_processed: int | Tensor,
+    output_path: str,
+    dtype: torch.dtype,
+) -> tuple[dict[str, Tensor], dict[str, Tensor]]:
+    """Eigendecompose the covariances this rank owns, in full, from memory.
+
+    ``owners`` maps every module in ``shapes`` to its rank, or is ``None`` in a
+    single process. Saves this rank's row shard of every module's eigenvectors
+    to ``output_path`` in ``dtype``, the layout :func:`compute_eigendecomposition`
+    writes.
+
+    Returns the eigenvectors of the modules this rank owns, in full, on this
+    rank's device, and this rank's row shard of every module's eigenvalues on
+    the CPU. Empties ``covariances`` as it goes, so each covariance is freed once
+    its eigenvectors exist.
+    """
+    rank = dist.get_rank() if owners is not None else 0
+    device = get_device(rank)
+    total_processed = torch.as_tensor(total_processed, device=device)
+
+    eigenvectors: dict[str, Tensor] = {}
+    eigenvalues: dict[str, Tensor] = {}
+    for key in tqdm(
+        list(covariances),
+        desc=f"Rank {rank}: Computing eigenvectors",
+        position=rank,
+        leave=False,
+    ):
+        values, vectors = _eigh(key, covariances.pop(key), total_processed, dtype)
+        eigenvectors[key] = vectors
+        eigenvalues[key] = values
+
+    shards = eigenvectors
+    value_shards = {key: values.cpu() for key, values in eigenvalues.items()}
+    if owners is not None:
+        shards = owned_to_row_shards(eigenvectors, shapes, owners, dtype, device)
+        value_shards = owned_to_row_shards(
+            eigenvalues,
+            {key: shape[:1] for key, shape in shapes.items()},
+            owners,
+            dtype,
+            device,
+        )
+
+    os.makedirs(output_path, exist_ok=True)
+    save_file(shards, os.path.join(output_path, f"shard_{rank}.safetensors"))
+    get_logger().info(f"Saved eigenvectors to {output_path}")
+    return eigenvectors, value_shards
+
+
 def compute_eigendecomposition(
     covariance_path: str,
     total_processed: int | Tensor,
@@ -383,30 +482,10 @@ def compute_eigendecomposition(
             world_size=world_size,
         )
 
-        # original_dtype = matrix.dtype
-        matrix_normalized = matrix.to(torch.float64) / total_processed
-        matrix_normalized = (matrix_normalized + matrix_normalized.T).div(2)
-
-        if not torch.isfinite(matrix_normalized).all():
-            raise ValueError(
-                f"Covariance matrix for {key} contains NaNs or Infs. "
-                "Consider using fp32."
-            )
-
-        try:
-            eigenvalues, eigenvectors = torch.linalg.eigh(matrix_normalized)
-        except Exception as e:
-            raise RuntimeError(f"Eigendecomposition failed for {key}") from e
-
-        # TODO: Maybe possible to avoid CPU transfer here?
-        eigenvectors = eigenvectors.to(original_dtype).to(device="cpu").contiguous()
-        covariance_eigenvectors[key] = eigenvectors
-        covariance_eigenvalues[key] = (
-            eigenvalues.to(original_dtype).to(device="cpu").contiguous()
-        )
-        covariance_eigenvalues[key] = (
-            eigenvalues.to(original_dtype).to(device="cpu").contiguous()
-        )
+        eigenvalues, eigenvectors = _eigh(key, matrix, total_processed, original_dtype)
+        del matrix
+        covariance_eigenvectors[key] = eigenvectors.cpu()
+        covariance_eigenvalues[key] = eigenvalues.cpu()
 
     covariance_eigenvectors = _gather_and_shard_along_dim_0(
         input_dict=covariance_eigenvectors,

{bergson-2.2.0 → bergson-2.2.2}/bergson/hessians/hessian_approximations.py

@@ -9,6 +9,7 @@ import torch.distributed as dist
 from datasets import Dataset
 from safetensors import safe_open
 from safetensors.torch import save_file
+from torch import Tensor
 from transformers import PreTrainedModel
 
 from bergson.collector.collector import (
@@ -27,6 +28,7 @@ from bergson.hessians.autocorrelation import (
 from bergson.hessians.eigenvectors import (
     LambdaCollector,
     compute_eigendecomposition,
+    eigendecompose_owned,
     save_uncorrected_eigenvalues,
 )
 from bergson.hessians.kfac import CovarianceCollector
@@ -315,7 +317,11 @@ def fit_factored_hessians(
         "path": path,
     }
 
-    collect_hessians(**kwargs)
+    collector = collect_hessians(**kwargs)
+    if not do_eigendecomposition and isinstance(collector, CovarianceCollector):
+        # Only the eigendecomposition reads the covariances kept in memory.
+        collector.A_cov_dict.clear()
+        collector.S_cov_dict.clear()
     _release_device_memory()
 
     dist.barrier() if dist.is_initialized() else None
@@ -332,14 +338,35 @@ def fit_factored_hessians(
         weights_only=False,
     )
 
-    eigenvalues_a = compute_eigendecomposition(
-        os.path.join(path, "activation_sharded"),
-        total_processed=total_processed,
-    )
-    eigenvalues_g = compute_eigendecomposition(
-        os.path.join(path, "gradient_sharded"),
-        total_processed=total_processed,
-    )
+    eigenvectors = None
+    if isinstance(collector, CovarianceCollector):
+        # Each rank still holds the covariances it owns, so skip the files.
+        eigenvectors_a, eigenvalues_a = eigendecompose_owned(
+            collector.A_cov_dict,
+            collector.A_shapes,
+            collector.owners,
+            total_processed,
+            os.path.join(path, "eigen_activation_sharded"),
+            collector.dtype,
+        )
+        eigenvectors_g, eigenvalues_g = eigendecompose_owned(
+            collector.S_cov_dict,
+            collector.S_shapes,
+            collector.owners,
+            total_processed,
+            os.path.join(path, "eigen_gradient_sharded"),
+            collector.dtype,
+        )
+        eigenvectors = (eigenvectors_a, eigenvectors_g)
+    else:
+        eigenvalues_a = compute_eigendecomposition(
+            os.path.join(path, "activation_sharded"),
+            total_processed=total_processed,
+        )
+        eigenvalues_g = compute_eigendecomposition(
+            os.path.join(path, "gradient_sharded"),
+            total_processed=total_processed,
+        )
 
     dist.barrier() if dist.is_initialized() else None
 
@@ -354,7 +381,12 @@ def fit_factored_hessians(
     )
 
     if hessian_cfg.ev_correction:
-        collect_hessians(**kwargs, ev_correction=True, num_documents=len(data))
+        collect_hessians(
+            **kwargs,
+            ev_correction=True,
+            num_documents=len(data),
+            eigenvectors=eigenvectors,
+        )
         _release_device_memory()
 
 
@@ -380,11 +412,13 @@ def collect_hessians(
     output_subdir: str = "eigenvalue_correction_sharded",
     path: str | None = None,
     num_documents: int = 1,
-):
+    eigenvectors: tuple[dict[str, Tensor], dict[str, Tensor]] | None = None,
+) -> HookCollectorBase:
     """
     Compute Hessian approximations using the hooks specified in the collector.
     If ev_correction is True, uses LambdaCollector to compute eigenvalue corrections.
-    ``path`` overrides where the collector writes.
+    ``path`` overrides where the collector writes. ``eigenvectors`` are passed to
+    the LambdaCollector. Returns the finished collector.
     """
 
     hessian_dtype = convert_precision_to_torch(hessian_cfg.hessian_dtype)
@@ -405,6 +439,7 @@ def collect_hessians(
             eigen_path=eigen_path,
             output_subdir=output_subdir,
             num_documents=num_documents,
+            eigenvectors=eigenvectors,
         )
         desc += " (eigenvalue correction)"
     else:
@@ -421,3 +456,4 @@ def collect_hessians(
     computer.forward_backward = fwd_bwd_hessian_factory(index_cfg, hessian_cfg)
 
     computer.run_with_collector_hooks(desc=desc)
+    return collector

{bergson-2.2.0 → bergson-2.2.2}/bergson/hessians/kfac.py

@@ -31,6 +31,8 @@ class CovarianceCollector(HookCollectorBase):
 
     Distributed, each module's covariances belong to one rank, which receives
     every rank's positions for that module; teardown saves the usual row shards.
+    Teardown leaves the covariances this rank owns, in full, on its device in
+    ``A_cov_dict`` and ``S_cov_dict`` for the eigendecomposition.
     """
 
     dtype: torch.dtype
@@ -152,4 +154,4 @@ class CovarianceCollector(HookCollectorBase):
             else:
                 shards = covariances
             save_file(shards, os.path.join(path, f"shard_{self.rank}.safetensors"))
-            covariances.clear()
+            del shards

{bergson-2.2.0 → bergson-2.2.2}/bergson/hessians/sharded_computation.py

@@ -74,19 +74,31 @@ def owned_to_row_shards(
     device: str | torch.device,
 ) -> dict[str, Tensor]:
     """This rank's row shard of every matrix, on the CPU, from ``owned``, the
-    matrices this rank owns in full. One matrix at a time is sent from its owner."""
+    matrices this rank owns in full. One matrix at a time is sent from its owner,
+    which sends each rank only its rows."""
     rank, world_size = dist.get_rank(), dist.get_world_size()
     shards = {}
     for name, shape in shapes.items():
         owner = owners[name]
+        # Rank 0's shard is the largest; the others are padded to its size so
+        # every rank receives the same shape.
+        _, rows = shard_bounds(shape[0], 0, world_size)
+        received = torch.empty((rows, *shape[1:]), device=device, dtype=dtype)
+        blocks = None
         if rank == owner:
-            full = owned[name].contiguous()
-        else:
-            full = torch.empty(shape, device=device, dtype=dtype)
-        dist.broadcast(full, src=owner)
+            full = owned[name].to(device, dtype).contiguous()
+            blocks = []
+            for r in range(world_size):
+                block = full[slice(*shard_bounds(shape[0], r, world_size))]
+                if len(block) < rows:
+                    block = torch.cat(
+                        [block, block.new_zeros(rows - len(block), *shape[1:])]
+                    )
+                blocks.append(block)
+        dist.scatter(received, blocks, src=owner)
         start, end = shard_bounds(shape[0], rank, world_size)
-        shards[name] = full[start:end].cpu()
-        del full
+        shards[name] = received[: end - start].cpu()
+        del received, blocks
     return shards