From 1e595d277115f1df9714155e845661fc0c48ec0c Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Thu, 6 Aug 2026 20:33:11 +0900 Subject: [PATCH 1/4] feat(magic): write scores as a score directory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MAGIC wrote a bare tensor to scores.pt, so every consumer had to recover from the shape what a score directory simply states. A 2-D tensor is [docs, seq_len] or [docs, queries] depending only on how the run was configured, which is why load_attribution_scores had to read config.yaml, fall back to the rank when there was none, and why scores_are_per_token needed the same treatment. The tensor was never self-contained anyway: it already leaned on config.yaml to disambiguate and on a doc_ids.pt sidecar to map shuffled chunks back to documents. Write the same score directory the scoring pipeline writes. info.json records attribute_tokens and num_scores, so both questions are answered by the store, and the reader collapses to: load it, negate if higher_is_better, report num_scores > 1. Per-token scores pack into the ragged token store exactly. The weight grid is [rows, seq_len] and weighted_causal_lm_ce reads example_weight[:, :-1], so a row's columns past length - 1 never receive gradient — precisely the rows compute_num_token_grads allocates. to_grid then restores the same width, since max(length - 1) + 1 is the seq_len the grid started with. The round-trip is exact in shape and value, tested both ways. doc_ids moves into the directory as doc_ids.npy, in the grid's shape, so a per-token store carries everything needed to aggregate per document. Deleted along the way: the .pt branch of load_attribution_scores, the config lookup and shape fallback behind it, scores_are_per_token's two non-directory branches, and read_first_step_config, whose only caller was that lookup. Co-Authored-By: Claude Opus 5 (1M context) --- bergson/config/config_io.py | 25 --------- bergson/magic/cli.py | 87 ++++++++++++++++++------------ bergson/validate.py | 56 +++++++------------ docs/magic.rst | 42 +++++++++------ tests/test_distributed_magic.py | 13 ++--- tests/test_magic.py | 52 +++++++++++++++--- tests/test_multi_query_validate.py | 40 -------------- tests/test_per_query_magic.py | 25 +++------ tests/test_per_token_lds.py | 46 +--------------- 9 files changed, 160 insertions(+), 226 deletions(-) diff --git a/bergson/config/config_io.py b/bergson/config/config_io.py index 105accc7..da50616d 100644 --- a/bergson/config/config_io.py +++ b/bergson/config/config_io.py @@ -120,31 +120,6 @@ def read_config(path: str | Path) -> dict[str, Any]: return config -def read_first_step_config(path: str | Path) -> dict[str, Any] | None: - """The first step's configuration from a run's ``config.yaml``. - - ``path`` may be the run directory or an artifact inside it (``scores.pt``), - so callers holding an output path need not resolve it. Returns ``None`` - when there is no readable config, which is the signal that a caller must - fall back to inspecting the artifact itself. - """ - p = Path(path) - cfg_path = (p if p.is_dir() else p.parent) / CONFIG_FILENAME - if not cfg_path.is_file(): - return None - - try: - steps = read_config(cfg_path)["steps"] - except (ValueError, yaml.YAMLError): - return None - - for step in steps: - for step_cfg in step.values(): - if isinstance(step_cfg, dict): - return step_cfg - return None - - T = TypeVar("T", bound="FromDict") diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index 47434c78..f837224d 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -6,6 +6,7 @@ from dataclasses import asdict from pathlib import Path +import numpy as np import torch import torch.distributed as dist from datasets import Dataset, concatenate_datasets @@ -28,8 +29,10 @@ ) from ..config.config import TrainingConfig, ValidationConfig -from ..config.config_io import read_first_step_config, save_run_config +from ..config.config_io import save_run_config +from ..data import compute_num_token_grads from ..distributed import launch_distributed_run +from ..score.score_writer import save_sequence_scores, save_token_scores from ..utils.load_from_optimizer import ( save_second_moments_as_optimizer_pt, ) @@ -244,20 +247,11 @@ def compute_per_query_magic_scores( def scores_are_per_token(score_path: str) -> bool: - if os.path.isdir(score_path): - info_path = os.path.join(score_path, "info.json") - if not os.path.isfile(info_path): - return False - with open(info_path) as f: - return bool(json.load(f).get("attribute_tokens", False)) - step_cfg = read_first_step_config(score_path) - if step_cfg is not None: - return bool(step_cfg.get("attribute_tokens") or step_cfg.get("per_token")) - - scores = torch.load(score_path, map_location="cpu") - return isinstance(scores, torch.Tensor) and ( - scores.ndim == 3 or (scores.ndim == 2 and scores.shape[1] > 1) - ) + info_path = os.path.join(score_path, "info.json") + if not os.path.isfile(info_path): + return False + with open(info_path) as f: + return bool(json.load(f).get("attribute_tokens", False)) def attach_doc_ids_if_missing(dataset: Dataset) -> Dataset: @@ -286,16 +280,46 @@ def attach_doc_ids_if_missing(dataset: Dataset) -> Dataset: ) -def save_doc_ids(run_path: str, train_dataset: Dataset, pad_count: int) -> str: - """Write ``doc_ids.pt`` beside a per-token ``scores.pt``.""" - doc_ids = torch.tensor(train_dataset["doc_ids"]) +def save_magic_scores( + run_path: str, + scores: torch.Tensor, + train_dataset: Dataset, + pad_count: int, + per_token: bool, +) -> str: + """Write MAGIC scores as a score directory under ``/scores``. + + Columns past a row's ``length - 1`` never receive gradient, so the dense + per-token grid packs into the ragged store without loss. + """ + path = Path(run_path) / "scores" + + if not per_token: + arr = scores.float().numpy() + save_sequence_scores(path, arr if arr.ndim > 1 else arr[:, None]) + print(f"Saved attribution scores to {path}") + return str(path) + + 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] - doc_ids_path = os.path.join(run_path, "doc_ids.pt") - torch.save(doc_ids, doc_ids_path) - print(f"Saved doc_ids to {doc_ids_path}") - return doc_ids_path + offsets = np.zeros(len(num_token_grads) + 1, dtype=np.int64) + np.cumsum(num_token_grads, out=offsets[1:]) + + grid = scores.float().numpy() + if grid.ndim == 2: + grid = grid[:, :, None] + flat = np.concatenate( + [grid[r, : num_token_grads[r]] for r in range(len(num_token_grads))] + ) + + save_token_scores(path, flat, offsets) + np.save(path / "doc_ids.npy", doc_ids) + print(f"Saved attribution scores to {path}") + return str(path) def shuffled_epochs(dataset: Dataset, seed: int, num_epochs: int) -> Dataset: @@ -512,9 +536,9 @@ def worker( if global_rank == 0: print(f"Baseline loss: {baseline}") print(f"Score summary: {describe(scores.flatten())}") - score_path = os.path.join(run_cfg.run_path, "scores.pt") - torch.save(scores, score_path) - print(f"Saved per-query attribution scores to {score_path}") + score_path = save_magic_scores( + run_cfg.run_path, scores, train_dataset, pad_count, bool(per_token) + ) elif not score_path: # Sanity check if not isinstance(run_cfg, MagicConfig): @@ -566,16 +590,11 @@ def worker( summ = describe(scores.flatten()) print(f"Score summary: {summ}") - score_path = os.path.join(run_cfg.run_path, "scores.pt") - torch.save(scores, score_path) - print(f"Saved attribution scores to {score_path}") - elif os.path.isdir(score_path): - scores, multi_query = load_attribution_scores(score_path) + score_path = save_magic_scores( + run_cfg.run_path, scores, train_dataset, pad_count, bool(per_token) + ) else: - scores = torch.load(score_path, map_location="cpu") - - if per_token and global_rank == 0: - save_doc_ids(run_cfg.run_path, train_dataset, pad_count) + scores, multi_query = load_attribution_scores(score_path) stream.requires_grad = False diff --git a/bergson/validate.py b/bergson/validate.py index 71c3805f..d071483e 100644 --- a/bergson/validate.py +++ b/bergson/validate.py @@ -30,8 +30,8 @@ ) from .config.config import ScoreConfig, ValidationConfig -from .config.config_io import load_subconfig, read_first_step_config, save_run_config -from .data import Scores, load_scores, pad_and_tensor +from .config.config_io import load_subconfig, save_run_config +from .data import load_scores, pad_and_tensor from .magic.data_stream import DataStream, pad_dataset_to_batch_size from .magic.trainer import TrainerState, prepare_trainer from .utils.csv_writer import CSVWriter @@ -40,47 +40,29 @@ def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]: - """Load attribution scores from a score directory or ``.pt`` file. - - Returns ``(scores, multi_query)``. Token score directories are per-token, - with a query dimension when ``num_scores > 1`` (``[docs, seq_len, - queries]``); plain score directories are per-document, with one column per - query. A 3-D ``.pt`` tensor is per-token per-query ``[docs, seq_len, - queries]``. A 2-D ``.pt`` tensor is ambiguous — per-token MAGIC scores are - ``[docs, seq_len]`` and per-query MAGIC scores are ``[docs, queries]`` — - so the run config next to it decides: it is per-query iff the run used - ``query_method: none`` without attributing tokens. - - Score directories are negated when their ``score_cfg.higher_is_better`` is - set, aligning them with the loss-diff convention. - """ - if os.path.isdir(score_path): - loaded = load_scores(Path(score_path)) - score_cfg = load_subconfig(score_path, "score_cfg", ScoreConfig) - negate = score_cfg is not None and score_cfg.higher_is_better + """Load attribution scores from a score directory. - if isinstance(loaded, Scores) and loaded.offsets is not None: - scores = loaded.to_grid() - if negate: - scores = -scores - return scores, loaded.num_scores > 1 + Returns ``(scores, multi_query)``. Token stores are per-token, with a + query dimension when ``num_scores > 1`` (``[docs, seq_len, queries]``); + plain stores are per-document, with one column per query. Scores are + negated when ``score_cfg.higher_is_better`` is set, aligning them with + the loss-diff convention. + """ + loaded = load_scores(Path(score_path)) + score_cfg = load_subconfig(score_path, "score_cfg", ScoreConfig) + negate = score_cfg is not None and score_cfg.higher_is_better + if loaded.offsets is not None: + scores = loaded.to_grid() + else: arr = np.asarray(loaded[:]) # Copy: the slice is a read-only view onto the memmap. out_dtype = arr.dtype if np.issubdtype(arr.dtype, np.floating) else np.float32 scores = torch.from_numpy(arr.astype(out_dtype, copy=True)) - if negate: - scores = -scores - return scores, loaded.num_scores > 1 - - scores = torch.load(score_path, map_location="cpu") - if not isinstance(scores, torch.Tensor) or scores.ndim not in (2, 3): - return scores, False - - step_cfg = read_first_step_config(score_path) - if step_cfg is None: - return scores, scores.ndim == 3 - return scores, step_cfg.get("query_method") == "none" + + if negate: + scores = -scores + return scores, loaded.num_scores > 1 def bank_loss_cache_key( diff --git a/docs/magic.rst b/docs/magic.rst index 41481d72..b6ff3171 100644 --- a/docs/magic.rst +++ b/docs/magic.rst @@ -32,25 +32,35 @@ Output files After a run completes, ``run_cfg.run_path`` contains: -* ``scores.pt`` — attribution scores tensor. Shape depends on the weight - parameterization: - - * Per-example (1D weights): ``(num_train_docs,)``, indexed directly by - ``doc_id``. - * Per-token (2D weights): ``(num_chunks, seq_len)``, indexed by - ``(chunk_idx, token_idx)`` in the *post-shuffle* order used during - training. Pad rows appended to make the dataset divisible by - ``batch_size`` are trimmed before saving. - -* ``doc_ids.pt`` — written alongside ``scores.pt`` for every per-token - run, shape ``(num_chunks, seq_len)`` matching ``scores.pt`` row-for-row. - Each entry is the original (pre-shuffle) document id for that token - position. Downstream aggregation is one line: +* ``scores/`` — a score directory, the same self-describing format the + scoring pipeline writes. ``info.json`` records ``attribute_tokens`` and + ``num_scores``, so consumers never infer the layout from the shape. + Read it with :func:`bergson.validate.load_attribution_scores`, which + returns ``(scores, multi_query)``: + + * Per-example: ``(num_train_docs, 1)``, indexed directly by ``doc_id``. + * Per-token: ``(num_chunks, seq_len)``, indexed by ``(chunk_idx, + token_idx)`` in the *post-shuffle* order used during training. + * Per-query (``query_method: none``) adds a trailing query axis, so + per-token per-query scores are ``(num_chunks, seq_len, + num_query_docs)``. + + Pad rows appended to make the dataset divisible by ``batch_size`` are + trimmed before saving. Per-token scores are stored ragged — a row holds + ``length - 1`` values, the positions ``weighted_causal_lm_ce`` can reach + — and are unpacked back into the dense grid on load. + +* ``scores/doc_ids.npy`` — written for every per-token run, shape + ``(num_chunks, seq_len)`` matching the loaded scores row-for-row. Each + entry is the original (pre-shuffle) document id for that token position. + Downstream aggregation is one line: .. code-block:: python - scores = torch.load("scores.pt") # (num_chunks, seq_len) - doc_ids = torch.load("doc_ids.pt") # (num_chunks, seq_len) + from bergson.validate import load_attribution_scores + + scores, _ = load_attribution_scores("runs/magic/scores") + doc_ids = torch.from_numpy(np.load("runs/magic/scores/doc_ids.npy")) num_docs = int(doc_ids.max()) + 1 per_doc = torch.zeros(num_docs, dtype=scores.dtype) per_doc.scatter_add_(0, doc_ids.flatten(), scores.flatten()) diff --git a/tests/test_distributed_magic.py b/tests/test_distributed_magic.py index b736d309..6efc895d 100644 --- a/tests/test_distributed_magic.py +++ b/tests/test_distributed_magic.py @@ -13,6 +13,7 @@ from bergson.config import DataConfig, DistributedConfig, LRScheduleConfig from bergson.magic.cli import MagicConfig, run_magic +from bergson.validate import load_attribution_scores # Both tests consume the module-scoped noclip_scores fixture, so they must run # on the same xdist worker or each worker recomputes the two no-clip runs. @@ -81,7 +82,7 @@ def noclip_scores(tmp_path_factory) -> dict[str, torch.Tensor]: scores = {} for mode, fsdp in [("ddp", False), ("fsdp", True)]: run_magic(magic_cfg(f"{tmpdir}/{mode}", fsdp=fsdp, clip=False)) - scores[mode] = torch.load(f"{tmpdir}/{mode}/scores.pt", weights_only=True) + scores[mode] = load_attribution_scores(f"{tmpdir}/{mode}/scores")[0] return scores @@ -130,8 +131,8 @@ def test_fsdp_ddp_scores_match_with_grad_clipping(noclip_scores, tmp_path): run_magic(magic_cfg(f"{tmp_path}/ddp", fsdp=False, clip=True)) run_magic(magic_cfg(f"{tmp_path}/fsdp", fsdp=True, clip=True)) - ddp_scores = torch.load(f"{tmp_path}/ddp/scores.pt", weights_only=True) - fsdp_scores = torch.load(f"{tmp_path}/fsdp/scores.pt", weights_only=True) + ddp_scores = load_attribution_scores(f"{tmp_path}/ddp/scores")[0] + fsdp_scores = load_attribution_scores(f"{tmp_path}/fsdp/scores")[0] assert fsdp_scores.shape == ddp_scores.shape @@ -167,8 +168,8 @@ def test_grad_accum_matches_full_batch(noclip_scores, tmp_path): run_magic(magic_cfg(f"{tmp_path}/ddp", fsdp=False, clip=False, grad_accum=2)) run_magic(magic_cfg(f"{tmp_path}/fsdp", fsdp=True, clip=False, grad_accum=2)) - ddp_scores = torch.load(f"{tmp_path}/ddp/scores.pt", weights_only=True) - fsdp_scores = torch.load(f"{tmp_path}/fsdp/scores.pt", weights_only=True) + ddp_scores = load_attribution_scores(f"{tmp_path}/ddp/scores")[0] + fsdp_scores = load_attribution_scores(f"{tmp_path}/fsdp/scores")[0] assert fsdp_scores.shape == ddp_noclip.shape @@ -226,7 +227,7 @@ def run(name: str, *, fsdp: bool, grad_accum: int) -> torch.Tensor: lr=1e-5, ) ) - return torch.load(f"{path}/scores.pt", weights_only=True) + return load_attribution_scores(f"{path}/scores")[0] ddp1 = run("ddp1", fsdp=False, grad_accum=1) ddp2 = run("ddp2", fsdp=False, grad_accum=2) diff --git a/tests/test_magic.py b/tests/test_magic.py index cffa807b..eeda4191 100644 --- a/tests/test_magic.py +++ b/tests/test_magic.py @@ -1358,16 +1358,20 @@ def test_worker_writes_doc_ids_for_fresh_per_token_run(tmp_path): worker(0, 0, 1, train_ds, query_ds, num_docs, 2, run_cfg) - scores_path = tmp_path / "scores.pt" - doc_ids_path = tmp_path / "doc_ids.pt" - assert scores_path.is_file(), "worker() did not write scores.pt" + import numpy as np + + from bergson.validate import load_attribution_scores + + score_dir = tmp_path / "scores" + doc_ids_path = score_dir / "doc_ids.npy" + assert score_dir.is_dir(), "worker() did not write a score directory" assert doc_ids_path.is_file(), ( - "worker() wrote scores.pt but no doc_ids.pt; per-token scores are " + "worker() wrote scores but no doc_ids.npy; per-token scores are " "indexed by shuffled chunk, so they are unaggregatable without it" ) - scores = torch.load(scores_path) - doc_ids = torch.load(doc_ids_path) + scores, _ = load_attribution_scores(str(score_dir)) + doc_ids = torch.from_numpy(np.load(doc_ids_path)) assert scores.ndim == 2, f"expected per-token scores, got {scores.shape}" assert doc_ids.shape == scores.shape @@ -1376,3 +1380,39 @@ def test_worker_writes_doc_ids_for_fresh_per_token_run(tmp_path): torch.testing.assert_close( agg.sum(), scores.sum().to(torch.float64), atol=1e-6, rtol=1e-5 ) + + +@pytest.mark.parametrize("num_scores", [1, 2]) +def test_save_magic_scores_round_trips_the_grid(tmp_path, num_scores): + """The ragged store returns the dense grid MAGIC computed. Columns past a + row's ``length - 1`` never receive gradient, so packing drops only zeros + and ``to_grid`` restores the same shape and values.""" + import numpy as np + from datasets import Dataset + + from bergson.magic.cli import save_magic_scores + from bergson.validate import load_attribution_scores + + rows, seq_len = 4, 6 + data = Dataset.from_dict( + { + "input_ids": [[1] * seq_len] * rows, + "labels": [[1] * seq_len] * rows, + "doc_ids": [[d] * seq_len for d in range(rows)], + "length": [seq_len] * rows, + } + ) + + shape = (rows, seq_len) if num_scores == 1 else (rows, seq_len, num_scores) + grid = torch.arange(int(np.prod(shape)), dtype=torch.float32).reshape(shape) + grid[:, seq_len - 1] = 0.0 # the column weighted_causal_lm_ce never reaches + + save_magic_scores(str(tmp_path), grid, data, pad_count=0, per_token=True) + loaded, multi_query = load_attribution_scores(str(tmp_path / "scores")) + + assert multi_query == (num_scores > 1) + assert loaded.shape == grid.shape + torch.testing.assert_close(loaded, grid) + + doc_ids = np.load(tmp_path / "scores" / "doc_ids.npy") + assert doc_ids.shape == (rows, seq_len) diff --git a/tests/test_multi_query_validate.py b/tests/test_multi_query_validate.py index ad81b007..d8e28fd3 100644 --- a/tests/test_multi_query_validate.py +++ b/tests/test_multi_query_validate.py @@ -71,14 +71,6 @@ def test_load_attribution_scores_single_column_dir(tmp_path): assert scores.shape == (4, 1) -def test_load_attribution_scores_pt_without_config(tmp_path): - # 2D tensors in .pt files are per-token MAGIC scores, never multi-query. - pt_path = tmp_path / "scores.pt" - torch.save(torch.zeros(5, 2), pt_path) - scores, multi_query = load_attribution_scores(str(pt_path)) - assert not multi_query - - def test_weighted_ce_sum_of_means_reduction(): """sum_of_means = per-sample token-mean, summed over batch (no /B) — the MAGIC/metagradients convention (arXiv 2503.13751 App. D).""" @@ -146,35 +138,3 @@ def test_metasmoothness_score(): # Zero movement -> defined as perfectly smooth. assert metasmoothness_score(theta0, theta0, theta0) == 1.0 - - -def test_load_attribution_scores_pt_per_query(tmp_path): - """A per-query [docs, queries] scores.pt is flagged multi-query when the - run config beside it says query_method: none; per-token and configless - tensors stay single-query.""" - import yaml - - values = torch.arange(12, dtype=torch.float32).reshape(4, 3) - torch.save(values, tmp_path / "scores.pt") - - # No config.yaml: ambiguous, keep the per-token interpretation. - _, multi_query = load_attribution_scores(str(tmp_path / "scores.pt")) - assert not multi_query - - cfg = {"steps": [{"magic": {"query_method": "none", "per_token": False}}]} - (tmp_path / "config.yaml").write_text(yaml.safe_dump(cfg)) - scores, multi_query = load_attribution_scores(str(tmp_path / "scores.pt")) - assert multi_query - torch.testing.assert_close(scores, values) - - # Attributing tokens as well makes it 3-D, and still per-query. - cfg = {"steps": [{"magic": {"query_method": "none", "attribute_tokens": True}}]} - (tmp_path / "config.yaml").write_text(yaml.safe_dump(cfg)) - torch.save(torch.zeros(4, 3, 2), tmp_path / "tok.pt") - _, multi_query = load_attribution_scores(str(tmp_path / "tok.pt")) - assert multi_query - - cfg = {"steps": [{"magic": {"query_method": "mean"}}]} - (tmp_path / "config.yaml").write_text(yaml.safe_dump(cfg)) - _, multi_query = load_attribution_scores(str(tmp_path / "scores.pt")) - assert not multi_query diff --git a/tests/test_per_query_magic.py b/tests/test_per_query_magic.py index a8910f8e..1c41960c 100644 --- a/tests/test_per_query_magic.py +++ b/tests/test_per_query_magic.py @@ -232,11 +232,13 @@ def ds(n): ) worker(0, 0, 1, ds(num_docs), ds(n_query), num_docs, n_query, run_cfg) - doc_ids = run_path / "doc_ids.pt" - return ( - torch.load(run_path / "scores.pt"), - torch.load(doc_ids) if doc_ids.exists() else None, - ) + import numpy as np + + from bergson.validate import load_attribution_scores + + scores, _ = load_attribution_scores(str(run_path / "scores")) + doc_ids = run_path / "scores" / "doc_ids.npy" + return scores, (torch.from_numpy(np.load(doc_ids)) if doc_ids.exists() else None) def test_per_query_per_token_aggregates_to_per_doc(tmp_path): @@ -258,16 +260,3 @@ def test_per_query_per_token_aggregates_to_per_doc(tmp_path): ) torch.testing.assert_close(agg, per_doc.to(torch.float64), atol=1e-5, rtol=1e-4) - - -def test_three_dim_scores_load_as_per_token_multi_query(tmp_path): - """A 3-D scores.pt is unambiguous: neither classifier needs the run config.""" - from bergson.magic.cli import scores_are_per_token - from bergson.validate import load_attribution_scores - - path = tmp_path / "scores.pt" - torch.save(torch.randn(4, 6, 3), path) - - assert scores_are_per_token(str(path)) - _, multi_query = load_attribution_scores(str(path)) - assert multi_query diff --git a/tests/test_per_token_lds.py b/tests/test_per_token_lds.py index f5fe33ec..2bd09340 100644 --- a/tests/test_per_token_lds.py +++ b/tests/test_per_token_lds.py @@ -16,20 +16,6 @@ from bergson.validate import load_attribution_scores -def test_load_attribution_scores_per_token_pt(tmp_path): - """A 2-D ``.pt`` tensor is treated as per-token MAGIC scores, never - multi-query -- so ``validate_scores`` keeps it 2-D for token re-weighting.""" - scores = torch.randn(6, 5) - path = tmp_path / "scores.pt" - torch.save(scores, path) - - loaded, multi_query = load_attribution_scores(str(path)) - - assert not multi_query - assert loaded.shape == (6, 5) - torch.testing.assert_close(loaded, scores) - - def test_datastream_serves_reweighted_per_token_weights(): """End of the pipeline: after a flat per-token re-weight, the DataStream serves the updated per-token ``example_weight`` for each batch, so the @@ -81,9 +67,8 @@ def test_load_token_dir_scatters_packed_to_grid(tmp_path): def test_scores_are_per_token_inference(tmp_path): - """``bergson validate`` infers per-token-ness from the scores themselves: - token dirs and 2-D ``.pt`` tensors are per-token; 1-D ``.pt`` and plain - score dirs are per-document.""" + """Per-token-ness comes from the store's info.json; a directory without + one is not a score store.""" from bergson.magic.cli import scores_are_per_token token_dir = tmp_path / "token_dir" @@ -91,18 +76,6 @@ def test_scores_are_per_token_inference(tmp_path): _write_token_dir(token_dir, ntg=[3, 2], values=[1, 2, 3, 4, 5]) assert scores_are_per_token(str(token_dir)) - pt_2d = tmp_path / "per_token.pt" - torch.save(torch.randn(6, 5), pt_2d) - assert scores_are_per_token(str(pt_2d)) - - pt_1d = tmp_path / "per_doc.pt" - torch.save(torch.randn(6), pt_1d) - assert not scores_are_per_token(str(pt_1d)) - - pt_col = tmp_path / "per_doc_col.pt" - torch.save(torch.randn(6, 1), pt_col) - assert not scores_are_per_token(str(pt_col)) - plain_dir = tmp_path / "plain_dir" plain_dir.mkdir() assert not scores_are_per_token(str(plain_dir)) @@ -120,18 +93,3 @@ def test_load_token_dir_keeps_query_dim_when_multiscore(tmp_path): torch.testing.assert_close(scores[0, 0], torch.tensor([1.0, 2.0])) torch.testing.assert_close(scores[1, 0], torch.tensor([5.0, 6.0])) torch.testing.assert_close(scores[1, 1], torch.tensor([0.0, 0.0])) - - -def test_scores_are_per_token_from_run_config(tmp_path): - """With a run config beside it, the flag decides rather than the shape -- - including ``attribute_tokens``, not just the deprecated ``per_token``.""" - import yaml - - from bergson.magic.cli import scores_are_per_token - - # A shape the fallback above would read as per-doc. - torch.save(torch.zeros(4, 1), tmp_path / "scores.pt") - cfg = {"steps": [{"magic": {"query_method": "mean", "attribute_tokens": True}}]} - (tmp_path / "config.yaml").write_text(yaml.safe_dump(cfg)) - - assert scores_are_per_token(str(tmp_path / "scores.pt")) From e6a62db0e14ed112fd59f2df95f971a740b34aed Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Thu, 6 Aug 2026 21:41:30 +0900 Subject: [PATCH 2/4] refactor: rename load_attribution_scores to load_scores_loss_signed The old name said what was loaded but not in which orientation, and the orientation is the part a caller gets wrong. Scores come back in the loss-diff convention: higher_is_better stores are negated on load, and validate correlates them against baseline - loss, so a proponent lands negative. The name now carries that, and the docstring states it outright. Co-Authored-By: Claude Opus 5 (1M context) --- bergson/magic/cli.py | 4 ++-- bergson/validate.py | 14 ++++---------- docs/magic.rst | 6 +++--- tests/test_distributed_magic.py | 14 +++++++------- tests/test_magic.py | 8 ++++---- tests/test_multi_query_validate.py | 10 +++++----- tests/test_per_query_magic.py | 4 ++-- tests/test_per_token_lds.py | 6 +++--- tests/test_score.py | 6 +++--- 9 files changed, 33 insertions(+), 39 deletions(-) diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index f837224d..052ba0d6 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -39,7 +39,7 @@ from ..utils.logging import wandb_log_fn from ..utils.utils import get_device, get_device_index from ..utils.worker_utils import setup_data_pipeline -from ..validate import load_attribution_scores, validate_scores +from ..validate import load_scores_loss_signed, validate_scores from .config import MagicConfig from .data_stream import DataStream, pad_dataset_to_batch_size from .grad_accum import accumulate_grads @@ -594,7 +594,7 @@ def worker( run_cfg.run_path, scores, train_dataset, pad_count, bool(per_token) ) else: - scores, multi_query = load_attribution_scores(score_path) + scores, multi_query = load_scores_loss_signed(score_path) stream.requires_grad = False diff --git a/bergson/validate.py b/bergson/validate.py index d071483e..ef444b55 100644 --- a/bergson/validate.py +++ b/bergson/validate.py @@ -39,15 +39,9 @@ from .utils.worker_utils import setup_data_pipeline -def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]: - """Load attribution scores from a score directory. - - Returns ``(scores, multi_query)``. Token stores are per-token, with a - query dimension when ``num_scores > 1`` (``[docs, seq_len, queries]``); - plain stores are per-document, with one column per query. Scores are - negated when ``score_cfg.higher_is_better`` is set, aligning them with - the loss-diff convention. - """ +def load_scores_loss_signed(score_path: str) -> tuple[torch.Tensor, bool]: + """Loads scores using the sign convention that negative scores reduce + query loss (proponents are negative).""" loaded = load_scores(Path(score_path)) score_cfg = load_subconfig(score_path, "score_cfg", ScoreConfig) negate = score_cfg is not None and score_cfg.higher_is_better @@ -554,7 +548,7 @@ def evaluate_retrained( subsets = [torch.tensor(s, dtype=torch.long) for s in subset_lists] # Load per-query attribution scores (mirrors run_magic's score loading). - scores, multi_query = load_attribution_scores(score_path) + scores, multi_query = load_scores_loss_signed(score_path) if not multi_query: assert ( scores.ndim == 1 or scores.shape[1] == 1 diff --git a/docs/magic.rst b/docs/magic.rst index b6ff3171..3eefbc98 100644 --- a/docs/magic.rst +++ b/docs/magic.rst @@ -35,7 +35,7 @@ After a run completes, ``run_cfg.run_path`` contains: * ``scores/`` — a score directory, the same self-describing format the scoring pipeline writes. ``info.json`` records ``attribute_tokens`` and ``num_scores``, so consumers never infer the layout from the shape. - Read it with :func:`bergson.validate.load_attribution_scores`, which + Read it with :func:`bergson.validate.load_scores_loss_signed`, which returns ``(scores, multi_query)``: * Per-example: ``(num_train_docs, 1)``, indexed directly by ``doc_id``. @@ -57,9 +57,9 @@ After a run completes, ``run_cfg.run_path`` contains: .. code-block:: python - from bergson.validate import load_attribution_scores + from bergson.validate import load_scores_loss_signed - scores, _ = load_attribution_scores("runs/magic/scores") + scores, _ = load_scores_loss_signed("runs/magic/scores") doc_ids = torch.from_numpy(np.load("runs/magic/scores/doc_ids.npy")) num_docs = int(doc_ids.max()) + 1 per_doc = torch.zeros(num_docs, dtype=scores.dtype) diff --git a/tests/test_distributed_magic.py b/tests/test_distributed_magic.py index 6efc895d..7c6a1af8 100644 --- a/tests/test_distributed_magic.py +++ b/tests/test_distributed_magic.py @@ -13,7 +13,7 @@ from bergson.config import DataConfig, DistributedConfig, LRScheduleConfig from bergson.magic.cli import MagicConfig, run_magic -from bergson.validate import load_attribution_scores +from bergson.validate import load_scores_loss_signed # Both tests consume the module-scoped noclip_scores fixture, so they must run # on the same xdist worker or each worker recomputes the two no-clip runs. @@ -82,7 +82,7 @@ def noclip_scores(tmp_path_factory) -> dict[str, torch.Tensor]: scores = {} for mode, fsdp in [("ddp", False), ("fsdp", True)]: run_magic(magic_cfg(f"{tmpdir}/{mode}", fsdp=fsdp, clip=False)) - scores[mode] = load_attribution_scores(f"{tmpdir}/{mode}/scores")[0] + scores[mode] = load_scores_loss_signed(f"{tmpdir}/{mode}/scores")[0] return scores @@ -131,8 +131,8 @@ def test_fsdp_ddp_scores_match_with_grad_clipping(noclip_scores, tmp_path): run_magic(magic_cfg(f"{tmp_path}/ddp", fsdp=False, clip=True)) run_magic(magic_cfg(f"{tmp_path}/fsdp", fsdp=True, clip=True)) - ddp_scores = load_attribution_scores(f"{tmp_path}/ddp/scores")[0] - fsdp_scores = load_attribution_scores(f"{tmp_path}/fsdp/scores")[0] + ddp_scores = load_scores_loss_signed(f"{tmp_path}/ddp/scores")[0] + fsdp_scores = load_scores_loss_signed(f"{tmp_path}/fsdp/scores")[0] assert fsdp_scores.shape == ddp_scores.shape @@ -168,8 +168,8 @@ def test_grad_accum_matches_full_batch(noclip_scores, tmp_path): run_magic(magic_cfg(f"{tmp_path}/ddp", fsdp=False, clip=False, grad_accum=2)) run_magic(magic_cfg(f"{tmp_path}/fsdp", fsdp=True, clip=False, grad_accum=2)) - ddp_scores = load_attribution_scores(f"{tmp_path}/ddp/scores")[0] - fsdp_scores = load_attribution_scores(f"{tmp_path}/fsdp/scores")[0] + ddp_scores = load_scores_loss_signed(f"{tmp_path}/ddp/scores")[0] + fsdp_scores = load_scores_loss_signed(f"{tmp_path}/fsdp/scores")[0] assert fsdp_scores.shape == ddp_noclip.shape @@ -227,7 +227,7 @@ def run(name: str, *, fsdp: bool, grad_accum: int) -> torch.Tensor: lr=1e-5, ) ) - return load_attribution_scores(f"{path}/scores")[0] + return load_scores_loss_signed(f"{path}/scores")[0] ddp1 = run("ddp1", fsdp=False, grad_accum=1) ddp2 = run("ddp2", fsdp=False, grad_accum=2) diff --git a/tests/test_magic.py b/tests/test_magic.py index eeda4191..030fb6c2 100644 --- a/tests/test_magic.py +++ b/tests/test_magic.py @@ -1360,7 +1360,7 @@ def test_worker_writes_doc_ids_for_fresh_per_token_run(tmp_path): import numpy as np - from bergson.validate import load_attribution_scores + from bergson.validate import load_scores_loss_signed score_dir = tmp_path / "scores" doc_ids_path = score_dir / "doc_ids.npy" @@ -1370,7 +1370,7 @@ def test_worker_writes_doc_ids_for_fresh_per_token_run(tmp_path): "indexed by shuffled chunk, so they are unaggregatable without it" ) - scores, _ = load_attribution_scores(str(score_dir)) + scores, _ = load_scores_loss_signed(str(score_dir)) doc_ids = torch.from_numpy(np.load(doc_ids_path)) assert scores.ndim == 2, f"expected per-token scores, got {scores.shape}" assert doc_ids.shape == scores.shape @@ -1391,7 +1391,7 @@ def test_save_magic_scores_round_trips_the_grid(tmp_path, num_scores): from datasets import Dataset from bergson.magic.cli import save_magic_scores - from bergson.validate import load_attribution_scores + from bergson.validate import load_scores_loss_signed rows, seq_len = 4, 6 data = Dataset.from_dict( @@ -1408,7 +1408,7 @@ def test_save_magic_scores_round_trips_the_grid(tmp_path, num_scores): grid[:, seq_len - 1] = 0.0 # the column weighted_causal_lm_ce never reaches save_magic_scores(str(tmp_path), grid, data, pad_count=0, per_token=True) - loaded, multi_query = load_attribution_scores(str(tmp_path / "scores")) + loaded, multi_query = load_scores_loss_signed(str(tmp_path / "scores")) assert multi_query == (num_scores > 1) assert loaded.shape == grid.shape diff --git a/tests/test_multi_query_validate.py b/tests/test_multi_query_validate.py index d8e28fd3..2576d46f 100644 --- a/tests/test_multi_query_validate.py +++ b/tests/test_multi_query_validate.py @@ -4,7 +4,7 @@ from bergson.magic.data_stream import DataStream from bergson.score.score_writer import MemmapSequenceScoreWriter -from bergson.validate import load_attribution_scores, per_doc_query_losses +from bergson.validate import load_scores_loss_signed, per_doc_query_losses def test_per_doc_query_losses_matches_hf_loss(model): @@ -47,26 +47,26 @@ def test_per_doc_query_losses_packed_rows(model): torch.testing.assert_close(losses[d], expected, rtol=1e-4, atol=1e-5) -def test_load_attribution_scores_score_dir(tmp_path): +def test_load_scores_loss_signed_score_dir(tmp_path): """Score dirs load all query columns and flag multi-query.""" writer = MemmapSequenceScoreWriter(tmp_path, num_items=6, num_scores=3) values = torch.arange(18, dtype=torch.float32).reshape(6, 3) writer(list(range(6)), values) writer.flush() - scores, multi_query = load_attribution_scores(str(tmp_path)) + scores, multi_query = load_scores_loss_signed(str(tmp_path)) assert multi_query assert scores.shape == (6, 3) # No score_cfg saved, so no higher_is_better negation is applied. torch.testing.assert_close(scores, values) -def test_load_attribution_scores_single_column_dir(tmp_path): +def test_load_scores_loss_signed_single_column_dir(tmp_path): writer = MemmapSequenceScoreWriter(tmp_path, num_items=4, num_scores=1) writer(list(range(4)), torch.ones(4, 1)) writer.flush() - scores, multi_query = load_attribution_scores(str(tmp_path)) + scores, multi_query = load_scores_loss_signed(str(tmp_path)) assert not multi_query assert scores.shape == (4, 1) diff --git a/tests/test_per_query_magic.py b/tests/test_per_query_magic.py index 1c41960c..981a69ef 100644 --- a/tests/test_per_query_magic.py +++ b/tests/test_per_query_magic.py @@ -234,9 +234,9 @@ def ds(n): import numpy as np - from bergson.validate import load_attribution_scores + from bergson.validate import load_scores_loss_signed - scores, _ = load_attribution_scores(str(run_path / "scores")) + scores, _ = load_scores_loss_signed(str(run_path / "scores")) doc_ids = run_path / "scores" / "doc_ids.npy" return scores, (torch.from_numpy(np.load(doc_ids)) if doc_ids.exists() else None) diff --git a/tests/test_per_token_lds.py b/tests/test_per_token_lds.py index 2bd09340..a649f13c 100644 --- a/tests/test_per_token_lds.py +++ b/tests/test_per_token_lds.py @@ -13,7 +13,7 @@ from datasets import Dataset from bergson.magic.data_stream import DataStream -from bergson.validate import load_attribution_scores +from bergson.validate import load_scores_loss_signed def test_datastream_serves_reweighted_per_token_weights(): @@ -58,7 +58,7 @@ def test_load_token_dir_scatters_packed_to_grid(tmp_path): ``[docs, seq_len]`` grid as MAGIC, with tokens at positions 0..len-2.""" _write_token_dir(tmp_path, ntg=[3, 2], values=[1, 2, 3, 4, 5]) - scores, multi_query = load_attribution_scores(str(tmp_path)) + scores, multi_query = load_scores_loss_signed(str(tmp_path)) assert not multi_query assert scores.shape == (2, 4) # width = max(len-1)+1 @@ -86,7 +86,7 @@ def test_load_token_dir_keeps_query_dim_when_multiscore(tmp_path): tmp_path, ntg=[2, 1], values=[[1, 2], [3, 4], [5, 6]], num_scores=2 ) - scores, multi_query = load_attribution_scores(str(tmp_path)) + scores, multi_query = load_scores_loss_signed(str(tmp_path)) assert multi_query assert scores.shape == (2, 3, 2) diff --git a/tests/test_score.py b/tests/test_score.py index 2795fd26..7997af4d 100644 --- a/tests/test_score.py +++ b/tests/test_score.py @@ -314,19 +314,19 @@ def test_memmap_score_writer_float32(tmp_path: Path): ) -def test_load_attribution_scores_bfloat16(tmp_path: Path): +def test_load_scores_loss_signed_bfloat16(tmp_path: Path): """A bf16 per-document store must be readable back. ``torch.from_numpy`` cannot ingest ``ml_dtypes.bfloat16`` directly, so the per-document branch has to cast the way the per-token one does. """ - from bergson.validate import load_attribution_scores + from bergson.validate import load_scores_loss_signed writer = MemmapSequenceScoreWriter(tmp_path, 4, 1, dtype=torch.bfloat16) writer([0, 1, 2, 3], torch.tensor([[1.0], [2.0], [3.0], [4.0]])) writer.flush() - scores, multi_query = load_attribution_scores(str(tmp_path)) + scores, multi_query = load_scores_loss_signed(str(tmp_path)) assert not multi_query assert scores.dtype == torch.float32 np.testing.assert_array_equal(scores.squeeze(-1).numpy(), [1.0, 2.0, 3.0, 4.0]) From ad09dba3a340bcec53ffd86a8a11161e5aef943d Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Thu, 6 Aug 2026 21:46:02 +0900 Subject: [PATCH 3/4] refactor: move load_scores_loss_signed into data MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit It is a thin sign-normalising wrapper over load_scores, which lives in data alongside the store format it reads — Scores, offsets, to_grid — so the two belong together. Nothing about it is validation-specific; magic imported it from validate only because that is where it happened to sit. No new dependency edge: data already imports from config.config, and the config layer does not import data, so ScoreConfig and load_subconfig come along without a cycle. Co-Authored-By: Claude Opus 5 (1M context) --- bergson/data.py | 23 ++++++++++++++++++++++- bergson/magic/cli.py | 4 ++-- bergson/validate.py | 26 +++----------------------- docs/magic.rst | 4 ++-- tests/test_distributed_magic.py | 2 +- tests/test_magic.py | 4 ++-- tests/test_multi_query_validate.py | 3 ++- tests/test_per_query_magic.py | 2 +- tests/test_per_token_lds.py | 2 +- tests/test_score.py | 2 +- 10 files changed, 37 insertions(+), 35 deletions(-) diff --git a/bergson/data.py b/bergson/data.py index caa21a58..618da13c 100644 --- a/bergson/data.py +++ b/bergson/data.py @@ -25,7 +25,8 @@ from numpy.typing import DTypeLike, NDArray from transformers import PreTrainedTokenizerFast, logging -from .config.config import DataConfig +from .config.config import DataConfig, ScoreConfig +from .config.config_io import load_subconfig from .utils.utils import ( assert_type, simple_parse_kwargs_string, @@ -662,6 +663,26 @@ def load_scores(path: Path) -> Scores: return Scores(mmap, info, offsets) +def load_scores_loss_signed(score_path: str) -> tuple[torch.Tensor, bool]: + """Loads scores using the sign convention that negative scores reduce + query loss (proponents are negative).""" + loaded = load_scores(Path(score_path)) + score_cfg = load_subconfig(score_path, "score_cfg", ScoreConfig) + negate = score_cfg is not None and score_cfg.higher_is_better + + if loaded.offsets is not None: + scores = loaded.to_grid() + else: + arr = np.asarray(loaded[:]) + # Copy: the slice is a read-only view onto the memmap. + out_dtype = arr.dtype if np.issubdtype(arr.dtype, np.floating) else np.float32 + scores = torch.from_numpy(arr.astype(out_dtype, copy=True)) + + if negate: + scores = -scores + return scores, loaded.num_scores > 1 + + def sorted_checkpoints(folder: str) -> list[tuple[int, str]]: """ Return a list of (step, filepath) sorted by step diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index 052ba0d6..e33087fb 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -30,7 +30,7 @@ from ..config.config import TrainingConfig, ValidationConfig from ..config.config_io import save_run_config -from ..data import compute_num_token_grads +from ..data import compute_num_token_grads, load_scores_loss_signed from ..distributed import launch_distributed_run from ..score.score_writer import save_sequence_scores, save_token_scores from ..utils.load_from_optimizer import ( @@ -39,7 +39,7 @@ from ..utils.logging import wandb_log_fn from ..utils.utils import get_device, get_device_index from ..utils.worker_utils import setup_data_pipeline -from ..validate import load_scores_loss_signed, validate_scores +from ..validate import validate_scores from .config import MagicConfig from .data_stream import DataStream, pad_dataset_to_batch_size from .grad_accum import accumulate_grads diff --git a/bergson/validate.py b/bergson/validate.py index ef444b55..efb8e290 100644 --- a/bergson/validate.py +++ b/bergson/validate.py @@ -29,9 +29,9 @@ set_verbosity_error as hf_set_verbosity_error, ) -from .config.config import ScoreConfig, ValidationConfig -from .config.config_io import load_subconfig, save_run_config -from .data import load_scores, pad_and_tensor +from .config.config import ValidationConfig +from .config.config_io import save_run_config +from .data import load_scores_loss_signed, pad_and_tensor from .magic.data_stream import DataStream, pad_dataset_to_batch_size from .magic.trainer import TrainerState, prepare_trainer from .utils.csv_writer import CSVWriter @@ -39,26 +39,6 @@ from .utils.worker_utils import setup_data_pipeline -def load_scores_loss_signed(score_path: str) -> tuple[torch.Tensor, bool]: - """Loads scores using the sign convention that negative scores reduce - query loss (proponents are negative).""" - loaded = load_scores(Path(score_path)) - score_cfg = load_subconfig(score_path, "score_cfg", ScoreConfig) - negate = score_cfg is not None and score_cfg.higher_is_better - - if loaded.offsets is not None: - scores = loaded.to_grid() - else: - arr = np.asarray(loaded[:]) - # Copy: the slice is a read-only view onto the memmap. - out_dtype = arr.dtype if np.issubdtype(arr.dtype, np.floating) else np.float32 - scores = torch.from_numpy(arr.astype(out_dtype, copy=True)) - - if negate: - scores = -scores - return scores, loaded.num_scores > 1 - - def bank_loss_cache_key( run_cfg: ValidationConfig, multi_query: bool, num_subsets: int ) -> str: diff --git a/docs/magic.rst b/docs/magic.rst index 3eefbc98..1de8aa6d 100644 --- a/docs/magic.rst +++ b/docs/magic.rst @@ -35,7 +35,7 @@ After a run completes, ``run_cfg.run_path`` contains: * ``scores/`` — a score directory, the same self-describing format the scoring pipeline writes. ``info.json`` records ``attribute_tokens`` and ``num_scores``, so consumers never infer the layout from the shape. - Read it with :func:`bergson.validate.load_scores_loss_signed`, which + Read it with :func:`bergson.data.load_scores_loss_signed`, which returns ``(scores, multi_query)``: * Per-example: ``(num_train_docs, 1)``, indexed directly by ``doc_id``. @@ -57,7 +57,7 @@ After a run completes, ``run_cfg.run_path`` contains: .. code-block:: python - from bergson.validate import load_scores_loss_signed + from bergson.data import load_scores_loss_signed scores, _ = load_scores_loss_signed("runs/magic/scores") doc_ids = torch.from_numpy(np.load("runs/magic/scores/doc_ids.npy")) diff --git a/tests/test_distributed_magic.py b/tests/test_distributed_magic.py index 7c6a1af8..f47aa906 100644 --- a/tests/test_distributed_magic.py +++ b/tests/test_distributed_magic.py @@ -12,8 +12,8 @@ import torch from bergson.config import DataConfig, DistributedConfig, LRScheduleConfig +from bergson.data import load_scores_loss_signed from bergson.magic.cli import MagicConfig, run_magic -from bergson.validate import load_scores_loss_signed # Both tests consume the module-scoped noclip_scores fixture, so they must run # on the same xdist worker or each worker recomputes the two no-clip runs. diff --git a/tests/test_magic.py b/tests/test_magic.py index 030fb6c2..b2d2f109 100644 --- a/tests/test_magic.py +++ b/tests/test_magic.py @@ -1360,7 +1360,7 @@ def test_worker_writes_doc_ids_for_fresh_per_token_run(tmp_path): import numpy as np - from bergson.validate import load_scores_loss_signed + from bergson.data import load_scores_loss_signed score_dir = tmp_path / "scores" doc_ids_path = score_dir / "doc_ids.npy" @@ -1390,8 +1390,8 @@ def test_save_magic_scores_round_trips_the_grid(tmp_path, num_scores): import numpy as np from datasets import Dataset + from bergson.data import load_scores_loss_signed from bergson.magic.cli import save_magic_scores - from bergson.validate import load_scores_loss_signed rows, seq_len = 4, 6 data = Dataset.from_dict( diff --git a/tests/test_multi_query_validate.py b/tests/test_multi_query_validate.py index 2576d46f..898a84b3 100644 --- a/tests/test_multi_query_validate.py +++ b/tests/test_multi_query_validate.py @@ -2,9 +2,10 @@ import torch.nn.functional as F from datasets import Dataset +from bergson.data import load_scores_loss_signed from bergson.magic.data_stream import DataStream from bergson.score.score_writer import MemmapSequenceScoreWriter -from bergson.validate import load_scores_loss_signed, per_doc_query_losses +from bergson.validate import per_doc_query_losses def test_per_doc_query_losses_matches_hf_loss(model): diff --git a/tests/test_per_query_magic.py b/tests/test_per_query_magic.py index 981a69ef..c7f927ef 100644 --- a/tests/test_per_query_magic.py +++ b/tests/test_per_query_magic.py @@ -234,7 +234,7 @@ def ds(n): import numpy as np - from bergson.validate import load_scores_loss_signed + from bergson.data import load_scores_loss_signed scores, _ = load_scores_loss_signed(str(run_path / "scores")) doc_ids = run_path / "scores" / "doc_ids.npy" diff --git a/tests/test_per_token_lds.py b/tests/test_per_token_lds.py index a649f13c..9ca6eae2 100644 --- a/tests/test_per_token_lds.py +++ b/tests/test_per_token_lds.py @@ -12,8 +12,8 @@ import torch from datasets import Dataset +from bergson.data import load_scores_loss_signed from bergson.magic.data_stream import DataStream -from bergson.validate import load_scores_loss_signed def test_datastream_serves_reweighted_per_token_weights(): diff --git a/tests/test_score.py b/tests/test_score.py index 7997af4d..94bdca8f 100644 --- a/tests/test_score.py +++ b/tests/test_score.py @@ -320,7 +320,7 @@ def test_load_scores_loss_signed_bfloat16(tmp_path: Path): ``torch.from_numpy`` cannot ingest ``ml_dtypes.bfloat16`` directly, so the per-document branch has to cast the way the per-token one does. """ - from bergson.validate import load_scores_loss_signed + from bergson.data import load_scores_loss_signed writer = MemmapSequenceScoreWriter(tmp_path, 4, 1, dtype=torch.bfloat16) writer([0, 1, 2, 3], torch.tensor([[1.0], [2.0], [3.0], [4.0]])) From ca5de0f3c97c9d8cbec175f4c167f47bfadbea92 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Thu, 6 Aug 2026 23:15:46 +0900 Subject: [PATCH 4/4] fix(magic): keep reading legacy scores.pt Deleting the .pt branch made every run finished before score directories unreadable, and not with a message: load_scores(Path("scores.pt")) opens scores.pt/info.json and raises NotADirectoryError. Restore the read path only, behind a suffix check, with the ambiguity resolution it had: a 2-D tensor is per-token unless the run config beside it records query_method: none. Marked for removal in December 2026. Co-Authored-By: Claude Opus 5 (1M context) --- bergson/data.py | 40 +++++++++++++++++++++++++++++- tests/test_multi_query_validate.py | 27 ++++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/bergson/data.py b/bergson/data.py index 618da13c..b61547fc 100644 --- a/bergson/data.py +++ b/bergson/data.py @@ -13,6 +13,7 @@ import pyarrow as pa import torch import torch.distributed as dist +import yaml from datasets import ( Dataset, DatasetDict, @@ -26,7 +27,7 @@ from transformers import PreTrainedTokenizerFast, logging from .config.config import DataConfig, ScoreConfig -from .config.config_io import load_subconfig +from .config.config_io import CONFIG_FILENAME, load_subconfig, read_config from .utils.utils import ( assert_type, simple_parse_kwargs_string, @@ -663,9 +664,46 @@ def load_scores(path: Path) -> Scores: return Scores(mmap, info, offsets) +def _load_legacy_pt_scores(score_path: str) -> tuple[torch.Tensor, bool]: + """Read a bare ``scores.pt`` written by a MAGIC run that predates score + directories. + + A 2-D tensor is ambiguous — per-token scores are ``[docs, seq_len]`` and + per-query scores are ``[docs, queries]`` — so the run config beside it + decides: per-query iff the run used ``query_method: none``. Already in the + loss-diff convention, so nothing is negated. + + TODO: Lucia Quirke remove December 2026 + """ + scores = torch.load(score_path, map_location="cpu") + if not isinstance(scores, torch.Tensor) or scores.ndim not in (2, 3): + return scores, False + if scores.ndim == 3: + return scores, True + + cfg_path = Path(score_path).parent / CONFIG_FILENAME + if not cfg_path.is_file(): + return scores, False + + try: + steps = read_config(cfg_path)["steps"] + except (ValueError, yaml.YAMLError): + return scores, False + + for step in steps: + for step_cfg in step.values(): + if isinstance(step_cfg, dict): + return scores, step_cfg.get("query_method") == "none" + return scores, False + + def load_scores_loss_signed(score_path: str) -> tuple[torch.Tensor, bool]: """Loads scores using the sign convention that negative scores reduce query loss (proponents are negative).""" + # TODO: Lucia Quirke remove December 2026 + if score_path.endswith(".pt"): + return _load_legacy_pt_scores(score_path) + loaded = load_scores(Path(score_path)) score_cfg = load_subconfig(score_path, "score_cfg", ScoreConfig) negate = score_cfg is not None and score_cfg.higher_is_better diff --git a/tests/test_multi_query_validate.py b/tests/test_multi_query_validate.py index 898a84b3..a3e889db 100644 --- a/tests/test_multi_query_validate.py +++ b/tests/test_multi_query_validate.py @@ -139,3 +139,30 @@ def test_metasmoothness_score(): # Zero movement -> defined as perfectly smooth. assert metasmoothness_score(theta0, theta0, theta0) == 1.0 + + +def test_load_scores_loss_signed_legacy_pt(tmp_path): + """Runs that predate score directories still load from ``scores.pt``. A + 2-D tensor is per-token unless the run config beside it says the run was + per-query; a 3-D tensor is unambiguous. + + TODO: Lucia Quirke remove December 2026 + """ + import yaml + + per_token = torch.randn(6, 5) + torch.save(per_token, tmp_path / "scores.pt") + + scores, multi_query = load_scores_loss_signed(str(tmp_path / "scores.pt")) + assert not multi_query + torch.testing.assert_close(scores, per_token) + + cfg = {"steps": [{"magic": {"query_method": "none"}}]} + (tmp_path / "config.yaml").write_text(yaml.safe_dump(cfg)) + _, multi_query = load_scores_loss_signed(str(tmp_path / "scores.pt")) + assert multi_query + + torch.save(torch.randn(4, 6, 3), tmp_path / "tok.pt") + scores, multi_query = load_scores_loss_signed(str(tmp_path / "tok.pt")) + assert multi_query + assert scores.shape == (4, 6, 3)