fix(magic): reject chunked query sets in per-query MAGIC

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."
+            )

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,

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