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

[luciaquirke]

Third take on the bug behind #414 and #417, and the smallest one: 6 lines of source instead of ~70. Don't support chunked query sets in per-query MAGIC at all.

The bug

bergson magic runs/x --model EleutherAI/pythia-14m \
    --data.dataset Salesforce/wikitext --data.subset wikitext-2-raw-v1 \
    --data.split "train[:96]" --data.chunk_length 32 \
    --query.dataset Salesforce/wikitext --query.subset wikitext-2-raw-v1 \
    --query.split "train[300:320]" --query.chunk_length 32 --batch_size 2

IndexError: index 2 is out of bounds for dimension 0 with size 2
  bergson/magic/data_stream.py:132   "example_weight": self.weights[indices]

Per-query MAGIC scores one column per query document, but a chunked query set has rows that pack several documents and documents that span several rows, so query i is not row i. The per-query stream sizes its weights by row while DataStream indexes them by document id whenever the batch carries doc_ids.

Why reject instead of support

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 splitting a document across rows only makes the per-query path reconstruct what chunking took apart.

The check is at config time, so it fires in ~4s on the repro above rather than after a run has trained for hours:

ValueError: query.chunk_length must be 0 for per-query MAGIC (query_method='none');
use query.truncation for long documents.

It's scoped to per-query mode. The aggregate-query backward (query_method: mean / sum) weights the query stream by document id consistently and is fine with chunked queries, so it's left alone — the test asserts both halves.

What it costs

Every shipped config already passes query.chunk_length: 0 (examples/magic/*.yaml, examples/compare_wikitext/*.yaml, examples/replicate_bae_approx_unrolling_source/*.yaml), so no run configuration changes.

One test did rely on the rejected combination: test_distributed_magic.py built its query from raw wikitext with chunk_length=32 under the default query_method="none", so the guard would have made every test in that file raise at config construction. Its query now comes from EleutherAI/bergson-wikitext-512-chunks at chunk_length 0 — the same pre-chunked dataset the example configs use, where a row is already a document. One ~650-token document either way, and the tests only assert FSDP-vs-DDP agreement, so the query content is immaterial. That file is @requires_multi_gpu and I have no GPU here, so it wants a run on a GPU box before merge.

Compared to the alternatives

	#414 dataset surgery	#417 DataStream	this
non-test lines	+72 −11	+86 −27	+6
chunked query sets	supported	supported	rejected with a clear error

#414 rebuilds a masked mini-dataset per query; #417 teaches DataStream about documents (rows= / doc_id=). Both produce bit-identical scores and both work; this one deletes the requirement instead of meeting it. Close whichever two lose.

Residual worth knowing

The check keys on query.chunk_length, which is the only way our pipeline puts a doc_ids column on a query set (setup_data_pipeline only chunks; attach_doc_ids_if_missing is train-side). A hand-built query dataset that ships its own doc_ids column would slip past it and hit the original IndexError. Happy to add a one-line assert on the per-query path too, but it can only fire after training, which is why the config check is the useful one.

Testing

• Repro command above now fails in 4.4s with the message, before downloading data or training.
• The supported path still runs end to end (pythia-14m, pre-chunked wikitext, --noskip_validation): scores written, leave-k-out validation completes.
• pytest tests/test_per_query_magic.py tests/test_config_runner.py tests/test_distributed_magic.py → 21 passed, 5 skipped (the skips are the multi-GPU tests).
• The per-token per-query fixture from feat(magic): per-token per-query MAGIC, and score-format cleanups #415 built its MagicConfig with a chunked query config while handing worker() pre-tokenized datasets directly; its query config is now chunk_length 0, which is what its data already was.

🤖 Generated with Claude Code