From 41bb037433a734704edae006dfdf36f78bbfbdc4 Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Thu, 6 Aug 2026 22:18:48 +0900 Subject: [PATCH] fix(magic): reject chunked query sets in per-query MAGIC MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-query MAGIC scores one column per query *document*, but a chunked query set (`query.chunk_length > 0`) has rows that pack several documents and documents that span several rows, so query `i` is not row `i`. `compute_per_query_magic_scores` selected row `qi` and sized the stream's weights by row, while `DataStream` indexes weights by document id whenever the batch carries `doc_ids`, so a run crashed with: IndexError: index 2 is out of bounds for dimension 0 with size 2 bergson/magic/data_stream.py:132 self.weights[indices] Chunking exists to pack a training set efficiently; a query set is small — 50 documents in the compare_wikitext runs — so packing it buys nothing, and every shipped config already passes query.chunk_length 0. Rather than teach the per-query path to split and repack documents, require the query rows to be documents and say so at config time, before a run trains for hours. Co-Authored-By: Claude Opus 5 (1M context) --- bergson/magic/config.py | 6 ++++++ tests/test_distributed_magic.py | 10 +++++----- tests/test_per_query_magic.py | 15 ++++++++++++++- 3 files changed, 25 insertions(+), 6 deletions(-) diff --git a/bergson/magic/config.py b/bergson/magic/config.py index 930d4d49..2b238609 100644 --- a/bergson/magic/config.py +++ b/bergson/magic/config.py @@ -28,3 +28,9 @@ def __post_init__(self): super().__post_init__() if self.per_token: self.attribute_tokens = True + # Per-query MAGIC needs one document per row. + if self.query_method == "none" and self.query.chunk_length > 0: + raise ValueError( + "query.chunk_length must be 0 for per-query MAGIC " + "(query_method='none'); use query.truncation for long documents." + ) diff --git a/tests/test_distributed_magic.py b/tests/test_distributed_magic.py index b736d309..0e386a81 100644 --- a/tests/test_distributed_magic.py +++ b/tests/test_distributed_magic.py @@ -46,12 +46,12 @@ def magic_cfg( split="train[:512]", chunk_length=32, ) - # Single query doc: query_method="none" runs one backward per query. + # Single query doc: query_method="none" runs one backward per query, which + # needs one document per row — so query the pre-chunked dataset the example + # configs use, where a row is a document already. query = DataConfig( - dataset="Salesforce/wikitext", - subset="wikitext-2-raw-v1", - split="train[9:10]", - chunk_length=32, + dataset="EleutherAI/bergson-wikitext-512-chunks", + split="test[0:1]", ) return MagicConfig( run_path=run_path, diff --git a/tests/test_per_query_magic.py b/tests/test_per_query_magic.py index a8910f8e..55d686c7 100644 --- a/tests/test_per_query_magic.py +++ b/tests/test_per_query_magic.py @@ -224,7 +224,7 @@ def ds(n): run_path=str(run_path), model="EleutherAI/pythia-14m", data=DataConfig(dataset="unused", chunk_length=seq_len), - query=DataConfig(dataset="unused", chunk_length=seq_len), + query=DataConfig(dataset="unused"), batch_size=4, attribute_tokens=attribute_tokens, query_method="none", @@ -271,3 +271,16 @@ def test_three_dim_scores_load_as_per_token_multi_query(tmp_path): assert scores_are_per_token(str(path)) _, multi_query = load_attribution_scores(str(path)) assert multi_query + + +def test_chunked_query_set_rejected_for_per_query(): + """Rejected at config time, before a run trains for hours.""" + from bergson.config.config import DataConfig + + with pytest.raises(ValueError, match="query.chunk_length must be 0"): + MagicConfig( + run_path="x", query_method="none", query=DataConfig(chunk_length=32) + ) + + # Chunked query sets are fine for the aggregate-query backward. + MagicConfig(run_path="x", query_method="mean", query=DataConfig(chunk_length=32))