bergson 2.2.3__tar.gz → 2.2.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bergson
-Version: 2.2.3
+Version: 2.2.4
 Summary: Tracing the memory of neural nets with data attribution
 License: MIT License
 Keywords: interpretability,explainable-ai
@@ -91,16 +91,16 @@ Indicative performance of data attribution methods in the finetuning regime - se
 | 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] |
+| EKFAC + 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] |
+| EKFAC | 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] |
+| Activation similarity (activation covariance preconditioner) | 0.055 [0.043, 0.069] | 0.157 [0.120, 0.192] | 0.073 [0.056, 0.092] |
 | 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) | 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] |
 
 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.
 

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

@@ -30,16 +30,16 @@ Indicative performance of data attribution methods in the finetuning regime - se
 | 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] |
+| EKFAC + 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] |
+| EKFAC | 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] |
+| Activation similarity (activation covariance preconditioner) | 0.055 [0.043, 0.069] | 0.157 [0.120, 0.192] | 0.073 [0.056, 0.092] |
 | 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) | 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] |
 
 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.
 

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

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

{bergson-2.2.3 → bergson-2.2.4}/bergson/build.py

@@ -84,12 +84,20 @@ def build_worker(
     model, target_modules = setup_model_and_peft(index_cfg)
     processor = create_processor(model, index_cfg, target_modules)
 
-    maybe_auto_batch_size(index_cfg, model, ds, processor, target_modules, rank)
-
     attention_cfgs = {
         module: index_cfg.attention for module in index_cfg.split_attention_modules
     }
 
+    maybe_auto_batch_size(
+        index_cfg,
+        model,
+        ds,
+        processor,
+        target_modules,
+        attention_cfgs,
+        preprocess_cfg,
+    )
+
     kwargs = {
         "model": model,
         "data": ds,

{bergson-2.2.3 → bergson-2.2.4}/bergson/hessians/shampoo.py

@@ -15,14 +15,15 @@ from bergson.utils.utils import assert_type
 @dataclass(kw_only=True)
 class ShampooCollector(HookCollectorBase):
     """
-    Collects activation and gradient covariances for TKFAC.
+    Collects activation and gradient covariances for Shampoo.
 
     Computes:
-        A_shampoo = sum over batches of Grad @ Grad.T  for activations
-        S_shampoo = sum over batches of Grad.T @ Grad for gradients * trace(A_cov)
+        A_shampoo = sum over batches of (Grad.T @ Grad) / trace  for activations
+        S_shampoo = sum over batches of (Grad @ Grad.T)  for gradients
 
-
-    where X is input activations [N*S, I] and G is output gradients [N*S, O].
+    where Grad = G.T @ X is the batch gradient [O, I], X is input activations
+    [N*S, I], G is output gradients [N*S, O], and trace is the trace of the
+    summed activation factor.
     """
 
     dtype: torch.dtype
@@ -42,7 +43,7 @@ class ShampooCollector(HookCollectorBase):
         )
 
     def forward_hook(self, module: nn.Module, a: Tensor) -> None:
-        """Compute activation covariance: A^T @ A."""
+        """Store the input activations for the backward hook."""
 
         mask = self.collection_mask(module)
         assert mask is not None, "Collection mask not set for forward hook."
@@ -60,7 +61,7 @@ class ShampooCollector(HookCollectorBase):
         module._inputs = a_bi
 
     def backward_hook(self, module: nn.Module, g: Tensor) -> None:
-        """Compute gradient covariance: G^T @ G."""
+        """Accumulate gradient and activation covariances, A.T @ A and G.T @ G."""
         name = assert_type(str, module._name)
         S_shampoo_po = self.S_shampoo_dict[name]
         A_shampoo_ki = self.A_shampoo_dict[name]

bergson-2.2.4/bergson/utils/batch_size.py

@@ -0,0 +1,260 @@
+import dataclasses
+import gc
+import hashlib
+import json
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from datasets import Dataset
+from peft import PeftModel
+from transformers import AutoTokenizer, PreTrainedModel
+
+from bergson.collector.collector import fwd_bwd_factory
+from bergson.collector.gradient_collectors import GradientCollector
+from bergson.config.config import AttentionConfig, IndexConfig, PreprocessConfig
+from bergson.data import _allocate_batches_world, pad_and_tensor
+from bergson.gradients import GradientProcessor
+from bergson.utils.worker_utils import max_tokens_for_model
+
+CACHE_PATH = Path(".cache") / "batch_size_cache.jsonl"
+
+PROBE_MEMORY_FRACTION = 0.9
+"""Share of device memory the probe batches may use, leaving the rest for the
+fragmentation a full run builds up."""
+
+UNKEYED_FIELDS = {
+    "run_path",
+    "overwrite",
+    "token_batch_size",
+    "auto_batch_size",
+    "distributed",
+    "debug",
+    "profile",
+}
+"""IndexConfig fields that leave a run's memory unchanged."""
+
+
+def maybe_auto_batch_size(
+    cfg: IndexConfig,
+    model: PreTrainedModel | PeftModel,
+    ds: Dataset,
+    processor: GradientProcessor,
+    target_modules: set[str] | None,
+    attention_cfgs: dict[str, AttentionConfig],
+    preprocess_cfg: PreprocessConfig,
+) -> None:
+    """Set ``cfg.token_batch_size`` to the largest budget whose batches fit, if
+    ``cfg.auto_batch_size``.
+
+    Every rank probes the same batches and the ranks agree on each result, so
+    sharded models and multi-node runs search together.
+    """
+    if not cfg.auto_batch_size:
+        return
+
+    lengths = [int(n) for n in ds["length"][:]]
+    key = _cache_key(cfg, preprocess_cfg, lengths)
+    rank = dist.get_rank() if dist.is_initialized() else 0
+
+    found = [_check_cache(CACHE_PATH, key) if rank == 0 else None]
+    if dist.is_initialized():
+        dist.broadcast_object_list(found, src=0)
+    if found[0] is not None:
+        cfg.token_batch_size = found[0]
+        print(f"Loaded token_batch_size from cache: {found[0]}")
+        return
+
+    tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer or cfg.model)
+    max_len = max_tokens_for_model(tokenizer, cfg.model, cfg.revision)
+    collector = GradientCollector(
+        model=model.base_model,  # type: ignore
+        cfg=cfg,
+        processor=processor,
+        target_modules=target_modules,
+        data=ds,
+        preprocess_cfg=preprocess_cfg,
+        attention_cfgs=attention_cfgs,
+        filter_modules=cfg.filter_modules,
+    )
+    try:
+        cfg.token_batch_size = determine_batch_size(
+            cfg, model, ds, collector, lengths, max_len
+        )
+    finally:
+        del collector
+        _clear_cache()
+
+    if rank == 0:
+        _append_to_cache(CACHE_PATH, key, cfg.token_batch_size)
+
+
+def determine_batch_size(
+    cfg: IndexConfig,
+    model: PreTrainedModel | PeftModel,
+    ds: Dataset,
+    collector: GradientCollector,
+    lengths: list[int],
+    max_len: int,
+) -> int:
+    """Find the largest token budget whose batches fit, doubling from the
+    longest document and then bisecting."""
+    # Documents longer than the budget can't be batched, and a budget past
+    # the context length is rejected when the data is tokenized
+    lo = max(lengths)
+    limit = max(lo, min(max_len, sum(lengths)))
+    print(f"Determining token_batch_size in [{lo}, {limit}]...")
+    if not _fits(cfg, model, ds, collector, lengths, lo):
+        raise RuntimeError(
+            f"token_batch_size={lo}, the longest document, does not fit in memory."
+        )
+
+    hi = None
+    while lo < limit:
+        candidate = min(2 * lo, limit)
+        if not _fits(cfg, model, ds, collector, lengths, candidate):
+            hi = candidate
+            break
+        lo = candidate
+
+    # lo fits and hi, when set, does not
+    while hi is not None and hi - lo > max(256, lo // 16):
+        mid = (lo + hi) // 2
+        if _fits(cfg, model, ds, collector, lengths, mid):
+            lo = mid
+        else:
+            hi = mid
+
+    print(f"Optimal token_batch_size: {lo}")
+    return lo
+
+
+def probe_batches(
+    lengths: list[int], N: int, world_size: int, max_batch_size: int | None
+) -> list[list[int]]:
+    """The batches of budget ``N`` with the most padded tokens, the most
+    documents and the most attention entries, over every rank's batches."""
+    batches = [
+        b
+        for rank in _allocate_batches_world(
+            lengths, N, world_size, max_batch_size=max_batch_size
+        )
+        for b in rank
+    ]
+
+    def shape(b):
+        return len(b), max(lengths[i] for i in b)
+
+    keys = (
+        lambda b: shape(b)[0] * shape(b)[1],
+        lambda b: shape(b)[0],
+        lambda b: shape(b)[0] * shape(b)[1] ** 2,
+    )
+    return [max(batches, key=key) for key in keys]
+
+
+def _fits(
+    cfg: IndexConfig,
+    model: PreTrainedModel | PeftModel,
+    ds: Dataset,
+    collector: GradientCollector,
+    lengths: list[int],
+    N: int,
+) -> bool:
+    """Whether every rank runs the probe batches of budget ``N`` in memory."""
+    world_size = dist.get_world_size() if dist.is_initialized() else 1
+    device = torch.device(model.device)  # type: ignore[attr-defined]
+    fwd_bwd = fwd_bwd_factory(cfg)
+
+    print(f"  Testing {N}...", end=" ", flush=True)
+    fits = True
+    _clear_cache()
+    if device.type == "cuda":
+        torch.cuda.set_per_process_memory_fraction(PROBE_MEMORY_FRACTION, device)
+    try:
+        for indices in probe_batches(lengths, N, world_size, cfg.max_batch_size):
+            batch = ds[indices]
+            x, y, _, collection_mask = pad_and_tensor(
+                batch["input_ids"],
+                labels=batch.get("labels"),
+                device=device,
+                sync_max_len=False,
+            )
+            with collector.with_batch(collection_mask):
+                losses = fwd_bwd(model, x, y, batch)
+            collector.process_batch(indices, losses=losses)
+    except torch.cuda.OutOfMemoryError:
+        fits = False
+    finally:
+        collector.mod_grads.clear()
+        model.zero_grad(set_to_none=True)
+        if device.type == "cuda":
+            torch.cuda.set_per_process_memory_fraction(1.0, device)
+        _clear_cache()
+
+    if dist.is_initialized():
+        flag = torch.tensor(int(fits), device=device)
+        dist.all_reduce(flag, op=dist.ReduceOp.MIN)
+        fits = bool(flag.item())
+    print("fits" if fits else "OOM")
+    return fits
+
+
+def _clear_cache() -> None:
+    """Aggressively clear memory."""
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def _cache_key(
+    cfg: IndexConfig, preprocess_cfg: PreprocessConfig, lengths: list[int]
+) -> dict[str, Any]:
+    """Everything a run's memory depends on: its configs, the document lengths,
+    the device and the number of ranks."""
+    index = {
+        k: v for k, v in dataclasses.asdict(cfg).items() if k not in UNKEYED_FIELDS
+    }
+    gpu_name, gpu_mem = "cpu", 0
+    if torch.cuda.is_available():
+        gpu_name = torch.cuda.get_device_name(0)
+        gpu_mem = torch.cuda.get_device_properties(0).total_memory
+    configs = json.dumps(
+        [index, dataclasses.asdict(preprocess_cfg)], sort_keys=True, default=str
+    )
+    return {
+        "configs": hashlib.sha256(configs.encode()).hexdigest(),
+        "lengths": hashlib.sha256(np.asarray(lengths, np.int64).tobytes()).hexdigest(),
+        "world_size": dist.get_world_size() if dist.is_initialized() else 1,
+        "gpu_name": gpu_name,
+        "gpu_memory": gpu_mem,
+        "torch": torch.__version__,
+    }
+
+
+def _check_cache(cache_file: Path, key: dict[str, Any]) -> int | None:
+    """The batch size cached under ``key``, the latest if there are several."""
+    if not cache_file.exists():
+        return None
+
+    found = None
+    with open(cache_file) as f:
+        for line in f:
+            try:
+                row = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            if row.get("key") == key:
+                found = row.get("token_batch_size")
+    return found
+
+
+def _append_to_cache(cache_file: Path, key: dict[str, Any], batch_size: int) -> None:
+    """Append a new row to the JSONL cache file."""
+    cache_file.parent.mkdir(parents=True, exist_ok=True)
+    with open(cache_file, "a") as f:
+        f.write(json.dumps({"key": key, "token_batch_size": batch_size}) + "\n")
+
+    print(f"Cached batch size {batch_size} to {cache_file}")

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bergson
-Version: 2.2.3
+Version: 2.2.4
 Summary: Tracing the memory of neural nets with data attribution
 License: MIT License
 Keywords: interpretability,explainable-ai
@@ -91,16 +91,16 @@ Indicative performance of data attribution methods in the finetuning regime - se
 | 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] |
+| EKFAC + 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] |
+| EKFAC | 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] |
+| Activation similarity (activation covariance preconditioner) | 0.055 [0.043, 0.069] | 0.157 [0.120, 0.192] | 0.073 [0.056, 0.092] |
 | 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) | 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] |
 
 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.
 

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

@@ -101,6 +101,7 @@ tests/test_approx_unrolling_checkpoint_parse.py
 tests/test_attention_head_normalizers.py
 tests/test_attribute_tokens.py
 tests/test_attributor.py
+tests/test_auto_batch_size.py
 tests/test_bank_loss_cache.py
 tests/test_banked_model_loading.py
 tests/test_batch_size_invariance.py

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

@@ -37,7 +37,7 @@ dependencies = [
     "bitsandbytes",
     "huggingface-hub>=1.13.0",
 ]
-version = "2.2.3"
+version = "2.2.4"
 [project.optional-dependencies]
 dev = [
     "pre-commit",

bergson-2.2.4/tests/test_auto_batch_size.py

@@ -0,0 +1,90 @@
+import subprocess
+from dataclasses import replace
+from pathlib import Path
+
+import pytest
+import torch
+
+from bergson.config import IndexConfig, PreprocessConfig
+from bergson.data import _allocate_batches_world
+from bergson.utils.batch_size import _cache_key, probe_batches
+
+from .cli_command import bergson_cmd, bergson_env
+
+
+def test_probe_batches_cover_every_rank():
+    lengths = [512] * 4 + [8] * 300 + [100] * 40
+    world_size = 4
+    probes = probe_batches(lengths, 1024, world_size, None)
+    batches = [
+        b for rank in _allocate_batches_world(lengths, 1024, world_size) for b in rank
+    ]
+
+    def docs(b):
+        return len(b)
+
+    def padded(b):
+        return len(b) * max(lengths[i] for i in b)
+
+    assert padded(probes[0]) == max(padded(b) for b in batches)
+    assert docs(probes[1]) == max(docs(b) for b in batches)
+    # The many-short-documents batch, which a probe of full-length rows misses
+    assert docs(probes[1]) == 1024 // 8
+
+
+def test_cache_key_tracks_memory():
+    cfg = IndexConfig(run_path="a", model="EleutherAI/pythia-14m")
+    pre = PreprocessConfig()
+    lengths = [10, 20, 30]
+    key = _cache_key(cfg, pre, lengths)
+
+    assert _cache_key(replace(cfg, run_path="b", overwrite=True), pre, lengths) == key
+    assert _cache_key(replace(cfg, token_batch_size=4096), pre, lengths) == key
+    for changed in (
+        replace(cfg, attribute_tokens=True),
+        replace(cfg, include_bias=True),
+        replace(cfg, filter_modules="*.mlp.*"),
+        replace(cfg, loss_fn="kl"),
+        replace(cfg, max_batch_size=4),
+    ):
+        assert _cache_key(changed, pre, lengths) != key
+    assert _cache_key(cfg, replace(pre, aggregation="mean"), lengths) != key
+    assert _cache_key(cfg, pre, [10, 20, 31]) != key
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@pytest.mark.parametrize("extra", [[], ["--attribute_tokens"]])
+def test_auto_batch_size_build(tmp_path: Path, extra: list[str]):
+    def build(run: str):
+        return subprocess.run(
+            bergson_cmd(
+                "build",
+                run,
+                "--model",
+                "EleutherAI/pythia-14m",
+                "--dataset",
+                "NeelNanda/pile-10k",
+                "--split",
+                "train[:100]",
+                "--truncation",
+                "--projection_dim",
+                "16",
+                "--token_batch_size",
+                "512",
+                "--auto_batch_size",
+                *extra,
+            ),
+            cwd=tmp_path,
+            env=bergson_env(),
+            capture_output=True,
+            text=True,
+        )
+
+    first = build("first")
+    assert first.returncode == 0, first.stderr
+    assert "Optimal token_batch_size" in first.stdout
+    assert (tmp_path / "first" / "gradients.bin").exists()
+
+    second = build("second")
+    assert second.returncode == 0, second.stderr
+    assert "Loaded token_batch_size from cache" in second.stdout