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))