fix(magic): score per-query MAGIC by document, not by row (DataStream variant)

[luciaquirke]

Alternative to #414 — same bug, same semantics, fixed inside DataStream instead of by rebuilding a dataset per query. Open both, merge one.

The bug (identical in both)

Per-query MAGIC crashes on a chunked query set (--query.chunk_length 32):

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

A chunked query set carries a per-token doc_ids column: a row can pack several documents and a document can span several rows, so query i is not row i. compute_per_query_magic_scores selected row qi and sized the per-query weights by row, while DataStream indexes weights by document id whenever doc_ids is present.

Repro:

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

How this differs from #414

DataStream is already the only thing in the codebase that interprets doc_ids, so the document logic goes there:

• doc_rows(dataset, num_docs) — which rows hold each document's tokens (moved here from cli.py).
• DataStream(..., rows=..., doc_id=...) — rows restricts the stream to those dataset rows, cycling them to fill whole batches; doc_id masks every other document's tokens out of the labels in __getitem__, three lines where pad_and_tensor has just produced them.

shift_loss_mask is rebuilt alongside the labels. That's load-bearing: it is the denominator in weighted_causal_lm_ce's mean reduction (bergson/utils/math.py), so leaving it stale would divide each document's loss by the whole row's token count and scale every score column differently. The same fact rules out the tempting "just use the existing 2-D per-token weights as the mask" approach — weights affect the numerator only.

Versus #414, which builds a masked mini-dataset per query (select → remove_columns → add_column):

• query_doc_batch disappears; the call site is one DataStream(...) construction
• no Arrow copy per query, and no labels-column collision to guard against
• batch_rows slices self.rows instead of computing a range — a no-op when rows defaults to every row, and its only external caller is per_doc_query_losses

Cost: DataStream is shared by every training and backward path, so the blast radius is wider than #414's, which is confined to the per-query function.

Shared with #414

Padding cycles through the document's own rows instead of appending dead ones: compute_query_gradients discards the data weights that would silence a pad row, so every row is a real term in the loss, and an all-pad batch on some rank would scale the query gradient by 1/world_size. With one row per document this repeats that row exactly as pad_dataset_to_batch_size did, so unchunked runs are unchanged.

Documents that chunking dropped entirely (blank wikitext lines) score zeros, matching the zero baseline loss per_doc_query_losses gives them; their correlation is undefined, so the reported mean is now Mean Spearman 1.0000 (9/20 nan) instead of nan.

Testing

• The repro above runs end to end — 20 query documents → 20 score columns, the 9 zero columns are exactly the empty documents — with --noskip_validation and via bergson validate afterwards, producing byte-identical scores to fix(magic): score per-query MAGIC by document, not by row #414.
• New tests: document↔row mapping (packed, split, dropped, pad rows), a doc-restricted stream's cycled rows + masked labels + the shift mask tracking them, and packed documents scoring the same as the equivalent unpacked pair with a dropped document scoring zeros.
• pytest tests/test_per_query_magic.py green; the wider sweep over everything that touches DataStream (test_magic, test_ddp, test_attribute_tokens, test_multi_query_validate, test_per_token_lds) was still running at PR time — I'll report in a comment. No GPU on this machine, so multi-GPU tests are unrun.

🤖 Generated with Claude Code

[luciaquirke]

Test sweep from the PR description, now finished — all green:

• pytest tests/test_per_query_magic.py tests/test_magic.py tests/test_ddp.py tests/test_attribute_tokens.py tests/test_multi_query_validate.py tests/test_per_token_lds.py → 63 passed, 2 skipped (the paths DataStream is shared by: training, DDP, per-token attribution, leave-k-out validation). No GPU on this machine, so multi-GPU tests are unrun.
• CI green, including the test job.

Also verified the "identical scores to #414" claim directly rather than by eye. Same job on both branches (pythia-14m, wikitext-2 chunked at 32, 20 query documents, --noskip_validation), then torch.equal on the two scores.pt:

--- 417 (DataStream) ---
[per-query MAGIC] 9 queries have no tokens
Score summary: mean=-3.662985e-06, variance=0.00020203233, minmax=(-0.15384506, 0.20646359)
Mean Spearman 1.0000 (9/20 nan)
--- 414 (dataset surgery) ---
[per-query MAGIC] 9 queries have no tokens
Score summary: mean=-3.662985e-06, variance=0.00020203233, minmax=(-0.15384506, 0.20646359)
Mean Spearman 1.0000 (9/20 nan)
shapes (96, 20) (96, 20)
identical: True

So the choice between this and #414 is purely about where the document logic lives, not about behaviour.