bergson 0.0.1__tar.gz

bergson-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 EleutherAI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

bergson-0.0.1/PKG-INFO

@@ -0,0 +1,158 @@
+Metadata-Version: 2.4
+Name: bergson
+Version: 0.0.1
+Summary: Tracing the memory of neural nets with data attribution
+License: MIT License
+Keywords: interpretability,explainable-ai
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: accelerate
+Requires-Dist: datasets
+Requires-Dist: natsort
+Requires-Dist: peft>=0.17.0
+Requires-Dist: simple-parsing
+Requires-Dist: torch
+Requires-Dist: transformers
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: trl; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Provides-Extra: example
+Requires-Dist: trl; extra == "example"
+Provides-Extra: faiss
+Requires-Dist: faiss-gpu-cu12; extra == "faiss"
+Dynamic: license-file
+
+# Bergson
+This library enables you to trace the memory of deep neural nets with gradient-based data attribution techniques. We currently focus on TrackStar, as described in [Scalable Influence and Fact Tracing for Large Language Model Pretraining](https://arxiv.org/abs/2410.17413v3) by Chang et al. (2024), although we plan to add support for other methods inspired by influence functions in the near future.
+
+We view attribution as a counterfactual question: **_If we "unlearned" this training sample, how would the model's behavior change?_** This formulation ties attribution to some notion of what it means to "unlearn" a training sample. Here we focus on a very simple notion of unlearning: taking a gradient _ascent_ step on the loss with respect to the training sample. To mimic the behavior of popular optimizers, we precondition the gradient using Adam or Adafactor-style estimates of the second moments of the gradient.
+
+# Announcements
+
+**September 2025**
+- Saving per-head gradients: https://github.com/EleutherAI/bergson/pull/40
+- Eigendecompositions of preconditioners: https://github.com/EleutherAI/bergson/pull/34
+- Dr. GRPO-based loss gradients: https://github.com/EleutherAI/bergson/pull/35
+- Choosing between summing and averaging losses across tokens: https://github.com/EleutherAI/bergson/pull/36
+- Saving the order training data is seen in while using the gradient collector callback for HF's Trainer/SFTTrainer: https://github.com/EleutherAI/bergson/pull/40
+  - Saving training gradients adds a ~17% wall clock overhead
+- Improved static index build ETA accuracy: https://github.com/EleutherAI/bergson/pull/41
+- Several small quality of life improvements for querying indexes: https://github.com/EleutherAI/bergson/pull/38
+
+# Installation
+
+We're not yet on PyPI, but you can `git clone` the repo and install it as a package using pip:
+
+```bash
+git clone https://github.com/EleutherAI/bergson.git
+cd bergson
+pip install .
+```
+
+# Usage
+The first step is to build an index of gradients for each training sample. You can do this from the command line, using `bergson` as a CLI tool:
+
+```bash
+bergson <output_path> --model <model_name> --dataset <dataset_name>
+```
+
+This will create a directory at `<output_path>` containing the gradients for each training sample in the specified dataset. The `--model` and `--dataset` arguments should be compatible with the Hugging Face `transformers` library. By default it assumes that the dataset has a `text` column, but you can specify other columns using `--prompt_column` and optionally `--completion_column`. The `--help` flag will show you all available options.
+
+You can also use the library programmatically to build the index. The `collect_gradients` function is just a bit lower level the CLI tool, and allows you to specify the model and dataset directly as arguments. The result is a HuggingFace dataset which contains a handful of new columns, including `gradients`, which contains the gradients for each training sample. You can then use this dataset to compute attributions.
+
+At the lowest level of abstraction, the `GradientCollector` context manager allows you to efficiently collect gradients for _each individual example_ in a batch during a backward pass, simultaneously randomly projecting the gradients to a lower-dimensional space to save memory. If you use Adafactor normalization, which is the default, we will do this in a very compute-efficient way which avoids computing the full gradient for each example before projecting it to the lower dimension. There are two main ways you can use `GradientCollector`:
+
+1. Using a `closure` argument, which enables you to make use of the per-example gradients immediately after they are computed, during the backward pass. If you're computing summary statistics or other per-example metrics, this is the most efficient way to do it.
+2. Without a `closure` argument, in which case the gradients are collected and returned as a dictionary mapping module names to batches of gradients. This is the simplest and most flexible approach but is a bit more memory-intensive.
+
+## Training Gradients
+
+Gradient collection during training is supported via an integration with HuggingFace's Trainer and SFTTrainer classes. Training gradients are saved in the original order corresponding to their dataset items, and when the `track_order` flag is set the training steps associated with each training item are separately saved.
+
+```python
+from bergson import GradientCollectorCallback, prepare_for_gradient_collection
+
+callback = GradientCollectorCallback(
+    path="runs/example",
+    track_order=True,
+    use_optimizer_state=False,
+)
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset,
+    eval_dataset=dataset,
+    callbacks=[callback],
+)
+trainer = prepare_for_gradient_collection(trainer)
+trainer.train()
+```
+
+## Attention Head Gradients
+
+By default Bergson collects gradients for named parameter matrices, but gradients for individual attention heads within a named matrix can be collected too. To collect head gradients add a `head_cfgs` dictionary to the training calllback or static index config.
+
+```python
+from bergson import HeadConfig, IndexConfig, DataConfig
+from transformers import AutoModelForCausalLM
+
+model = AutoModelForCausalLM.from_pretrained("RonenEldan/TinyStories-1M", trust_remote_code=True, use_safetensors=True)
+
+collect_gradients(
+    model=model,
+    data=data,
+    processor=processor,
+    path="runs/example_with_heads",
+    head_cfgs={
+        # Head configuration for the TinyStories-1M transformer
+        "h.0.attn.attention.out_proj": HeadConfig(num_heads=16, head_size=4, head_dim=2),
+    },
+)
+```
+
+## GRPO
+
+Where a reward signal is available we compute gradients using a weighted advantage estimate based on Dr. GRPO:
+
+```bash
+bergson <output_path> --model <model_name> --dataset <dataset_name> --reward_column <reward_column_name>
+```
+
+## Queries
+
+We provide a query Attributor which supports unit normalized gradients and KNN search out of the box.
+
+```
+from bergson import Attributor, FaissConfig
+
+attr = Attributor(args.index, device="cuda")
+
+...
+query_tokens = tokenizer(query, return_tensors="pt").to("cuda:0")["input_ids"]
+
+# Query the index
+with attr.trace(model.base_model, 5) as result:
+    model(query_tokens, labels=query_tokens).loss.backward()
+    model.zero_grad()
+```
+
+To efficiently query on-disk indexes, perform ANN searches, and explore many other scalability features add a FAISS config:
+
+```
+attr = Attributor(args.index, device="cuda", faiss_cfg=FaissConfig("IVF1,SQfp16", mmap_index=True))
+
+with attr.trace(model.base_model, 5) as result:
+    model(query_tokens, labels=query_tokens).loss.backward()
+    model.zero_grad()
+```
+
+# Development
+
+```bash
+pip install -e .[dev]
+pytest
+```
+
+We use [conventional commits](https://www.conventionalcommits.org/en/v1.0.0/) for releases.

bergson-0.0.1/README.md

@@ -0,0 +1,132 @@
+# Bergson
+This library enables you to trace the memory of deep neural nets with gradient-based data attribution techniques. We currently focus on TrackStar, as described in [Scalable Influence and Fact Tracing for Large Language Model Pretraining](https://arxiv.org/abs/2410.17413v3) by Chang et al. (2024), although we plan to add support for other methods inspired by influence functions in the near future.
+
+We view attribution as a counterfactual question: **_If we "unlearned" this training sample, how would the model's behavior change?_** This formulation ties attribution to some notion of what it means to "unlearn" a training sample. Here we focus on a very simple notion of unlearning: taking a gradient _ascent_ step on the loss with respect to the training sample. To mimic the behavior of popular optimizers, we precondition the gradient using Adam or Adafactor-style estimates of the second moments of the gradient.
+
+# Announcements
+
+**September 2025**
+- Saving per-head gradients: https://github.com/EleutherAI/bergson/pull/40
+- Eigendecompositions of preconditioners: https://github.com/EleutherAI/bergson/pull/34
+- Dr. GRPO-based loss gradients: https://github.com/EleutherAI/bergson/pull/35
+- Choosing between summing and averaging losses across tokens: https://github.com/EleutherAI/bergson/pull/36
+- Saving the order training data is seen in while using the gradient collector callback for HF's Trainer/SFTTrainer: https://github.com/EleutherAI/bergson/pull/40
+  - Saving training gradients adds a ~17% wall clock overhead
+- Improved static index build ETA accuracy: https://github.com/EleutherAI/bergson/pull/41
+- Several small quality of life improvements for querying indexes: https://github.com/EleutherAI/bergson/pull/38
+
+# Installation
+
+We're not yet on PyPI, but you can `git clone` the repo and install it as a package using pip:
+
+```bash
+git clone https://github.com/EleutherAI/bergson.git
+cd bergson
+pip install .
+```
+
+# Usage
+The first step is to build an index of gradients for each training sample. You can do this from the command line, using `bergson` as a CLI tool:
+
+```bash
+bergson <output_path> --model <model_name> --dataset <dataset_name>
+```
+
+This will create a directory at `<output_path>` containing the gradients for each training sample in the specified dataset. The `--model` and `--dataset` arguments should be compatible with the Hugging Face `transformers` library. By default it assumes that the dataset has a `text` column, but you can specify other columns using `--prompt_column` and optionally `--completion_column`. The `--help` flag will show you all available options.
+
+You can also use the library programmatically to build the index. The `collect_gradients` function is just a bit lower level the CLI tool, and allows you to specify the model and dataset directly as arguments. The result is a HuggingFace dataset which contains a handful of new columns, including `gradients`, which contains the gradients for each training sample. You can then use this dataset to compute attributions.
+
+At the lowest level of abstraction, the `GradientCollector` context manager allows you to efficiently collect gradients for _each individual example_ in a batch during a backward pass, simultaneously randomly projecting the gradients to a lower-dimensional space to save memory. If you use Adafactor normalization, which is the default, we will do this in a very compute-efficient way which avoids computing the full gradient for each example before projecting it to the lower dimension. There are two main ways you can use `GradientCollector`:
+
+1. Using a `closure` argument, which enables you to make use of the per-example gradients immediately after they are computed, during the backward pass. If you're computing summary statistics or other per-example metrics, this is the most efficient way to do it.
+2. Without a `closure` argument, in which case the gradients are collected and returned as a dictionary mapping module names to batches of gradients. This is the simplest and most flexible approach but is a bit more memory-intensive.
+
+## Training Gradients
+
+Gradient collection during training is supported via an integration with HuggingFace's Trainer and SFTTrainer classes. Training gradients are saved in the original order corresponding to their dataset items, and when the `track_order` flag is set the training steps associated with each training item are separately saved.
+
+```python
+from bergson import GradientCollectorCallback, prepare_for_gradient_collection
+
+callback = GradientCollectorCallback(
+    path="runs/example",
+    track_order=True,
+    use_optimizer_state=False,
+)
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset,
+    eval_dataset=dataset,
+    callbacks=[callback],
+)
+trainer = prepare_for_gradient_collection(trainer)
+trainer.train()
+```
+
+## Attention Head Gradients
+
+By default Bergson collects gradients for named parameter matrices, but gradients for individual attention heads within a named matrix can be collected too. To collect head gradients add a `head_cfgs` dictionary to the training calllback or static index config.
+
+```python
+from bergson import HeadConfig, IndexConfig, DataConfig
+from transformers import AutoModelForCausalLM
+
+model = AutoModelForCausalLM.from_pretrained("RonenEldan/TinyStories-1M", trust_remote_code=True, use_safetensors=True)
+
+collect_gradients(
+    model=model,
+    data=data,
+    processor=processor,
+    path="runs/example_with_heads",
+    head_cfgs={
+        # Head configuration for the TinyStories-1M transformer
+        "h.0.attn.attention.out_proj": HeadConfig(num_heads=16, head_size=4, head_dim=2),
+    },
+)
+```
+
+## GRPO
+
+Where a reward signal is available we compute gradients using a weighted advantage estimate based on Dr. GRPO:
+
+```bash
+bergson <output_path> --model <model_name> --dataset <dataset_name> --reward_column <reward_column_name>
+```
+
+## Queries
+
+We provide a query Attributor which supports unit normalized gradients and KNN search out of the box.
+
+```
+from bergson import Attributor, FaissConfig
+
+attr = Attributor(args.index, device="cuda")
+
+...
+query_tokens = tokenizer(query, return_tensors="pt").to("cuda:0")["input_ids"]
+
+# Query the index
+with attr.trace(model.base_model, 5) as result:
+    model(query_tokens, labels=query_tokens).loss.backward()
+    model.zero_grad()
+```
+
+To efficiently query on-disk indexes, perform ANN searches, and explore many other scalability features add a FAISS config:
+
+```
+attr = Attributor(args.index, device="cuda", faiss_cfg=FaissConfig("IVF1,SQfp16", mmap_index=True))
+
+with attr.trace(model.base_model, 5) as result:
+    model(query_tokens, labels=query_tokens).loss.backward()
+    model.zero_grad()
+```
+
+# Development
+
+```bash
+pip install -e .[dev]
+pytest
+```
+
+We use [conventional commits](https://www.conventionalcommits.org/en/v1.0.0/) for releases.

bergson-0.0.1/bergson/__init__.py

@@ -0,0 +1,21 @@
+__version__ = "0.0.1"
+
+from .attributor import Attributor
+from .collection import collect_gradients
+from .data import DataConfig, IndexConfig, load_gradients
+from .faiss_index import FaissConfig
+from .gradcheck import FiniteDiff
+from .gradients import GradientCollector, GradientProcessor, HeadConfig
+
+__all__ = [
+    "collect_gradients",
+    "load_gradients",
+    "Attributor",
+    "FaissConfig",
+    "FiniteDiff",
+    "GradientCollector",
+    "GradientProcessor",
+    "IndexConfig",
+    "DataConfig",
+    "HeadConfig",
+]

bergson-0.0.1/bergson/__main__.py

@@ -0,0 +1,12 @@
+from simple_parsing import parse
+
+from .build import build_gradient_dataset
+from .data import IndexConfig
+
+
+def main():
+    build_gradient_dataset(parse(IndexConfig))
+
+
+if __name__ == "__main__":
+    main()

bergson-0.0.1/bergson/attributor.py

@@ -0,0 +1,159 @@
+from collections import defaultdict
+from contextlib import contextmanager
+from typing import Generator
+
+import torch
+from torch import Tensor, nn
+
+from .data import load_gradients
+from .faiss_index import FaissConfig, FaissIndex
+from .gradients import GradientCollector, GradientProcessor
+
+
+class TraceResult:
+    """Result of a .trace() call."""
+
+    def __init__(self):
+        # Should be set by the Attributor after a search
+        self._indices: Tensor | None = None
+        self._scores: Tensor | None = None
+
+    @property
+    def indices(self) -> Tensor:
+        """The indices of the top-k examples."""
+        if self._indices is None:
+            raise ValueError("No indices available. Exit the context manager first.")
+
+        return self._indices
+
+    @property
+    def scores(self) -> Tensor:
+        """The attribution scores of the top-k examples."""
+        if self._scores is None:
+            raise ValueError("No scores available. Exit the context manager first.")
+
+        return self._scores
+
+
+class Attributor:
+    def __init__(
+        self,
+        index_path: str,
+        device: str = "cpu",
+        dtype: torch.dtype = torch.float32,
+        unit_norm: bool = False,
+        faiss_cfg: FaissConfig | None = None,
+    ):
+        self.device = device
+        self.dtype = dtype
+        self.unit_norm = unit_norm
+        self.faiss_index = None
+
+        # Load the gradient processor
+        self.processor = GradientProcessor.load(index_path, map_location=device)
+
+        # Load the gradient index
+        if faiss_cfg:
+            self.faiss_index = FaissIndex(index_path, faiss_cfg, device, unit_norm)
+            self.N = self.faiss_index.ntotal
+        else:
+            mmap = load_gradients(index_path)
+
+            # Copy gradients into device memory
+            self.grads = {
+                name: torch.tensor(mmap[name], device=device, dtype=dtype)
+                for name in mmap.dtype.names
+            }
+            self.N = mmap[mmap.dtype.names[0]].shape[0]
+
+            if unit_norm:
+                norm = torch.cat([grad for grad in self.grads.values()], dim=1).norm(
+                    dim=1, keepdim=True
+                )
+                for name in self.grads:
+                    self.grads[name] /= norm
+
+    def search(
+        self, queries: dict[str, Tensor], k: int, modules: list[str] | None = None
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Search for the `k` nearest examples in the index based on the query or queries.
+
+        Args:
+            queries: The query tensor of shape [..., d].
+            k: The number of nearest examples to return for each query.
+            module: The name of the module to search for. If `None`,
+                all modules will be searched.
+
+        Returns:
+            A namedtuple containing the top `k` indices and inner products for each
+            query. Both have shape [..., k].
+        """
+        q = {name: item.to(self.device, self.dtype) for name, item in queries.items()}
+
+        if self.unit_norm:
+            norm = torch.cat(list(q.values()), dim=1).norm(dim=1, keepdim=True)
+
+            for name in q:
+                q[name] /= norm + 1e-8
+
+        if self.faiss_index:
+            if modules:
+                raise NotImplementedError(
+                    "FAISS index does not implement module-specific search."
+                )
+
+            q = torch.cat([q[name] for name in q], dim=1).cpu().numpy()
+
+            distances, indices = self.faiss_index.search(q, k)
+
+            return torch.from_numpy(distances.squeeze()), torch.from_numpy(
+                indices.squeeze()
+            )
+
+        modules = modules or list(q.keys())
+        k = min(k, self.N)
+
+        scores = torch.stack(
+            [q[name] @ self.grads[name].mT for name in modules], dim=-1
+        ).sum(-1)
+
+        return torch.topk(scores, k)
+
+    @contextmanager
+    def trace(
+        self, module: nn.Module, k: int, *, precondition: bool = False
+    ) -> Generator[TraceResult, None, None]:
+        """
+        Context manager to trace the gradients of a module and return the
+        corresponding Attributor instance.
+        """
+        mod_grads = defaultdict(list)
+        result = TraceResult()
+
+        def callback(name: str, g: Tensor):
+            # Precondition the gradient using Cholesky solve
+            if precondition:
+                eigval, eigvec = self.processor.preconditioners_eigen[name]
+                eigval_inverse_sqrt = 1.0 / (eigval).sqrt()
+                P = eigvec * eigval_inverse_sqrt @ eigvec.mT
+                g = g.flatten(1).type_as(P)
+                g = g @ P
+            else:
+                g = g.flatten(1)
+
+            # Store the gradient for later use
+            mod_grads[name].append(g.to(self.device, self.dtype, non_blocking=True))
+
+        with GradientCollector(module, callback, self.processor):
+            yield result
+
+        if not mod_grads:
+            raise ValueError("No grads collected. Did you forget to call backward?")
+
+        queries = {name: torch.cat(g, dim=1) for name, g in mod_grads.items()}
+
+        if any(q.isnan().any() for q in queries.values()):
+            raise ValueError("NaN found in queries.")
+
+        result._scores, result._indices = self.search(queries, k)