diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py index c5f75095..3bdcc130 100644 --- a/bergson/magic/cli.py +++ b/bergson/magic/cli.py @@ -285,6 +285,18 @@ 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"]) + if 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 + + def shuffled_epochs(dataset: Dataset, seed: int, num_epochs: int) -> Dataset: """Concatenate `num_epochs` independently shuffled copies of `dataset`. @@ -562,16 +574,8 @@ def worker( else: scores = torch.load(score_path, map_location="cpu") - # Per-token scores are indexed by (shuffled_chunk_idx, token_idx). - # Save doc_ids alongside so downstream can aggregate per-doc with - # one scatter_add and no reference to the raw dataset or seed. - if scores.ndim == 2: - doc_ids = torch.tensor(train_dataset["doc_ids"]) - if pad_count: - doc_ids = doc_ids[:-pad_count] - doc_ids_path = os.path.join(run_cfg.run_path, "doc_ids.pt") - torch.save(doc_ids, doc_ids_path) - print(f"Saved doc_ids to {doc_ids_path}") + if per_token and global_rank == 0: + save_doc_ids(run_cfg.run_path, train_dataset, pad_count) stream.requires_grad = False diff --git a/tests/test_magic.py b/tests/test_magic.py index 783006d2..cffa807b 100644 --- a/tests/test_magic.py +++ b/tests/test_magic.py @@ -1309,3 +1309,70 @@ def test_weighted_ce_preserves_fp64(): ).dtype == torch.float32 ) + + +def _tiny_magic_dataset(num_docs: int, seq_len: int): + """A dataset shaped the way run_magic hands it to worker().""" + from datasets import Dataset + + return Dataset.from_dict( + { + "input_ids": [ + [(d * seq_len + t) % 50 + 1 for t in range(seq_len)] + for d in range(num_docs) + ], + "labels": [ + [(d * seq_len + t) % 50 + 1 for t in range(seq_len)] + for d in range(num_docs) + ], + "doc_ids": [[d] * seq_len for d in range(num_docs)], + "length": [seq_len] * num_docs, + } + ) + + +def test_worker_writes_doc_ids_for_fresh_per_token_run(tmp_path): + """A per-token ``bergson magic`` run must write doc_ids.pt next to scores.pt. + + Calls worker() rather than reimplementing it, so it covers what actually + lands on disk. + """ + from bergson.config.config import DataConfig + from bergson.magic.cli import worker + from bergson.magic.config import MagicConfig + + num_docs, seq_len = 4, 8 + train_ds = _tiny_magic_dataset(num_docs, seq_len) + query_ds = _tiny_magic_dataset(2, seq_len) + + run_cfg = MagicConfig( + run_path=str(tmp_path), + model="EleutherAI/pythia-14m", + data=DataConfig(dataset="unused", chunk_length=seq_len), + query=DataConfig(dataset="unused", chunk_length=seq_len), + batch_size=2, + attribute_tokens=True, + query_method="mean", + skip_validation=True, + ) + + 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" + assert doc_ids_path.is_file(), ( + "worker() wrote scores.pt but no doc_ids.pt; 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) + assert scores.ndim == 2, f"expected per-token scores, got {scores.shape}" + assert doc_ids.shape == scores.shape + + agg = torch.zeros(num_docs, dtype=torch.float64) + agg.scatter_add_(0, doc_ids.reshape(-1), scores.reshape(-1).to(torch.float64)) + torch.testing.assert_close( + agg.sum(), scores.sum().to(torch.float64), atol=1e-6, rtol=1e-5 + )