bergson 2.2.3__tar.gz → 2.3.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bergson
-Version: 2.2.3
+Version: 2.3.0
 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.3.0}/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.3.0}/bergson/__init__.py

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

{bergson-2.2.3 → bergson-2.3.0}/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.3.0}/bergson/config/config.py

@@ -52,8 +52,13 @@ class DataConfig(Serializable):
     When specified, gradients are calculated using the policy
     gradient loss from Dr. GRPO. https://arxiv.org/abs/2503.20783"""
 
+    advantage_column: str = ""
+    """Optional column in the dataset that contains per-example advantages, e.g.
+    those logged by a GRPO trainer. When specified, gradients are calculated using
+    the Dr. GRPO policy gradient loss."""
+
     skip_nan_rewards: bool = False
-    """Whether to skip examples with NaN rewards."""
+    """Whether to skip examples with NaN rewards or advantages."""
 
     truncation: bool = False
     """Whether to truncate long documents to fit the token budget."""
@@ -73,6 +78,10 @@ class DataConfig(Serializable):
     sequences of this length. Incompatible with truncation and format_template."""
 
     def __post_init__(self):
+        if self.reward_column and self.advantage_column:
+            raise ValueError(
+                "reward_column and advantage_column cannot both be specified"
+            )
         if self.chunk_length > 0:
             if self.truncation:
                 raise ValueError("chunk_length and truncation cannot both be True")

{bergson-2.2.3 → bergson-2.3.0}/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.3.0/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.3.0}/bergson/utils/worker_utils.py

@@ -347,7 +347,9 @@ def filter_by_max_tokens(ds: Dataset, cfg: AttributionConfig) -> Dataset:
     return ds
 
 
-def max_tokens_for_model(tokenizer, model_str: str, revision: str | None) -> int:
+def max_tokens_for_model(
+    tokenizer, model_str: str, revision: str | None, trust_remote_code: bool = False
+) -> int:
     # You might think model_max_length should always be the same as
     # max_position_embeddings, but some models (e.g. Pythia!) have a smaller
     # max_position_embeddings than model_max_length, so we need to check both.
@@ -360,7 +362,9 @@ def max_tokens_for_model(tokenizer, model_str: str, revision: str | None) -> int
     except ValueError:
         pass
 
-    model_cfg = AutoConfig.from_pretrained(model_str, revision=revision)
+    model_cfg = AutoConfig.from_pretrained(
+        model_str, revision=revision, trust_remote_code=trust_remote_code
+    )
     model_max_length = getattr(tokenizer, "model_max_length", BIG_NUM)
     max_pos_emb = getattr(model_cfg, "max_position_embeddings", BIG_NUM)
     return min(model_max_length, max_pos_emb)
@@ -377,7 +381,12 @@ def setup_data_pipeline(
         data_cfg.dataset, data_cfg.split, data_cfg.subset, data_cfg.data_kwargs
     )
     tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer or cfg.model)
-    max_model_length = max_tokens_for_model(tokenizer, cfg.model, cfg.revision)
+    trust_remote_code = bool(
+        simple_parse_kwargs_string(cfg.model_kwargs).get("trust_remote_code", False)
+    )
+    max_model_length = max_tokens_for_model(
+        tokenizer, cfg.model, cfg.revision, trust_remote_code
+    )
 
     if data_cfg.chunk_length > 0:
         # Sanity check
@@ -449,23 +458,28 @@ def setup_data_pipeline(
                 f"Consider increasing --token_batch_size or using --truncation."
             )
 
-    if data_cfg.reward_column:
+    value_column = data_cfg.reward_column or data_cfg.advantage_column
+    if value_column:
         assert isinstance(ds, Dataset), "Dataset required for advantage estimation"
 
-        rewards = np.array(ds[data_cfg.reward_column], dtype=np.float64)
-        nan_mask = np.isnan(rewards)
+        values = np.array(ds[value_column], dtype=np.float64)
+        nan_mask = np.isnan(values)
         if nan_mask.any():
             if data_cfg.skip_nan_rewards:
-                print(f"Warning: Filtering out {nan_mask.sum()} rows with NaN rewards")
+                print(f"Warning: Filtering out {nan_mask.sum()} rows with NaN values")
                 ds = ds.filter(lambda _, idx: not nan_mask[idx], with_indices=True)
             else:
-                raise ValueError(
-                    f"Reward column '{data_cfg.reward_column}' contains NaN values"
-                )
+                raise ValueError(f"Column '{value_column}' contains NaN values")
 
+        if data_cfg.reward_column:
+            advantages = estimate_advantage(ds, data_cfg)
+        else:
+            advantages = ds[data_cfg.advantage_column][:]
+        if "advantage" in ds.column_names:
+            ds = ds.remove_columns("advantage")
         ds = ds.add_column(
             "advantage",
-            estimate_advantage(ds, data_cfg),
+            advantages,
             new_fingerprint="advantage",  # type: ignore
         )
 
@@ -475,6 +489,8 @@ def setup_data_pipeline(
 
     # Remove extraneous columns
     keep = {"length", "input_ids", "labels"}
+    if value_column:
+        keep.add("advantage")
     remove_columns -= keep
     remove_columns &= set(ds.column_names)
     if remove_columns:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bergson
-Version: 2.2.3
+Version: 2.3.0
 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.3.0}/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.3.0}/pyproject.toml

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

{bergson-2.2.3 → bergson-2.3.0}/tests/test_advantages.py

@@ -1,5 +1,6 @@
 import json
 
+import pytest
 from datasets import Dataset
 
 from bergson.config import DataConfig, IndexConfig
@@ -83,3 +84,57 @@ def test_advantages_computed_with_drop_columns(tmp_path):
     # "foo bar" group mean = 15.0: advantages = [-5.0, 5.0]
     assert abs(advantages[2] - (-5.0)) < 1e-9
     assert abs(advantages[3] - 5.0) < 1e-9
+
+
+def _pipeline_advantages(tmp_path, rows, **data_kwargs):
+    path = tmp_path / "rows.json"
+    path.write_text("\n".join(json.dumps(r) for r in rows))
+    cfg = IndexConfig(
+        run_path=str(tmp_path / "run"),
+        model="EleutherAI/pythia-14m",
+        token_batch_size=2048,
+        data=DataConfig(dataset=str(path), **data_kwargs),
+    )
+    ds, _ = setup_data_pipeline(cfg)
+    return ds["advantage"]
+
+
+def test_advantage_column_used_as_given(tmp_path):
+    rows = [
+        {"text": "hello world", "adv": 0.5},
+        {"text": "hello world", "adv": 2.0},
+        {"text": "foo bar", "adv": -1.5},
+    ]
+    assert _pipeline_advantages(tmp_path, rows, advantage_column="adv") == [
+        0.5,
+        2.0,
+        -1.5,
+    ]
+
+
+def test_existing_advantage_column(tmp_path):
+    """A dataset column named "advantage" can be the advantage column, and is
+    replaced by the estimate when rewards are given."""
+    rows = [
+        {"text": "hello world", "advantage": 0.5, "reward": 1.0},
+        {"text": "hello world", "advantage": 2.0, "reward": 3.0},
+    ]
+    assert _pipeline_advantages(tmp_path, rows, advantage_column="advantage") == [
+        0.5,
+        2.0,
+    ]
+    assert _pipeline_advantages(tmp_path, rows, reward_column="reward") == [-1.0, 1.0]
+
+
+def test_nan_advantages(tmp_path):
+    rows = [{"text": "hello world", "adv": 1.0}, {"text": "foo bar", "adv": None}]
+    with pytest.raises(ValueError, match="Column 'adv' contains NaN"):
+        _pipeline_advantages(tmp_path, rows, advantage_column="adv")
+    assert _pipeline_advantages(
+        tmp_path, rows, advantage_column="adv", skip_nan_rewards=True
+    ) == [1.0]
+
+
+def test_conflicting_advantage_config():
+    with pytest.raises(ValueError, match="cannot both be specified"):
+        DataConfig(reward_column="reward", advantage_column="adv")