bergson 2.2.1__tar.gz → 2.2.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bergson
-Version: 2.2.1
+Version: 2.2.2
 Summary: Tracing the memory of neural nets with data attribution
 License: MIT License
 Keywords: interpretability,explainable-ai
@@ -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
 
@@ -120,7 +120,7 @@ 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`.
 
@@ -132,33 +132,27 @@ Per-module and per-attention head gradients can be extracted from the store.
 
 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.1 → bergson-2.2.2}/README.md

@@ -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.
-
-| 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] |
-
-Results for GPT-2 finetuned on 4 epochs of the WikiText corpus. Held-out loss dropped from 3.545 to 3.111 over training.
+[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 | 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 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
 
@@ -59,7 +59,7 @@ 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`.
 
@@ -71,33 +71,27 @@ Per-module and per-attention head gradients can be extracted from the store.
 
 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.1 → bergson-2.2.2}/bergson/__init__.py

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

{bergson-2.2.1 → 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.1 → bergson-2.2.2}/bergson/magic/cli.py

@@ -52,7 +52,12 @@ from ..utils.utils import (
 from ..utils.worker_utils import setup_data_pipeline
 from ..validate import build_contrast_stream, validate_scores
 from .config import MagicConfig
-from .data_stream import DataStream, mask_padded_rows, pad_dataset_to_batch_size
+from .data_stream import (
+    DataStream,
+    Padding,
+    mask_padded_rows,
+    pad_dataset_to_batch_size,
+)
 from .grad_accum import accumulate_grads
 from .score_plot import plot_score_trajectory
 from .trainer import BackwardState, TrainerState, prepare_trainer, write_lr_history
@@ -171,8 +176,7 @@ def compute_per_query_magic_scores(
     run_cfg: "MagicConfig",
     world_size: int,
     global_rank: int,
-    pad_count: int,
-    weight_pad_count: int,
+    padding: Padding,
 ) -> torch.Tensor:
     """Per-query MAGIC scores: one backward per query, sharing the forward.
 
@@ -243,7 +247,7 @@ def compute_per_query_magic_scores(
         one = query_dataset.select([qi])
         if "doc_ids" in one.column_names:
             one = one.remove_columns("doc_ids")
-        one, n_one, one_pad, one_wpad = pad_dataset_to_batch_size(
+        one, n_one, one_padding = pad_dataset_to_batch_size(
             one, world_size, 1, f"Query {qi}", global_rank
         )
         qstream = DataStream(
@@ -253,8 +257,7 @@ def compute_per_query_magic_scores(
             input_key=run_cfg.query.prompt_column,
             weight_shape=(n_one,),
         )
-        if one_pad:
-            qstream.weights.data[-one_wpad:] = 0.0
+        one_padding.zero_weights(qstream.weights.data)
         assert_ckpts_exist()
         qgrads, _ = compute_query_gradients(
             fwd_state, model, qstream, "mean", run_cfg.fsdp, run_cfg.grad_accum_steps
@@ -288,9 +291,7 @@ def compute_per_query_magic_scores(
         if world_size > 1:
             dist.all_reduce(bwd_state.weight_grads, op=dist.ReduceOp.SUM)
 
-        s = bwd_state.weight_grads.detach().cpu()
-        if pad_count:
-            s = s[:-weight_pad_count] if s.ndim == 1 else s[:-pad_count]
+        s = padding.trim(bwd_state.weight_grads.detach().cpu())
         if main:
             # Atomic write
             torch.save(s, qpath + ".tmp")
@@ -357,7 +358,7 @@ def save_magic_scores(
     run_path: str,
     scores: torch.Tensor,
     train_dataset: Dataset,
-    pad_count: int,
+    padding: Padding,
     per_token: bool,
 ) -> str:
     """Write MAGIC scores as a score directory under ``<run_path>/scores``.
@@ -375,9 +376,19 @@ def save_magic_scores(
 
     num_token_grads = compute_num_token_grads(train_dataset)
     doc_ids = np.asarray(train_dataset["doc_ids"], dtype=np.int64)
-    if pad_count:
-        num_token_grads = num_token_grads[:-pad_count]
-        doc_ids = doc_ids[:-pad_count]
+    if "example_ids" in train_dataset.column_names:
+        # The grid has one row per example; the stream's rows (shuffled
+        # epochs, then padding past the grid) fold onto it.
+        example_ids = np.asarray(train_dataset["example_ids"], dtype=np.int64)
+        keep = example_ids < len(scores)
+        per_example_ntg = np.zeros(len(scores), dtype=np.int64)
+        per_example_ntg[example_ids[keep]] = num_token_grads[keep]
+        per_example_doc_ids = np.zeros((len(scores), doc_ids.shape[1]), np.int64)
+        per_example_doc_ids[example_ids[keep]] = doc_ids[keep]
+        num_token_grads, doc_ids = per_example_ntg, per_example_doc_ids
+    elif padding:
+        num_token_grads = num_token_grads[: -padding.num_rows]
+        doc_ids = doc_ids[: -padding.num_rows]
 
     offsets = np.zeros(len(num_token_grads) + 1, dtype=np.int64)
     np.cumsum(num_token_grads, out=offsets[1:])
@@ -395,6 +406,14 @@ def save_magic_scores(
     return str(path)
 
 
+def attach_example_ids(dataset: Dataset) -> Dataset:
+    """Number the rows so per-token weights and scores keep one row per
+    example across the shuffled epochs. No-op if ``example_ids`` is present."""
+    if "example_ids" in dataset.column_names:
+        return dataset
+    return dataset.add_column("example_ids", list(range(len(dataset))))
+
+
 def shuffled_epochs(dataset: Dataset, seed: int, num_epochs: int) -> Dataset:
     """Concatenate `num_epochs` independently shuffled copies of `dataset`.
 
@@ -444,24 +463,26 @@ def worker(
     # Ensure total effective batch size is divisible by world size
     assert run_cfg.batch_size % world_size == 0
 
-    # Pad train dataset to be divisible by batch_size (weight=0 for padding)
-    train_dataset, num_train_docs, pad_count, weight_pad_count = (
-        pad_dataset_to_batch_size(
-            train_dataset, run_cfg.batch_size, num_train_docs, "Train", global_rank
-        )
-    )
-
     # Plain magic runs enter with score_path="" (scores are computed below).
     per_token = (isinstance(run_cfg, MagicConfig) and run_cfg.attribute_tokens) or (
         score_path and scores_are_per_token(score_path)
     )
+    if per_token:
+        train_dataset = attach_example_ids(train_dataset)
+
+    # Pad train dataset to be divisible by batch_size (weight=0 for padding)
+    train_dataset, num_train_docs, padding = pad_dataset_to_batch_size(
+        train_dataset, run_cfg.batch_size, num_train_docs, "Train", global_rank
+    )
+
     if per_token:
         seq_len = run_cfg.data.chunk_length
         if seq_len <= 0:
             seq_len = max(train_dataset["length"])
             print(f"Using max sequence length {seq_len} for per-token attribution")
 
-        w_shape = (len(train_dataset), seq_len)
+        # One weight row per example; the pad rows share one synthetic row.
+        w_shape = (max(train_dataset["example_ids"]) + 1, seq_len)
     else:
         w_shape = (num_train_docs,)
 
@@ -472,11 +493,7 @@ def worker(
         input_key=run_cfg.data.prompt_column,
         weight_shape=w_shape,
     )
-    if pad_count:
-        if stream.weights.ndim == 1:
-            stream.weights.data[-weight_pad_count:] = 0.0
-        else:
-            stream.weights.data[-pad_count:] = 0.0
+    padding.zero_weights(stream.weights.data)
 
     log_fn = None
     if run_cfg.wandb_project and global_rank == 0:
@@ -564,10 +581,8 @@ def worker(
 
     # Pad query dataset to be divisible by batch_size (weight=0 for padding)
     num_real_query_docs = num_query_docs
-    query_dataset, num_query_docs, query_pad_count, query_weight_pad_count = (
-        pad_dataset_to_batch_size(
-            query_dataset, run_cfg.batch_size, num_query_docs, "Query", global_rank
-        )
+    query_dataset, num_query_docs, query_padding = pad_dataset_to_batch_size(
+        query_dataset, run_cfg.batch_size, num_query_docs, "Query", global_rank
     )
     if len(query_dataset) < run_cfg.batch_size:
         raise ValueError(
@@ -584,9 +599,8 @@ def worker(
         input_key=run_cfg.query.prompt_column,
         weight_shape=(num_query_docs,),
     )
-    if query_pad_count:
-        # query_stream.weights is always 1D (weight_shape=(num_query_docs,))
-        query_stream.weights.data[-query_weight_pad_count:] = 0.0
+    # query_stream.weights is always 1D (weight_shape=(num_query_docs,))
+    query_padding.zero_weights(query_stream.weights.data)
 
     query_grads, baseline = compute_query_gradients(
         fwd_state,
@@ -635,15 +649,14 @@ def worker(
             run_cfg,
             world_size,
             global_rank,
-            pad_count,
-            weight_pad_count,
+            padding,
         )
         multi_query = True
         if global_rank == 0:
             print(f"Baseline loss: {baseline}")
             print(f"Score summary: {describe(scores.flatten())}")
             score_path = save_magic_scores(
-                run_cfg.run_path, scores, train_dataset, pad_count, bool(per_token)
+                run_cfg.run_path, scores, train_dataset, padding, bool(per_token)
             )
             plot_score_trajectory(
                 scores.float().numpy(),
@@ -688,12 +701,7 @@ def worker(
         if world_size > 1:
             dist.all_reduce(bwd_state.weight_grads, op=dist.ReduceOp.SUM)
 
-        scores = bwd_state.weight_grads.cpu()
-        if pad_count:
-            if scores.ndim == 1:
-                scores = scores[:-weight_pad_count]
-            else:
-                scores = scores[:-pad_count]
+        scores = padding.trim(bwd_state.weight_grads.cpu())
 
         if global_rank == 0:
             print(f"Baseline loss: {baseline}")
@@ -702,7 +710,7 @@ def worker(
             print(f"Score summary: {summ}")
 
             score_path = save_magic_scores(
-                run_cfg.run_path, scores, train_dataset, pad_count, bool(per_token)
+                run_cfg.run_path, scores, train_dataset, padding, bool(per_token)
             )
             plot_score_trajectory(
                 scores.float().numpy(),
@@ -731,9 +739,8 @@ def worker(
         model=model,
         baseline=baseline,
         num_query_docs=num_query_docs,
-        query_weight_pad_count=query_weight_pad_count,
-        pad_count=pad_count,
-        weight_pad_count=weight_pad_count,
+        query_padding=query_padding,
+        padding=padding,
     )
 
 
@@ -766,15 +773,18 @@ def run_magic(
         save_run_config(run_cfg, run_path)
 
     # HF datasets caches are not safe for concurrent writers, so the main node
-    # must finish populating the cache before others read from it.
-    barrier = run_path / ".preprocess_done" if multi_node else None
+    # must finish populating the cache before others read from it. The job id
+    # keeps a stale barrier from releasing the wait early.
+    job_id = os.environ.get("SLURM_JOB_ID", "")
+    barrier = run_path / f".preprocess_done{job_id}" if multi_node else None
     if barrier is not None and not is_main_node:
-        run_path.mkdir(parents=True, exist_ok=True)
+        # Don't create run_path here to avoid multi-node hang.
         while not barrier.exists():
             time.sleep(0.5)
 
     train_ds, train_n = setup_data_pipeline(run_cfg)
     train_ds = attach_doc_ids_if_missing(train_ds)
+    train_ds = attach_example_ids(train_ds)
 
     train_ds = shuffled_epochs(train_ds, run_cfg.seed, max(1, run_cfg.num_epochs))
 

{bergson-2.2.1 → bergson-2.2.2}/bergson/magic/data_stream.py

@@ -1,6 +1,9 @@
+from dataclasses import dataclass
+
 import torch
 import torch.distributed as dist
-from datasets import Dataset
+from datasets import Dataset, concatenate_datasets
+from torch import Tensor
 
 from ..data import pad_and_tensor
 
@@ -18,61 +21,88 @@ def mask_padded_rows(batch: dict) -> tuple[dict, int]:
     return batch, int((batch["labels"][:, 1:] != -100).sum())
 
 
+@dataclass(frozen=True)
+class Padding:
+    """How much of the padded dataset is padding, and how to keep padding from
+    affecting influence scores.
+
+    A run weights either documents, as a vector with one entry per document,
+    or tokens, as a grid with one row per example. ``num_docs`` counts the
+    vector entries that are pure padding and ``num_examples`` the grid rows.
+    Each is one if the pad rows share a synthetic id, ``doc_ids`` for the
+    first and ``example_ids`` for the second, and ``num_rows`` if they do not;
+    the two are decided separately. Pass the weights to :meth:`zero_weights`
+    and the scores to :meth:`trim`, which use whichever count matches.
+    """
+
+    num_rows: int = 0
+    num_docs: int = 0
+    num_examples: int = 0
+
+    def __bool__(self) -> bool:
+        return bool(self.num_rows)
+
+    def _trailing_pad_len(self, t: Tensor) -> int:
+        return self.num_docs if t.ndim == 1 else self.num_examples
+
+    def zero_weights(self, weights: Tensor) -> None:
+        """Zero the pad rows' weights in place."""
+        if n := self._trailing_pad_len(weights):
+            weights.data[-n:] = 0.0
+
+    def trim(self, scores: Tensor) -> Tensor:
+        """Return ``scores`` without the entries the pad rows produced."""
+        n = self._trailing_pad_len(scores)
+        return scores[:-n] if n else scores
+
+
 def pad_dataset_to_batch_size(
     dataset: Dataset,
     batch_size: int,
     num_docs: int,
     label: str,
     global_rank: int,
-) -> tuple[Dataset, int, int, int]:
+) -> tuple[Dataset, int, Padding]:
     """Pad dataset to be divisible by batch_size by repeating the last example.
 
-    Returns (padded_dataset, num_docs, pad_count, weight_pad_count).
-
-    `pad_count` is the number of rows appended to the dataset (0 if unchanged).
-    `weight_pad_count` is the number of trailing entries of a *1D* per-doc
-    weight tensor that should be zeroed to silence the pad rows' training
-    contribution.
-
-    - If the dataset has a "doc_ids" column, `.select(total - 1, ...)` copies
-      the last doc's doc_ids into every pad row. Zeroing the last `pad_count`
-      entries of a weights-indexed-by-doc_id tensor would silence real docs,
-      so we instead route pad rows to a fresh synthetic doc id (=num_docs),
-      bump num_docs by 1, and set `weight_pad_count = 1`.
-    - Otherwise rows are self-identified docs: num_docs becomes the padded
-      length and `weight_pad_count = pad_count` zeros the pad rows directly.
-
-    In per-token (2D) mode callers should zero `weights[-pad_count:]` instead
-    — `weight_pad_count` applies only to 1D per-doc weights.
+    Repeating a row copies its ``doc_ids``, which would add the pad rows'
+    scores to the last document; they are given a separate document id
+    instead, so one zeroed weight entry covers all of them. An ``example_ids``
+    column, which per-token weights index by, is re-pointed the same way.
     """
     remainder = len(dataset) % batch_size
     if not remainder:
-        return dataset, num_docs, 0, 0
+        return dataset, num_docs, Padding()
 
     pad_count = batch_size - remainder
     total = len(dataset)
-    pad_indices = list(range(total)) + [total - 1] * pad_count
-    dataset = dataset.select(pad_indices)
+    last = dataset[total - 1]
+
+    if "example_ids" in dataset.column_names:
+        last = {**last, "example_ids": max(dataset["example_ids"]) + 1}
+        num_examples = 1
+    else:
+        num_examples = pad_count
 
     if "doc_ids" in dataset.column_names:
-        synthetic_doc_id = num_docs
-        new_doc_ids = [
-            row if i < total else [synthetic_doc_id] * len(row)
-            for i, row in enumerate(dataset["doc_ids"])
-        ]
-        dataset = dataset.remove_columns("doc_ids").add_column("doc_ids", new_doc_ids)
+        last = {**last, "doc_ids": [num_docs] * len(last["doc_ids"])}
         num_docs += 1
-        weight_pad_count = 1
+        padding = Padding(pad_count, num_docs=1, num_examples=num_examples)
     else:
-        num_docs = len(dataset)
-        weight_pad_count = pad_count
-
+        num_docs = total + pad_count
+        padding = Padding(pad_count, num_docs=pad_count, num_examples=num_examples)
+
+    # Built in memory: the pad rows are a handful of copies of the last row.
+    pad_rows = Dataset.from_dict(
+        {k: [v] * pad_count for k, v in last.items()}, features=dataset.features
+    )
+    dataset = concatenate_datasets([dataset, pad_rows])
     if global_rank == 0:
         print(
             f"{label}: padded {pad_count}/{total + pad_count} examples "
             f"(weight=0) to fill last batch"
         )
-    return dataset, num_docs, pad_count, weight_pad_count
+    return dataset, num_docs, padding
 
 
 class DataStream:
@@ -130,8 +160,15 @@ class DataStream:
         # If the weights are 1D, we assume they correspond to documents and look for
         # "doc_ids" in the batch to index them. If they're 2D, they correspond to tokens
         if self.weights.ndim == 2:
-            # Truncate to the max sequence length in the batch to avoid indexing errors
-            indices = (indices, slice(None, x.shape[1]))
+            # One weight row per example: a shuffled multi-epoch stream reaches
+            # its rows through "example_ids". Truncate to the max sequence
+            # length in the batch to avoid indexing errors.
+            rows = (
+                torch.tensor(batch["example_ids"], device=self.device)
+                if "example_ids" in batch
+                else indices
+            )
+            indices = (rows, slice(None, x.shape[1]))
         elif "doc_ids" in batch:
             indices = torch.tensor(batch["doc_ids"], device=self.device)
             # doc_ids may be longer than the per-batch padded seq_len (unpacked