bergson 2.2.2__tar.gz → 2.2.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bergson
-Version: 2.2.2
+Version: 2.2.3
 Summary: Tracing the memory of neural nets with data attribution
 License: MIT License
 Keywords: interpretability,explainable-ai
@@ -97,7 +97,7 @@ Indicative performance of data attribution methods in the finetuning regime - se
 | 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] |
+| TrackStar (no optimizer) | 0.045 [0.036, 0.055] | 0.276 [0.252, 0.299] | 0.492 [0.467, 0.517] |
 | 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] |

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

@@ -36,7 +36,7 @@ Indicative performance of data attribution methods in the finetuning regime - se
 | 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] |
+| TrackStar (no optimizer) | 0.045 [0.036, 0.055] | 0.276 [0.252, 0.299] | 0.492 [0.467, 0.517] |
 | 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] |

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

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

{bergson-2.2.2 → bergson-2.2.3}/bergson/process_autocorrelation.py

@@ -4,6 +4,14 @@ import torch.distributed as dist
 from bergson.gradients import GradientProcessor
 
 
+def _eigh(prec: torch.Tensor, device, dtype) -> tuple[torch.Tensor, torch.Tensor]:
+    eigvals, eigvecs = torch.linalg.eigh(prec.to(dtype=torch.float64, device=device))
+    return (
+        eigvals.to(dtype=dtype).contiguous().cpu(),
+        eigvecs.to(dtype=dtype).contiguous().cpu(),
+    )
+
+
 def process_autocorrelation_matrices(
     processor: GradientProcessor,
     hessians: dict[str, torch.Tensor],
@@ -20,9 +28,11 @@ def process_autocorrelation_matrices(
     count when ``attribute_tokens`` produced per-token rows — so the
     normalized matrix is the second moment over whichever unit the rows
     represent.
-    """
-    hessians_eigen = {}
 
+    A rank may hold a partial Gram over its own rows for every module, or the
+    full Gram for the modules it owns. Each module's Gram is summed onto one
+    of the ranks that hold it, which eigendecomposes the total.
+    """
     device = next(iter(hessians.values())).device
     dtype = next(iter(hessians.values())).dtype
 
@@ -32,57 +42,58 @@ def process_autocorrelation_matrices(
     for name, prec in hessians.items():
         hessians[name] = (prec / num_rows).cpu()
 
-    if rank == 0:
-        print("Computing hessian eigen decompositions...")
-
-    for name in hessians.keys():
-        prec = hessians[name].to(dtype=torch.float64, device=device)
-        eigvals, eigvecs = torch.linalg.eigh(prec)
-        hessians_eigen[name] = (
-            eigvals.to(dtype=dtype).contiguous().cpu(),
-            eigvecs.to(dtype=dtype).contiguous().cpu(),
-        )
-
     if not dist.is_initialized():
+        if rank == 0:
+            print("Computing hessian eigen decompositions...")
         processor.hessians = hessians
-        processor.hessians_eigen = hessians_eigen
+        processor.hessians_eigen = {
+            name: _eigh(prec, device, dtype) for name, prec in hessians.items()
+        }
         return
 
     if rank == 0:
-        print("Gathering hessians...")
+        print("Reducing hessians and computing eigen decompositions...")
 
     cpu_group: dist.ProcessGroup = dist.new_group(backend="gloo")  # type: ignore[assignment]
-
-    for name, grad_size in grad_sizes.items():
+    world_size = dist.get_world_size()
+    held: list[list[str]] = [[] for _ in range(world_size)]
+    dist.all_gather_object(held, sorted(hessians), group=cpu_group)
+
+    totals, hessians_eigen = {}, {}
+    for i, (name, grad_size) in enumerate(grad_sizes.items()):
+        holders = [r for r in range(world_size) if name in held[r]] or [0]
+        owner = holders[i % len(holders)]
         if name in hessians:
-            local_prec = hessians[name]
-            del hessians[name]
+            local_prec = hessians.pop(name)
         else:
             local_prec = torch.zeros([grad_size, grad_size], dtype=dtype, device="cpu")
 
-        dist.reduce(local_prec, dst=0, op=dist.ReduceOp.SUM, group=cpu_group)
+        dist.reduce(local_prec, dst=owner, op=dist.ReduceOp.SUM, group=cpu_group)
 
-        if rank == 0:
-            hessians[name] = local_prec
+        if rank == owner:
+            totals[name] = local_prec
+            hessians_eigen[name] = _eigh(local_prec, device, dtype)
 
     if rank == 0:
-        processor.hessians = hessians
-
-        print("Gathering eigen decompositions...")
+        print("Gathering hessians and eigen decompositions...")
 
     for name, grad_size in grad_sizes.items():
         prec_size = torch.Size([grad_size, grad_size])
-        if name not in hessians_eigen:
+        prec = totals.pop(name, None)
+        if name in hessians_eigen:
+            eigval, eigvec = hessians_eigen[name]
+        else:
+            prec = torch.zeros(prec_size, dtype=dtype)
             eigval = torch.zeros(prec_size[0], dtype=dtype)
             eigvec = torch.zeros(prec_size, dtype=dtype)
-        else:
-            eigval, eigvec = hessians_eigen[name]
 
-        dist.reduce(eigval, dst=0, op=dist.ReduceOp.SUM, group=cpu_group)
-        dist.reduce(eigvec, dst=0, op=dist.ReduceOp.SUM, group=cpu_group)
+        for t in (prec, eigval, eigvec):
+            dist.reduce(t, dst=0, op=dist.ReduceOp.SUM, group=cpu_group)
 
         if rank == 0:
+            hessians[name] = prec
             hessians_eigen[name] = (eigval, eigvec)
 
     if rank == 0:
+        processor.hessians = hessians
         processor.hessians_eigen = hessians_eigen

{bergson-2.2.2 → bergson-2.2.3}/bergson/utils/worker_utils.py

@@ -217,13 +217,11 @@ def setup_model_and_peft(
     model_kwargs.update(simple_parse_kwargs_string(cfg.model_kwargs))
 
     if meta_init:
-        config = AutoConfig.from_pretrained(base_model_path, revision=cfg.revision)
+        config = AutoConfig.from_pretrained(
+            base_model_path, revision=cfg.revision, **model_kwargs
+        )
         with torch.device("meta"):
-            model = AutoModelForCausalLM.from_config(
-                config,
-                dtype=dtype,
-                **model_kwargs,
-            )
+            model = AutoModelForCausalLM.from_config(config, dtype=dtype)
     else:
         model = AutoModelForCausalLM.from_pretrained(
             base_model_path,

{bergson-2.2.2 → bergson-2.2.3}/bergson.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bergson
-Version: 2.2.2
+Version: 2.2.3
 Summary: Tracing the memory of neural nets with data attribution
 License: MIT License
 Keywords: interpretability,explainable-ai
@@ -97,7 +97,7 @@ Indicative performance of data attribution methods in the finetuning regime - se
 | 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] |
+| TrackStar (no optimizer) | 0.045 [0.036, 0.055] | 0.276 [0.252, 0.299] | 0.492 [0.467, 0.517] |
 | 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] |

{bergson-2.2.2 → bergson-2.2.3}/bergson.egg-info/SOURCES.txt

@@ -145,6 +145,7 @@ tests/test_padding.py
 tests/test_per_query_magic.py
 tests/test_per_token_lds.py
 tests/test_pretokenized.py
+tests/test_process_autocorrelation.py
 tests/test_projection_inner_products.py
 tests/test_projection_matrix.py
 tests/test_projection_settings.py

{bergson-2.2.2 → bergson-2.2.3}/pyproject.toml

@@ -37,7 +37,7 @@ dependencies = [
     "bitsandbytes",
     "huggingface-hub>=1.13.0",
 ]
-version = "2.2.2"
+version = "2.2.3"
 [project.optional-dependencies]
 dev = [
     "pre-commit",
@@ -93,6 +93,9 @@ reportPrivateImportUsage = false
 [tool.setuptools.packages.find]
 include = ["bergson*"]
 
+[tool.uv.pip]
+torch-backend = "auto"
+
 [tool.black]
 target-version = ["py310", "py311", "py312"]
 

bergson-2.2.3/tests/test_process_autocorrelation.py

@@ -0,0 +1,62 @@
+"""The eigendecomposition saved across ranks must be that of the summed Gram."""
+
+import socket
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from bergson.gradients import GradientProcessor
+from bergson.process_autocorrelation import process_autocorrelation_matrices
+
+SIZES = {"a": 5, "b": 3, "c": 4}
+NUM_ROWS = 12
+
+
+def rank_rows(rank: int, name: str) -> torch.Tensor:
+    gen = torch.Generator().manual_seed(100 * rank + ord(name))
+    return torch.randn(NUM_ROWS // 2, SIZES[name], generator=gen, dtype=torch.float64)
+
+
+def full_gram(name: str, world_size: int) -> torch.Tensor:
+    rows = torch.cat([rank_rows(r, name) for r in range(world_size)])
+    return rows.mT @ rows / NUM_ROWS
+
+
+def run(rank: int, world_size: int, port: int, owned: bool, out: dict):
+    dist.init_process_group(
+        "gloo", init_method=f"tcp://localhost:{port}", rank=rank, world_size=world_size
+    )
+    if owned:
+        # Each rank holds the full Gram of the modules it owns.
+        names = [n for i, n in enumerate(SIZES) if i % world_size == rank]
+        hessians = {n: full_gram(n, world_size) * NUM_ROWS for n in names}
+    else:
+        # Each rank holds a partial Gram over its own rows for every module.
+        hessians = {n: rank_rows(rank, n).mT @ rank_rows(rank, n) for n in SIZES}
+
+    processor = GradientProcessor()
+    process_autocorrelation_matrices(processor, hessians, NUM_ROWS, SIZES, rank)
+    if rank == 0:
+        out.update(
+            hessians=dict(processor.hessians), eigen=dict(processor.hessians_eigen)
+        )
+    dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("owned", [False, True])
+def test_eigen_matches_summed_gram(owned: bool):
+    world_size = 2
+    with socket.socket() as s:
+        s.bind(("", 0))
+        port = s.getsockname()[1]
+    out = mp.Manager().dict()
+    mp.spawn(run, args=(world_size, port, owned, out), nprocs=world_size)
+
+    for name in SIZES:
+        expected = full_gram(name, world_size)
+        torch.testing.assert_close(out["hessians"][name], expected)
+        eigvals, eigvecs = out["eigen"][name]
+        torch.testing.assert_close(eigvals, torch.linalg.eigvalsh(expected))
+        torch.testing.assert_close(eigvecs @ torch.diag(eigvals) @ eigvecs.mT, expected)